Update src/frontend/visualizations.py

src/frontend/visualizations.py

@@ -5,17 +5,23 @@ import matplotlib.pyplot as plt
 import numpy as np
 from streamlit_agraph import agraph, Node, Edge, Config
 
-def plot_compatibility(plants, compatibility_matrix, is_mini=False):
 
+def plot_compatibility(plants, compatibility_matrix, is_mini=False):
     # Create the graph
     G = nx.Graph()
     G.add_nodes_from(plants)
     for i in range(len(plants)):
         for j in range(i + 1, len(plants)):
             if compatibility_matrix[i][j] == 0:
-                G.add_edge(plants[i], plants[j], color='dimgrey')
+                G.add_edge(plants[i], plants[j], color="dimgrey")
             else:
-                G.add_edge(plants[i], plants[j], color='green' if compatibility_matrix[i][j] == 1 else 'mediumvioletred')
+                G.add_edge(
+                    plants[i],
+                    plants[j],
+                    color="green"
+                    if compatibility_matrix[i][j] == 1
+                    else "mediumvioletred",
+                )
 
     # Generate positions for the nodes
     pos = nx.spring_layout(G)
@@ -25,31 +31,27 @@ def plot_compatibility(plants, compatibility_matrix, is_mini=False):
         x=[pos[node][0] for node in G.nodes()],
         y=[pos[node][1] for node in G.nodes()],
         text=list(G.nodes()),
-        mode='markers+text',
-        textposition='top center',
-        hoverinfo='text',
+        mode="markers+text",
+        textposition="top center",
+        hoverinfo="text",
         marker=dict(
             size=40,
-            color='lightblue',
+            color="lightblue",
             line_width=2,
-        )
+        ),
     )
 
     # Create edge trace
     edge_trace = go.Scatter(
-        x=[],
-        y=[],
-        line=dict(width=1, color='dimgrey'),
-        hoverinfo='none',
-        mode='lines'
+        x=[], y=[], line=dict(width=1, color="dimgrey"), hoverinfo="none", mode="lines"
     )
 
     # Add coordinates to edge trace
     for edge in G.edges():
         x0, y0 = pos[edge[0]]
         x1, y1 = pos[edge[1]]
-        edge_trace['x'] += tuple([x0, x1, None])
-        edge_trace['y'] += tuple([y0, y1, None])
+        edge_trace["x"] += tuple([x0, x1, None])
+        edge_trace["y"] += tuple([y0, y1, None])
 
     # Create edge traces for colored edges
     edge_traces = []
@@ -57,13 +59,13 @@ def plot_compatibility(plants, compatibility_matrix, is_mini=False):
     for edge in G.edges(data=True):
         x0, y0 = pos[edge[0]]
         x1, y1 = pos[edge[1]]
-        color = edge[2]['color']
+        color = edge[2]["color"]
         trace = go.Scatter(
             x=[x0, x1],
             y=[y0, y1],
-            mode='lines',
+            mode="lines",
             line=dict(width=2, color=color),
-            hoverinfo='none'
+            hoverinfo="none",
         )
         edge_traces.append(trace)
         edge_legend.add(color)  # Add edge color to the set
@@ -71,57 +73,55 @@ def plot_compatibility(plants, compatibility_matrix, is_mini=False):
     # Create layout
     layout = go.Layout(
         showlegend=False,
-        hovermode='closest',
+        hovermode="closest",
         margin=dict(b=20, l=5, r=5, t=40),
         xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
+        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
     )
 
     # Create figure
     fig = go.Figure(data=[edge_trace, *edge_traces, node_trace], layout=layout)
 
-
     # Create custom legend for edge colors
     custom_legend = []
-    legend_names = ['Neutral', 'Negative', 'Positive']
-    legend_colors = ['dimgrey', 'mediumvioletred', 'green']
+    legend_names = ["Neutral", "Negative", "Positive"]
+    legend_colors = ["dimgrey", "mediumvioletred", "green"]
 
     for name, color in zip(legend_names, legend_colors):
         custom_legend.append(
             go.Scatter(
                 x=[None],
                 y=[None],
-                mode='markers',
+                mode="markers",
                 marker=dict(color=color),
-                name=f'{name}',
+                name=f"{name}",
                 showlegend=True,
-                hoverinfo='none'
+                hoverinfo="none",
             )
         )
     if is_mini == False:
         # Create layout for custom legend figure
         legend_layout = go.Layout(
-            title='Plant Compatibility Network Graph',
+            title="Plant Compatibility Network Graph",
             showlegend=True,
             margin=dict(b=1, t=100),
             xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
             yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
             height=120,
             legend=dict(
-                title='Edge Colors',
-                orientation='h',
+                title="Edge Colors",
+                orientation="h",
                 x=-1,
                 y=1.1,
-                bgcolor='rgba(0,0,0,0)'
-            )
+                bgcolor="rgba(0,0,0,0)",
+            ),
         )
     else:
         fig.update_layout(
-        autosize=False,
-        width=300,
-        height=300,)
-
-
+            autosize=False,
+            width=300,
+            height=300,
+        )
 
     if is_mini == False:
         # Create figure for custom legend
