Merge pull request #1931 from jamesmkrieger/parallel_savePDBwbt

Parallel savePDBWaterBridgesTrajectory
prody · Aug 7, 2024 · 34fa7be · 34fa7be
2 parents 922b426 + fd727cf
commit 34fa7be
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 5 deletions.
diff --git a/docs/release/index.rst b/docs/release/index.rst
@@ -10,6 +10,10 @@ Release Notes
    :maxdepth: 2
    :glob:
 
+   v2.4_series
+   v2.3_series
+   v2.2_series
+   v2.1_series
    v2.0_series
    v1.11_series
    v1.10_series

diff --git a/prody/proteins/waterbridges.py b/prody/proteins/waterbridges.py
@@ -530,7 +530,8 @@ def calcWaterBridgesTrajectory(atoms, trajectory, **kwargs):
     combining selections satifying the criteria in each.
 
     :arg return_selection: whether to return the combined common selection
-        Default is **False**
+        Default is **False** to keep expected behaviour.
+        However, this output is required when using selstr.
     :type return_selection: bool    
     """
     start_frame = kwargs.pop('start_frame', 0)
@@ -554,7 +555,9 @@ def calcWaterBridgesTrajectory(atoms, trajectory, **kwargs):
             traj = trajectory[start_frame:stop_frame+1]
 
         indices = None
+        selection = atoms
         if selstr is not None:
+            LOGGER.info('Finding common selection')
             indices = []
             for frame0 in traj:
                 atoms_copy = atoms.copy()
@@ -567,7 +570,10 @@ def calcWaterBridgesTrajectory(atoms, trajectory, **kwargs):
                 indices.extend(list(selection.getIndices()))
 
             indices = np.unique(indices)
+            selection = atoms_copy[indices]
 
+            LOGGER.info('Common selection found with {0} atoms and {1} protein chains'.format(selection.numAtoms(),
+                                                                                              len(list(selection.protein.getHierView()))))
 
         def analyseFrame(j0, start_frame, frame0, interactions_all):
             LOGGER.info('Frame: {0}'.format(j0))
@@ -599,7 +605,7 @@ def analyseFrame(j0, start_frame, frame0, interactions_all):
                 while j0 < traj.numConfs()+start_frame:
 
                     processes = []
-                    for i in range(max_proc):
+                    for _ in range(max_proc):
                         frame0 = traj[j0-start_frame]
 
                         p = mp.Process(target=analyseFrame, args=(j0, start_frame,
@@ -644,7 +650,7 @@ def analyseFrame(i, interactions_all):
                     i = start_frame
                     while i < len(atoms.getCoordsets()[start_frame:stop_frame]):
                         processes = []
-                        for i in range(max_proc):
+                        for _ in range(max_proc):
                             p = mp.Process(target=analyseFrame, args=(i, interactions_all))
                             p.start()
                             processes.append(p)
@@ -1108,7 +1114,7 @@ def savePDBWaterBridges(bridges, atoms, filename):
     return writePDB(filename, atomsToSave)
 
 
-def savePDBWaterBridgesTrajectory(bridgeFrames, atoms, filename, trajectory=None):
+def savePDBWaterBridgesTrajectory(bridgeFrames, atoms, filename, trajectory=None, max_proc=1):
     """Saves one PDB per frame with occupancy and beta on protein atoms and waters forming bridges in frame.
 
     :arg bridgeFrames: atomic output from calcWaterBridgesTrajectory
@@ -1131,7 +1137,8 @@ def savePDBWaterBridgesTrajectory(bridgeFrames, atoms, filename, trajectory=None
     atoms = atoms.copy()
     mofifyBeta(bridgeFrames, atoms)
 
-    for frameIndex, frame in enumerate(bridgeFrames):
+    def saveBridgesFrame(trajectory, atoms, frameIndex, frame):
+        LOGGER.info('Frame: {0}'.format(frameIndex))
         if trajectory:
             coords = trajectory[frameIndex].getCoords()
             atoms.setCoords(coords)
@@ -1157,6 +1164,26 @@ def savePDBWaterBridgesTrajectory(bridgeFrames, atoms, filename, trajectory=None
             writePDB(f'{filename}_{frameIndex}.pdb',
                      atomsToSave, csets=frameIndex)
 
+    if max_proc == 1:
+        for frameIndex, frame in enumerate(bridgeFrames):
+            saveBridgesFrame(trajectory, atoms, frameIndex, frame)
+    else:
+        frameIndex = 0
+        numFrames = len(bridgeFrames)
+        while frameIndex < numFrames:
+            processes = []
+            for _ in range(max_proc):
+                p = mp.Process(target=saveBridgesFrame, args=(trajectory, atoms, frameIndex,
+                                                              bridgeFrames[frameIndex]))
+                p.start()
+                processes.append(p)
+
+                frameIndex += 1
+                if frameIndex >= numFrames:
+                    break
+
+            for p in processes:
+                p.join()
 
 def getBridgeIndicesString(bridge):
     return ' '.join(map(lambda a: str(a.getIndex()), bridge.proteins))\