scripts from other repo

Dockerfile

@@ -41,6 +41,15 @@ ENV AMBERHOME="/usr/bin/amber22/"
 ENV PATH="$AMBERHOME/bin:$PATH"
 ENV PYTHONPATH="$AMBERHOME/lib/python3.10/site-packages"
 
+
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+WORKDIR $HOME/app
+COPY --chown=user . $HOME/app
+
+
 RUN adduser -u 5678 --disabled-password --gecos "" appuser && chown -R appuser .
 USER appuser
 CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

MDmodel.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import GCNConv
+
+class GNN_MD(torch.nn.Module):
+    def __init__(self, num_features, hidden_dim):
+        super(GNN_MD, self).__init__()
+        self.conv1 = GCNConv(num_features, hidden_dim)
+        self.bn1 = nn.BatchNorm1d(hidden_dim)
+        self.conv2 = GCNConv(hidden_dim, hidden_dim*2)
+        self.bn2 = nn.BatchNorm1d(hidden_dim*2)
+        self.conv3 = GCNConv(hidden_dim*2, hidden_dim*4)
+        self.bn3 = nn.BatchNorm1d(hidden_dim*4)
+        self.conv4 = GCNConv(hidden_dim*4, hidden_dim*4)
+        self.bn4 = nn.BatchNorm1d(hidden_dim*4)
+        self.conv5 = GCNConv(hidden_dim*4, hidden_dim*8)
+        self.bn5 = nn.BatchNorm1d(hidden_dim*8)
+        self.fc1 = nn.Linear(hidden_dim*8, hidden_dim*4)
+        self.fc2 = nn.Linear(hidden_dim*4, 1)
+
+
+    def forward(self, data):
+        x = self.conv1(data.x, data.edge_index, data.edge_attr.view(-1))
+        x = F.relu(x)
+        x = self.bn1(x)
+        x = self.conv2(x, data.edge_index, data.edge_attr.view(-1))
+        x = F.relu(x)
+        x = self.bn2(x)
+        x = self.conv3(x, data.edge_index, data.edge_attr.view(-1))
+        x = F.relu(x)
+        x = self.bn3(x)
+        x = self.conv4(x, data.edge_index, data.edge_attr.view(-1))
+        x = self.bn4(x)
+        x = F.relu(x)
+        x = self.conv5(x, data.edge_index, data.edge_attr.view(-1))
+        x = self.bn5(x)
+        #x = global_add_pool(x, x.batch)
+        x = F.relu(x)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, p=0.25)
+        return self.fc2(x).view(-1)

app.py

