Add protein visualization capabilities

app.py

@@ -1,6 +1,7 @@
 
 import gradio as gr
 import numpy as np
+import os, tempfile
 import torch
 import py3Dmol
 from huggingface_hub import login
@@ -12,6 +13,9 @@ from esm.sdk.api import (
     GenerationConfig,
 )
 
+from gradio_molecule3d import Molecule3D
+
+
 theme = gr.themes.Monochrome(
     primary_hue="gray",
 )
@@ -28,26 +32,51 @@ def get_model(model_name, token):
     # model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
     return model
 
-## Function to render 3D structure using py3Dmol
-def render_pdb(pdb_string, motif_start=None, motif_end=None):
-    view = py3Dmol.view(width=800, height=800)
-    view.addModel(pdb_string, "pdb")
-    view.setStyle({"cartoon": {"color": "spectrum"}})
-    if motif_start is not None and motif_end is not None:
-        motif_inds = np.arange(motif_start, motif_end)
-        view.setStyle({"cartoon": {"color": "lightgrey"}})
-        motif_res_inds = (motif_inds + 1).tolist()
-        view.addStyle({"resi": motif_res_inds}, {"cartoon": {"color": "cyan"}})
-    view.zoomTo()
-    return view
-
 ## Function to get PDB data
 def get_pdb(pdb_id, chain_id):
     pdb = ProteinChain.from_rcsb(pdb_id, chain_id)
     # return [pdb.sequence, render_pdb(pdb.to_pdb_string())]
     return pdb
 
+## Function to generate rep for 3D structure
+def make_reps(res_start=None, res_end=None, main_color="whiteCarbon", highlight_color="redCarbon", main_style="cartoon", highlight_style="cartoon"):
+    residue_range = f"{res_start}-{res_end}" if res_start != res_end else ""
 
+    return [
+     {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": main_style,
+        "color": main_color,
+        "residue_range": "",
+        "around": 0,
+        "byres": False,
+        "visible": True
+    },
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": highlight_style,
+        "color": highlight_color,
+        "residue_range": residue_range,
+        "around": 0,
+        "byres": False,
+        "visible": True
+    }]
+
+## Function to render 3D structure
+def render_pdb(pdb_id, chain_id, res_start, res_end, pdb_string=None):
+    if pdb_string is None:
+        pdb_string = get_pdb(pdb_id, chain_id).to_pdb_string()
+    ## Write to temporary file and read back in to get the 3D structure
+    tmp_pdb = tempfile.NamedTemporaryFile(delete=False, prefix=f"{pdb_id}_chain{chain_id}_", suffix=".pdb")
+    tmp_pdb.write(str.encode(pdb_string))
+
+    return Molecule3D(tmp_pdb.name, reps=make_reps(res_start=res_start, res_end=res_end))
+
+## Function for Scaffolding
 def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt_length, insert_size):
     pdb = get_pdb(pdb_id, chain_id)
 
@@ -75,16 +104,39 @@ def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt
     sequence_generation = model.generate(protein_prompt, sequence_generation_config)
     generated_sequence = sequence_generation.sequence
 
+    ## Generate structure
+    structure_prediction_config = GenerationConfig(
+        track="structure", # We want ESM3 to generate tokens for the structure track
+        num_steps=len(sequence_generation) // 8,
+        temperature=0.7, 
+    )
+    structure_prediction_prompt = ESMProtein(sequence=sequence_generation.sequence)
+    structure_prediction = model.generate(structure_prediction_prompt, structure_prediction_config)
+    ## Convert the generated structure to a back into a ProteinChain object
+    structure_prediction_chain = structure_prediction.to_protein_chain()
+    motif_inds_in_generation = np.arange(insert_size, insert_size+len(motif_sequence))
+    structure_prediction_chain.align(pdb, mobile_inds=motif_inds_in_generation, target_inds=motif_inds)
+    # crmsd = structure_prediction_chain.rmsd(renal_dipep_chain, mobile_inds=motif_inds_in_generation, target_inds=motif_inds)
+
+    structure_orig_highlight = render_pdb(pdb_id, chain_id, res_start=motif_start, res_end=motif_end)
+    structure_new_highlight = render_pdb(pdb_id, chain_id, res_start=insert_size, res_end=insert_size+len(motif_sequence), 
+                                         pdb_string=structure_prediction_chain.to_pdb_string())
+
     return [
         pdb.sequence,
         motif_sequence,
+        structure_orig_highlight,
         # motif_atom37_positions,
         sequence_prompt,
         # structure_prompt,
         # protein_prompt
-        generated_sequence
+        generated_sequence,
+        # structure_prediction,
+        # structure_prediction_chain,
+        structure_new_highlight
     ]
 
+## Function for Secondary Structure Editing
 def ss_edit(model_name, token, pdb_id, chain_id, region_start, region_end, shortened_region_length, shortening_ss8):
     pdb = get_pdb(pdb_id, chain_id)
     edit_region = np.arange(region_start, region_end)
