Implement basic 3DMol.js

app.py

@@ -1,17 +1,22 @@
-import shutil
+import copy
 import tempfile
 from pathlib import Path
 from functools import partial
 
 import gradio as gr
-import numpy as np
 import torch
+import py3Dmol
+from biopandas.pdb import PandasPdb
+from colour import Color
 
 from mutils.pdb import download_pdb
+from mutils.mutations import Mutation
 from ppiref.extraction import PPIExtractor
 from ppiref.utils.ppi import PPIPath
+from ppiref.utils.residue import Residue
 from ppiformer.tasks.node import DDGPPIformer
 from ppiformer.utils.api import predict_ddg
+from ppiformer.utils.torch import fill_diagonal
 from ppiformer.definitions import PPIFORMER_WEIGHTS_DIR
 
 
@@ -63,26 +68,149 @@ def process_inputs(inputs, temp_dir):
     
     muts = list(map(lambda x: x.strip(), muts.split(';')))
 
-    return ppi_path, muts
+    return pdb_path, ppi_path, muts
+
+
+def plot_3dmol(pdb_path, ppi_path, muts, attn):
+    # 3DMol.js adapted from https://huggingface.co/spaces/huhlim/cg2all/blob/main/app.py
+
+    # Read PDB for 3Dmol.js
+    with open(pdb_path, "r") as fp:
+        lines = fp.readlines()
+    mol = ""
+    for l in lines:
+        mol += l
+    mol = mol.replace("OT1", "O  ")
+    mol = mol.replace("OT2", "OXT")
+
+    # Read PPI to customize 3Dmol.js visualization
+    ppi_df = PandasPdb().read_pdb(ppi_path).df['ATOM']
+    ppi_df = ppi_df.groupby(list(Residue._fields)).apply(lambda df: df[df['atom_name'] == 'CA'].iloc[0]).reset_index(drop=True)
+    chains = ppi_df['chain_id'].unique()
+    ppi_df['id'] = ppi_df.apply(lambda row: ':'.join([row['residue_name'], row['chain_id'], str(row['residue_number']), row['insertion']]), axis=1)
+    ppi_df['id'] = ppi_df['id'].apply(lambda x: x[:-1] if x[-1] == ':' else x)
+    muts_id = sum([Mutation(mut).wt_to_graphein() for mut in muts], start=[])  # flatten ids of all sp muts 
+    ppi_df['mutated'] = ppi_df.apply(lambda row: row['id'] in muts_id, axis=1)
+
+    # Prepare attention coeffictients per residue (normalized sum of direct attention from mutated residues)
+    attn = torch.nan_to_num(attn, nan=1e-10)
+    attn_sub = attn[:, 0, :, 0, :, :, :]  # models, layers, heads, tokens, tokens
+    idx_mutated = torch.from_numpy(ppi_df.index[ppi_df['mutated']].to_numpy())
+    attn_sub = fill_diagonal(attn_sub, 1e-10)
+    attn_mutated = attn_sub[..., idx_mutated, :]
+    attn_mutated.shape
+    attns_per_token = torch.sum(attn_mutated, dim=(0, 1, 2, 3))
+    attns_per_token = (attns_per_token - attns_per_token.min()) / (attns_per_token.max() - attns_per_token.min())
+    attns_per_token += 1e-10
+    ppi_df['attn'] = attns_per_token.numpy()
+
+    # Customize 3Dmol.js visualization https://3dmol.csb.pitt.edu/doc/
+    styles = []
+    zoom_atoms = []
+
+    # Cartoon chains
+    colors = [Color(c) for c in ['LimeGreen', 'HotPink', 'RoyalBlue']]
+    chain_to_color = dict(zip(chains, colors))
+    for chain in chains:
+        styles.append([{"chain": chain}, {"cartoon": {"color": chain_to_color[chain].hex_l, "opacity": 0.6}}])
+
+    # Stick PPI and atoms for zoom
+    # TODO Insertions
+    for _, row in ppi_df.iterrows():
+        color = copy.deepcopy(chain_to_color[row['chain_id']])
+        color.saturation = row['attn']
+        color = color.hex_l
+        if row['mutated']:
+            styles.append([{'chain': row['chain_id'], 'resi': str(row['residue_number'])}, {'stick': {'color': 'red', 'radius': 0.2, 'opacity': 1.0}}])
+            zoom_atoms.append(row['atom_number'])
+        else:
+            styles.append([{'chain': row['chain_id'], 'resi': str(row['residue_number'])}, {'stick': {'color': color, 'radius': row['attn'] / 5, 'opacity': row['attn']}}])
+
+    # Convert style dicts to JS lines
+    styles = '\n'.join(['viewer.addStyle(' + ', '.join([str(s).replace("'", '"') for s in dcts]) + ');' for dcts in styles])
+
+    # Connert zoom atoms to 3DMol.js selection
+    zoom = 'viewer.zoomTo({\"or\": [' + ', '.join(["{\"serial\": " + str(a) + "}" for a in zoom_atoms]) + ']}, 1000);'
+
+    # Construct 3Dmol.js visualization script in HTML
+    html = (
+        """<!DOCTYPE html>
+        <html>
+        <head>    
+    <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+    <style>
+    body{
+        font-family:sans-serif
+    }
+    .mol-container {
+    width: 100%;
+    height: 600px;
+    position: relative;
+    }
+    .mol-container select{
+        background-image:None;
+    }
+    </style>
+     <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+    <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>  
+    <div id="container" class="mol-container"></div>
+  
+            <script>
+               let pdb = `"""
+        + mol
+        + """`  
+      
+             $(document).ready(function () {
+                let element = $("#container");
+                let config = { backgroundColor: "white" };
+                let viewer = $3Dmol.createViewer(element, config);
+                viewer.addModel(pdb, "pdb");
+                viewer.setBackgroundColor("black");
+                viewer.setStyle({"model": 0}, {"ray_opaque_background": "off"}, {"stick": {"color": "lightgrey", "opacity": 0.5}});
+          """
+        + styles
+        + zoom
+        + """
+                viewer.render();
+              })
+        </script>
+        </body></html>"""
+    )
+    print(html)
+
+    return f"""<iframe style="width: 100%; height: 600px" 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='{html}'></iframe>"""
 
 
 def predict(models, temp_dir, *inputs):
     # Process input