@@ -0,0 +1,150 @@
+import gradio as gr
+
+import py3Dmol
+
+from Bio.PDB import *
+
+import numpy as np
+from Bio.PDB import PDBParser
+import pandas as pd
+import torch
+import os
+from MDmodel import GNN_MD
+import h5py
+from transformMD import GNNTransformMD
+
+# JavaScript functions
+resid_hover = """function(atom,viewer) {{
+    if(!atom.label) {{
+        atom.label = viewer.addLabel('{0}:'+atom.atom+atom.serial,
+            {{position: atom, backgroundColor: 'mintcream', fontColor:'black'}});
+    }}
+}}"""
+hover_func = """
+function(atom,viewer) {
+    if(!atom.label) {
+        atom.label = viewer.addLabel(atom.interaction,
+            {position: atom, backgroundColor: 'black', fontColor:'white'});
+    }
+}"""
+unhover_func = """
+function(atom,viewer) {
+    if(atom.label) {
+        viewer.removeLabel(atom.label);
+        delete atom.label;
+    }
+}"""
+atom_mapping = {0:'H', 1:'C', 2:'N', 3:'O', 4:'F', 5:'P', 6:'S', 7:'CL', 8:'BR', 9:'I', 10: 'UNK'}
+
+model = GNN_MD(11, 64)
+state_dict = torch.load(
+    "best_weights_rep0.pt",
+    map_location=torch.device("cpu"),
+)["model_state_dict"]
+model.load_state_dict(state_dict)
+model = model.to('cpu')
+model.eval()
+
+
+
+def get_pdb(pdb_code="", filepath=""):
+    try:
+        return filepath.name
+    except AttributeError as e:
+        if pdb_code is None or pdb_code == "":
+            return None
+        else:
+            os.system(f"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb")
+            return f"{pdb_code}.pdb"
+
+
+def get_offset(pdb):
+    pdb_multiline = pdb.split("\n")
+    for line in pdb_multiline:
+        if line.startswith("ATOM"):
+            return int(line[22:27])
+
+
+def predict(pdb_code, pdb_file):
+    #path_to_pdb = get_pdb(pdb_code=pdb_code, filepath=pdb_file)
+
+    #pdb = open(path_to_pdb, "r").read()
+    # switch to misato env if not running from container
+    mdh5_file = "inference_for_md.hdf5"
+    md_H5File = h5py.File(mdh5_file)
+
+    column_names = ["x", "y", "z", "element"]
+    atoms_protein = pd.DataFrame(columns = column_names)
+    cutoff = md_H5File["11GS"]["molecules_begin_atom_index"][:][-1] # cutoff defines protein atoms
+
+    atoms_protein["x"] = md_H5File["11GS"]["atoms_coordinates_ref"][:][:cutoff, 0]
+    atoms_protein["y"] = md_H5File["11GS"]["atoms_coordinates_ref"][:][:cutoff, 1]
+    atoms_protein["z"] = md_H5File["11GS"]["atoms_coordinates_ref"][:][:cutoff, 2]
+
+    atoms_protein["element"] = md_H5File["11GS"]["atoms_element"][:][:cutoff]  
+
+    item = {}
+    item["scores"] = 0
+    item["id"] = "11GS"
+    item["atoms_protein"] = atoms_protein
+
+    transform = GNNTransformMD()
+    data_item = transform(item)
+    adaptability = model(data_item)
+    adaptability = adaptability.detach().numpy()
+    
+    data = []
+
+
+    for i in range(adaptability.shape[0]):
+        data.append([i, atom_mapping(atoms_protein.iloc[i, atoms_protein.columns.get_loc("element")] - 1), atoms_protein.iloc[i, atoms_protein.columns.get_loc("x")],atoms_protein.iloc[i, atoms_protein.columns.get_loc("y")],atoms_protein.iloc[i, atoms_protein.columns.get_loc("z")],adaptability[i]])
+
+    topN = 100
+    topN_ind = np.argsort(adaptability)[::-1][:topN]    
+
+    pdb = open(pdb_file.name, "r").read()
+
+    view = py3Dmol.view(width=600, height=400)
+    view.setBackgroundColor('black')
+    view.addModel(pdb, "pdb")
+    view.setStyle({'stick': {'colorscheme': {'prop': 'resi', 'C': 'turquoise'}}})
+   
+    for i in range(topN):
+        view.addSphere({'center':{'x':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("x")], 'y':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("y")],'z':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("z")]},'radius':adaptability[topN_ind[i]]/1.5,'color':'orange','alpha':0.75})    
+
+    view.zoomTo()
+
+    output = view._make_html().replace("'", '"')
+
+    x = f"""<!DOCTYPE html><html> {output} </html>"""  # do not use ' in this input
+    return f"""<iframe  style="width: 100%; height:420px" name="result" allow="midi; geolocation; microphone; camera; 
+    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
+    allow-scripts allow-same-origin allow-popups 
+    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
+    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>""", pd.DataFrame(data, columns=['index','element','x','y','z','Adaptability'])
+
+
+callback = gr.CSVLogger()
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Protein Adaptability Prediction")
+    
+    #text_input = gr.Textbox()
+    #text_output = gr.Textbox()
+    #text_button = gr.Button("Flip")
+    inp = gr.Textbox(placeholder="PDB Code or upload file below", label="Input structure")
+    pdb_file = gr.File(label="PDB File Upload")
+    #with gr.Row():
+    #    helix = gr.ColorPicker(label="helix")
+    #    sheet = gr.ColorPicker(label="sheet")
+    #    loop = gr.ColorPicker(label="loop")
+    single_btn = gr.Button(label="Run")
+    with gr.Row():
+        html = gr.HTML()
+    with gr.Row():
+        dataframe = gr.Dataframe()
+            
+    single_btn.click(fn=predict, inputs=[inp, pdb_file], outputs=[html, dataframe])
+
+
+demo.launch(debug=True)