@@ -109,17 +161,28 @@ def ss_edit(model_name, token, pdb_id, chain_id, region_start, region_end, short
     model = get_model(model_name, token)
     sequence_generation = model.generate(protein_prompt, GenerationConfig(track="sequence", num_steps=protein_prompt.sequence.count("_") // 2, temperature=0.5))
     
+    ## Generate structure
+    structure_prediction = model.generate(ESMProtein(sequence=sequence_generation.sequence), GenerationConfig(track="structure", num_steps=len(protein_prompt) // 4, temperature=0))
+    structure_prediction_chain = structure_prediction.to_protein_chain()
+
+    structure_orig_highlight = render_pdb(pdb_id, chain_id, res_start=region_start, res_end=region_end)
+    structure_new_highlight = render_pdb(pdb_id, chain_id, res_start=region_start, res_end=region_end,
+                                         pdb_string=structure_prediction_chain.to_pdb_string())
+
     return [
         original_sequence,
         original_ss8,
         original_ss8_region,
+        structure_orig_highlight,
         sequence_prompt,
         ss8_prompt,
         proposed_ss8_region,
         # protein_prompt,
-        sequence_generation
+        sequence_generation,
+        structure_new_highlight
         ]
 
+## Function for SASA Editing
 def sasa_edit(model_name, token, pdb_id, chain_id, span_start, span_end, n_samples):
     pdb = get_pdb(pdb_id, chain_id)
 
@@ -144,10 +207,16 @@ def sasa_edit(model_name, token, pdb_id, chain_id, span_start, span_end, n_sampl
     ## Sort generations by ptm
     generated_proteins = sorted(generated_proteins, key=lambda x: x.ptm.item(), reverse=True)
 
+    structure_orig_highlight = render_pdb(pdb_id, chain_id, res_start=span_start, res_end=span_end)
+    structure_new_highlight = render_pdb(pdb_id, chain_id, res_start=span_start, res_end=span_end,
+                                         pdb_string=generated_proteins[0].to_protein_chain().to_pdb_string())
+
     return [
-        protein_prompt,
-        sequence_generation,
-        generated_proteins
+        protein_prompt.sequence,
+        structure_orig_highlight,
+        [seq.sequence for seq in sequence_generation],
+        # [pro.sequence for pro in generated_proteins]
+        structure_new_highlight
     ]
 
 
@@ -166,13 +235,14 @@ scaffold_app = gr.Interface(
         ],
     outputs=[
         gr.Textbox(label="Sequence"),
-        # gr.Plot(label="3D Structure")
         gr.Textbox(label="Motif Sequence"),
+        Molecule3D(label="Original Structure"),
         # gr.Textbox(label="Motif Positions")
         gr.Textbox(label="Sequence Prompt"),
         # gr.Textbox(label="Structure Prompt"),
         # gr.Textbox(label="Protein Prompt"),
-        gr.Textbox(label="Generated Sequence")
+        gr.Textbox(label="Generated Sequence"),
+        Molecule3D(label="Generated Structure")
     ]
     )
 
@@ -193,11 +263,13 @@ ss_app = gr.Interface(
         gr.Textbox(label="Original Sequence"),
         gr.Textbox(label="Original SS8"),
         gr.Textbox(label="Original SS8 Edit Region"),
+        Molecule3D(label="Original Structure"),
         gr.Textbox(label="Sequence Prompt"),
         gr.Textbox(label="Edited SS8 Prompt"),
         gr.Textbox(label="Proposed SS8 of Edit Region"),
         # gr.Textbox(label="Protein Prompt"),
-        gr.Textbox(label="Generated Sequence")
+        gr.Textbox(label="Generated Sequence"),
+        Molecule3D(label="Generated Structure")
     ]
     )
 
@@ -212,15 +284,32 @@ sasa_app = gr.Interface(
         gr.Number(value=105, label="Span Start"),
         gr.Number(value=116, label="Span End"),
         # gr.Textbox(value="CCSSCCCCSSCHHHHHHTEEETTBBTTBCSSEEEEECCTTCCHHHHHTTTHHHHHHHTTCEEEEECCTTTTCSCHHHHHHHHHHHHHHHHHHTTSCCEEEEEETHHHHHHHHHHHHCGGGGGTEEEEEEESCCTTCBGGGHHHHHTTCBCHHHHHTBTTCHHHHHHHHTTTTBCSSCEEEEECTTCSSSCCCCSSSTTSTTCCBTSEEEEHHHHHCTTCCCCSHHHHHBHHHHHHHHHHHHCTTSSCCGGGCCSTTCCCSBCTTSCHHHHHHHHSTHHHHHHHHHHSCCBSSCCCCCGGGGGGSTTCEETTEECCC", label="SS8 String")
-        gr.Number(value=4, label="Number of Samples")
+        gr.Number(value=1, label="Number of Samples")
     ],
     outputs = [
         gr.Textbox(label="Protein Prompt"),
+        Molecule3D(label="Original Structure"),
         gr.Textbox(label="Generated Sequences"),
-        gr.Textbox(label="Generated Proteins")
+        # gr.Textbox(label="Generated Proteins")
+        Molecule3D(label="Best Generated Structure")
     ]
 )
 
+
+protein_viewer = gr.Interface(
+    fn=render_pdb,
+    inputs=[
+        gr.Textbox(value = "1LBS", label="PDB ID"),
+        gr.Textbox(value = "A", label="Chain ID"),
+        gr.Number(value=10, label="Residue Highlight Start"),
+        gr.Number(value=20, label="Residue Highlight End")
+    ],
+    outputs=[
+        Molecule3D(label="3D Structure")
+    ]
+)
+
+
 ## Main Interface
 with gr.Blocks(theme=theme) as esm_app:
     with gr.Row():
@@ -239,12 +328,14 @@ with gr.Blocks(theme=theme) as esm_app:
         gr.TabbedInterface([
             scaffold_app,
             ss_app,
-            sasa_app
+            sasa_app,
+            protein_viewer
             ],
             [
                 "Scaffolding Example",
                 "Secondary Structure Editing Example",
-                "SASA Editing Example"
+                "SASA Editing Example",
+                "PDB Viewer"
             ])
 
 if __name__ == "__main__":

requirements.txt

@@ -2,4 +2,5 @@ esm
 numpy
 torch>=2.3.0
 py3Dmol
-huggingface_hub
+huggingface_hub
+gradio_molecule3d