evaluate_results can now also print the SCORE (dG)

eval/docking.py

 import os
 from vina import Vina
-from eval.chemutils import kd
+from src.chemutils import kd
 import time
 
-def docking(
-  receptor_file:str, 
-  ligand_file:str, 
-  center:"tuple[float, float, float]"=(0,0,0), 
-  box_size:"tuple[float, float, float]"=(20,20,20),
-  n_dockings:int=64, 
-  n_poses:int=32,
-  score=False,
-  write=False) -> "list[float] | float":
 
+def docking(receptor_file: str,
+            ligand_file: str,
+            center: "tuple[float, float, float]" = (0, 0, 0),
+            box_size: "tuple[float, float, float]" = (20, 20, 20),
+            n_dockings: int = 32,
+            n_poses: int = 20) -> "dict[str, list[float]]":
   """
   Docking simulation function : returns ...
   @param receptor_file: protein (pdbqt file)
@@ -20,47 +17,39 @@ def docking(
   @param center: docking window center
   @param box_size: docking window size
   @param n_dockings: number of docking simulations
-  @param n_poses: number of pose attempts per simulation
-  @param score (bool): wether or not the output should be a single score
-  @param write (bool): wether or not the poses should be saved to a file
-  @return (list[float] | float): a list of scores or a single score
+  @param n_poses: number of pose attempts per simulation 
+  @return: dockings (pdbqt files), delta_G
   """
 
   receptor_name = os.path.splitext(receptor_file)[-1].split('.')[0]
   ligand_name = os.path.splitext(ligand_file)[-1].split('.')[0]
 
-  # initialises vina
-  v = Vina(sf_name='vina', verbosity=1)
+  #On initialise vina
+  v = Vina(sf_name='vina', verbosity=0)
   v.set_receptor(receptor_file)
   v.set_ligand_from_file(ligand_file)
 
-  # set the docking frame
-  v.compute_vina_maps(center=center,box_size=box_size)
+  #On pose la box de docking
+  v.compute_vina_maps(center=center, box_size=box_size)
 
-  # scores the current pose
+  # Score the current pose
   energy = v.score()
-  print('Score before minimization: %.3f (kcal/mol)' % energy[0])
+  #print('Score before minimization: %.3f (kcal/mol)' % energy[0])
 
-  # minimizes locally the current pose
+  # Minimized locally the current pose
   energy_minimized = v.optimize()
-  print('Score after minimization : %.3f (kcal/mol)' % energy_minimized[0])
+  #print('Score after minimization : %.3f (kcal/mol)' % energy_minimized[0])
   # v.write_pose(f'{ligand_name}_minimized.pdbqt', overwrite=True)
 
-  # docks the ligand
+  # Dock the ligand
   v.dock(exhaustiveness=n_dockings, n_poses=20)
+  v.write_poses(f'results/docked/{ligand_name}_docked_{time.time()}.pdbqt',
+                n_poses=n_poses)
 
-  if write:
-    v.write_poses(
-      f'results/docked/{ligand_name}_docked_{time.time()}.pdbqt', 
-      n_poses=n_poses)
+  results = v.energies(n_poses=n_poses)
 
-  output = None
-
-  # single hi-score or the list of all poses energies
-  if not score:
-    results = v.energies(n_poses=n_poses)
-    output = [energies[0] for energies in results]
-  else:
-    output = v.energies(n_poses=1)[0][0]
-
-  return output
+  return {
+      "poses": v.poses(coordinates_only=True),
+      "Kd": [kd(energies[0]) for energies in results],
+      "dG": [energies[0] for energies in results],
+  }
\ No newline at end of file

evaluate_results.py

@@ -18,14 +18,15 @@ def get_energy(receptor_path, ligand_path):
   receptor_path_preped = prepare(receptor_path)
   ligand_path_preped = prepare(ligand_path)
   try:
-    energy = docking(receptor_path_preped,
+    dic = docking(receptor_path_preped,
                      ligand_path_preped,
                      center=docking_box["center"],
                      box_size=docking_box["size"],
                      n_dockings=40,
-                     n_poses=20)["Kd"]
+                     n_poses=20)
+    energy, dG =  dic["Kd"], dic["dG"]
     print("Kd:", energy)
-    return energy
+    return energy, dG
   except Exception as e:
     print(f"[!]Error: {e}.")
     print("------")
@@ -33,9 +34,9 @@ def get_energy(receptor_path, ligand_path):
 
 
 output_path = "./summary.tsv"
-res_path = "./data/PG_res/"
+res_path = "./results/mutants/"
 
-DATA = [["ID", "KD", "NUM_MUTATIONS"]]
+DATA = [["ID","SCORE" ,"KD", "NUM_MUTATIONS"]]
 
 for mol in os.listdir(res_path):
   if mol.endswith(".sdf"):
@@ -46,15 +47,15 @@ for mol in os.listdir(res_path):
       if rec.startswith(prefix) and rec.endswith(".pdb"):
         # We have a receptor that starts with the prefix of the molecule
         print("Doing:", rec, "with ligand:", mol)
-        energy = get_energy(os.path.join(res_path, rec),
+        energy,dG = get_energy(os.path.join(res_path, rec),
                             os.path.join(res_path, mol))
         
         numbers = "".join([s for s in rec.split() if s.isdigit()])
-
+        print("Numbers:", numbers)
         num_mutations = get_num_mutations(
             os.path.join(res_path, rec),
             os.path.join(res_path, numbers + "_whole.pdb"))
-        DATA.append([rec, energy, num_mutations])
+        DATA.append([rec, dG,energy, num_mutations])
         if int(numbers) >= 3:
           break