Allow selecting individual chains

hexviz/app.py

@@ -3,7 +3,7 @@ import stmol
 import streamlit as st
 from stmol import showmol
 
-from hexviz.attention import get_attention_pairs, get_structure
+from hexviz.attention import get_attention_pairs, get_chains, get_structure
 from hexviz.models import Model, ModelType
 
 st.title("Attention Visualization on proteins")
@@ -21,7 +21,20 @@ models = [
 selected_model_name = st.selectbox("Select a model", [model.name.value for model in models], index=0)
 selected_model = next((model for model in models if model.name.value == selected_model_name), None)
 
-pdb_id = st.text_input("PDB ID", "1I60")
+st.sidebar.title("Settings")
+
+pdb_id = st.sidebar.text_input(
+        label="PDB ID",
+        value="4RW0",
+    )
+structure = get_structure(pdb_id)
+chains = get_chains(structure)
+selected_chains = st.sidebar.multiselect(label="Chain(s)", options=chains, default=chains)
+
+hl_resi_list = st.sidebar.multiselect(label="Highlight Residues",options=list(range(1,5000)))
+
+label_resi = st.sidebar.checkbox(label="Label Residues", value=True)
+
 
 left, right = st.columns(2)
 with left:
@@ -35,21 +48,31 @@ with right:
 with st.expander("Configure parameters", expanded=False):
     min_attn = st.slider("Minimum attention", min_value=0.0, max_value=0.4, value=0.1)
     try:
-        structure = get_structure(pdb_id)
         ec_class = structure.header["compound"]["1"]["ec"]
     except KeyError:
         ec_class = None
     if ec_class and selected_model.name == ModelType.ZymCTRL:
         ec_class = st.text_input("Enzyme classification number fetched from PDB", ec_class)
 
-attention_pairs = get_attention_pairs(pdb_id, layer, head, min_attn, model_type=selected_model.name)
+attention_pairs = get_attention_pairs(pdb_id, chain_ids=selected_chains, layer=layer, head=head, threshold=min_attn, model_type=selected_model.name)
 
 def get_3dview(pdb):
     xyzview = py3Dmol.view(query=f"pdb:{pdb}")
     xyzview.setStyle({"cartoon": {"color": "spectrum"}})
     stmol.add_hover(xyzview, backgroundColor="black", fontColor="white")
+
+
+    hidden_chains = [x for x in chains if x not in selected_chains]
+    for chain in hidden_chains:
+        xyzview.setStyle({"chain": chain},{"cross":{"hidden":"true"}})
+
     for att_weight, first, second in attention_pairs:
         stmol.add_cylinder(xyzview, start=first, end=second, cylradius=att_weight, cylColor='red', dashed=False)
+
+    if label_resi:
+        for hl_resi in hl_resi_list:
+            xyzview.addResLabels({"chain": chain,"resi": hl_resi},
+            {"backgroundColor": "lightgray","fontColor": "black","backgroundOpacity": 0.5})
     return xyzview
 
 

hexviz/attention.py

@@ -1,6 +1,5 @@
-from enum import Enum
 from io import StringIO
-from typing import List, Tuple
+from typing import List, Optional
 from urllib import request
 
 import streamlit as st
@@ -21,20 +20,27 @@ def get_structure(pdb_code: str) -> Structure:
     structure = parser.get_structure(pdb_code, file)
     return structure
 
-def get_sequences(structure: Structure) -> List[str]:
+def get_chains(structure: Structure) -> List[str]:
     """
-    Get list of sequences with residues on a single letter format
+    Get list of chains in a structure
+    """
+    chains = []
+    for model in structure:
+        for chain in model.get_chains():
+            chains.append(chain.id)
+    return chains
+
+def get_sequence(chain) -> str:
+    """
+    Get sequence from a chain
 
     Residues not in the standard 20 amino acids are replaced with X
     """
-    sequences = []
-    for seq in structure.get_chains():
-        residues = [residue.get_resname() for residue in seq.get_residues()]
-        # TODO ask if using protein_letters_3to1_extended makes sense
-        residues_single_letter = map(lambda x: Polypeptide.protein_letters_3to1.get(x, "X"), residues)
+    residues = [residue.get_resname() for residue in chain.get_residues()]
+    # TODO ask if using protein_letters_3to1_extended makes sense
+    residues_single_letter = map(lambda x: Polypeptide.protein_letters_3to1.get(x, "X"), residues)
 
-        sequences.append("".join(list(residues_single_letter)))
-    return sequences
+    return "".join(list(residues_single_letter))
 
 @st.cache
 def get_attention(
@@ -100,17 +106,19 @@ def unidirectional_avg_filtered(attention, layer, head, threshold):
     return unidirectional_avg_for_head
  
 @st.cache
-def get_attention_pairs(pdb_code: str, layer: int, head: int, threshold: int = 0.2, model_type: ModelType = ModelType.TAPE_BERT):
-    # fetch structure
+def get_attention_pairs(pdb_code: str, layer: int, head: int, chain_ids: Optional[str] = None ,threshold: int = 0.2, model_type: ModelType = ModelType.TAPE_BERT):
     structure = get_structure(pdb_code=pdb_code)
-    # Get list of sequences
-    sequences = get_sequences(structure)
+
+    if chain_ids:
+        chains = [ch for ch in structure.get_chains() if ch.id in chain_ids]
+    else:
+        chains = list(structure.get_chains())
 
     attention_pairs = []
-    for i, sequence in enumerate(sequences):
+    for chain in chains:
+        sequence = get_sequence(chain)
         attention = get_attention(sequence=sequence, model_type=model_type)
         attention_unidirectional = unidirectional_avg_filtered(attention, layer, head, threshold)
-        chain = list(structure.get_chains())[i]
         for attn_value, res_1, res_2 in attention_unidirectional:
             try:
                 coord_1 = chain[res_1]["CA"].coord.tolist()

tests/test_attention.py

@@ -70,7 +70,7 @@ def test_get_unidirection_avg_filtered():
     assert result is not None
     assert len(result) == 10
 
-    attention= torch.tensor([[[[1, 2, 3],
+    attention = torch.tensor([[[[1, 2, 3],
                                [2, 5, 6],
                                [4, 7, 91]]]], dtype=torch.float32)