@@ -129,21 +129,22 @@ def plot_compatibility(plants, compatibility_matrix, is_mini=False):
         # Render the custom legend using Plotly in Streamlit
         st.plotly_chart(legend_fig, use_container_width=True)
 
-    
     # Render the graph using Plotly in Streamlit
     st.plotly_chart(fig)
 
 
-
 # this is not used as it needs to be refactored and is not working as intended
 def show_plant_tips():
     tips_string = st.session_state.plant_care_tips
 
     tips_list = tips_string.split("\n")
     num_tips = len(tips_list)
-    st.markdown("## Plant Care Tips for your plants: " + str(st.session_state.input_plants_raw) + "\n\n" + st.session_state.plant_care_tips)
-   
-
+    st.markdown(
+        "## Plant Care Tips for your plants: "
+        + str(st.session_state.input_plants_raw)
+        + "\n\n"
+        + st.session_state.plant_care_tips
+    )
 
 
 def visualize_groupings_sankey():
@@ -169,33 +170,34 @@ def visualize_groupings_sankey():
                     compatibility = compatibility_matrix[species1_index][species2_index]
 
                     if compatibility == 1:
-                        color = 'green'
+                        color = "green"
                     elif compatibility == -1:
-                        color = 'pink'
+                        color = "pink"
                     else:
-                        color = 'grey'
+                        color = "grey"
 
-                    links.append(dict(source=j, target=k, value=compatibility, color=color))
+                    links.append(
+                        dict(source=j, target=k, value=compatibility, color=color)
+                    )
 
         # Create the Sankey diagram
-        fig = go.Figure(data=[go.Sankey(
-            node=dict(
-                label=nodes,
-                color="lightblue"
-            ),
-            link=dict(
-                source=[link['source'] for link in links],
-                target=[link['target'] for link in links],
-                value=[link['value'] for link in links],
-                color=[link['color'] for link in links]
-            )
-        )])
+        fig = go.Figure(
+            data=[
+                go.Sankey(
+                    node=dict(label=nodes, color="lightblue"),
+                    link=dict(
+                        source=[link["source"] for link in links],
+                        target=[link["target"] for link in links],
+                        value=[link["value"] for link in links],
+                        color=[link["color"] for link in links],
+                    ),
+                )
+            ]
+        )
 
         # Set the layout properties
         layout = go.Layout(
-            plot_bgcolor='black',
-            paper_bgcolor='black',
-            title_font=dict(color='white')
+            plot_bgcolor="black", paper_bgcolor="black", title_font=dict(color="white")
         )
 
         # Set the figure layout
@@ -209,7 +211,7 @@ def visualize_groupings():
     groupings = st.session_state.grouping
     compatibility_matrix = st.session_state.extracted_mat
     plant_list = st.session_state.input_plants_raw
-    
+
     def generate_grouping_matrices(groupings, compatibility_matrix, plant_list):
         grouping_matrices = []
         for grouping in groupings:
@@ -217,17 +219,20 @@ def visualize_groupings():
             submatrix = [[compatibility_matrix[i][j] for j in indices] for i in indices]
             grouping_matrices.append(submatrix)
         return grouping_matrices
-    
-    grouping_matrices = generate_grouping_matrices(groupings, compatibility_matrix, plant_list)
+
+    grouping_matrices = generate_grouping_matrices(
+        groupings, compatibility_matrix, plant_list
+    )
     for i, submatrix in enumerate(grouping_matrices):
-        col1, col2= st.columns([1,3])
-        with col1: 
+        col1, col2 = st.columns([1, 3])
+        with col1:
             st.write(f"Plant Bed {i + 1}")
             st.write("Plant List")
             st.write(groupings[i])
         with col2:
-            plot_compatibility_with_agraph(groupings[i], st.session_state.full_mat, is_mini=True)
-
+            plot_compatibility_with_agraph(
+                groupings[i], st.session_state.full_mat, is_mini=True
+            )
 
 
 def plot_compatibility_with_agraph(plants, compatibility_matrix, is_mini=False):
@@ -245,18 +250,22 @@ def plot_compatibility_with_agraph(plants, compatibility_matrix, is_mini=False):
     size_n = 32 if not is_mini else 24
     # Create nodes with images
     for plant in plants:
-        nodes.append(Node(id=plant, 
-                          label=plant, 
-                            # make text bigger
-                            font={'size': 20},
-                            # spread nodes out
-                            scaling={'label': {'enabled': True}},
-                          size=size_n, 
-                          shape="circularImage",
-                          image=get_image_url(plant)))
+        nodes.append(
+            Node(
+                id=plant,
+                label=plant,
+                # make text bigger
+                font={"size": 20},
+                # spread nodes out
+                scaling={"label": {"enabled": True}},
+                size=size_n,
+                shape="circularImage",
+                image=get_image_url(plant),
+            )
+        )
 
     # Create edges based on compatibility
-    #for i in range(len(st.session_state.plant_list)):
+    # for i in range(len(st.session_state.plant_list)):
     # loop through all plants in raw long list and find the index of the plant in the plant list to get relevant metadata. skip if we are looking at the same plant
     for i, i_p in enumerate(st.session_state.plant_list):
         for j, j_p in enumerate(st.session_state.plant_list):
