Update inference_app.py

inference_app.py

@@ -5,15 +5,80 @@ import gradio as gr
 
 from gradio_molecule3d import Molecule3D
 
+import numpy as np
+from biotite.structure.io.pdb import PDBFile
+from rdkit import Chem
+from rdkit.Chem import AllChem
 
 
+def generate_input_conformer(
+    ligand_smiles: str,
+    addHs: bool = False,
+    minimize_maxIters: int = -1,
+) -> Chem.Mol:
 
-def predict (input_sequence, input_ligand,input_msa, input_protein):
+    _mol = Chem.MolFromSmiles(ligand_smiles)
+    # need to add Hs to generate sensible conformers
+    _mol = Chem.AddHs(_mol)
+
+    # try embedding molecule using ETKDGv2 (default)
+    confid = AllChem.EmbedMolecule(
+        _mol,
+        useRandomCoords=True,
+        useBasicKnowledge=True,
+        maxAttempts=100,
+        randomSeed=42,
+    )
+    if confid != -1:
+        if minimize_maxIters > 0:
+            # molecule successfully embedded - minimize
+            success = AllChem.MMFFOptimizeMolecule(_mol, maxIters=minimize_maxIters)
+            # 0 if the optimization converged,
+            # -1 if the forcefield could not be set up,
+            # 1 if more iterations are required.
+            if success == 1:
+                # extend optimization to double the steps (extends by the same amount)
+                AllChem.MMFFOptimizeMolecule(_mol, maxIters=minimize_maxIters)
+    else:
+        # this means EmbedMolecule failed
+        # try less optimal approach
+        confid = AllChem.EmbedMolecule(
+            _mol,
+            useRandomCoords=True,
+            useBasicKnowledge=False,
+            maxAttempts=100,
+            randomSeed=42,
+        )
+    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.array([new_coord] * len(structure.coord))
+    file = PDBFile()
+    file.set_structure(structure)
+    file.write(output_file)
+
+
+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)
+    with Chem.SDWriter("test_docking_pose.sdf") as writer:
+        writer.write(mol)
+    mol_coords = mol.GetConformer().GetPositions()
+    
+    # new_coord = [0, 0, 0]
+    new_coord = np.mean(mol_coords, axis=1)
+    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 = {"mean_plddt": 80, "binding_affinity": -2}
+    metrics = {}
+    
     end_time = time.time()
     run_time = end_time - start_time
     return ["test_out.pdb", "test_docking_pose.sdf"], metrics, run_time
@@ -43,12 +108,13 @@ with gr.Blocks() as app:
     gr.Examples(
         [
             [
-                "SVKSEYAEAAAVGQEAVAVFNTMKAAFQNGDKEAVAQYLARLASLYTRHEELLNRILEKARREGNKEAVTLMNEFTATFQTGKSIFNAMVAAFKNGDDDSFESYLQALEKVTAKGETLADQIAKAL:SVKSEYAEAAAVGQEAVAVFNTMKAAFQNGDKEAVAQYLARLASLYTRHEELLNRILEKARREGNKEAVTLMNEFTATFQTGKSIFNAMVAAFKNGDDDSFESYLQALEKVTAKGETLADQIAKAL",
+                None,
                 "COc1ccc(cc1)n2c3c(c(n2)C(=O)N)CCN(C3=O)c4ccc(cc4)N5CCCCC5=O",
+                None,
                 "test_out.pdb"
             ],
         ],
-        [input_sequence, input_ligand, input_protein],
+        [input_sequence, input_ligand, input_msa, input_protein],
     )
     reps =    [
     {
@@ -68,6 +134,6 @@ with gr.Blocks() as app:
     metrics = gr.JSON(label="Metrics")
     run_time = gr.Textbox(label="Runtime")
 
-    btn.click(predict, inputs=[input_sequence, input_ligand, input_msa, input_protein], outputs=[out,metrics, run_time])
+    btn.click(predict, inputs=[input_sequence, input_ligand, input_protein], outputs=[out,metrics, run_time])
 
 app.launch()