Update inference_app.py

inference_app.py

@@ -10,6 +10,7 @@ from scipy.optimize import differential_evolution, NonlinearConstraint
 from biotite.structure.io.pdb import PDBFile
 from rdkit import Chem
 from rdkit.Chem import AllChem
+from biotite.structure import AtomArrayStack
 
 
 def generate_input_conformer(
@@ -53,14 +54,6 @@ def generate_input_conformer(
     return _mol
 
 
-def set_protein_to_new_coord(input_pdb_file, new_coord, output_file):
-    structure = PDBFile.read(input_pdb_file).get_structure()
-    structure.coord = np.ones_like(structure.coord) * np.array(new_coord)
-    file = PDBFile()
-    file.set_structure(structure)
-    file.write(output_file)
-
-
 def optimize_coordinate(points, bound_buffer=15, dmin=6.05):
 
     bounds = list(
@@ -83,6 +76,44 @@ def optimize_coordinate(points, bound_buffer=15, dmin=6.05):
     return result.x, result.fun
 
 
+def optimize_decoy_coordinate(points, bound_buffer=15, dmin=6.05, decoy_min=4.0, decoy_max=4.99):
+    bounds = list(
+        zip(
+            np.average(points, axis=0) - [bound_buffer]*3,
+            np.average(points, axis=0) + [bound_buffer]*3
+            )
+        )
+    # Define the constraint function (ensure dmin distance for all but one atom)
+    con1 = NonlinearConstraint(lambda x: np.sum(np.linalg.norm(points - x, axis=1) < dmin), 1, 1)
+    con2 = NonlinearConstraint(lambda x: np.min(np.linalg.norm(points - x, axis=1)), decoy_min, decoy_max)
+    # Define the objective function (maximize pairwise distance)
+    def objective(x):
+        return - np.sum(np.linalg.norm(points - x, axis=1))
+    # Perform differential evolution to find the optimal coordinate
+    result = differential_evolution(objective, bounds, constraints=(con1, con2))
+    return result.x, result.fun
+
+
+def add_decoy_atom(structure, decoy_pos):
+    decoy = AtomArrayStack(length=1, depth=1)
+    decoy.coord = np.ones_like(struct.coord) * decoy_pos
+    decoy.chain_id = ["q"]
+    decoy.element = ["C"]
+    decoy.atom_name = ["C"]
+    decoy.res_name = ["GLY"]
+    return structure + struct
+
+
+def set_protein_to_new_coord_plus_decoy_atom(input_pdb_file, new_coord, decoy_coord, output_file):
+    structure = PDBFile.read(input_pdb_file).get_structure()
+    structure.coord = np.ones_like(structure.coord) * np.array(new_coord)
+    # add decoy 
+    structure = add_decoy_atom(structure, decoy_pos)
+    file = PDBFile()
+    file.set_structure(structure)
+    file.write(output_file)
+
+    
 def predict(input_sequence, input_ligand, input_msa, input_protein):
     start_time = time.time()
     
@@ -93,18 +124,21 @@ def predict(input_sequence, input_ligand, input_msa, input_protein):
     molwriter.write(mol)
 
     mol_coords = mol.GetConformer().GetPositions()    
-    # new_coord = [0, 0, 0]
-    # new_coord = np.mean(mol_coords, axis=0) + [3.5, 3.5, 3.5]
+    # get opt coords
     new_coord, min_dist_sum = optimize_coordinate(mol_coords)
     # get mindist to protein
     min_dist = np.min(np.linalg.norm(mol_coords - new_coord, axis=1))
+    # decoy coord
+    decoy_coord = optimize_decoy_coordinate(mol_coords)
+    decoy_min_dist = np.min(np.linalg.norm(mol_coords - decoy_coord, axis=1))
 
+    # save protein
     output_file = "test_out.pdb"
-    set_protein_to_new_coord(input_protein, new_coord, output_file)
+    set_protein_to_new_coord_plus_decoy_atom(input_protein, new_coord, decoy_coord, output_file)
     
     # return an output pdb file with the protein and ligand with resname LIG or UNK. 
     # also return any metrics you want to log, metrics will not be used for evaluation but might be useful for users
-    metrics = {"min_dist": min_dist, "min_dist_sum": min_dist_sum}
+    metrics = {"min_dist": min_dist, "min_dist_sum": min_dist_sum, "decoy_min_dist": decoy_min_dist}
     
     end_time = time.time()
     run_time = end_time - start_time