best_weights_rep0.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6b9499f2dc7b16eb3c669d2c6c26e11a53e21047264d3eb0cdda6bbc1d17f91
+size 4517600

graph.py

@@ -0,0 +1,121 @@
+import numpy as np
+import scipy.spatial as ss
+import torch
+import torch.nn.functional as F
+from torch_geometric.utils import to_undirected
+from torch_sparse import coalesce
+
+atom_mapping = {0:'H', 1:'C', 2:'N', 3:'O', 4:'F', 5:'P', 6:'S', 7:'CL', 8:'BR', 9:'I', 10: 'UNK'}
+residue_mapping = {0:'ALA', 1:'ARG', 2:'ASN', 3:'ASP', 4:'CYS', 5:'CYX', 6:'GLN', 7:'GLU', 8:'GLY', 9:'HIE', 10:'ILE', 11:'LEU', 12:'LYS', 13:'MET', 14:'PHE', 15:'PRO', 16:'SER', 17:'THR', 18:'TRP', 19:'TYR', 20:'VAL', 21:'UNK'}
+        
+ligand_atoms_mapping = {8: 0, 16: 1, 6: 2, 7: 3, 1: 4, 15: 5, 17: 6, 9: 7, 53: 8, 35: 9, 5: 10, 33: 11, 26: 12, 14: 13, 34: 14, 44: 15, 12: 16, 23: 17, 77: 18, 27: 19, 52: 20, 30: 21, 4: 22, 45: 23}
+
+
+def prot_df_to_graph(item, df, edge_dist_cutoff, feat_col='element'):
+    r"""
+    Converts protein in dataframe representation to a graph compatible with Pytorch-Geometric, where each node is an atom.
+
+    :param df: Protein structure in dataframe format.
+    :type df: pandas.DataFrame
+    :param node_col: Column of dataframe to find node feature values. For example, for atoms use ``feat_col="element"`` and for residues use ``feat_col="resname"``
+    :type node_col: str, optional
+    :param allowable_feats: List containing all possible values of node type, to be converted into 1-hot node features. 
+        Any elements in ``feat_col`` that are not found in ``allowable_feats`` will be added to an appended "unknown" bin (see :func:`atom3d.util.graph.one_of_k_encoding_unk`).
+    :type allowable_feats: list, optional
+    :param edge_dist_cutoff: Maximum distance cutoff (in Angstroms) to define an edge between two atoms, defaults to 4.5.
+    :type edge_dist_cutoff: float, optional
+
+    :return: tuple containing
+
+        - node_feats (torch.FloatTensor): Features for each node, one-hot encoded by values in ``allowable_feats``.
+
+        - edges (torch.LongTensor): Edges in COO format
+
+        - edge_weights (torch.LongTensor): Edge weights, defined as a function of distance between atoms given by :math:`w_{i,j} = \frac{1}{d(i,j)}`, where :math:`d(i, j)` is the Euclidean distance between node :math:`i` and node :math:`j`.
+
+        - node_pos (torch.FloatTensor): x-y-z coordinates of each node
+    :rtype: Tuple
+    """ 
+
+    allowable_feats = atom_mapping
+    
+    try : 
+        node_pos = torch.FloatTensor(df[['x', 'y', 'z']].to_numpy())
+        kd_tree = ss.KDTree(node_pos)
+        edge_tuples = list(kd_tree.query_pairs(edge_dist_cutoff))
+        edges = torch.LongTensor(edge_tuples).t().contiguous()
+        edges = to_undirected(edges)
+    except:
+        print(f"Problem with PDB Id is {item['id']}")
+        
+    node_feats = torch.FloatTensor([one_of_k_encoding_unk_indices(e-1, allowable_feats) for e in df[feat_col]])
+    edge_weights = torch.FloatTensor(
+        [1.0 / (np.linalg.norm(node_pos[i] - node_pos[j]) + 1e-5) for i, j in edges.t()]).view(-1)
+
+    
+    return node_feats, edges, edge_weights, node_pos
+
+
+def mol_df_to_graph_for_qm(df, bonds=None, allowable_atoms=None, edge_dist_cutoff=4.5, onehot_edges=True):
+    """
+    Converts molecule in dataframe to a graph compatible with Pytorch-Geometric
+    :param df: Molecule structure in dataframe format
+    :type mol: pandas.DataFrame
+    :param bonds: Molecule structure in dataframe format
+    :type bonds: pandas.DataFrame
+    :param allowable_atoms: List containing allowable atom types
+    :type allowable_atoms: list[str], optional
+    :return: Tuple containing \n
+        - node_feats (torch.FloatTensor): Features for each node, one-hot encoded by atom type in ``allowable_atoms``.
+        - edge_index (torch.LongTensor): Edges from chemical bond graph in COO format.
+        - edge_feats (torch.FloatTensor): Edge features given by bond type. Single = 1.0, Double = 2.0, Triple = 3.0, Aromatic = 1.5.
+        - node_pos (torch.FloatTensor): x-y-z coordinates of each node.
+    """
+    if allowable_atoms is None:
+        allowable_atoms = ligand_atoms_mapping
+    node_pos = torch.FloatTensor(df[['x', 'y', 'z']].to_numpy())
+    
+    if bonds is not None:
+        N = df.shape[0]
+        bond_mapping = {1.0: 0, 2.0: 1, 3.0: 2, 1.5: 3}
+        bond_data = torch.FloatTensor(bonds)
+        edge_tuples = torch.cat((bond_data[:, :2], torch.flip(bond_data[:, :2], dims=(1,))), dim=0)
+        edge_index = edge_tuples.t().long().contiguous()
+        
+        if onehot_edges:
+            bond_idx = list(map(lambda x: bond_mapping[x], bond_data[:,-1].tolist())) + list(map(lambda x: bond_mapping[x], bond_data[:,-1].tolist()))
+            edge_attr = F.one_hot(torch.tensor(bond_idx), num_classes=4).to(torch.float)
+            edge_index, edge_attr = coalesce(edge_index, edge_attr, N, N)
+  
+        else:
+            edge_attr = torch.cat((torch.FloatTensor(bond_data[:,-1]).view(-1), torch.FloatTensor(bond_data[:,-1]).view(-1)), dim=0)
+    else:
+        kd_tree = ss.KDTree(node_pos)
+        edge_tuples = list(kd_tree.query_pairs(edge_dist_cutoff))
+        edge_index = torch.LongTensor(edge_tuples).t().contiguous()
+        edge_index = to_undirected(edge_index)
+        edge_attr = torch.FloatTensor([1.0 / (np.linalg.norm(node_pos[i] - node_pos[j]) + 1e-5) for i, j in edge_index.t()]).view(-1)
+        edge_attr = edge_attr.unsqueeze(1)
+    
+    node_feats = torch.FloatTensor([one_of_k_encoding_unk_indices_qm(e, allowable_atoms) for e in df['element']])    
+    
+    return node_feats, edge_index, edge_attr, node_pos
+
+
+def one_of_k_encoding_unk_indices(x, allowable_set):
+    """Converts input to 1-hot encoding given a set of allowable values. Additionally maps inputs not in the allowable set to the last element."""
+    one_hot_encoding = [0] * len(allowable_set)
+    if x in allowable_set:
+        one_hot_encoding[x] = 1
+    else:
+        one_hot_encoding[-1] = 1
+    return one_hot_encoding
+
+def one_of_k_encoding_unk_indices_qm(x, allowable_set):
+    """Converts input to 1-hot encoding given a set of allowable values. Additionally maps inputs not in the allowable set to the last element."""
+    one_hot_encoding = [0] * (len(allowable_set)+1)
+    if x in allowable_set:
+        one_hot_encoding[allowable_set[x]] = 1
+    else:
+        one_hot_encoding[-1] = 1
+    return one_hot_encoding

