Persist selected PDB, chain and model in session state

Selected chain(s) are not shared between views.
Because one of the views allows multiselect of chains and the
other does not.

hexviz/pages/🗺️Identify_Interesting_Heads.py

@@ -3,6 +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
 
 st.set_page_config(layout="wide")
 st.subheader("Find interesting heads and layers")
@@ -13,13 +14,15 @@ models = [
     Model(name=ModelType.ZymCTRL, layers=36, heads=16),
 ]
 
-selected_model_name = st.sidebar.selectbox("Select a model", [model.name.value for model in models], index=0)
+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="1AKE",
+        value=st.session_state.get("pdb_id", "2FZ5"),
     )
+st.session_state.pdb_id = pdb_id
 
 
 structure = get_structure(pdb_id)
@@ -29,7 +32,9 @@ chain_ids = [chain.id for chain in chains]
 chain_selection = st.sidebar.selectbox(
     label="Select Chain",
     options=chain_ids,
+    index=st.session_state.get("selected_chain_index", 0)
 )
+st.session_state.selected_chain_index = chain_ids.index(chain_selection)
 
 selected_chain = next(chain for chain in chains if chain.id == chain_selection)
 sequence = get_sequence(selected_chain)

hexviz/view.py

@@ -0,0 +1,9 @@
+import streamlit as st
+
+
+def get_selecte_model_index(models):
+    selected_model_name = st.session_state.get("selected_model_name", None)
+    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)

hexviz/🧬Attention_Visualization.py

@@ -6,6 +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
 
 st.title("Attention Visualization on proteins")
 
@@ -22,11 +23,13 @@ st.sidebar.markdown(
     """)
 pdb_id = st.sidebar.text_input(
         label="PDB ID",
-        value="2FZ5",
-    )
+        value=st.session_state.get("pdb_id", "2FZ5"))
+st.session_state.pdb_id = pdb_id
 structure = get_structure(pdb_id)
 chains = get_chains(structure)
-selected_chains = st.sidebar.multiselect(label="Chain(s)", options=chains, default=chains)
+
+selected_chains = st.sidebar.multiselect(label="Chain(s)", options=chains, default=st.session_state.get("selected_chains", None) or chains)
+st.session_state.selected_chains = selected_chains
 
 
 st.sidebar.markdown(
@@ -56,7 +59,8 @@ 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=0)
+    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)
 with mid:
     layer_one = st.number_input("Layer", value=5, min_value=1, max_value=selected_model.layers)