Persist query

pages/input.py

@@ -37,20 +37,37 @@ with st.spinner('Loading knowledge graph...'):
     if not allow_reverse_edges:
        edge_types = edge_types[edge_types.direction == 'forward']
 
+# If query is not in session state, initialize it
+if "query" not in st.session_state:
+    source_node_type_index = 0
+    source_node_index = 0
+    target_node_type_index = 0
+    relation_index = 0
+else:
+    source_node_type_index = st.session_state.query_options['source_node_type'].index(st.session_state.query['source_node_type'])
+    source_node_index = st.session_state.query_options['source_node'].index(st.session_state.query['source_node'])
+    target_node_type_index = st.session_state.query_options['target_node_type'].index(st.session_state.query['target_node_type'])
+    relation_index = st.session_state.query_options['relation'].index(st.session_state.query['relation'])
+
 # Select source node type
-source_node_type = st.selectbox("Source Node Type", node_types['node_type'],
-                                format_func = lambda x: x.replace("_", " "))
+source_node_type_options = node_types['node_type']
+source_node_type = st.selectbox("Source Node Type", source_node_type_options,
+                                format_func = lambda x: x.replace("_", " "), index = source_node_type_index)
 
 # Select source node
-source_node = st.selectbox("Source Node", kg_nodes[kg_nodes['node_type'] == source_node_type]['node_name'])
+source_node_options = kg_nodes[kg_nodes['node_type'] == source_node_type]['node_name']
+source_node = st.selectbox("Source Node", source_node_options,
+                           index = source_node_index)
 
 # Select target node type
-target_node_type = st.selectbox("Target Node Type", edge_types[edge_types.x_type == source_node_type].y_type.unique(),
-                                format_func = lambda x: x.replace("_", " "))
+target_node_type_options = edge_types[edge_types.x_type == source_node_type].y_type.unique()
+target_node_type = st.selectbox("Target Node Type", target_node_type_options,
+                                format_func = lambda x: x.replace("_", " "), index = target_node_type_index)
 
 # Select relation
-relation = st.selectbox("Edge Type", edge_types[(edge_types.x_type == source_node_type) & (edge_types.y_type == target_node_type)].relation.unique(),
-                        format_func = lambda x: x.replace("_", "-"))
+relation_options = edge_types[(edge_types.x_type == source_node_type) & (edge_types.y_type == target_node_type)].relation.unique()
+relation = st.selectbox("Edge Type", relation_options,
+                        format_func = lambda x: x.replace("_", "-"), index = relation_index)
 
 # Button to submit query
 if st.button("Submit Query"):
@@ -63,6 +80,14 @@ if st.button("Submit Query"):
         "relation": relation
     }
 
+    # Save query options to session state
+    st.session_state.query_options = {
+        "source_node_type": list(source_node_type_options),
+        "source_node": list(source_node_options),
+        "target_node_type": list(target_node_type_options),
+        "relation": list(relation_options)
+    }
+
     # # Write query to console
     # st.write("Current Query:")
     # st.write(st.session_state.query)

pages/validate.py

@@ -43,6 +43,9 @@ predictions = st.session_state.predictions
 kg_nodes = load_kg()
 kg_edges = load_kg_edges()
 
+# Convert tuple to hex
+def rgba_to_hex(rgba):
+    return mcolors.to_hex(rgba[:3])
 
 with st.spinner('Searching known relationships...'):
 
@@ -59,48 +62,53 @@ with st.spinner('Searching known relationships...'):
     edges_in_kg = edges_in_kg[['relation'] + [col for col in edges_in_kg.columns if col != 'relation']]
     edges_in_kg = edges_in_kg.rename(columns = {'relation': 'Known Relation'})
 
+# If there exist edges in KG
+if len(edges_in_kg) > 0:
 
