fix inference_app.py

inference_app.py

@@ -104,28 +104,11 @@ def get_system(system_id: str) -> PinderSystem:
     return PinderSystem(system_id)
 from Bio import PDB
 from Bio.PDB.PDBIO import PDBIO
+from pinder.core.structure.atoms import atom_array_from_pdb_file
+from pathlib import Path
+from pinder.eval.dockq.biotite_dockq import BiotiteDockQ
 
-def extract_coordinates_from_pdb(filename):
-    r"""
-    Extracts atom coordinates from a PDB file and returns them as a list of tuples.
-    Each tuple contains (x, y, z) coordinates of an atom.
-    """
-    parser = PDB.PDBParser(QUIET=True)
-    structure = parser.get_structure("structure", filename)
-
-    coordinates = []
-
-    # Loop through each model, chain, residue, and atom to collect coordinates
-    for model in structure:
-        for chain in model:
-            for residue in chain:
-                # Retrieve atoms and their coordinates
-                for atom in residue:
-                    xyz = atom.coord  # Coordinates are in a numpy array
-                    # Append the coordinates (x, y, z) as a tuple
-                    coordinates.append((xyz[0], xyz[1], xyz[2]))
 
-    return coordinates
 log = setup_logger(__name__)
 
 try:
@@ -302,8 +285,8 @@ def create_graph(pdb1, pdb2, k=5):
         HeteroData: A PyG HeteroData object containing ligand and receptor data.
     """
     # Extract coordinates from PDB files
-    coords1 = torch.tensor(extract_coordinates_from_pdb(pdb1),dtype=torch.float)  
-    coords2 = torch.tensor(extract_coordinates_from_pdb(pdb2),dtype=torch.float)  
+    coords1 = torch.tensor(atom_array_from_pdb_file(pdb1),dtype=torch.float)  
+    coords2 = torch.tensor(atom_array_from_pdb_file(pdb2),dtype=torch.float)  
     # coords3 = torch.tensor(extract_coordinates_from_pdb(pdb3),dtype=torch.float)
     # Create the HeteroData object
     data = HeteroData()
@@ -422,6 +405,7 @@ def merge_pdb_files(file1, file2, output_file):
     
     print(f"Merged PDB saved to {output_file}")   
     return output_file
+
 class MPNNLayer(MessagePassing):
     def __init__(self, emb_dim=64, edge_dim=4, aggr='add'):
         r"""Message Passing Neural Network Layer
@@ -892,21 +876,34 @@ def predict (input_seq_1, input_msa_1, input_protein_1, input_seq_2,input_msa_2,
     start_time = time.time()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     data = create_graph(input_protein_1, input_protein_2, k=10)
-    
+    R_chain, L_chain = ["R"], ["L"]
     with torch.no_grad():
         mat, vect = model(data)
     mat = mat.to(device)
     vect = vect.to(device)
-    ligand1 = torch.tensor(extract_coordinates_from_pdb(input_protein_1),dtype=torch.float).to(device)
+    ligand1 = torch.tensor(atom_array_from_pdb_file(input_protein_1),dtype=torch.float).to(device)
     # receptor1 = torch.tensor(extract_coordinates_from_pdb(input_protein_2),dtype=torch.float).to(device)
     transformed_ligand = torch.matmul(ligand1, mat) + vect
     # transformed_receptor = torch.matmul(receptor1, mat) + vect
     file1 = update_pdb_coordinates_from_tensor(input_protein_1, "holo_ligand.pdb", transformed_ligand)
     # file2 = update_pdb_coordinates_from_tensor(input_protein_2, "holo_receptor.pdb", transformed_receptor)
     out_pdb = merge_pdb_files(file1,input_protein_2,"output.pdb")
+    
     # return an output pdb file with the protein and two chains A and B.  
     # also return a JSON with any metrics you want to report
     metrics = {"mean_plddt": 80, "binding_affinity": 2}
+    native = './test_out (1)'
+    decoys = out_pdb
+    bdq = BiotiteDockQ(
+    native=native, decoys=decoys,
+    # These are optional and if not specified will be assigned based on number of atoms (receptor > ligand)
+    native_receptor_chain=R_chain,
+    native_ligand_chain=L_chain,
+    decoy_receptor_chain=R_chain,
+    decoy_ligand_chain=L_chain,
+)
+    dockq = bdq.calculate()
+    metrics['DockQ'] = dockq
     end_time = time.time()
     run_time = end_time - start_time