Share pdb selector between views

hexviz/pages/🗺️Identify_Interesting_Heads.py

@@ -3,7 +3,7 @@ import streamlit as st
 from hexviz.attention import get_attention, get_sequence, get_structure
 from hexviz.models import Model, ModelType
 from hexviz.plot import plot_tiled_heatmap
-from hexviz.view import get_selecte_model_index
+from hexviz.view import select_model, select_pdb
 
 st.set_page_config(layout="wide")
 st.subheader("Find interesting heads and layers")
@@ -14,16 +14,9 @@ models = [
     Model(name=ModelType.ZymCTRL, layers=36, heads=16),
 ]
 
-selected_model_name = st.selectbox("Select a model", [model.name.value for model in models], index=get_selecte_model_index(models))
-st.session_state.selected_model_name = selected_model_name
-selected_model = next((model for model in models if model.name.value == selected_model_name), None)
-
-pdb_id = st.sidebar.text_input(
-        label="PDB ID",
-        value=st.session_state.get("pdb_id", "2FZ5"),
-    )
-st.session_state.pdb_id = pdb_id
+selected_model = select_model(models)
 
+pdb_id = select_pdb()
 
 structure = get_structure(pdb_id)
 

hexviz/view.py

@@ -6,4 +6,27 @@ def get_selecte_model_index(models):
     if selected_model_name is None:
         return 0
     else:
-        return next((i for i, model in enumerate(models) if model.name.value == selected_model_name), None)
+        return next((i for i, model in enumerate(models) if model.name.value == selected_model_name), None)
+
+def select_model(models):
+    """
+    Select model, prefil selector with selected model from session storage
+
+    Saves the selected model in session storage.
+    """
+    selected_model_name = st.selectbox("Select model", [model.name.value for model in models], index=get_selecte_model_index(models))
+    st.session_state.selected_model_name = selected_model_name
+    select_model = next((model for model in models if model.name.value == selected_model_name), None)
+    return select_model
+
+def select_pdb():
+    st.sidebar.markdown(
+        """
+        Select Protein 
+        ---
+        """)
+    pdb_id = st.sidebar.text_input(
+            label="PDB ID",
+            value=st.session_state.get("pdb_id", "2FZ5"))
+    st.session_state.pdb_id = pdb_id
+    return pdb_id

hexviz/🧬Attention_Visualization.py

@@ -6,7 +6,7 @@ from stmol import showmol
 
 from hexviz.attention import get_attention_pairs, get_chains, get_structure
 from hexviz.models import Model, ModelType
-from hexviz.view import get_selecte_model_index
+from hexviz.view import select_model, select_pdb
 
 st.title("Attention Visualization on proteins")
 
@@ -16,19 +16,11 @@ models = [
     Model(name=ModelType.ZymCTRL, layers=36, heads=16),
 ]
 
-st.sidebar.markdown(
-    """
-    Select Protein 
-    ---
-    """)
-pdb_id = st.sidebar.text_input(
-        label="PDB ID",
-        value=st.session_state.get("pdb_id", "2FZ5"))
-st.session_state.pdb_id = pdb_id
+pdb_id = select_pdb()
 structure = get_structure(pdb_id)
 chains = get_chains(structure)
 
-selected_chains = st.sidebar.multiselect(label="Chain(s)", options=chains, default=st.session_state.get("selected_chains", None) or chains)
+selected_chains = st.sidebar.multiselect(label="Select Chain(s)", options=chains, default=st.session_state.get("selected_chains", None) or chains)
 st.session_state.selected_chains = selected_chains
 
 
@@ -59,9 +51,7 @@ label_highest = st.sidebar.checkbox("Label highest attention pairs", value=True)
 
 left, mid, right = st.columns(3)
 with left:
-    selected_model_name = st.selectbox("Select a model", [model.name.value for model in models], index=get_selecte_model_index(models))
-    st.session_state.selected_model_name = selected_model_name
-    selected_model = next((model for model in models if model.name.value == selected_model_name), None)
+    selected_model = select_model(models)
 with mid:
     layer_one = st.number_input("Layer", value=5, min_value=1, max_value=selected_model.layers)
     layer = layer_one - 1