'MDAnalysis.analysis.diffusionmap' parallelization (#4745)

`MDAnalysis.analysis.diffusionmap` parallelization

- fix #4679
- enable parallelization for `diffusionmap` with all backends
- update documentation and tests with appropriate fixture

---------

Co-authored-by: Egor Marin <[email protected]>

Author: talagayev

package/CHANGELOG

@@ -14,7 +14,7 @@ The rules for this file:
 
 
 -------------------------------------------------------------------------------
-??/??/?? IAlibay, orbeckst, marinegor, tylerjereddy, ljwoods2
+??/??/?? IAlibay, orbeckst, marinegor, tylerjereddy, ljwoods2, talagayev
 
  * 2.11.0
 
@@ -26,6 +26,10 @@ Fixes
  * Fixes incorrect assignment of secondary structure to proline residues in
    DSSP by porting upstream PyDSSP 0.9.1 fix (Issue #4913)
 
+Enhancements
+ * Enables parallelization for analysis.diffusionmap.DistanceMatrix
+   (Issue #4679, PR #4745)
+
 Changes
 
 Deprecations

package/MDAnalysis/analysis/diffusionmap.py

@@ -143,7 +143,7 @@
 from MDAnalysis.core.universe import Universe
 from MDAnalysis.core.groups import AtomGroup, UpdatingAtomGroup
 from .rms import rmsd
-from .base import AnalysisBase
+from .base import AnalysisBase, ResultsGroup
 
 logger = logging.getLogger("MDAnalysis.analysis.diffusionmap")
 
@@ -234,8 +234,22 @@ class DistanceMatrix(AnalysisBase):
     .. versionchanged:: 2.8.0
          :class:`DistanceMatrix` is now correctly works with `frames=...`
          parameter (#4432) by iterating over `self._sliced_trajectory`
+    .. versionchanged:: 2.11.0
+       Enabled **parallel execution** with the ``multiprocessing`` and ``dask``
+       backends; use the new method :meth:`get_supported_backends` to see all
+       supported backends.
     """
 
+    _analysis_algorithm_is_parallelizable = True
+
+    @classmethod
+    def get_supported_backends(cls):
+        return (
+            "serial",
+            "multiprocessing",
+            "dask",
+        )
+
     def __init__(
         self,
         universe,
@@ -265,27 +279,16 @@ def __init__(
         self._calculated = False
 
     def _prepare(self):
-        self.results.dist_matrix = np.zeros((self.n_frames, self.n_frames))
+        # Perpare for parallelization workers
+        n_atoms = self.atoms.n_atoms
+        n_dim = self.atoms.positions.shape[1]
+        self.results._positions = np.zeros(
+            (self.n_frames, n_atoms, n_dim), dtype=np.float64
+        )
 
     def _single_frame(self):
-        iframe = self._frame_index
-        i_ref = self.atoms.positions
-        # diagonal entries need not be calculated due to metric(x,x) == 0 in
-        # theory, _ts not updated properly. Possible savings by setting a
-        # cutoff for significant decimal places to sparsify matrix
-        for j, ts in enumerate(self._sliced_trajectory[iframe:]):
-            self._ts = ts
-            j_ref = self.atoms.positions
-            dist = self._metric(i_ref, j_ref, weights=self._weights)
-            self.results.dist_matrix[
-                self._frame_index, j + self._frame_index
-            ] = (dist if dist > self._cutoff else 0)
-            self.results.dist_matrix[
-                j + self._frame_index, self._frame_index
-            ] = self.results.dist_matrix[
-                self._frame_index, j + self._frame_index
-            ]
-        self._ts = self._sliced_trajectory[iframe]
+        # Store current frame positions
+        self.results._positions[self._frame_index] = self.atoms.positions
 
     @property
     def dist_matrix(self):
@@ -298,8 +301,38 @@ def dist_matrix(self):
         return self.results.dist_matrix
 
     def _conclude(self):
+        # Build the full pairwise distance matrix from stored positions
+        # Calculate and store results
+        pos = np.asarray(
+            self.results._positions, dtype=np.float64
+        )  # (n_frames, n_atoms, n_dim)
+        n = pos.shape[0]
+
+        D = np.zeros((n, n), dtype=np.float64)
+
+        metric = self._metric
+        cutoff = self._cutoff
+        weights = self._weights
+
+        for i in range(n):
+            pi = pos[i]
+            for j in range(i + 1, n):
+                pj = pos[j]
+                d = metric(pi, pj, weights=weights)
+                if d > cutoff:
+                    D[i, j] = d
+                    D[j, i] = d
+
+        self.results.dist_matrix = D
         self._calculated = True
 
+    def _get_aggregator(self):
+        return ResultsGroup(
+            lookup={
+                "_positions": ResultsGroup.ndarray_vstack,  # Get positions
+            }
+        )
+
 
 class DiffusionMap(object):
     """Non-linear dimension reduction method

testsuite/MDAnalysisTests/analysis/conftest.py

@@ -10,6 +10,7 @@
 from MDAnalysis.analysis.dihedrals import Dihedral, Ramachandran, Janin
 from MDAnalysis.analysis.bat import BAT
 from MDAnalysis.analysis.gnm import GNMAnalysis
+from MDAnalysis.analysis.diffusionmap import DistanceMatrix
 from MDAnalysis.analysis.dssp.dssp import DSSP
 from MDAnalysis.analysis.hydrogenbonds.hbond_analysis import (
     HydrogenBondAnalysis,
@@ -208,3 +209,11 @@ def client_InterRDF(request):
 @pytest.fixture(scope="module", params=params_for_cls(InterRDF_s))
 def client_InterRDF_s(request):
     return request.param
+
+
+# MDAnalysis.analysis.diffusionmap
+
+
+@pytest.fixture(scope="module", params=params_for_cls(DistanceMatrix))
+def client_DistanceMatrix(request):
+    return request.param

testsuite/MDAnalysisTests/analysis/test_diffusionmap.py

@@ -52,13 +52,13 @@ def test_eg(dist, dmap):
     # makes no sense to test values here, no physical meaning
 
 
-def test_dist_weights(u):
+def test_dist_weights(u, client_DistanceMatrix):
     backbone = u.select_atoms("backbone")
     weights_atoms = np.ones(len(backbone.atoms))
     dist = diffusionmap.DistanceMatrix(
         u, select="backbone", weights=weights_atoms
     )
-    dist.run(step=3)
+    dist.run(**client_DistanceMatrix, step=3)
     dmap = diffusionmap.DiffusionMap(dist)
     dmap.run()
     assert_array_almost_equal(dmap.eigenvalues, [1, 1, 1, 1], 4)
@@ -76,11 +76,11 @@ def test_dist_weights(u):
     )
 
 
-def test_dist_weights_frames(u):
+def test_dist_weights_frames(u, client_DistanceMatrix):
     backbone = u.select_atoms("backbone")
     weights_atoms = np.ones(len(backbone.atoms))
     dist = diffusionmap.DistanceMatrix(
-        u, select="backbone", weights=weights_atoms
+        u, **client_DistanceMatrix, select="backbone", weights=weights_atoms
     )
     frames = np.arange(len(u.trajectory))
     dist.run(frames=frames[::3])
@@ -101,19 +101,25 @@ def test_dist_weights_frames(u):
     )
 
 
-def test_distvalues_ag_universe(u):
-    dist_universe = diffusionmap.DistanceMatrix(u, select="backbone").run()
+def test_distvalues_ag_universe(u, client_DistanceMatrix):
+    dist_universe = diffusionmap.DistanceMatrix(u, select="backbone").run(
+        **client_DistanceMatrix
+    )
     ag = u.select_atoms("backbone")
-    dist_ag = diffusionmap.DistanceMatrix(ag).run()
+    dist_ag = diffusionmap.DistanceMatrix(ag).run(**client_DistanceMatrix)
     assert_allclose(
         dist_universe.results.dist_matrix, dist_ag.results.dist_matrix
     )
 
 
-def test_distvalues_ag_select(u):
-    dist_universe = diffusionmap.DistanceMatrix(u, select="backbone").run()
+def test_distvalues_ag_select(u, client_DistanceMatrix):
+    dist_universe = diffusionmap.DistanceMatrix(u, select="backbone").run(
+        **client_DistanceMatrix
+    )
     ag = u.select_atoms("protein")
-    dist_ag = diffusionmap.DistanceMatrix(ag, select="backbone").run()
+    dist_ag = diffusionmap.DistanceMatrix(ag, select="backbone").run(
+        **client_DistanceMatrix
+    )
     assert_allclose(
         dist_universe.results.dist_matrix, dist_ag.results.dist_matrix
     )
@@ -156,8 +162,10 @@ def test_not_universe_atomgroup_error(u):
         diffusionmap.DiffusionMap(trj_only)
 
 
-def test_DistanceMatrix_attr_warning(u):
-    dist = diffusionmap.DistanceMatrix(u, select="backbone").run(step=3)
+def test_DistanceMatrix_attr_warning(u, client_DistanceMatrix):
+    dist = diffusionmap.DistanceMatrix(u, select="backbone").run(
+        **client_DistanceMatrix, step=3
+    )
     wmsg = f"The `dist_matrix` attribute was deprecated in MDAnalysis 2.0.0"
     with pytest.warns(DeprecationWarning, match=wmsg):
         assert getattr(dist, "dist_matrix") is dist.results.dist_matrix