Merge branch 'formalcharge-pdb' of github.com:MDAnalysis/mdanalysis i…

…nto formalcharge-pdb

package/AUTHORS

@@ -196,6 +196,7 @@ Chronological list of authors
   - Ricky Sexton
   - Rafael R. Pappalardo
   - Tengyu Xie
+  - Raymond Zhao
 
 
 External code

package/CHANGELOG

@@ -13,7 +13,7 @@ The rules for this file:
   * release numbers follow "Semantic Versioning" http://semver.org
 
 ------------------------------------------------------------------------------
-??/??/?? IAlibay, PicoCentauri, orbeckst, hmacdope, rmeli, miss77jun
+??/??/?? IAlibay, PicoCentauri, orbeckst, hmacdope, rmeli, miss77jun, rzhao271
 
  * 2.3.0
 
@@ -24,6 +24,8 @@ Fixes
 Enhancements
   * Add a new `formalcharge` attribute for storing formal charges (PR #3755)
   * Add the ability to read & write formal charges from PDB files (PR #3755)
+  * Change functions in `lib.distances` to accept AtomGroups as well as NumPy
+    arrays (CZI Performance track PR #3730)
   * Add `norm` parameter to InterRDF, InterRDF_s to normalize as rdf,
     number density or do not normalize at all. (Issue #3687)
   * Additional logger.info output when per-frame analysis starts (PR #3710)

package/MDAnalysis/lib/util.py

@@ -214,6 +214,7 @@
 from .picklable_file_io import pickle_open, bz2_pickle_open, gzip_pickle_open
 
 from ..exceptions import StreamWarning, DuplicateWarning
+
 try:
     from ._cutil import unique_int_1d
 except ImportError:
@@ -1937,12 +1938,17 @@ def check_coords(*coord_names, **options):
     :mod:`MDAnalysis.lib.distances`.
     It takes an arbitrary number of positional arguments which must correspond
     to names of positional arguments of the decorated function.
-    It then checks if the corresponding values are valid coordinate arrays.
-    If all these arrays are single coordinates (i.e., their shape is ``(3,)``),
-    the decorated function can optionally return a single coordinate (or angle)
-    instead of an array of coordinates (or angles). This can be used to enable
-    computations of single observables using functions originally designed to
-    accept only 2-d coordinate arrays.
+    It then checks if the corresponding values are valid coordinate arrays or
+    an :class:`~MDAnalysis.core.groups.AtomGroup`.
+    If the input is an array and all these arrays are single coordinates
+    (i.e., their shape is ``(3,)``), the decorated function can optionally
+    return a single coordinate (or angle) instead of an array of coordinates
+    (or angles). This can be used to enable computations of single observables
+    using functions originally designed to accept only 2-d coordinate arrays.
+
+    If the input is an :class:`~MDAnalysis.core.groups.AtomGroup` it is
+    converted into its corresponding position array via a call to
+    `AtomGroup.positions`.
 
     The checks performed on each individual coordinate array are:
 
@@ -1979,6 +1985,9 @@ def check_coords(*coord_names, **options):
         * **check_lengths_match** (:class:`bool`, optional) -- If ``True``, a
           :class:`ValueError` is raised if not all coordinate arrays contain the
           same number of coordinates. Default: ``True``
+        * **allow_atomgroup** (:class:`bool`, optional) -- If ``False``, a
+          :class:`TypeError` is raised if an :class:`AtomGroup` is supplied
+          Default: ``False``
 
     Raises
     ------
@@ -1992,15 +2001,16 @@ def check_coords(*coord_names, **options):
 
         If any of the coordinate arrays has a wrong shape.
     TypeError
-        If any of the coordinate arrays is not a :class:`numpy.ndarray`.
+        If any of the coordinate arrays is not a :class:`numpy.ndarray` or an
+        :class:`~MDAnalysis.core.groups.AtomGroup`.
 
         If the dtype of any of the coordinate arrays is not convertible to
           ``numpy.float32``.
 
     Example
     -------
 
-    >>> @check_coords('coords1', 'coords2')
+    >>> @check_coords('coords1', 'coords2', allow_atomgroup=True)
     ... def coordsum(coords1, coords2):
     ...     assert coords1.dtype == np.float32
     ...     assert coords2.flags['C_CONTIGUOUS']
@@ -2014,12 +2024,20 @@ def check_coords(*coord_names, **options):
     >>> coordsum(np.zeros(3), np.ones(6)[::2])
     array([1., 1., 1.], dtype=float32)
     >>>
+    >>> # automatic handling of AtomGroups
+    >>> u = mda.Universe(PSF, DCD)
+    >>> coordsum(u.atoms, u.select_atoms("index 1 to 10"))
+    ...
+    >>>
     >>> # automatic shape checking:
     >>> coordsum(np.zeros(3), np.ones(6))
     ValueError: coordsum(): coords2.shape must be (3,) or (n, 3), got (6,).
 
 
     .. versionadded:: 0.19.0
+    .. versionchanged:: 2.3.0
+       Can now accept an :class:`AtomGroup` as input, and added option
+       allow_atomgroup with default False to retain old behaviour
     """
     enforce_copy = options.get('enforce_copy', True)
     enforce_dtype = options.get('enforce_dtype', True)
@@ -2028,6 +2046,7 @@ def check_coords(*coord_names, **options):
     reduce_result_if_single = options.get('reduce_result_if_single', True)
     check_lengths_match = options.get('check_lengths_match',
                                       len(coord_names) > 1)
+    allow_atomgroup = options.get('allow_atomgroup', False)
     if not coord_names:
         raise ValueError("Decorator check_coords() cannot be used without "
                          "positional arguments.")
@@ -2051,32 +2070,49 @@ def check_coords_decorator(func):
                                  "".format(name, func.__name__))
 
         def _check_coords(coords, argname):
-            if not isinstance(coords, np.ndarray):
-                raise TypeError("{}(): Parameter '{}' must be a numpy.ndarray, "
-                                "got {}.".format(fname, argname, type(coords)))
             is_single = False
-            if allow_single:
-                if (coords.ndim not in (1, 2)) or (coords.shape[-1] != 3):
-                    raise ValueError("{}(): {}.shape must be (3,) or (n, 3), "
-                                     "got {}.".format(fname, argname,
-                                                      coords.shape))
-                if coords.ndim == 1:
-                    is_single = True
-                    if convert_single:
-                        coords = coords[None, :]
+            if isinstance(coords, np.ndarray):
+                if allow_single:
+                    if (coords.ndim not in (1, 2)) or (coords.shape[-1] != 3):
+                        errmsg = (f"{fname}(): {argname}.shape must be (3,) "
+                                  f"or (n, 3), got {coords.shape}")
+                        raise ValueError(errmsg)
+                    if coords.ndim == 1:
+                        is_single = True
+                        if convert_single:
+                            coords = coords[None, :]
+                else:
+                    if (coords.ndim != 2) or (coords.shape[1] != 3):
+                        errmsg = (f"{fname}(): {argname}.shape must be (n, 3) "
+                                  f"got {coords.shape}")
+                        raise ValueError(errmsg)
+                if enforce_dtype:
+                    try:
+                        coords = coords.astype(
+                            np.float32, order='C', copy=enforce_copy)
+                    except ValueError:
+                        errmsg = (f"{fname}(): {argname}.dtype must be"
+                                  f"convertible to float32, got"
+                                  f" {coords.dtype}.")
+                        raise TypeError(errmsg) from None
+                # coordinates should now be the right shape
+                ncoord = coords.shape[0]
             else:
-                if (coords.ndim != 2) or (coords.shape[1] != 3):
-                    raise ValueError("{}(): {}.shape must be (n, 3), got {}."
-                                     "".format(fname, argname, coords.shape))
-            if enforce_dtype:
                 try:
-                    coords = coords.astype(
-                        np.float32, order='C', copy=enforce_copy)
-                except ValueError:
-                    errmsg = (f"{fname}(): {argname}.dtype must be convertible to "
-                              f"float32, got {coords.dtype}.")
-                    raise TypeError(errmsg) from None
-            return coords, is_single
+                    coords = coords.positions  # homogenise to a numpy array
+                    ncoord = coords.shape[0]
+                    if not allow_atomgroup:
+                        err = TypeError("AtomGroup or other class with a"
+                                        "`.positions` method supplied as an"
+                                        "argument, but allow_atomgroup is"
+                                        " False")
+                        raise err
+                except AttributeError:
+                    raise TypeError(f"{fname}(): Parameter '{argname}' must be"
+                                    f" a numpy.ndarray or an AtomGroup,"
+                                    f" got {type(coords)}.")
+
+            return coords, is_single, ncoord
 
         @wraps(func)
         def wrapper(*args, **kwargs):
@@ -2107,14 +2143,15 @@ def wrapper(*args, **kwargs):
             for name in coord_names:
                 idx = posargnames.index(name)
                 if idx < len(args):
-                    args[idx], is_single = _check_coords(args[idx], name)
+                    args[idx], is_single, ncoord = _check_coords(args[idx],
+                                                                 name)
                     all_single &= is_single
-                    ncoords.append(args[idx].shape[0])
+                    ncoords.append(ncoord)
                 else:
-                    kwargs[name], is_single = _check_coords(kwargs[name],
-                                                            name)
+                    kwargs[name], is_single, ncoord = _check_coords(kwargs[name],
+                                                                    name)
                     all_single &= is_single
-                    ncoords.append(kwargs[name].shape[0])
+                    ncoords.append(ncoord)
             if check_lengths_match and ncoords:
                 if ncoords.count(ncoords[0]) != len(ncoords):
                     raise ValueError("{}(): {} must contain the same number of "

testsuite/MDAnalysisTests/converters/test_openmm.py

@@ -22,7 +22,7 @@
 #
 import pytest
 import numpy as np
-from numpy.testing import assert_almost_equal
+from numpy.testing import assert_allclose
 
 from MDAnalysisTests.datafiles import PDBX
 from MDAnalysisTests.coordinates.reference import RefAdKSmall
@@ -65,18 +65,19 @@ def omm_sim_uni(self):
         return mda.Universe(simulation)
 
     def test_dimensions(self, omm_sim_uni):
-        assert_almost_equal(
+        assert_allclose(
             omm_sim_uni.trajectory.ts.dimensions,
             np.array([20., 20., 20., 90., 90., 90.]),
-            3,
-            "OpenMMBasicSimulationReader failed to get unitcell dimensions " +
-            "from OpenMM Simulation Object",
+            rtol=0,
+            atol=1e-3,
+            err_msg=("OpenMMBasicSimulationReader failed to get unitcell "
+                     "dimensions from OpenMM Simulation Object"),
         )
 
     def test_coordinates(self, omm_sim_uni):
         up = omm_sim_uni.atoms.positions
         reference = np.ones((5, 3))
-        assert_almost_equal(up, reference, decimal=3)
+        assert_allclose(up, reference, rtol=0, atol=1e-3)
 
     def test_basic_topology(self, omm_sim_uni):
         assert omm_sim_uni.atoms.n_atoms == 5
@@ -96,18 +97,19 @@ def setUp(self):
         self.prec = 3
 
     def test_dimensions(self):
-        assert_almost_equal(
+        assert_allclose(
             self.universe.trajectory.ts.dimensions,
             self.ref.trajectory.ts.dimensions,
-            self.prec,
-            "OpenMMPDBFileReader failed to get unitcell dimensions " +
-            "from OpenMMPDBFile",
+            rtol=0,
+            atol=1e-3,
+            err_msg=("OpenMMPDBFileReader failed to get unitcell dimensions "
+                     "from OpenMMPDBFile"),
         )
 
     def test_coordinates(self):
         up = self.universe.atoms.positions
         rp = self.ref.atoms.positions
-        assert_almost_equal(up, rp, decimal=3)
+        assert_allclose(up, rp, rtol=0, atol=1e-3)
 
 
 class TestOpenMMModellerReader(_SingleFrameReader):
@@ -121,18 +123,19 @@ def setUp(self):
         self.prec = 3
 
     def test_dimensions(self):
-        assert_almost_equal(
+        assert_allclose(
             self.universe.trajectory.ts.dimensions,
             self.ref.trajectory.ts.dimensions,
-            self.prec,
-            "OpenMMModellerReader failed to get unitcell dimensions " +
-            "from OpenMMModeller",
+            rtol=0,
+            atol=1e-3,
+            err_msg=("OpenMMModellerReader failed to get unitcell dimensions "
+                     "from OpenMMModeller"),
         )
 
     def test_coordinates(self):
         up = self.universe.atoms.positions
         rp = self.ref.atoms.positions
-        assert_almost_equal(up, rp, decimal=3)
+        assert_allclose(up, rp, rtol=0, atol=1e-3)
 
 
 class TestOpenMMSimulationReader(_SingleFrameReader):
@@ -154,18 +157,19 @@ def setUp(self):
         self.prec = 3
 
     def test_dimensions(self):
-        assert_almost_equal(
+        assert_allclose(
             self.universe.trajectory.ts.dimensions,
             self.ref.trajectory.ts.dimensions,
-            self.prec,
-            "OpenMMSimulationReader failed to get unitcell dimensions " +
-            "from OpenMMSimulation",
+            rtol=0,
+            atol=1e-3,
+            err_msg=("OpenMMSimulationReader failed to get unitcell "
+                     "dimensions from OpenMMSimulation"),
         )
 
     def test_coordinates(self):
         up = self.universe.atoms.positions
         rp = self.ref.atoms.positions
-        assert_almost_equal(up, rp, decimal=3)
+        assert_allclose(up, rp, rtol=0, atol=1e-3)
 
     @pytest.mark.xfail(reason='OpenMM pickling not supported yet')
     def test_pickle_singleframe_reader(self):
@@ -246,4 +250,4 @@ def test_pdbx_coordinates(PDBX_U):
         ]
     )
     rp = PDBX_U.atoms.positions
-    assert_almost_equal(ref_pos, rp, decimal=3)
+    assert_allclose(ref_pos, rp, rtol=0, atol=1e-3)