WIP Add option to upload PDB files
Browse filesThis 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 +12 -3
- hexviz/view.py +27 -0
- hexviz/🧬Attention_Visualization.py +12 -5
hexviz/attention.py
CHANGED
@@ -20,6 +20,16 @@ def get_structure(pdb_code: str) -> Structure:
|
|
20 |
structure = parser.get_structure(pdb_code, file)
|
21 |
return structure
|
22 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
23 |
def get_chains(structure: Structure) -> List[str]:
|
24 |
"""
|
25 |
Get list of chains in a structure
|
@@ -107,9 +117,8 @@ def unidirectional_avg_filtered(attention, layer, head, threshold):
|
|
107 |
return unidirectional_avg_for_head
|
108 |
|
109 |
@st.cache
|
110 |
-
def get_attention_pairs(
|
111 |
-
structure = get_structure(
|
112 |
-
|
113 |
if chain_ids:
|
114 |
chains = [ch for ch in structure.get_chains() if ch.id in chain_ids]
|
115 |
else:
|
|
|
20 |
structure = parser.get_structure(pdb_code, file)
|
21 |
return structure
|
22 |
|
23 |
+
def get_pdb_file(pdb_code: str) -> Structure:
|
24 |
+
"""
|
25 |
+
Get structure from PDB
|
26 |
+
"""
|
27 |
+
pdb_url = f"https://files.rcsb.org/download/{pdb_code}.pdb"
|
28 |
+
pdb_data = request.urlopen(pdb_url).read().decode("utf-8")
|
29 |
+
file = StringIO(pdb_data)
|
30 |
+
return file
|
31 |
+
|
32 |
+
|
33 |
def get_chains(structure: Structure) -> List[str]:
|
34 |
"""
|
35 |
Get list of chains in a structure
|
|
|
117 |
return unidirectional_avg_for_head
|
118 |
|
119 |
@st.cache
|
120 |
+
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):
|
121 |
+
structure = PDBParser().get_structure("pdb", StringIO(pdb_str))
|
|
|
122 |
if chain_ids:
|
123 |
chains = [ch for ch in structure.get_chains() if ch.id in chain_ids]
|
124 |
else:
|
hexviz/view.py
CHANGED
@@ -1,4 +1,9 @@
|
|
|
|
|
|
1 |
import streamlit as st
|
|
|
|
|
|
|
2 |
|
3 |
menu_items = {
|
4 |
"Get Help": "https://huggingface.co/spaces/aksell/hexviz/discussions/new",
|
@@ -51,9 +56,31 @@ def select_pdb():
|
|
51 |
st.session_state.selected_chain_index = 0
|
52 |
if "sequence_slice" in st.session_state:
|
53 |
del st.session_state.sequence_slice
|
|
|
|
|
54 |
st.session_state.pdb_id = pdb_id
|
55 |
return pdb_id
|
56 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
57 |
def select_heads_and_layers(sidebar, model):
|
58 |
sidebar.markdown(
|
59 |
"""
|
|
|
1 |
+
from io import StringIO
|
2 |
+
|
3 |
import streamlit as st
|
4 |
+
from Bio.PDB import PDBParser
|
5 |
+
|
6 |
+
from hexviz.attention import get_pdb_file
|
7 |
|
8 |
menu_items = {
|
9 |
"Get Help": "https://huggingface.co/spaces/aksell/hexviz/discussions/new",
|
|
|
56 |
st.session_state.selected_chain_index = 0
|
57 |
if "sequence_slice" in st.session_state:
|
58 |
del st.session_state.sequence_slice
|
59 |
+
if "uploaded_pdb_str" in st.session_state:
|
60 |
+
del st.session_state.uploaded_pdb_str
|
61 |
st.session_state.pdb_id = pdb_id
|
62 |
return pdb_id
|
63 |
|
64 |
+
def select_protein(pdb_code, uploaded_file):
|
65 |
+
# We get the pdb from 1 of 3 places:
|
66 |
+
# 1. Cached pdb from session storage
|
67 |
+
# 2. PDB file from uploaded file
|
68 |
+
# 3. PDB file fetched based on the pdb_code input
|
69 |
+
parser = PDBParser()
|
70 |
+
if uploaded_file is not None:
|
71 |
+
if "pdb_str" in st.session_state:
|
72 |
+
del st.session_state.pdb_str
|
73 |
+
pdb_str = uploaded_file.read().decode("utf-8")
|
74 |
+
st.session_state["uploaded_pdb_str"] = pdb_str
|
75 |
+
if "uploaded_pdb_str" in st.session_state:
|
76 |
+
pdb_str = st.session_state.uploaded_pdb_str
|
77 |
+
else:
|
78 |
+
file = get_pdb_file(pdb_code)
|
79 |
+
pdb_str = file.read()
|
80 |
+
|
81 |
+
structure = parser.get_structure(pdb_code, StringIO(pdb_str))
|
82 |
+
return pdb_str, structure
|
83 |
+
|
84 |
def select_heads_and_layers(sidebar, model):
|
85 |
sidebar.markdown(
|
86 |
"""
|
hexviz/🧬Attention_Visualization.py
CHANGED
@@ -1,12 +1,15 @@
|
|
|
|
|
|
1 |
import pandas as pd
|
2 |
import py3Dmol
|
3 |
import stmol
|
4 |
import streamlit as st
|
|
|
5 |
from stmol import showmol
|
6 |
|
7 |
-
from hexviz.attention import get_attention_pairs, get_chains
|
8 |
from hexviz.models import Model, ModelType
|
9 |
-
from hexviz.view import menu_items, select_model, select_pdb
|
10 |
|
11 |
st.set_page_config(layout="centered", menu_items=menu_items)
|
12 |
st.title("Attention Visualization on proteins")
|
@@ -18,7 +21,10 @@ models = [
|
|
18 |
]
|
19 |
|
20 |
pdb_id = select_pdb()
|
21 |
-
|
|
|
|
|
|
|
22 |
chains = get_chains(structure)
|
23 |
|
24 |
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:
|
|
68 |
if ec_class and selected_model.name == ModelType.ZymCTRL:
|
69 |
ec_class = st.sidebar.text_input("Enzyme classification number fetched from PDB", ec_class)
|
70 |
|
71 |
-
attention_pairs, top_residues = get_attention_pairs(
|
72 |
|
73 |
sorted_by_attention = sorted(attention_pairs, key=lambda x: x[0], reverse=True)
|
74 |
|
75 |
def get_3dview(pdb):
|
76 |
-
xyzview = py3Dmol.view(
|
|
|
77 |
xyzview.setStyle({"cartoon": {"color": "spectrum"}})
|
78 |
stmol.add_hover(xyzview, backgroundColor="black", fontColor="white")
|
79 |
|
|
|
1 |
+
from io import StringIO
|
2 |
+
|
3 |
import pandas as pd
|
4 |
import py3Dmol
|
5 |
import stmol
|
6 |
import streamlit as st
|
7 |
+
from Bio.PDB import PDBParser
|
8 |
from stmol import showmol
|
9 |
|
10 |
+
from hexviz.attention import get_attention_pairs, get_chains
|
11 |
from hexviz.models import Model, ModelType
|
12 |
+
from hexviz.view import menu_items, select_model, select_pdb, select_protein
|
13 |
|
14 |
st.set_page_config(layout="centered", menu_items=menu_items)
|
15 |
st.title("Attention Visualization on proteins")
|
|
|
21 |
]
|
22 |
|
23 |
pdb_id = select_pdb()
|
24 |
+
with st.expander("Input sequence or upload PDB file"):
|
25 |
+
uploaded_file = st.file_uploader("Upload PDB", type=["pdb"])
|
26 |
+
|
27 |
+
pdb_str, structure = select_protein(pdb_id, uploaded_file)
|
28 |
chains = get_chains(structure)
|
29 |
|
30 |
selected_chains = st.sidebar.multiselect(label="Select Chain(s)", options=chains, default=st.session_state.get("selected_chains", None) or chains)
|
|
|
74 |
if ec_class and selected_model.name == ModelType.ZymCTRL:
|
75 |
ec_class = st.sidebar.text_input("Enzyme classification number fetched from PDB", ec_class)
|
76 |
|
77 |
+
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)
|
78 |
|
79 |
sorted_by_attention = sorted(attention_pairs, key=lambda x: x[0], reverse=True)
|
80 |
|
81 |
def get_3dview(pdb):
|
82 |
+
xyzview = py3Dmol.view()
|
83 |
+
xyzview.addModel(pdb_str, "pdb")
|
84 |
xyzview.setStyle({"cartoon": {"color": "spectrum"}})
|
85 |
stmol.add_hover(xyzview, backgroundColor="black", fontColor="white")
|
86 |
|