Implement text labels and more examples

app.py

@@ -5,10 +5,11 @@ from functools import partial
 
 import gradio as gr
 import torch
-import py3Dmol
+from Bio.PDB.Polypeptide import protein_letters_3to1
 from biopandas.pdb import PandasPdb
 from colour import Color
 
+from mutils.proteins import AMINO_ACID_CODES_1
 from mutils.pdb import download_pdb
 from mutils.mutations import Mutation
 from ppiref.extraction import PPIExtractor
@@ -121,16 +122,29 @@ def plot_3dmol(pdb_path, ppi_path, muts, attn):
         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}}])
+            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']}}])
+            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);'
+    styles = ''.join(['viewer.addStyle(' + ', '.join([str(s).replace("'", '"') for s in dcts]) + ');\n' for dcts in styles])
+
+    # Connert zoom atoms to 3DMol.js selection and add labels for mutated residues
+    zoom_animation_duration = 500
+    sel = '{\"or\": [' + ', '.join(["{\"serial\": " + str(a) + "}" for a in zoom_atoms]) + ']}'
+    zoom = 'viewer.zoomTo(' + sel + ',' + f'{zoom_animation_duration});'
+    for atom in zoom_atoms:
+        sel = '{\"serial\": ' + str(atom) + '}'
+        row = ppi_df[ppi_df['atom_number'] == atom].iloc[0]
+        label = protein_letters_3to1[row['residue_name']] + row['chain_id'] + str(row['residue_number']) + row['insertion']
+        styles += 'viewer.addLabel(' + f"\"{label}\"," + "{fontSize:16, fontColor:\"red\", backgroundOpacity: 0.0}," + sel + ');\n'
 
     # Construct 3Dmol.js visualization script in HTML
     html = (
@@ -164,10 +178,9 @@ def plot_3dmol(pdb_path, ppi_path, muts, attn):
       
              $(document).ready(function () {
                 let element = $("#container");
-                let config = { backgroundColor: "white" };
+                let config = { backgroundColor: "black" };
                 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
@@ -227,7 +240,10 @@ with app:
             muts_path = gr.File(file_count="single", label="Or file with mutations")
 
     examples = gr.Examples(
-        examples=[["1BUI", "A,B,C", "SC16A;FC47A;SC16A,FC47A"]],
+        examples=[
+            ["1BUI", "A,B,C", "SC16A;FC47A;SC16A,FC47A"],
+            ["1KNE", "A,P", ';'.join([f"TP6{a}" for a in AMINO_ACID_CODES_1])]
+        ],
         inputs=[pdb_code, partners, muts],
         label="Examples (press line to fill)"
     )