Store selected sequence slice

hexviz/pages/🗺️Identify_Interesting_Heads.py

@@ -4,7 +4,7 @@ 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 (menu_items, select_heads_and_layers, select_model,
-                         select_pdb)
+                         select_pdb, select_sequence_slice)
 
 st.set_page_config(layout="wide", menu_items=menu_items)
 st.subheader("Find interesting heads and layers")
@@ -34,12 +34,7 @@ selected_chain = next(chain for chain in chains if chain.id == chain_selection)
 sequence = get_sequence(selected_chain)
 
 l = len(sequence)
-st.sidebar.markdown("Sequence segment to plot")
-if "sequence-slice" not in st.session_state:
-    st.session_state["sequence-slice"] = (1, min(50, l))
-slice_start, slice_end = st.sidebar.slider("Sequence", key="sequence-slice", min_value=1, max_value=l, step=1)
-# slice_start= st.sidebar.number_input(f"Section start(1-{l})",value=1, min_value=1, max_value=l)
-# slice_end = st.sidebar.number_input(f"Section end(1-{l})",value=50, min_value=1, max_value=l)
+slice_start, slice_end = select_sequence_slice(l)
 truncated_sequence = sequence[slice_start-1:slice_end]
 
 

hexviz/view.py

@@ -41,9 +41,12 @@ def select_pdb():
     pdb_id = st.sidebar.text_input(
             label="PDB ID",
             value=stored_pdb or "2FZ5")
-    if pdb_id != stored_pdb:
+    pdb_changed = stored_pdb != pdb_id
+    if pdb_changed:
         st.session_state.selected_chains = None
         st.session_state.selected_chain_index = 0
+        if "sequence_slice" in st.session_state:
+            del st.session_state.sequence_slice
     st.session_state.pdb_id = pdb_id
     return pdb_id
 
@@ -69,4 +72,12 @@ def select_heads_and_layers(sidebar, model):
     layer_sequence = list(range(layer_range[0]-1, layer_range[1], step_size))
     head_sequence = list(range(head_range[0]-1, head_range[1], step_size))
 
-    return layer_sequence, head_sequence
+    return layer_sequence, head_sequence
+
+def select_sequence_slice(sequence_length):
+    st.sidebar.markdown("Sequence segment to plot")
+    if "sequence_slice" not in st.session_state:
+        st.session_state["sequence_slice"] = (1, min(50, sequence_length))
+    slice = st.sidebar.slider("Sequence", value=st.session_state.sequence_slice, min_value=1, max_value=sequence_length, step=1)
+    st.session_state.sequence_slice = slice
+    return slice