feat: plot plddt

Dockerfile

@@ -27,6 +27,7 @@ RUN cd /app/folding-studio && uv pip install -e .
 COPY app.py /app/app.py
 COPY molecule.py /app/molecule.py
 # COPY pred.cif /app/boltz_results/pred_model_0.cif
+# COPY plddt_0.npz /app/boltz_results/plddt_0.npz
 
 EXPOSE 7860
 ENV GRADIO_SERVER_NAME="0.0.0.0"

app.py

@@ -1,12 +1,13 @@
-import hashlib
-import gradio as gr
-from pathlib import Path
 from Bio.PDB import MMCIFParser, PDBIO
-import logging
-import os
 from folding_studio.client import Client
 from folding_studio.query.boltz import BoltzQuery, BoltzParameters
-
+from pathlib import Path
+import gradio as gr
+import hashlib
+import logging
+import numpy as np
+import os
+import plotly.graph_objects as go
 
 from molecule import molecule
 
@@ -21,7 +22,7 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-def convert_cif_to_pdb(cif_path, pdb_path):
+def convert_cif_to_pdb(cif_path: str, pdb_path: str) -> None:
     """Convert a .cif file to .pdb format using Biopython.
     
     Args:
@@ -117,10 +118,13 @@ def predict(sequence: str, api_key: str) -> str:
     mol = _create_molecule_visualization(
         converted_pdb_path,
         sequence,
-        output_dir
-    )
-    
-    return _wrap_in_iframe(mol)
+    )    
+
+    plddt_file = list(pred_cif.parent.glob("plddt_*.npz"))[0]
+    logger.info("plddt file: %s", plddt_file)
+    plddt_vals = np.load(plddt_file)["plddt"]
+
+    return _wrap_in_iframe(mol), add_plddt_plot(plddt_vals=plddt_vals)
 
 
 def _write_fasta_file(filepath: Path, sequence: str) -> None:
@@ -129,7 +133,7 @@ def _write_fasta_file(filepath: Path, sequence: str) -> None:
         f.write(f">A|protein\n{sequence}")
 
 
-def _create_molecule_visualization(pdb_path: Path, sequence: str, output_dir: Path) -> str:
+def _create_molecule_visualization(pdb_path: Path, sequence: str) -> str:
     """Create molecular visualization using molecule module."""
     return molecule(
         str(pdb_path),
@@ -160,6 +164,28 @@ def _wrap_in_iframe(content: str) -> str:
         srcdoc='{content}'
     ></iframe>"""
 
+def add_plddt_plot(plddt_vals: list[float]) -> str:
+    """Create a plot of metrics."""
+    visible = True
+    plddt_trace = go.Scatter(
+            x=np.arange(len(plddt_vals)),
+            y=plddt_vals,
+            hovertemplate="<i>pLDDT</i>: %{y:.2f} <br><i>Residue index:</i> %{x}<br>",
+            name="seq",
+            visible=visible,
+        )
+    
+    plddt_fig = go.Figure(data=[plddt_trace])
+    plddt_fig.update_layout(
+        title="pLDDT",
+        xaxis_title="Residue index",
+        yaxis_title="pLDDT",
+        height=500,
+        template="simple_white",
+        legend=dict(yanchor="bottom", y=0.01, xanchor="left", x=0.99),
+    )
+    return plddt_fig
+
 demo = gr.Blocks(title="Folding Studio: structure prediction with Boltz-1")
 
 with demo:
@@ -172,12 +198,15 @@ with demo:
     with gr.Row():
         predict_btn = gr.Button("Predict")
     with gr.Row():
-        mol_output = gr.HTML()
-    
+        with gr.Column():
+            mol_output = gr.HTML()
+        with gr.Column():
+            metrics_plot = gr.Plot(label="pLDDT")
+           
     predict_btn.click(
         fn=predict,
         inputs=[sequence, api_key],
-        outputs=[mol_output]
+        outputs=[mol_output, metrics_plot]
     )
 
 demo.launch()