RMSD with multicore and timeout functionality

spyrmsd/rmsd.py

@@ -1,3 +1,7 @@
+import os
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
+from multiprocessing import Process, Queue
 from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
@@ -307,7 +311,7 @@ def rmsdwrapper(
     cache: bool = True,
 ) -> Any:
     """
-    Compute RMSD between two molecule.
+    Compute RMSD between two molecules.
 
     Parameters
     ----------
@@ -373,3 +377,188 @@ def rmsdwrapper(
             )
 
     return RMSDlist
+
+
+def _rmsd_queue(
+    molref: molecule.Molecule,
+    mols: Union[molecule.Molecule, List[molecule.Molecule]],
+    queue: Queue,
+    symmetry: bool = True,
+    center: bool = False,
+    minimize: bool = False,
+    strip: bool = True,
+    cache: bool = True,
+) -> Any:
+    """
+    Compute RMSD between two molecules and put it in a queue.
+
+    Parameters
+    ----------
+    molref: molecule.Molecule
+        Reference molecule
+    mols: Union[molecule.Molecule, List[molecule.Molecule]]
+        Molecules to compare to reference molecule
+    queue: Queue
+        The queue where to put the result
+    symmetry: bool, optional
+        Symmetry-corrected RMSD (using graph isomorphism)
+    center: bool, optional
+        Center molecules at origin
+    minimize: bool, optional
+        Minimised RMSD (using the quaternion polynomial method)
+    strip: bool, optional
+        Strip hydrogen atoms
+
+    Returns
+    -------
+    None
+        The RMSD result is put in the queue.
+    """
+    queue.put(
+        rmsdwrapper(
+            molref=molref,
+            mols=mols,
+            symmetry=symmetry,
+            center=center,
+            minimize=minimize,
+            strip=strip,
+            cache=cache,
+        )[0]
+    )
+
+
+def _rmsd_timeout(
+    molref: molecule.Molecule,
+    mols: Union[molecule.Molecule, List[molecule.Molecule]],
+    symmetry: bool = True,
+    center: bool = False,
+    minimize: bool = False,
+    strip: bool = True,
+    cache: bool = True,
+    timeout: Optional[float] = None,
+) -> Any:
+    """
+    Compute RMSD between two molecules with a timeout.
+
+    Parameters
+    ----------
+    molref: molecule.Molecule
+        Reference molecule
+    mols: Union[molecule.Molecule, List[molecule.Molecule]]
+        Molecules to compare to reference molecule
+    symmetry: bool, optional
+        Symmetry-corrected RMSD (using graph isomorphism)
+    center: bool, optional
+        Center molecules at origin
+    minimize: bool, optional
+        Minimised RMSD (using the quaternion polynomial method)
+    strip: bool, optional
+        Strip hydrogen atoms
+    timeout: float, optional
+        After how many seconds to stop the RMSD calculations
+
+    Returns
+    -------
+    List[float]
+        RMSDs
+
+    Notes
+    -----
+    Timeout implemenation inspired by https://superfastpython.com/task-with-timeout-child-process/
+    """
+
+    if not isinstance(mols, list):
+        mols = [mols]
+
+    queue = Queue()
+    process = Process(
+        target=_rmsd_queue,
+        args=(molref, mols, queue, symmetry, center, minimize, strip, cache),
+    )
+
+    process.start()
+    process.join(timeout=timeout)
+
+    # Check if the process finished running successfully
+    if process.exitcode != 0:
+        # Actually terminate the process if it still running
+        if process.is_alive():
+            process.terminate()
+
+        return [np.nan] * len(mols)
+    else:
+        # Retrieve the result from the finished job
+        return queue.get()
+
+
+def rmsd_parallel(
+    # Should we change this to reflect to possibility of also allowing just 1 molecule as input?
+    molrefs: Union[molecule.Molecule, List[molecule.Molecule]],
+    mols: Union[molecule.Molecule, List[molecule.Molecule]],
+    num_workers: int = 1,
+    symmetry: bool = True,
+    center: bool = False,
+    minimize: bool = False,
+    strip: bool = True,
+    cache: bool = True,
+    timeout: Optional[float] = None,
+) -> Any:
+    """
+    Compute RMSD between two molecules with a timeout.
+
+    Parameters
+    ----------
+    molrefs: Union[molecule.Molecule, List[molecule.Molecule]]
+        Reference molecule
+    mols: Union[molecule.Molecule, List[molecule.Molecule]]
+        Molecules to compare to reference molecule
+    num_workers: int
+        Amount of processor to use for the parallel calculations
+    symmetry: bool, optional
+        Symmetry-corrected RMSD (using graph isomorphism)
+    center: bool, optional
+        Center molecules at origin
+    minimize: bool, optional
+        Minimised RMSD (using the quaternion polynomial method)
+    strip: bool, optional
+        Strip hydrogen atoms
+    timeout: float, optional
+        After how many seconds to stop the RMSD calculations
+
+    Returns
+    -------
+    List[float]
+        RMSDs
+    """
+
+    # Ensure the num_workers is less or equal than the max number of CPUs
+    # Makes MyPy unhappy because os.cpu_count() can return None in some cases (and num_workers is defined as an int)
+    num_workers = min(num_workers, os.cpu_count())
+
+    # Cast the molecules to lists if they aren't already
+    if not isinstance(molrefs, list):
+        molrefs = [molrefs]
+    if not isinstance(mols, list):
+        mols = [mols]
+
+    # Match the length of the molref
+    if len(molrefs) == 1 and len(molrefs) < len(mols):
+        molrefs = molrefs * len(mols)
+
+    # Ensure molrefs and mols have the same len
+    if not len(molrefs) == len(mols):
+        raise ValueError("The input mol lists have different lengths")
+
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        rsmd_partial = partial(
+            _rmsd_timeout,
+            symmetry=symmetry,
+            center=center,
+            minimize=minimize,
+            strip=strip,
+            cache=cache,
+            timeout=timeout,
+        )
+        results = executor.map(rsmd_partial, molrefs, mols)
+
+    return list(results)