Parallel IFP (#36)

* minimize prints * MP for IFP calculation
drewnutt · Oct 27, 2023 · 463d90c · 463d90c
1 parent e40b0d3
commit 463d90c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/open_combind/features/features.py b/open_combind/features/features.py
@@ -283,14 +283,14 @@ def load_single_features(self, pvs, ligands=None, center_ligand=None):
                             print(f"skipped for {pv}")
                             continue
                     _poses.append(st)
-                print(len(poses))
+                # print(len(poses))
 
             keep = []
             for i in range(len(_names)):
                 if ((ligands == None or (_names[i] in ligands))
                     and sum(_names[:i] == _names[i]) < self.max_poses):
                     keep += [i]
-            print(keep)
+            # print(keep)
             rmsds += [_rmsds[keep]]
             if self.cnn_scores:
                 gscores += [_gscores[keep]]
@@ -316,7 +316,7 @@ def load_single_features(self, pvs, ligands=None, center_ligand=None):
             gscores = None
         return rmsds, gscores, gaffs, vaffs, poses, names, ifps
 
-    def compute_single_features(self, pvs, native_poses):
+    def compute_single_features(self, pvs, native_poses, processes=1):
         """
         Compute all of the single pose features (e.g. GNINA scores, RMSD to native, etc.) for the provided poses
 
@@ -372,11 +372,18 @@ def compute_single_features(self, pvs, native_poses):
                 self.compute_rmsd(bundle, native_poses, out)
 
         print('Computing interaction fingerprints.')
+        comp_ifp = Features.compute_ifp_wrapper
+        if processes != 1:
+            ifp_unfinished = []
+            comp_ifp = ifp_unfinished.append
         for pv in molbundles.keys():
             # print(pv)
             out = self.path('ifp', pv=pv)
             if not os.path.exists(out):
-                self.compute_ifp(pv, out)
+                # self.compute_ifp(pv, out)
+                comp_ifp((self,pv,out))
+        if processes != 1:
+            mp(Features.compute_ifp_wrapper, ifp_unfinished, processes=processes)
 
     def compute_pair_features(self, pvs, pvs2=None, ifp=True, shape=True, mcss=True, processes=1):
         """
@@ -562,6 +569,10 @@ def compute_ifp(self, pv, out):
         settings = IFP[self.ifp_version]
         ifp(settings, pv, out, self.max_poses)
 
+    @staticmethod
+    def compute_ifp_wrapper(self, pv, out):
+        self.compute_ifp(pv, out)
+
     def compute_ifp_pair(self, ifps1, ifps2, feature, out, processes=1):
         """
         Compute the pseudo-Tanimoto similarity between the given IFPs for the given features

diff --git a/open_combind/open_combind.py b/open_combind/open_combind.py
@@ -303,7 +303,7 @@ def featurize(root, poseviewers, native='structures/ligands/*_lig.sdf',
                         check_center_ligs=check_center_ligs, newscore=newscore)
 
     print(poseviewers)
-    features.compute_single_features(poseviewers, native_poses=native_poses)
+    features.compute_single_features(poseviewers, native_poses=native_poses, processes=processes)
 
     if screen:
         assert len(poseviewers) == 2