-with st.spinner('Plotting known relationships...'):
-
-    # Define a color map for different relations
-    color_map = plt.get_cmap('tab10')
-
-    # Group by relation and create separate plots
-    relations = edges_in_kg['Known Relation'].unique()
-    for idx, relation in enumerate(relations):
-        fig, ax = plt.subplots(figsize=(10, 3))
-        ax.plot(predictions['Rank'], predictions['Score'])
-        ax.set_xlabel('Rank', fontsize=12)
-        ax.set_ylabel('Score', fontsize=12)
-        ax.set_xlim(1, predictions['Rank'].max())
-
-        relation_data = edges_in_kg[edges_in_kg['Known Relation'] == relation]
-        color = color_map(idx % color_map.N)  # Get a color from the color map
-        
-        for i, node in relation_data.iterrows():
-            ax.axvline(node['Rank'], color=color, linestyle='--', label=node['Name'])
-            # ax.text(node['Rank'] + 100, node['Score'], node['Name'], fontsize=10, color=color)
-        
-        # ax.set_title(f'{relation.replace("_", "-")}')
-        # ax.legend()
-
-        # Convert tuple to hex
-        def rgba_to_hex(rgba):
-            return mcolors.to_hex(rgba[:3])
-        color_hex = rgba_to_hex(color)
-
-        # Write header in color of relation
-        st.markdown(f"<h3 style='color:{color_hex}'>{relation.replace('_', ' ').title()}</h2>", unsafe_allow_html=True)
-
-        # Show plot
-        st.pyplot(fig)
-
-        # Drop known relation column
-        relation_data = relation_data.drop(columns = 'Known Relation')
-        if target_node_type not in ['disease', 'anatomy']:
-            st.dataframe(relation_data, use_container_width=True,
-                         column_config={"Database": st.column_config.LinkColumn(width = "small",
-                                                                                help = "Click to visit external database.",
-                                                                                display_text = st.session_state.display_database)})
-        else:
-            st.dataframe(relation_data, use_container_width=True)
+    with st.spinner('Plotting known relationships...'):
+
+        # Define a color map for different relations
+        color_map = plt.get_cmap('tab10')
+
+        # Group by relation and create separate plots
+        relations = edges_in_kg['Known Relation'].unique()
+        for idx, relation in enumerate(relations):
+
+            relation_data = edges_in_kg[edges_in_kg['Known Relation'] == relation]
+
+            # Get a color from the color map
+            color = color_map(idx % color_map.N)
+
+            fig, ax = plt.subplots(figsize=(10, 3))
+            ax.plot(predictions['Rank'], predictions['Score'])
+            ax.set_xlabel('Rank', fontsize=12)
+            ax.set_ylabel('Score', fontsize=12)
+            ax.set_xlim(1, predictions['Rank'].max())
+            
+            for i, node in relation_data.iterrows():
+                ax.axvline(node['Rank'], color=color, linestyle='--', label=node['Name'])
+                # ax.text(node['Rank'] + 100, node['Score'], node['Name'], fontsize=10, color=color)
+            
+            # ax.set_title(f'{relation.replace("_", "-")}')
+            # ax.legend()
+            color_hex = rgba_to_hex(color)
+
+            # Write header in color of relation
+            st.markdown(f"<h3 style='color:{color_hex}'>{relation.replace('_', ' ').title()}</h2>", unsafe_allow_html=True)
+
+            # Show plot
+            st.pyplot(fig)
+
+            # Drop known relation column
+            relation_data = relation_data.drop(columns = 'Known Relation')
+            if target_node_type not in ['disease', 'anatomy']:
+                st.dataframe(relation_data, use_container_width=True,
+                            column_config={"Database": st.column_config.LinkColumn(width = "small",
+                                                                                    help = "Click to visit external database.",
+                                                                                    display_text = st.session_state.display_database)})
+            else:
+                st.dataframe(relation_data, use_container_width=True)
+
+else:
+
+    st.error('No ground truth relationships found for the given query in the knowledge graph.', icon="✖️")