Accept AtomGroup alongwith Universe in the DistanceMatrix

package/AUTHORS

@@ -184,6 +184,7 @@ Chronological list of authors
   - Tamandeep Singh
   - Mingyi Xue
   - Meghan Osato
+  - Anirvinya G
 
 External code
 -------------

package/CHANGELOG

@@ -14,7 +14,7 @@ The rules for this file:
 
 ------------------------------------------------------------------------------
 ??/??/?? IAlibay, BFedder, inomag, Agorfa, aya9aladdin, shudipto-amin, cbouy,
-         HenokB, umak1106, tamandeeps, Mrqeoqqt, megosato
+         HenokB, umak1106, tamandeeps, Mrqeoqqt, megosato, AnirG
 
  * 2.2.0
 
@@ -36,6 +36,8 @@ Fixes
   * Fixed FrameIteratorIndices doesn't rewind. (Issue #3416)
 
 Enhancements
+ * `DistanceMatrix` class in MDAnalysis.analysis.diffusionmap now also
+    accepts `AtomGroup` as a valid argument (Issue #3486, PR #3541)
   * Improves the RDKitConverter's accuracy (PR #3044): AtomGroups containing
     monatomic ion charges or edge cases with nitrogen, sulfur, phosphorus and
     conjugated systems should now have correctly assigned bond orders and

package/MDAnalysis/analysis/diffusionmap.py

@@ -156,6 +156,7 @@
 import numpy as np
 
 from MDAnalysis.core.universe import Universe
+from MDAnalysis.core.groups import AtomGroup, UpdatingAtomGroup
 from .rms import rmsd
 from .base import AnalysisBase
 
@@ -171,7 +172,7 @@ class DistanceMatrix(AnalysisBase):
 
     Parameters
     ----------
-    universe : `~MDAnalysis.core.universe.Universe`
+    universe : `~MDAnalysis.core.universe.Universe` or `~MDAnalysis.core.groups.AtomGroup`
         The MD Trajectory for dimension reduction, remember that
         computational cost of eigenvalue decomposition
         scales at O(N^3) where N is the number of frames.
@@ -243,12 +244,20 @@ class DistanceMatrix(AnalysisBase):
     .. versionchanged:: 2.0.0
          :attr:`dist_matrix` is now stored in a
          :class:`MDAnalysis.analysis.base.Results` instance.
-
+    .. versionchanged:: 2.2.0
+         :class:`DistanceMatrix` now also accepts `AtomGroup`.
     """
     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
-        super(DistanceMatrix, self).__init__(universe.trajectory, **kwargs)
+        super(DistanceMatrix, self).__init__(universe.universe.trajectory,
+                                             **kwargs)
+
+        if isinstance(universe, UpdatingAtomGroup):
+            wmsg = ("U must be a static AtomGroup. Parsing an updating AtomGroup "
+                    "will result in a static AtomGroup with the current frame "
+                    "atom selection.")
+            warnings.warn(wmsg)
 
         self.atoms = universe.select_atoms(select)
         self._metric = metric
@@ -309,14 +318,18 @@ class DiffusionMap(object):
     transform(n_eigenvectors, time)
         Perform an embedding of a frame into the eigenvectors representing
         the collective coordinates.
+
+
+    .. versionchanged:: 2.2.0
+         :class:`DiffusionMap` now also accepts `AtomGroup`.
     """
 
     def __init__(self, u, epsilon=1, **kwargs):
         """
         Parameters
         -------------
-        u : MDAnalysis Universe or DistanceMatrix object
-            Can be a Universe, in which case one must supply kwargs for the
+        u : MDAnalysis Universe or AtomGroup or DistanceMatrix object.
+            Can be a Universe or AtomGroup, in which case one must supply kwargs for the
             initialization of a DistanceMatrix. Otherwise, this can be a
             DistanceMatrix already initialized. Either way, this will be made
             into a diffusion kernel.
@@ -328,12 +341,12 @@ def __init__(self, u, epsilon=1, **kwargs):
             Parameters to be passed for the initialization of a
             :class:`DistanceMatrix`.
         """
-        if isinstance(u, Universe):
+        if isinstance(u, AtomGroup) or isinstance(u, Universe):
             self._dist_matrix = DistanceMatrix(u, **kwargs)
         elif isinstance(u, DistanceMatrix):
             self._dist_matrix = u
         else:
-            raise ValueError("U is not a Universe or DistanceMatrix and"
+            raise ValueError("U is not a Universe or AtomGroup or DistanceMatrix and"
                              " so the DiffusionMap has no data to work with.")
         self._epsilon = epsilon
 

testsuite/MDAnalysisTests/analysis/test_diffusionmap.py

@@ -25,7 +25,7 @@
 import numpy as np
 import pytest
 from MDAnalysisTests.datafiles import PDB, XTC
-from numpy.testing import assert_array_almost_equal
+from numpy.testing import assert_array_almost_equal, assert_allclose
 
 
 @pytest.fixture(scope='module')
@@ -69,6 +69,22 @@ def test_dist_weights(u):
                                 [.707, -.707, 0, 0]]), 2)
 
 
+def test_distvalues_ag_universe(u):
+    dist_universe = diffusionmap.DistanceMatrix(u, select='backbone').run()
+    ag = u.select_atoms('backbone')
+    dist_ag = diffusionmap.DistanceMatrix(ag).run()
+    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()
+    ag = u.select_atoms('protein')
+    dist_ag = diffusionmap.DistanceMatrix(ag, select='backbone').run()
+    assert_allclose(dist_universe.results.dist_matrix,
+                    dist_ag.results.dist_matrix)
+
+
 def test_different_steps(u):
     dmap = diffusionmap.DiffusionMap(u, select='backbone')
     dmap.run(step=3)
@@ -92,9 +108,16 @@ def test_long_traj(u):
         dmap.run()
 
 
-def test_not_universe_error(u):
+def test_updating_atomgroup(u):
+    with pytest.warns(UserWarning, match='U must be a static AtomGroup'):
+        resid_select = 'around 5 resname ALA'
+        ag = u.select_atoms(resid_select, updating=True)
+        dmap = diffusionmap.DiffusionMap(ag)
+        dmap.run()
+
+def test_not_universe_atomgroup_error(u):
     trj_only = u.trajectory
-    with pytest.raises(ValueError, match='U is not a Universe'):
+    with pytest.raises(ValueError, match='U is not a Universe or AtomGroup'):
         diffusionmap.DiffusionMap(trj_only)