inference_app.py

import random
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 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, 
    ligand_pdb: Path, 
    receptor_fasta: Path | None = None, 
    ligand_fasta: Path | None = None,
) -> tuple[str, float]:
    start_time = time.time()
    # Do inference here
    # return an output pdb file with the protein and two chains R and L.  
    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
    # Inititalize an empty stack with shape (m x n x 3) 
    stack = AtomArrayStack(M, ligand.shape[0])
    
    # copy annotations from ligand
    for annot in ligand.get_annotation_categories():
        stack.set_annotation(annot, np.copy(ligand.get_annotation(annot)))
    
    # Random translations sampled along 0-50 angstroms per axis
    translation_magnitudes = np.linspace(
        0, 26,
        num=26,
        endpoint=False
    )
    # generate one pose at a time
    for i in range(M):
        q = R.random()
        translation_vec = [
            random.choice(translation_magnitudes), # x
            random.choice(translation_magnitudes), # y
            random.choice(translation_magnitudes), # z
        ]
        # transform the ligand chain
        stack.coord[i, ...] = q.apply(ligand.coord) + translation_vec
        
    # Find clashes (1.2 A contact radius)
    stack_conts = get_stack_contacts(receptor, stack, threshold=1.2)
    
    # Keep the "best" pose based on pose w/fewest clashes
    pose_clashes = []
    for i in range(stack_conts.shape[0]):
        pose_conts = stack_conts[i]
        pose_clashes.append((i, np.argwhere(pose_conts != -1).shape[0]))
    
    best_pose_idx = sorted(pose_clashes, key=lambda x: x[1])[0][0]
    best_pose = receptor + stack[best_pose_idx]
    
    output_dir = Path(receptor_pdb).parent
    # System ID
    pdb_name = Path(receptor_pdb).stem + "--" + Path(ligand_pdb).name
    output_pdb = output_dir / pdb_name
    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

      # 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.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")

      eval_btn.click(evaluate, inputs=[input_system_id, input_prediction_pdb], outputs=[metric_table, pred_native, eval_run_time])

app.launch()