Update inference_app.py

inference_app.py

@@ -87,7 +87,28 @@ def run_smina(
     )
     return output_text
 
-def predict (input_sequence, input_ligand, input_protein):
+def predict (input_sequence, input_ligand, input_protein, exhaustiveness):
+    """
+    Main prediction function that calls ligsite and smina
+
+    Parameters
+    ----------
+    input_sequence: str
+        monomer sequence
+    input_ligand: str
+        ligand as SMILES string
+    protein_path: gradio.File
+        Gradio file object to monomer protein structure as PDB
+    exhaustiveness: int
+        SMINA parameter
+
+    Returns
+    -------
+    output_structures: tuple
+        (output_protein, output_ligand_sdf)
+    run_time: float
+        run time of the program
+    """
     start_time = time.time()
 
     if input_protein==None:
@@ -102,6 +123,7 @@ def predict (input_sequence, input_ligand, input_protein):
 
     os.system(f"obabel {input_protein.name} -xr -O /usr/src/app/receptor.pdbqt")
     os.system("obabel -isdf /usr/src/app/ligand.sdf -O /usr/src/app/ligand.pdbqt")
+
     #Find pocket
     pdb = md.load(input_protein.name)
     # run ligsite
@@ -110,23 +132,15 @@ def predict (input_sequence, input_ligand, input_protein):
     min_samples_value = 5
     dbscan = DBSCAN(eps=eps_value, min_samples=min_samples_value)
     labels = dbscan.fit_predict(pockets_xyz)
-    
     # Find the unique clusters and their sizes
     unique_labels, counts = np.unique(labels, return_counts=True)
-    
     # Exclude noise points
     valid_clusters = unique_labels[unique_labels != -1]
     valid_counts = counts[unique_labels != -1]
-    
     # Find the cluster with the most points (highest density)
     densest_cluster_label = valid_clusters[np.argmax(valid_counts)]
     densest_cluster_points = pockets_xyz[labels == densest_cluster_label]
-    
-    pocket_center = np.mean(densest_cluster_points, axis=0)
-    
-    
-    import pandas as pd
-    
+    # write cluster to PDB   
     top_df = pd.DataFrame()
     top_df['serial'] = list(range(densest_cluster_points.shape[0]))
     top_df['name'] = 'PK'
@@ -134,24 +148,22 @@ def predict (input_sequence, input_ligand, input_protein):
     top_df['resSeq'] = list(range(densest_cluster_points.shape[0]))
     top_df['resName'] = 'PCK'
     top_df['chainID'] = 0
-    
     pocket_top = md.Topology.from_dataframe(top_df, np.array([]))
     pocket_trj = md.Trajectory(xyz=densest_cluster_points, topology=pocket_top)
     pocket_trj.save('/usr/src/app/pockets_dense.pdb')
-    
     parser = PDBParser()
     struc = parser.get_structure("X", "/usr/src/app/pockets_dense.pdb")
     coords = [x.coord for x in struc.get_atoms()]
     pocket_center = np.mean(coords, axis=0)
+    # run smina
     output_text = run_smina(
         "/usr/src/app/ligand.pdbqt",
         "/usr/src/app/receptor.pdbqt",
         "/usr/src/app/docking_pose.sdf",
         pocket_center,
         [10,10,10],
-    )
-    os.system("cat /usr/src/app/docking_pose.sdf")
-      
+        exhaustiveness=exhaustiveness
+    )    
     end_time = time.time()
     run_time = end_time - start_time
     return [input_protein.name,"/usr/src/app/docking_pose.sdf"], run_time
@@ -170,7 +182,7 @@ with gr.Blocks() as app:
     # define any options here
 
     # for automated inference the default options are used
-    # slider_option = gr.Slider(0,10, label="Slider Option")
+    exhaustiveness = gr.Slider(1,10,value=1, label="Slider Option")
     # checkbox_option = gr.Checkbox(label="Checkbox Option")
     # dropdown_option = gr.Dropdown(["Option 1", "Option 2", "Option 3"], label="Radio Option")
 
@@ -181,7 +193,7 @@ with gr.Blocks() as app:
             [
                 "SVKSEYAEAAAVGQEAVAVFNTMKAAFQNGDKEAVAQYLARLASLYTRHEELLNRILEKARREGNKEAVTLMNEFTATFQTGKSIFNAMVAAFKNGDDDSFESYLQALEKVTAKGETLADQIAKAL:SVKSEYAEAAAVGQEAVAVFNTMKAAFQNGDKEAVAQYLARLASLYTRHEELLNRILEKARREGNKEAVTLMNEFTATFQTGKSIFNAMVAAFKNGDDDSFESYLQALEKVTAKGETLADQIAKAL",
                 "COc1ccc(cc1)n2c3c(c(n2)C(=O)N)CCN(C3=O)c4ccc(cc4)N5CCCCC5=O",
-                "test_out.pdb"
+                "input_test.pdb"
             ],
         ],
         [input_sequence, input_ligand, input_protein],
@@ -215,6 +227,6 @@ with gr.Blocks() as app:
     out = Molecule3D(reps=reps)
     run_time = gr.Textbox(label="Runtime")
 
-    btn.click(predict, inputs=[input_sequence, input_ligand, input_protein], outputs=[out, run_time])
+    btn.click(predict, inputs=[input_sequence, input_ligand, input_protein, exhaustiveness], outputs=[out, run_time])
 
 app.launch()