Update inference_app.py

inference_app.py

@@ -6,11 +6,10 @@ import torch
 from torch_geometric.data import HeteroData
 import numpy as np
 from loguru import logger
-from Bio import PDB
-from Bio.PDB.PDBIO import PDBIO
 from pinder.core.loader.geodata import structure2tensor
 from pinder.core.loader.structure import Structure
 from src.models.pinder_module import PinderLitModule
+from pathlib import Path
 
 try:
     from torch_cluster import knn_graph
@@ -37,13 +36,13 @@ def get_props_pdb(pdb_file):
         residue_types=calpha.res_name,
         residue_ids=calpha.res_id,
     )
-    return props
+    return structure, props
 
 
 def create_graph(pdb_1, pdb_2, k=5, device: torch.device = torch.device("cpu")):
 
-    props_ligand = get_props_pdb(pdb_1)
-    props_receptor = get_props_pdb(pdb_2)
+    ligand_structure, props_ligand = get_props_pdb(pdb_1)
+    receptor_structure, props_receptor = get_props_pdb(pdb_2)
 
     data = HeteroData()
 
@@ -56,74 +55,7 @@ def create_graph(pdb_1, pdb_2, k=5, device: torch.device = torch.device("cpu")):
     data["receptor", "receptor"].edge_index = knn_graph(data["receptor"].pos, k=k)
 
     data = data.to(device)
-    return data
-
-
-def update_pdb_coordinates_from_tensor(
-    input_filename, output_filename, coordinates_tensor
-):
-    r"""
-    Updates atom coordinates in a PDB file with new transformed coordinates provided in a tensor.
-
-    Parameters:
-    - input_filename (str): Path to the original PDB file.
-    - output_filename (str): Path to the new PDB file to save updated coordinates.
-    - coordinates_tensor (torch.Tensor): Tensor of shape (1, N, 3) with transformed coordinates.
-    """
-    # Convert the tensor to a list of tuples
-    new_coordinates = coordinates_tensor.squeeze(0).tolist()
-
-    # Create a parser and parse the structure
-    parser = PDB.PDBParser(QUIET=True)
-    structure = parser.get_structure("structure", input_filename)
-
-    # Flattened iterator for atoms to update coordinates
-    atom_iterator = (
-        atom
-        for model in structure
-        for chain in model
-        for residue in chain
-        for atom in residue
-    )
-
-    # Update each atom's coordinates
-    for atom, (new_x, new_y, new_z) in zip(atom_iterator, new_coordinates):
-        original_anisou = atom.get_anisou()
-        original_uij = atom.get_siguij()
-        original_tm = atom.get_sigatm()
-        original_occupancy = atom.get_occupancy()
-        original_bfactor = atom.get_bfactor()
-        original_altloc = atom.get_altloc()
-        original_serial_number = atom.get_serial_number()
-        original_element = atom.get_charge()
-        original_parent = atom.get_parent()
-        original_radius = atom.get_radius()
-
-        # Update only the atom coordinates, keep other fields intact
-        atom.coord = np.array([new_x, new_y, new_z])
-
-        # Reapply the preserved properties
-        atom.set_anisou(original_anisou)
-        atom.set_siguij(original_uij)
-        atom.set_sigatm(original_tm)
-        atom.set_occupancy(original_occupancy)
-        atom.set_bfactor(original_bfactor)
-        atom.set_altloc(original_altloc)
-        # atom.set_fullname(original_fullname)
-        atom.set_serial_number(original_serial_number)
-        atom.set_charge(original_element)
-        atom.set_radius(original_radius)
-        atom.set_parent(original_parent)
-        # atom.set_name(original_name)
-        # atom.set_leve
-
-    # Save the updated structure to a new PDB file
-    io = PDBIO()
-    io.set_structure(structure)
-    io.save(output_filename)
-
-    # Return the path to the updated PDB file
-    return output_filename
+    return data, receptor_structure, ligand_structure
 
 
 def merge_pdb_files(file1, file2, output_file):
@@ -156,7 +88,9 @@ def predict(
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     logger.info(f"Using device: {device}")
 
-    data = create_graph(input_protein_1, input_protein_2, k=10, device=device)
+    data, receptor_structure, ligand_structure = create_graph(
+        input_protein_1, input_protein_2, k=10, device=device
+    )
     logger.info("Created graph data")
 
     model = PinderLitModule.load_from_checkpoint("./checkpoints/epoch_010.ckpt")
@@ -167,13 +101,22 @@ def predict(
     with torch.no_grad():
         receptor_coords, ligand_coords = model(data)
 
-    file1 = update_pdb_coordinates_from_tensor(
-        input_protein_1, "holo_ligand.pdb", ligand_coords
+    receptor_structure.coord = receptor_coords.squeeze(0).cpu().numpy()
+    ligand_structure.coord = ligand_coords.squeeze(0).cpu().numpy()
+
+    receptor_pinder = Structure(
+        filepath=Path("./holo_receptor.pdb"), atom_array=receptor_structure
     )
-    file2 = update_pdb_coordinates_from_tensor(
-        input_protein_2, "holo_receptor.pdb", receptor_coords
+    ligand_pinder = Structure(
+        filepath=Path("./holo_ligand.pdb"), atom_array=ligand_structure
+    )
+
+    receptor_pinder.to_pdb()
+    ligand_pinder.to_pdb()
+
+    out_pdb = merge_pdb_files(
+        "./holo_receptor.pdb", "./holo_ligand.pdb", "./output.pdb"
     )
-    out_pdb = merge_pdb_files(file1, file2, "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
@@ -267,4 +210,4 @@ with gr.Blocks() as app:
         outputs=[out, metrics, run_time],
     )
 
-app.launch()
+app.launch()