WIP Add option to upload PDB files

This feel a bit messy. Passing the pdb_string to get_attention_pairs
is only done because if I pass the structure it is hashed and re-calculated
every run.
With newer streamlit this would be easy to fix just add a underscore
infront of the variable name.

It is also confusing now with 3 sources of the pdb data, either
the pdb code, an uploaded file or a pdb string that I store in
session storage to persist uploaded files across page navigation.
I think finding a way to persist the uploaded file would be
easier to understand for the users.

hexviz/attention.py

@@ -20,6 +20,16 @@ def get_structure(pdb_code: str) -> Structure:
     structure = parser.get_structure(pdb_code, file)
     return structure
 
+def get_pdb_file(pdb_code: str) -> Structure:
+    """
+    Get structure from PDB
+    """
+    pdb_url = f"https://files.rcsb.org/download/{pdb_code}.pdb"
+    pdb_data = request.urlopen(pdb_url).read().decode("utf-8")
+    file = StringIO(pdb_data)
+    return file
+
+
 def get_chains(structure: Structure) -> List[str]:
     """
     Get list of chains in a structure
@@ -107,9 +117,8 @@ 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, chain_ids: Optional[str] = None ,threshold: int = 0.2, model_type: ModelType = ModelType.TAPE_BERT, top_n: int = 2):
-    structure = get_structure(pdb_code=pdb_code)
-
+def get_attention_pairs(pdb_str: str, layer: int, head: int, chain_ids: Optional[str] = None ,threshold: int = 0.2, model_type: ModelType = ModelType.TAPE_BERT, top_n: int = 2):
+    structure = PDBParser().get_structure("pdb", StringIO(pdb_str))
     if chain_ids:
         chains = [ch for ch in structure.get_chains() if ch.id in chain_ids]
     else:

hexviz/view.py

@@ -1,4 +1,9 @@
+from io import StringIO
+
 import streamlit as st
+from Bio.PDB import PDBParser
+
+from hexviz.attention import get_pdb_file
 
 menu_items = {
     "Get Help": "https://huggingface.co/spaces/aksell/hexviz/discussions/new", 
@@ -51,9 +56,31 @@ def select_pdb():
         st.session_state.selected_chain_index = 0
         if "sequence_slice" in st.session_state:
             del st.session_state.sequence_slice
+        if "uploaded_pdb_str" in st.session_state:
+            del st.session_state.uploaded_pdb_str
     st.session_state.pdb_id = pdb_id
     return pdb_id
 
+def select_protein(pdb_code, uploaded_file):
+    # We get the pdb from 1 of 3 places:
+    # 1. Cached pdb from session storage
+    # 2. PDB file from uploaded file
+    # 3. PDB file fetched based on the pdb_code input
+    parser = PDBParser()
+    if uploaded_file is not None:
+        if "pdb_str" in st.session_state:
+            del st.session_state.pdb_str
+        pdb_str = uploaded_file.read().decode("utf-8")
+        st.session_state["uploaded_pdb_str"] = pdb_str
+    if "uploaded_pdb_str" in st.session_state:
+        pdb_str = st.session_state.uploaded_pdb_str
+    else:
+        file = get_pdb_file(pdb_code)
+        pdb_str = file.read()
+
+    structure = parser.get_structure(pdb_code, StringIO(pdb_str))
+    return pdb_str, structure
+
 def select_heads_and_layers(sidebar, model):
     sidebar.markdown(
         """

hexviz/🧬Attention_Visualization.py

@@ -1,12 +1,15 @@
+from io import StringIO
+
 import pandas as pd
 import py3Dmol
 import stmol
 import streamlit as st
+from Bio.PDB import PDBParser
 from stmol import showmol
 
-from hexviz.attention import get_attention_pairs, get_chains, get_structure
+from hexviz.attention import get_attention_pairs, get_chains
 from hexviz.models import Model, ModelType
-from hexviz.view import menu_items, select_model, select_pdb
+from hexviz.view import menu_items, select_model, select_pdb, select_protein
 
 st.set_page_config(layout="centered", menu_items=menu_items)
 st.title("Attention Visualization on proteins")
@@ -18,7 +21,10 @@ models = [
 ]
 
 pdb_id = select_pdb()
-structure = get_structure(pdb_id)
+with st.expander("Input sequence or upload PDB file"):
+    uploaded_file = st.file_uploader("Upload PDB", type=["pdb"])
+
+pdb_str, structure = select_protein(pdb_id, uploaded_file)
 chains = get_chains(structure)
 
 selected_chains = st.sidebar.multiselect(label="Select Chain(s)", options=chains, default=st.session_state.get("selected_chains", None) or chains)
@@ -68,12 +74,13 @@ if selected_model.name == ModelType.ZymCTRL:
     if ec_class and selected_model.name == ModelType.ZymCTRL:
         ec_class = st.sidebar.text_input("Enzyme classification number fetched from PDB", ec_class)
 
-attention_pairs, top_residues = get_attention_pairs(pdb_id, chain_ids=selected_chains, layer=layer, head=head, threshold=min_attn, model_type=selected_model.name, top_n=n_highest_resis)
+attention_pairs, top_residues = get_attention_pairs(pdb_str=pdb_str, chain_ids=selected_chains, layer=layer, head=head, threshold=min_attn, model_type=selected_model.name, top_n=n_highest_resis)
 
 sorted_by_attention = sorted(attention_pairs, key=lambda x: x[0], reverse=True) 
 
 def get_3dview(pdb):
-    xyzview = py3Dmol.view(query=f"pdb:{pdb}")
+    xyzview = py3Dmol.view()
+    xyzview.addModel(pdb_str, "pdb")
     xyzview.setStyle({"cartoon": {"color": "spectrum"}})
     stmol.add_hover(xyzview, backgroundColor="black", fontColor="white")