-    ppi_path, muts = process_inputs(inputs, temp_dir)
+    pdb_path, ppi_path, muts = process_inputs(inputs, temp_dir)
 
     print(ppi_path, muts)
 
     # Predict
     ddg, attn = predict_ddg(models, ppi_path, muts, return_attn=True)
 
+    # Create dataframe
     ddg = ddg.detach().numpy().tolist()
     df = list(zip(muts, ddg))
-    
-    return df
+
+    # Create 3DMol plot
+    plot = plot_3dmol(pdb_path, ppi_path, muts, attn)
+
+    return df, plot
 
 
 app = gr.Blocks()
 with app:
+    # print('app.theme.background_fill_primary', app.theme.background_fill_primary, type(app.theme.background_fill_primary))
+    # print('app.theme.background_fill_primary', app.theme.background_fill_primary_dark, type(app.theme.background_fill_primary))
+    # print(app.theme.to_dict())
 
     # Input GUI
     with gr.Row():
@@ -114,6 +242,7 @@ with app:
         datatype=["str", "number"],
         col_count=(2, "fixed"),
     )
+    plot = gr.HTML()
 
     # Load models
     models = [
@@ -130,7 +259,8 @@ with app:
 
     # Main logic
     inputs = [pdb_code, pdb_path, partners, muts, muts_path]
+    outputs = [df, plot]
     predict = partial(predict, models, temp_dir)
-    predict_button.click(predict, inputs=inputs, outputs=df)
+    predict_button.click(predict, inputs=inputs, outputs=outputs)
 
 app.launch()