'MDAnalysis.analysis.diffusionmap' parallelization

package/CHANGELOG

@@ -25,6 +25,7 @@ Fixes
    the function to prevent shared state. (Issue #4655)
 
 Enhancements
+ * Enables parallelization for analysis.diffusionmap.DistanceMatrix (Issue #4679, PR #4745)
  * Enables parallelization for analysis.density.DensityAnalysis (Issue #4677, PR #4729)
  * Enables parallelization for analysis.contacts.Contacts (Issue #4660)
  * Enable parallelization for analysis.nucleicacids.NucPairDist (Issue #4670)

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,7 +234,18 @@ 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.9.0
+       Introduced :meth:`get_supported_backends` allowing
+       for parallel execution on :mod:`multiprocessing`
+       and :mod:`dask` backends.
     """
+
+    _analysis_algorithm_is_parallelizable = True
+
+    @classmethod
+    def get_supported_backends(cls):
+        return ('serial', 'multiprocessing', 'dask')
+
     def __init__(self, universe, select='all', metric=rmsd, cutoff=1E0-5,
                  weights=None, **kwargs):
         # remember that this must be called before referencing self.n_frames
@@ -286,6 +297,12 @@ def dist_matrix(self):
     def _conclude(self):
         self._calculated = True
 
+    def _get_aggregator(self):
+        return ResultsGroup(
+            lookup={
+                "dist_matrix": ResultsGroup.ndarray_sum,
+            }
+        )
 
 class DiffusionMap(object):
     """Non-linear dimension reduction method

testsuite/MDAnalysisTests/analysis/conftest.py

@@ -17,6 +17,7 @@
 from MDAnalysis.analysis.nucleicacids import NucPairDist
 from MDAnalysis.analysis.contacts import Contacts
 from MDAnalysis.analysis.density import DensityAnalysis
+from MDAnalysis.analysis.diffusionmap import DistanceMatrix
 from MDAnalysis.lib.util import is_installed
 
 
@@ -165,3 +166,10 @@ def client_Contacts(request):
 @pytest.fixture(scope='module', params=params_for_cls(DensityAnalysis))
 def client_DensityAnalysis(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)
@@ -69,14 +69,14 @@ def test_dist_weights(u):
                                 [.707, -.707, 0, 0]]), 2)
 
 
-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)
     frames = np.arange(len(u.trajectory))
-    dist.run(frames=frames[::3])
+    dist.run(**client_DistanceMatrix, frames=frames[::3])
     dmap = diffusionmap.DiffusionMap(dist)
     dmap.run()
     assert_array_almost_equal(dmap.eigenvalues, [1, 1, 1, 1], 4)
@@ -86,18 +86,25 @@ def test_dist_weights_frames(u):
                                 [-.707, -.707, 0, 0],
                                 [.707, -.707, 0, 0]]), 2)
 
-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)