transformMD.py

@@ -0,0 +1,21 @@
+from transforms import prot_graph_transform
+        
+class GNNTransformMD(object):
+    """
+    Transform the dict returned by the ProtDataset class to a pyTorch Geometric graph
+    """
+
+    def __init__(self, edge_dist_cutoff=4.5):
+        """
+
+        Args:
+            edge_dist_cutoff (float, optional): distence between the edges. Defaults to 4.5.
+        """
+        self.edge_dist_cutoff = edge_dist_cutoff 
+
+    def __call__(self, item):
+        item = prot_graph_transform(item, atom_keys=['atoms_protein'], label_key='scores', edge_dist_cutoff=self.edge_dist_cutoff)
+        return item['atoms_protein']
+  
+
+    

transforms.py

@@ -0,0 +1,46 @@
+import torch
+from torch_geometric.data import Data
+from graph import prot_df_to_graph, mol_df_to_graph_for_qm
+
+def prot_graph_transform(item, atom_keys, label_key, edge_dist_cutoff):
+    """Transform for converting dataframes to Pytorch Geometric graphs, to be applied when defining a :mod:`Dataset <atom3d.datasets.datasets>`.
+    Operates on Dataset items, assumes that the item contains all keys specified in ``keys`` and ``labels`` arguments.
+
+    :param item: Dataset item to transform
+    :type item: dict
+    :param atom_keys: list of keys to transform, where each key contains a dataframe of atoms, defaults to ['atoms']
+    :type atom_keys: list, optional
+    :param label_key: name of key containing labels, defaults to ['scores']
+    :type label_key: str, optional
+    :return: Transformed Dataset item
+    :rtype: dict
+    """    
+
+    for key in atom_keys:
+        node_feats, edge_index, edge_feats, pos = prot_df_to_graph(item, item[key], edge_dist_cutoff)
+        item[key] = Data(node_feats, edge_index, edge_feats, y=torch.FloatTensor(item[label_key]), pos=pos, ids=item["id"])
+        
+    return item
+
+def mol_graph_transform_for_qm(item, atom_key, label_key, allowable_atoms, use_bonds, onehot_edges, edge_dist_cutoff):
+    """Transform for converting dataframes to Pytorch Geometric graphs, to be applied when defining a :mod:`Dataset <atom3d.datasets.datasets>`.
+    Operates on Dataset items, assumes that the item contains all keys specified in ``keys`` and ``labels`` arguments.
+
+    :param item: Dataset item to transform
+    :type item: dict
+    :param atom_key: name of key containing molecule structure as a dataframe, defaults to 'atoms'
+    :type atom_keys: list, optional
+    :param label_key: name of key containing labels, defaults to 'scores'
+    :type label_key: str, optional
+    :param use_bonds: whether to use molecular bond information for edges instead of distance. Assumes bonds are stored under 'bonds' key, defaults to False
+    :type use_bonds: bool, optional
+    :return: Transformed Dataset item
+    :rtype: dict
+    """    
+
+    bonds = item['bonds'] if use_bonds else  None
+     
+    node_feats, edge_index, edge_feats, pos = mol_df_to_graph_for_qm(item[atom_key], bonds=bonds, onehot_edges=onehot_edges, allowable_atoms=allowable_atoms, edge_dist_cutoff=edge_dist_cutoff)
+    item[atom_key] = Data(node_feats, edge_index, edge_feats, y=item[label_key], pos=pos)
+
+    return item