feat: update inference app template to implement a dummy poc predict endpoint that follows the suggested input specs

inference_app.py

@@ -1,20 +1,81 @@
-
+import random
 import time
+from pathlib import Path
+
+import numpy as np
+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.structure.contacts import get_stack_contacts
 
 import gradio as gr
 
 from gradio_molecule3d import Molecule3D
 
 
-
-
-def predict (input_seq_1, input_protein_1, input_seq_2,  input_protein_2):
+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 A and B.  
+    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)
+    
+    # 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, M + 1,
+        num=M + 1,
+        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 = "--".join([
+        Path(receptor_pdb).stem.rstrip("-R"), Path(ligand_pdb).stem.rstrip("-L")
+    ]) + ".pdb"
+    output_pdb = output_dir / pdb_name
+    write_pdb(best_pose, output_pdb)
     end_time = time.time()
     run_time = end_time - start_time
-    return "test_out.pdb", run_time
+    return str(output_pdb), run_time
+
 
 with gr.Blocks() as app:
 
@@ -23,11 +84,11 @@ with gr.Blocks() as app:
     gr.Markdown("Title, description, and other information about the model")   
     with gr.Row():
         with gr.Column():
-            input_seq_1 = gr.Textbox(lines=3, label="Input Protein 1 sequence (FASTA)")
-            input_protein_1 = gr.File(label="Input Protein 2 monomer (PDB)")
+            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_seq_2 = gr.Textbox(lines=3, label="Input Protein 1 sequence (FASTA)")
-            input_protein_2 = gr.File(label="Input Protein 2 structure (PDB)")
+            input_protein_2 = gr.File(label="Input Protein 2 monomer (PDB)")
+            input_fasta_2 = gr.File(label="Input Protein 2 monomer sequence (FASTA)")
         
         
     
@@ -43,38 +104,37 @@ with gr.Blocks() as app:
     gr.Examples(
         [
             [
-                "GSGSPLAQQIKNIHSFIHQAKAAGRMDEVRTLQENLHQLMHEYFQQSD",
-                "3v1c_A.pdb",
-                "GSGSPLAQQIKNIHSFIHQAKAAGRMDEVRTLQENLHQLMHEYFQQSD",
-                "3v1c_B.pdb",
-                
+                "8i5w_R.pdb",
+                "8i5w_R.fasta",
+                "8i5w_L.pdb",
+                "8i5w_L.fasta",
             ],
         ],
-        [input_seq_1, input_protein_1, input_seq_2,  input_protein_2],
+        [input_protein_1, input_fasta_1, input_protein_2, input_fasta_2],
     )
     reps =    [
     {
       "model": 0,
       "style": "cartoon",
-      "chain": "A",
+      "chain": "R",
       "color": "whiteCarbon",
     },
     {
       "model": 0,
       "style": "cartoon",
-       "chain": "B",
+       "chain": "L",
       "color": "greenCarbon",
     },
     {
       "model": 0,
-      "chain": "A",
+      "chain": "R",
       "style": "stick",
       "sidechain": True,
       "color": "whiteCarbon",
     },
     {
       "model": 0,
-      "chain": "B",
+      "chain": "L",
       "style": "stick",
       "sidechain": True,
       "color": "greenCarbon"
@@ -84,6 +144,6 @@ with gr.Blocks() as app:
     out = Molecule3D(reps=reps)
     run_time = gr.Textbox(label="Runtime")
 
-    btn.click(predict, inputs=[input_seq_1, input_protein_1, input_seq_2,  input_protein_2], outputs=[out, run_time])
+    btn.click(predict, inputs=[input_protein_1, input_protein_2, input_fasta_1,  input_fasta_2], outputs=[out, run_time])
 
 app.launch()