Update inference_app.py

inference_app.py

@@ -6,6 +6,7 @@ import gradio as gr
 from gradio_molecule3d import Molecule3D
 
 import numpy as np
+from scipy.optimize import differential_evolution, NonlinearConstraint
 from biotite.structure.io.pdb import PDBFile
 from rdkit import Chem
 from rdkit.Chem import AllChem
@@ -60,24 +61,48 @@ def set_protein_to_new_coord(input_pdb_file, new_coord, output_file):
     file.write(output_file)
 
 
+def optimize_coordinate(points, bound_buffer=15, dmin=6.05):
+
+    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)
+    con = NonlinearConstraint(lambda x: np.min(np.linalg.norm(points - x, axis=1)), dmin, 10)
+
+    # Define the objective function (minimize 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=con)
+
+    return result.x, result.fun
+
+
 def predict(input_sequence, input_ligand, input_msa, input_protein):
     start_time = time.time()
     
     # Do inference here
-    mol = generate_input_conformer(input_ligand, minimize_maxIters=100)
+    mol = generate_input_conformer(input_ligand, minimize_maxIters=500)
     
     molwriter = Chem.SDWriter("test_docking_pose.sdf")
     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]
+    # new_coord = np.mean(mol_coords, axis=0) + [3.5, 3.5, 3.5]
+    new_coord, min_dist = optimize_coordinate(mol_coords)
+
     output_file = "test_out.pdb"
     set_protein_to_new_coord(input_protein, new_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 = {"mean_plddt": 80, "binding_affinity": -2}
+    metrics = {"min_dist": min_dist}
     metrics = {}
     
     end_time = time.time()