@@ -268,75 +277,103 @@ def plot_compatibility_with_agraph(plants, compatibility_matrix, is_mini=False):
                 else:
                     length_e = 150
 
-                if i_p in st.session_state.input_plants_raw and j_p in st.session_state.input_plants_raw:
+                if (
+                    i_p in st.session_state.input_plants_raw
+                    and j_p in st.session_state.input_plants_raw
+                ):
                     # use the compatibility matrix and the plant to index mapping to determine the color of the edge
                     if compatibility_matrix[i][j] == 1:
-                        color = 'green'
-                        edges.append(Edge(source=i_p, target=j_p,width = 3.5, type="CURVE_SMOOTH", color=color, length=length_e))
-                        print(i,j,i_p,j_p,color)
+                        color = "green"
+                        edges.append(
+                            Edge(
+                                source=i_p,
+                                target=j_p,
+                                width=3.5,
+                                type="CURVE_SMOOTH",
+                                color=color,
+                                length=length_e,
+                            )
+                        )
+                        print(i, j, i_p, j_p, color)
                     elif compatibility_matrix[i][j] == -1:
-                        color = 'mediumvioletred'
-                        edges.append(Edge(source=i_p, target=j_p,width = 3.5, type="CURVE_SMOOTH", color=color, length=length_e))
-                        print(i,j,i_p,j_p,color)
+                        color = "mediumvioletred"
+                        edges.append(
+                            Edge(
+                                source=i_p,
+                                target=j_p,
+                                width=3.5,
+                                type="CURVE_SMOOTH",
+                                color=color,
+                                length=length_e,
+                            )
+                        )
+                        print(i, j, i_p, j_p, color)
 
                     else:
-                        color = 'dimgrey'
-                        edges.append(Edge(source=i_p, target=j_p,width = .2, type="CURVE_SMOOTH", color=color, length=length_e))
-                        print(i,j,i_p,j_p,color)
+                        color = "dimgrey"
+                        edges.append(
+                            Edge(
+                                source=i_p,
+                                target=j_p,
+                                width=0.2,
+                                type="CURVE_SMOOTH",
+                                color=color,
+                                length=length_e,
+                            )
+                        )
+                        print(i, j, i_p, j_p, color)
 
-
-        
     # Configuration for the graph
-    config = Config(width=650 if not is_mini else 400,
-                    height=400 if not is_mini else 400,
-                    directed=False, 
-                    physics=True, 
-                    hierarchical=False,
-                    nodeHighlightBehavior=True,
-                    highlightColor="#F7A7A6",
-                    collapsible=True,
-                    maxZoom=5,
-                    minZoom=0.2,
-                    initialZoom=4,
-                    )
-
+    config = Config(
+        width=650 if not is_mini else 400,
+        height=400 if not is_mini else 400,
+        directed=False,
+        physics=True,
+        hierarchical=False,
+        nodeHighlightBehavior=True,
+        highlightColor="#F7A7A6",
+        collapsible=True,
+        maxZoom=5,
+        minZoom=0.2,
+        initialZoom=4,
+    )
 
     # Handling for non-mini version
     if not is_mini:
         # Create custom legend for edge colors at the top of the page
         custom_legend = []
-        legend_names = ['Neutral', 'Negative', 'Positive']
-        legend_colors = ['dimgrey', 'mediumvioletred', 'green']
+        legend_names = ["Neutral", "Negative", "Positive"]
+        legend_colors = ["dimgrey", "mediumvioletred", "green"]
 
         for name, color in zip(legend_names, legend_colors):
             custom_legend.append(
                 go.Scatter(
                     x=[None],
                     y=[None],
-                    mode='markers',
+                    mode="markers",
                     marker=dict(color=color),
                     name=name,
                     showlegend=True,
-                    hoverinfo='none'
+                    hoverinfo="none",
                 )
             )
 
         # Create layout for custom legend figure
         legend_layout = go.Layout(
-            title='Plant Compatibility Network Graph',
+            title="Plant Compatibility Network Graph",
             showlegend=True,
             margin=dict(b=1, t=100),
             xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
             yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
             height=120,
             legend=dict(
-                title='Edge Colors',
-                orientation='h',
+                title="Edge Colors",
+                orientation="h",
                 # make it appear above the graph
                 x=-1,
                 y=1.1,
-                bgcolor='rgba(0,0,0,0)'
-            )
+                bgcolor="rgba(0,0,0,0)",
+            ),
         )
 
         # Create figure for custom legend
@@ -345,10 +382,7 @@ def plot_compatibility_with_agraph(plants, compatibility_matrix, is_mini=False):
         # Render the custom legend using Plotly in Streamlit
         st.plotly_chart(legend_fig, use_container_width=True)
 
-
     # Render the graph using streamlit-agraph
-    return_value = agraph(nodes=nodes, 
-                        edges=edges, 
-                        config=config)
-    
+    return_value = agraph(nodes=nodes, edges=edges, config=config)
+