feat: add example evaluate endpoint to fetch metrics for a single prediction given model and system ID

inference_app.py

@@ -3,15 +3,20 @@ import time
 from pathlib import Path
 
 import numpy as np
+import pandas as pd
 from biotite.structure.atoms import AtomArrayStack
 from scipy.spatial.transform import Rotation as R
-from pinder.core.structure.atoms import atom_array_from_pdb_file, normalize_orientation, write_pdb
+from pinder.core import PinderSystem
+from pinder.core.structure import atoms
 from pinder.core.structure.contacts import get_stack_contacts
+from pinder.core.loader.structure import Structure
+from pinder.eval.dockq import BiotiteDockQ
 
 import gradio as gr
 
 from gradio_molecule3d import Molecule3D
 
+EVAL_METRICS = ["system", "L_rms", "I_rms", "F_nat", "DOCKQ", "CAPRI_class"]
 
 def predict(
     receptor_pdb: Path, 
@@ -22,10 +27,10 @@ def predict(
     start_time = time.time()
     # Do inference here
     # return an output pdb file with the protein and two chains R and L.  
-    receptor = atom_array_from_pdb_file(receptor_pdb, extra_fields=["b_factor"])
-    ligand = atom_array_from_pdb_file(ligand_pdb, extra_fields=["b_factor"])
-    receptor = normalize_orientation(receptor)
-    ligand = normalize_orientation(ligand)
+    receptor = atoms.atom_array_from_pdb_file(receptor_pdb, extra_fields=["b_factor"])
+    ligand = atoms.atom_array_from_pdb_file(ligand_pdb, extra_fields=["b_factor"])
+    receptor = atoms.normalize_orientation(receptor)
+    ligand = atoms.normalize_orientation(ligand)
     
     # Number of random poses to generate
     M = 50
@@ -69,79 +74,145 @@ def predict(
     # System ID
     pdb_name = Path(receptor_pdb).stem + "--" + Path(ligand_pdb).name
     output_pdb = output_dir / pdb_name
-    write_pdb(best_pose, output_pdb)
+    atoms.write_pdb(best_pose, output_pdb)
     end_time = time.time()
     run_time = end_time - start_time
     return str(output_pdb), run_time
 
 
+def evaluate(
+    system_id: str, 
+    prediction_pdb: Path, 
+) -> tuple[pd.DataFrame, float]:
+    start_time = time.time()
+    system = PinderSystem(system_id)
+    native = system.native.filepath
+    bdq = BiotiteDockQ(native, Path(prediction_pdb), parallel_io=False)
+    metrics = bdq.calculate()
+    metrics = metrics[["system", "LRMS", "iRMS", "Fnat", "DockQ", "CAPRI"]].copy()
+    metrics.rename(columns={"LRMS": "L_rms", "iRMS": "I_rms", "Fnat": "F_nat", "DockQ": "DOCKQ", "CAPRI": "CAPRI_class"}, inplace=True)
+    end_time = time.time()
+    run_time = end_time - start_time
+    pred = Structure(Path(prediction_pdb))
+    nat = Structure(Path(native))
+    pred, _, _ = pred.superimpose(nat)
+    pred.to_pdb(Path(prediction_pdb))
+    return metrics, [str(prediction_pdb), str(native)], run_time
+
+
 with gr.Blocks() as app:
+    with gr.Tab("🧬 PINDER inference template"):
+      gr.Markdown("Title, description, and other information about the model")   
+      with gr.Row():
+          with gr.Column():
+              input_protein_1 = gr.File(label="Input Protein 1 monomer (PDB)")
+              input_fasta_1 = gr.File(label="Input Protein 1 monomer sequence (FASTA)")
+          with gr.Column():
+              input_protein_2 = gr.File(label="Input Protein 2 monomer (PDB)")
+              input_fasta_2 = gr.File(label="Input Protein 2 monomer sequence (FASTA)")
+          
+          
+      
+      # define any options here
 
-    gr.Markdown("# Template for inference")
+      # for automated inference the default options are used
+      # slider_option = gr.Slider(0,10, label="Slider Option")
+      # checkbox_option = gr.Checkbox(label="Checkbox Option")
+      # dropdown_option = gr.Dropdown(["Option 1", "Option 2", "Option 3"], label="Radio Option")
 
-    gr.Markdown("Title, description, and other information about the model")   
-    with gr.Row():
-        with gr.Column():
-            input_protein_1 = gr.File(label="Input Protein 1 monomer (PDB)")
-            input_fasta_1 = gr.File(label="Input Protein 1 monomer sequence (FASTA)")
-        with gr.Column():
-            input_protein_2 = gr.File(label="Input Protein 2 monomer (PDB)")
-            input_fasta_2 = gr.File(label="Input Protein 2 monomer sequence (FASTA)")
-        
-        
-    
-    # define any options here
-
-    # for automated inference the default options are used
-    # slider_option = gr.Slider(0,10, label="Slider Option")
-    # checkbox_option = gr.Checkbox(label="Checkbox Option")
-    # dropdown_option = gr.Dropdown(["Option 1", "Option 2", "Option 3"], label="Radio Option")
-
-    btn = gr.Button("Run Inference")
-
-    gr.Examples(
-        [
-            [
-                "8i5w_R.pdb",
-                "8i5w_R.fasta",
-                "8i5w_L.pdb",
-                "8i5w_L.fasta",
-            ],
-        ],
-        [input_protein_1, input_fasta_1, input_protein_2, input_fasta_2],
-    )
-    reps =    [
-    {
-      "model": 0,
-      "style": "cartoon",
-      "chain": "R",
-      "color": "whiteCarbon",
-    },
-    {
-      "model": 0,
-      "style": "cartoon",
-       "chain": "L",
-      "color": "greenCarbon",
-    },
-    {
-      "model": 0,
-      "chain": "R",
-      "style": "stick",
-      "sidechain": True,
-      "color": "whiteCarbon",
-    },
-    {
-      "model": 0,
-      "chain": "L",
-      "style": "stick",
-      "sidechain": True,
-      "color": "greenCarbon"
-    }      
-  ]
-    
-    out = Molecule3D(reps=reps)
-    run_time = gr.Textbox(label="Runtime")
+      btn = gr.Button("Run Inference")
+
+      gr.Examples(
+          [
+              [
+                  "8i5w_R.pdb",
+                  "8i5w_R.fasta",
+                  "8i5w_L.pdb",
+                  "8i5w_L.fasta",
+              ],
+          ],
+          [input_protein_1, input_fasta_1, input_protein_2, input_fasta_2],
+      )
+      reps =    [
+      {
+        "model": 0,
+        "style": "cartoon",
+        "chain": "R",
+        "color": "whiteCarbon",
+      },
+      {
+        "model": 0,
+        "style": "cartoon",
+        "chain": "L",
+        "color": "greenCarbon",
+      },
+      {
+        "model": 0,
+        "chain": "R",
+        "style": "stick",
+        "sidechain": True,
+        "color": "whiteCarbon",
+      },
+      {
+        "model": 0,
+        "chain": "L",
+        "style": "stick",
+        "sidechain": True,
+        "color": "greenCarbon"
+      }      
+    ]
+      
+      out = Molecule3D(reps=reps)
+      run_time = gr.Textbox(label="Runtime")
+
+      btn.click(predict, inputs=[input_protein_1, input_protein_2, input_fasta_1,  input_fasta_2], outputs=[out, run_time])
+    with gr.Tab("⚖️ PINDER evaluation template"):
+      with gr.Row():
+          with gr.Column():
+              input_system_id = gr.Textbox(label="PINDER system ID")
+              input_prediction_pdb = gr.File(label="Top ranked prediction (PDB with chains R and L)")
+
+      eval_btn = gr.Button("Run Evaluation")
+      gr.Examples(
+          [
+              [
+                  "3g9w__A1_Q71LX4--3g9w__D1_P05556",
+                  "3g9w_R--3g9w_L.pdb",
+              ],
+          ],
+          [input_system_id, input_prediction_pdb],
+      )
+      reps = [
+      {
+        "model": 0,
+        "style": "cartoon",
+        "chain": "R",
+        "color": "greenCarbon",
+      },
+      {
+        "model": 0,
+        "style": "cartoon",
+        "chain": "L",
+        "color": "cyanCarbon",
+      },
+      {
+        "model": 1,
+        "style": "cartoon",
+        "chain": "R",
+        "color": "grayCarbon",
+      },
+      {
+        "model": 1,
+        "style": "cartoon",
+        "chain": "L",
+        "color": "blueCarbon",
+      },
+    ]
+      
+      pred_native = Molecule3D(reps=reps, config={"backgroundColor": "black"})
+      eval_run_time = gr.Textbox(label="Evaluation runtime")
+      metric_table = gr.DataFrame(pd.DataFrame([], columns=EVAL_METRICS),label="Evaluation metrics")
 
-    btn.click(predict, inputs=[input_protein_1, input_protein_2, input_fasta_1,  input_fasta_2], outputs=[out, run_time])
+      eval_btn.click(evaluate, inputs=[input_system_id, input_prediction_pdb], outputs=[metric_table, pred_native, eval_run_time])
 
 app.launch()