From dbda8ebe8bf3789016c1a62fbe0620b36ba1083e Mon Sep 17 00:00:00 2001 From: Hugo MacDermott-Opeskin Date: Fri, 22 Jul 2022 17:42:22 +1000 Subject: [PATCH 1/2] Allow distance routines to accept an NumPy Array or AtomGroup [Core] (#3730) Allow distance routines to accept an NumPy Array or AtomGroup in one or both positions. @check_distances can now also extract coordinates from an AtomGroup. Co-authored-by: Irfan Alibay --- package/CHANGELOG | 2 + package/MDAnalysis/lib/distances.py | 166 +++++--- package/MDAnalysis/lib/util.py | 109 +++-- .../MDAnalysisTests/lib/test_distances.py | 383 ++++++++++++++++-- testsuite/MDAnalysisTests/lib/test_util.py | 73 +++- 5 files changed, 605 insertions(+), 128 deletions(-) diff --git a/package/CHANGELOG b/package/CHANGELOG index 679a3157b40..3770922ee6c 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -22,6 +22,8 @@ Fixes * add a 0.5 for correct midpoints in hole analysis (Issue #3715) Enhancements + * 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) diff --git a/package/MDAnalysis/lib/distances.py b/package/MDAnalysis/lib/distances.py index c529142e99c..8836b8725bd 100644 --- a/package/MDAnalysis/lib/distances.py +++ b/package/MDAnalysis/lib/distances.py @@ -26,10 +26,11 @@ =================================================================== Fast C-routines to calculate arrays of distances or angles from coordinate -arrays. Many of the functions also exist in parallel versions, which typically -provide higher performance than the serial code. -The boolean attribute `MDAnalysis.lib.distances.USED_OPENMP` can be checked to -see if OpenMP was used in the compilation of MDAnalysis. +arrays. Distance functions can accept a NumPy :class:`np.ndarray` or an +:class:`~MDAnalysis.core.groups.AtomGroup`. Many of the functions also exist +in parallel versions, which typically provide higher performance than the +serial code. The boolean attribute `MDAnalysis.lib.distances.USED_OPENMP` can +be checked to see if OpenMP was used in the compilation of MDAnalysis. Selection of acceleration ("backend") ------------------------------------- @@ -50,6 +51,9 @@ ========== ========================= ====================================== .. versionadded:: 0.13.0 +.. versionchanged:: 2.3.0 + Distance functions can now accept an + :class:`~MDAnalysis.core.groups.AtomGroup` or an :class:`np.ndarray` Functions --------- @@ -70,6 +74,10 @@ import numpy as np from numpy.lib.utils import deprecate +from typing import Union, Optional, Callable +from typing import TYPE_CHECKING +if TYPE_CHECKING: # pragma: no cover + from ..core.groups import AtomGroup from .util import check_coords, check_box from .mdamath import triclinic_vectors from ._augment import augment_coordinates, undo_augment @@ -91,7 +99,9 @@ pass del importlib -def _run(funcname, args=None, kwargs=None, backend="serial"): + +def _run(funcname: Callable, args: Optional[tuple] = None, + kwargs: Optional[dict] = None, backend: str = "serial") -> Callable: """Helper function to select a backend function `funcname`.""" args = args if args is not None else tuple() kwargs = kwargs if kwargs is not None else dict() @@ -129,7 +139,8 @@ def _run(funcname, args=None, kwargs=None, backend="serial"): from .c_distances_openmp import OPENMP_ENABLED as USED_OPENMP -def _check_result_array(result, shape): +# typing: numpy +def _check_result_array(result: np.ndarray, shape: tuple) -> np.ndarray: """Check if the result array is ok to use. The `result` array must meet the following requirements: @@ -170,10 +181,14 @@ def _check_result_array(result, shape): return result +# typing: numpy @check_coords('reference', 'configuration', reduce_result_if_single=False, - check_lengths_match=False) -def distance_array(reference, configuration, box=None, result=None, - backend="serial"): + check_lengths_match=False, allow_atomgroup=True) +def distance_array(reference: Union[np.ndarray, 'AtomGroup'], + configuration: Union[np.ndarray, 'AtomGroup'], + box: Optional[np.ndarray] = None, + result: Optional[np.ndarray] = None, + backend: str = "serial") -> np.ndarray: """Calculate all possible distances between a reference set and another configuration. @@ -190,12 +205,14 @@ def distance_array(reference, configuration, box=None, result=None, Parameters ---------- - reference : numpy.ndarray + reference :numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Reference coordinate array of shape ``(3,)`` or ``(n, 3)`` (dtype is - arbitrary, will be converted to ``numpy.float32`` internally). - configuration : numpy.ndarray + arbitrary, will be converted to ``numpy.float32`` internally). Also + accepts an :class:`~MDAnalysis.core.groups.AtomGroup`. + configuration : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Configuration coordinate array of shape ``(3,)`` or ``(m, 3)`` (dtype is - arbitrary, will be converted to ``numpy.float32`` internally). + arbitrary, will be converted to ``numpy.float32`` internally). Also + accepts an :class:`~MDAnalysis.core.groups.AtomGroup`. box : array_like, optional The unitcell dimensions of the system, which can be orthogonal or triclinic and must be provided in the same format as returned by @@ -221,6 +238,9 @@ def distance_array(reference, configuration, box=None, result=None, .. versionchanged:: 0.19.0 Internal dtype conversion of input coordinates to ``numpy.float32``. Now also accepts single coordinates as input. + .. versionchanged:: 2.3.0 + Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an + argument in any position and checks inputs using type hinting. """ confnum = configuration.shape[0] refnum = reference.shape[0] @@ -251,8 +271,12 @@ def distance_array(reference, configuration, box=None, result=None, return distances -@check_coords('reference', reduce_result_if_single=False) -def self_distance_array(reference, box=None, result=None, backend="serial"): +# typing: numpy +@check_coords('reference', reduce_result_if_single=False, allow_atomgroup=True) +def self_distance_array(reference: Union[np.ndarray, 'AtomGroup'], + box: Optional[np.ndarray] = None, + result: Optional[np.ndarray] = None, + backend: str = "serial") -> np.ndarray: """Calculate all possible distances within a configuration `reference`. If the optional argument `box` is supplied, the minimum image convention is @@ -265,9 +289,10 @@ def self_distance_array(reference, box=None, result=None, backend="serial"): Parameters ---------- - reference : numpy.ndarray + reference : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Reference coordinate array of shape ``(3,)`` or ``(n, 3)`` (dtype is - arbitrary, will be converted to ``numpy.float32`` internally). + arbitrary, will be converted to ``numpy.float32`` internally). Also + accepts an :class:`~MDAnalysis.core.groups.AtomGroup`. box : array_like, optional The unitcell dimensions of the system, which can be orthogonal or triclinic and must be provided in the same format as returned by @@ -298,6 +323,9 @@ def self_distance_array(reference, box=None, result=None, backend="serial"): Added *backend* keyword. .. versionchanged:: 0.19.0 Internal dtype conversion of input coordinates to ``numpy.float32``. + .. versionchanged:: 2.3.0 + Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an + argument in any position and checks inputs using type hinting. """ refnum = reference.shape[0] distnum = refnum * (refnum - 1) // 2 @@ -1241,8 +1269,13 @@ def transform_StoR(coords, box, backend="serial"): return coords -@check_coords('coords1', 'coords2') -def calc_bonds(coords1, coords2, box=None, result=None, backend="serial"): +# typing: numpy +@check_coords('coords1', 'coords2', allow_atomgroup=True) +def calc_bonds(coords1: Union[np.ndarray, 'AtomGroup'], + coords2: Union[np.ndarray, 'AtomGroup'], + box: Optional[np.ndarray] = None, + result: Optional[np.ndarray] = None, + backend: str = "serial") -> np.ndarray: """Calculates the bond lengths between pairs of atom positions from the two coordinate arrays `coords1` and `coords2`, which must contain the same number of coordinates. ``coords1[i]`` and ``coords2[i]`` represent the @@ -1266,14 +1299,16 @@ def calc_bonds(coords1, coords2, box=None, result=None, backend="serial"): Parameters ---------- - coords1 : numpy.ndarray + coords1 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Coordinate array of shape ``(3,)`` or ``(n, 3)`` for one half of a single or ``n`` bonds, respectively (dtype is arbitrary, will be - converted to ``numpy.float32`` internally). - coords2 : numpy.ndarray + converted to ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. + coords2 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Coordinate array of shape ``(3,)`` or ``(n, 3)`` for the other half of a single or ``n`` bonds, respectively (dtype is arbitrary, will be - converted to ``numpy.float32`` internally). + converted to ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. box : numpy.ndarray, optional The unitcell dimensions of the system, which can be orthogonal or triclinic and must be provided in the same format as returned by @@ -1300,6 +1335,9 @@ def calc_bonds(coords1, coords2, box=None, result=None, backend="serial"): .. versionchanged:: 0.19.0 Internal dtype conversion of input coordinates to ``numpy.float32``. Now also accepts single coordinates as input. + .. versionchanged:: 2.3.0 + Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an + argument in any position and checks inputs using type hinting. """ numatom = coords1.shape[0] bondlengths = _check_result_array(result, (numatom,)) @@ -1323,9 +1361,14 @@ def calc_bonds(coords1, coords2, box=None, result=None, backend="serial"): return bondlengths -@check_coords('coords1', 'coords2', 'coords3') -def calc_angles(coords1, coords2, coords3, box=None, result=None, - backend="serial"): +# typing: numpy +@check_coords('coords1', 'coords2', 'coords3', allow_atomgroup=True) +def calc_angles(coords1: Union[np.ndarray, 'AtomGroup'], + coords2: Union[np.ndarray, 'AtomGroup'], + coords3: Union[np.ndarray, 'AtomGroup'], + box: Optional[np.ndarray] = None, + result: Optional[np.ndarray] = None, + backend: str = "serial") -> np.ndarray: """Calculates the angles formed between triplets of atom positions from the three coordinate arrays `coords1`, `coords2`, and `coords3`. All coordinate arrays must contain the same number of coordinates. @@ -1351,18 +1394,21 @@ def calc_angles(coords1, coords2, coords3, box=None, result=None, Parameters ---------- - coords1 : numpy.ndarray + coords1 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Array of shape ``(3,)`` or ``(n, 3)`` containing the coordinates of one side of a single or ``n`` angles, respectively (dtype is arbitrary, will - be converted to ``numpy.float32`` internally) - coords2 : numpy.ndarray + be converted to ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. + coords2 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Array of shape ``(3,)`` or ``(n, 3)`` containing the coordinates of the apices of a single or ``n`` angles, respectively (dtype is arbitrary, - will be converted to ``numpy.float32`` internally) - coords3 : numpy.ndarray + will be converted to ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. + coords3 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Array of shape ``(3,)`` or ``(n, 3)`` containing the coordinates of the other side of a single or ``n`` angles, respectively (dtype is - arbitrary, will be converted to ``numpy.float32`` internally) + arbitrary, will be converted to ``numpy.float32`` internally). + Also accepts an :class:`~MDAnalysis.core.groups.AtomGroup`. box : numpy.ndarray, optional The unitcell dimensions of the system, which can be orthogonal or triclinic and must be provided in the same format as returned by @@ -1392,6 +1438,9 @@ def calc_angles(coords1, coords2, coords3, box=None, result=None, .. versionchanged:: 0.19.0 Internal dtype conversion of input coordinates to ``numpy.float32``. Now also accepts single coordinates as input. + .. versionchanged:: 2.3.0 + Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an + argument in any position and checks inputs using type hinting. """ numatom = coords1.shape[0] angles = _check_result_array(result, (numatom,)) @@ -1415,9 +1464,15 @@ def calc_angles(coords1, coords2, coords3, box=None, result=None, return angles -@check_coords('coords1', 'coords2', 'coords3', 'coords4') -def calc_dihedrals(coords1, coords2, coords3, coords4, box=None, result=None, - backend="serial"): +# typing: numpy +@check_coords('coords1', 'coords2', 'coords3', 'coords4', allow_atomgroup=True) +def calc_dihedrals(coords1: Union[np.ndarray, 'AtomGroup'], + coords2: Union[np.ndarray, 'AtomGroup'], + coords3: Union[np.ndarray, 'AtomGroup'], + coords4: Union[np.ndarray, 'AtomGroup'], + box: Optional[np.ndarray] = None, + result: Optional[np.ndarray] = None, + backend: str = "serial") -> np.ndarray: r"""Calculates the dihedral angles formed between quadruplets of positions from the four coordinate arrays `coords1`, `coords2`, `coords3`, and `coords4`, which must contain the same number of coordinates. @@ -1449,22 +1504,26 @@ def calc_dihedrals(coords1, coords2, coords3, coords4, box=None, result=None, Parameters ---------- - coords1 : numpy.ndarray + coords1 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Coordinate array of shape ``(3,)`` or ``(n, 3)`` containing the 1st positions in dihedrals (dtype is arbitrary, will be converted to - ``numpy.float32`` internally) - coords2 : numpy.ndarray + ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. + coords2 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Coordinate array of shape ``(3,)`` or ``(n, 3)`` containing the 2nd positions in dihedrals (dtype is arbitrary, will be converted to - ``numpy.float32`` internally) - coords3 : numpy.ndarray + ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. + coords3 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Coordinate array of shape ``(3,)`` or ``(n, 3)`` containing the 3rd positions in dihedrals (dtype is arbitrary, will be converted to - ``numpy.float32`` internally) - coords4 : numpy.ndarray + ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. + coords4 : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Coordinate array of shape ``(3,)`` or ``(n, 3)`` containing the 4th positions in dihedrals (dtype is arbitrary, will be converted to - ``numpy.float32`` internally) + ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. box : numpy.ndarray, optional The unitcell dimensions of the system, which can be orthogonal or triclinic and must be provided in the same format as returned by @@ -1497,6 +1556,9 @@ def calc_dihedrals(coords1, coords2, coords3, coords4, box=None, result=None, .. versionchanged:: 0.19.0 Internal dtype conversion of input coordinates to ``numpy.float32``. Now also accepts single coordinates as input. + .. versionchanged:: 2.3.0 + Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an + argument in any position and checks inputs using type hinting. """ numatom = coords1.shape[0] dihedrals = _check_result_array(result, (numatom,)) @@ -1520,15 +1582,19 @@ def calc_dihedrals(coords1, coords2, coords3, coords4, box=None, result=None, return dihedrals -@check_coords('coords') -def apply_PBC(coords, box, backend="serial"): +# typing: numpy +@check_coords('coords', allow_atomgroup=True) +def apply_PBC(coords: Union[np.ndarray, 'AtomGroup'], + box: Optional[np.ndarray] = None, + backend: str = "serial") -> np.ndarray: """Moves coordinates into the primary unit cell. Parameters ---------- - coords : numpy.ndarray + coords : numpy.ndarray or :class:`~MDAnalysis.core.groups.AtomGroup` Coordinate array of shape ``(3,)`` or ``(n, 3)`` (dtype is arbitrary, - will be converted to ``numpy.float32`` internally). + will be converted to ``numpy.float32`` internally). Also accepts an + :class:`~MDAnalysis.core.groups.AtomGroup`. box : numpy.ndarray The unitcell dimensions of the system, which can be orthogonal or triclinic and must be provided in the same format as returned by @@ -1550,6 +1616,9 @@ def apply_PBC(coords, box, backend="serial"): .. versionchanged:: 0.19.0 Internal dtype conversion of input coordinates to ``numpy.float32``. Now also accepts (and, likewise, returns) single coordinates. + .. versionchanged:: 2.3.0 + Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an + argument in any position and checks inputs using type hinting. """ if len(coords) == 0: return coords @@ -1562,8 +1631,9 @@ def apply_PBC(coords, box, backend="serial"): return coords +# typing: numpy @check_coords('vectors', enforce_copy=False, enforce_dtype=False) -def minimize_vectors(vectors, box): +def minimize_vectors(vectors: np.ndarray, box: np.ndarray) -> np.ndarray: """Apply minimum image convention to an array of vectors This function is required for calculating the correct vectors between two diff --git a/package/MDAnalysis/lib/util.py b/package/MDAnalysis/lib/util.py index 2b0cb03c060..22c73dbd10a 100644 --- a/package/MDAnalysis/lib/util.py +++ b/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,7 +2001,8 @@ 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``. @@ -2000,7 +2010,7 @@ def check_coords(*coord_names, **options): 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 " diff --git a/testsuite/MDAnalysisTests/lib/test_distances.py b/testsuite/MDAnalysisTests/lib/test_distances.py index 7168f74c87a..c8b3ee6b009 100644 --- a/testsuite/MDAnalysisTests/lib/test_distances.py +++ b/testsuite/MDAnalysisTests/lib/test_distances.py @@ -256,6 +256,16 @@ def ref_system(): return box, points, ref, conf +@pytest.fixture() +def ref_system_universe(ref_system): + box, points, ref, conf = ref_system + u = MDAnalysis.Universe.empty(points.shape[0], trajectory=True) + u.atoms.positions = points + u.trajectory.ts.dimensions = box + return (box, u.atoms, u.select_atoms("index 0"), + u.select_atoms("index 1 to 3")) + + @pytest.mark.parametrize('backend', ['serial', 'openmp']) class TestDistanceArray(object): @staticmethod @@ -264,24 +274,61 @@ def _dist(x, ref): r = x - ref return np.sqrt(np.dot(r, r)) - def test_noPBC(self, backend, ref_system): - box, points, ref, conf = ref_system - - d = distances.distance_array(ref, points, backend=backend) + # test both AtomGroup and numpy array + @pytest.mark.parametrize('pos', ['ref_system', 'ref_system_universe']) + def test_noPBC(self, backend, ref_system, pos, request): + _, points, reference, _ = ref_system # reference values + _, all, ref, _ = request.getfixturevalue(pos) + d = distances.distance_array(ref, all, backend=backend) assert_almost_equal(d, np.array([[ - self._dist(points[0], ref[0]), - self._dist(points[1], ref[0]), - self._dist(points[2], ref[0]), - self._dist(points[3], ref[0])] + self._dist(points[0], reference[0]), + self._dist(points[1], reference[0]), + self._dist(points[2], reference[0]), + self._dist(points[3], reference[0])] ])) - def test_PBC(self, backend, ref_system): - box, points, ref, conf = ref_system - - d = distances.distance_array(ref, points, box=box, backend=backend) + # cycle through combinations of numpy array and AtomGroup + @pytest.mark.parametrize('pos0', ['ref_system', 'ref_system_universe']) + @pytest.mark.parametrize('pos1', ['ref_system', 'ref_system_universe']) + def test_noPBC_mixed_combinations(self, backend, ref_system, pos0, pos1, + request): + _, points, reference, _ = ref_system # reference values + _, _, ref_val, _ = request.getfixturevalue(pos0) + _, points_val, _, _ = request.getfixturevalue(pos1) + d = distances.distance_array(ref_val, points_val, + backend=backend) + assert_almost_equal(d, np.array([[ + self._dist(points[0], reference[0]), + self._dist(points[1], reference[0]), + self._dist(points[2], reference[0]), + self._dist(points[3], reference[0])] + ])) - assert_almost_equal(d, np.array([[0., 0., 0., self._dist(points[3], ref=[1, 1, 2])]])) + # test both AtomGroup and numpy array + @pytest.mark.parametrize('pos', ['ref_system', 'ref_system_universe']) + def test_PBC(self, backend, ref_system, pos, request): + box, points, _, _ = ref_system + _, all, ref, _ = request.getfixturevalue(pos) + + d = distances.distance_array(ref, all, box=box, backend=backend) + + assert_almost_equal(d, np.array([[0., 0., 0., self._dist(points[3], + ref=[1, 1, 2])]])) + + # cycle through combinations of numpy array and AtomGroup + @pytest.mark.parametrize('pos0', ['ref_system', 'ref_system_universe']) + @pytest.mark.parametrize('pos1', ['ref_system', 'ref_system_universe']) + def test_PBC_mixed_combinations(self, backend, ref_system, pos0, pos1, + request): + box, points, _, _ = ref_system + _, _, ref_val, _ = request.getfixturevalue(pos0) + _, points_val, _, _ = request.getfixturevalue(pos1) + d = distances.distance_array(ref_val, points_val, + box=box, + backend=backend) + assert_almost_equal( + d, np.array([[0., 0., 0., self._dist(points[3], ref=[1, 1, 2])]])) def test_PBC2(self, backend): a = np.array([7.90146923, -13.72858524, 3.75326586], dtype=np.float32) @@ -321,11 +368,26 @@ class FakeArray(np.ndarray): def DCD_Universe(): universe = MDAnalysis.Universe(PSF, DCD) trajectory = universe.trajectory + return universe, trajectory + +# second independent universe required for +# TestDistanceArrayDCD_TRIC.test_atomgroup_simple +@pytest.fixture() +def DCD_Universe2(): + universe = MDAnalysis.Universe(PSF, DCD) + trajectory = universe.trajectory + return universe, trajectory + + +@pytest.fixture() +def Triclinic_Universe(): + universe = MDAnalysis.Universe(TRIC) + trajectory = universe.trajectory return universe, trajectory @pytest.mark.parametrize('backend', ['serial', 'openmp']) -class TestDistanceArrayDCD(object): +class TestDistanceArrayDCD_TRIC(object): # reasonable precision so that tests succeed on 32 and 64 bit machines # (the reference values were obtained on 64 bit) # Example: @@ -341,7 +403,8 @@ def test_simple(self, DCD_Universe, backend): trajectory[10] x1 = U.atoms.positions d = distances.distance_array(x0, x1, backend=backend) - assert_equal(d.shape, (3341, 3341), "wrong shape (should be (Natoms,Natoms))") + assert_equal(d.shape, (3341, 3341), "wrong shape (should be" + "(Natoms,Natoms))") assert_almost_equal(d.min(), 0.11981228170520701, self.prec, err_msg="wrong minimum distance value") assert_almost_equal(d.max(), 53.572192429459619, self.prec, @@ -356,7 +419,8 @@ def test_outarray(self, DCD_Universe, backend): natoms = len(U.atoms) d = np.zeros((natoms, natoms), np.float64) distances.distance_array(x0, x1, result=d, backend=backend) - assert_equal(d.shape, (natoms, natoms), "wrong shape, shoud be (Natoms,Natoms) entries") + assert_equal(d.shape, (natoms, natoms), "wrong shape, should be" + " (Natoms,Natoms) entries") assert_almost_equal(d.min(), 0.11981228170520701, self.prec, err_msg="wrong minimum distance value") assert_almost_equal(d.max(), 53.572192429459619, self.prec, @@ -371,16 +435,71 @@ def test_periodic(self, DCD_Universe, backend): x1 = U.atoms.positions d = distances.distance_array(x0, x1, box=U.coord.dimensions, backend=backend) - assert_equal(d.shape, (3341, 3341), "should be square matrix with Natoms entries") + assert_equal(d.shape, (3341, 3341), "should be square matrix with" + " Natoms entries") assert_almost_equal(d.min(), 0.11981228170520701, self.prec, err_msg="wrong minimum distance value with PBC") assert_almost_equal(d.max(), 53.572192429459619, self.prec, err_msg="wrong maximum distance value with PBC") + def test_atomgroup_simple(self, DCD_Universe, DCD_Universe2, backend): + # need two copies as moving ts updates underlying array on atomgroup + U1, trajectory1 = DCD_Universe + U2, trajectory2 = DCD_Universe2 + trajectory1.rewind() + trajectory2.rewind() + x0 = U1.select_atoms("all") + trajectory2[10] + x1 = U2.select_atoms("all") + d = distances.distance_array(x0, x1, backend=backend) + assert_equal(d.shape, (3341, 3341), "wrong shape (should be" + " (Natoms,Natoms))") + assert_almost_equal(d.min(), 0.11981228170520701, self.prec, + err_msg="wrong minimum distance value") + assert_almost_equal(d.max(), 53.572192429459619, self.prec, + err_msg="wrong maximum distance value") + + # check no box and ortho box types and some slices + @pytest.mark.parametrize('box', [None, [50, 50, 50, 90, 90, 90]]) + @pytest.mark.parametrize("sel, np_slice", [("all", np.s_[:, :]), + ("index 0 to 8 ", np.s_[0:9, :]), + ("index 9", np.s_[8, :])]) + def test_atomgroup_matches_numpy(self, DCD_Universe, backend, sel, + np_slice, box): + U, _ = DCD_Universe + x0_ag = U.select_atoms(sel) + x0_arr = U.atoms.positions[np_slice] + x1_ag = U.select_atoms(sel) + x1_arr = U.atoms.positions[np_slice] + d_ag = distances.distance_array(x0_ag, x1_ag, box=box, + backend=backend) + d_arr = distances.distance_array(x0_arr, x1_arr, box=box, + backend=backend) + assert_allclose(d_ag, d_arr, + err_msg="AtomGroup and NumPy distances do not match") + + # check triclinic box and some slices + @pytest.mark.parametrize("sel, np_slice", [("all", np.s_[:, :]), + ("index 0 to 8 ", np.s_[0:9, :]), + ("index 9", np.s_[8, :])]) + def test_atomgroup_matches_numpy_tric(self, Triclinic_Universe, backend, + sel, np_slice): + U, _ = Triclinic_Universe + x0_ag = U.select_atoms(sel) + x0_arr = U.atoms.positions[np_slice] + x1_ag = U.select_atoms(sel) + x1_arr = U.atoms.positions[np_slice] + d_ag = distances.distance_array(x0_ag, x1_ag, box=U.coord.dimensions, + backend=backend) + d_arr = distances.distance_array(x0_arr, x1_arr, + box=U.coord.dimensions, + backend=backend) + assert_allclose(d_ag, d_arr, + err_msg="AtomGroup and NumPy distances do not match") @pytest.mark.parametrize('backend', ['serial', 'openmp']) -class TestSelfDistanceArrayDCD(object): +class TestSelfDistanceArrayDCD_TRIC(object): prec = 5 def test_simple(self, DCD_Universe, backend): @@ -424,6 +543,54 @@ def test_periodic(self, DCD_Universe, backend): assert_almost_equal(d.max(), 52.4702570624190590, self.prec, err_msg="wrong maximum distance value with PBC") + def test_atomgroup_simple(self, DCD_Universe, backend): + U, trajectory = DCD_Universe + trajectory.rewind() + x0 = U.select_atoms("all") + d = distances.self_distance_array(x0, backend=backend) + N = 3341 * (3341 - 1) / 2 + assert_equal(d.shape, (N,), "wrong shape (should be" + " (Natoms*(Natoms-1)/2,))") + assert_almost_equal(d.min(), 0.92905562402529318, self.prec, + err_msg="wrong minimum distance value") + assert_almost_equal(d.max(), 52.4702570624190590, self.prec, + err_msg="wrong maximum distance value") + + # check no box and ortho box types and some slices + @pytest.mark.parametrize('box', [None, [50, 50, 50, 90, 90, 90]]) + @pytest.mark.parametrize("sel, np_slice", [("all", np.s_[:, :]), + ("index 0 to 8 ", np.s_[0:9, :]), + ("index 9", np.s_[8, :])]) + def test_atomgroup_matches_numpy(self, DCD_Universe, backend, + sel, np_slice, box): + U, _ = DCD_Universe + + x0_ag = U.select_atoms(sel) + x0_arr = U.atoms.positions[np_slice] + d_ag = distances.self_distance_array(x0_ag, box=box, + backend=backend) + d_arr = distances.self_distance_array(x0_arr, box=box, + backend=backend) + assert_allclose(d_ag, d_arr, + err_msg="AtomGroup and NumPy distances do not match") + + # check triclinic box and some slices + @pytest.mark.parametrize("sel, np_slice", [ + ("index 0 to 8 ", np.s_[0:9, :]), + ("index 9", np.s_[8, :])]) + def test_atomgroup_matches_numpy_tric(self, Triclinic_Universe, backend, + sel, np_slice): + U, _ = Triclinic_Universe + x0_ag = U.select_atoms(sel) + x0_arr = U.atoms.positions[np_slice] + d_ag = distances.self_distance_array(x0_ag, box=U.coord.dimensions, + backend=backend) + d_arr = distances.self_distance_array(x0_arr, box=U.coord.dimensions, + backend=backend) + assert_allclose(d_ag, d_arr, + err_msg="AtomGroup and NumPy distances do not match") + + @pytest.mark.parametrize('backend', ['serial', 'openmp']) class TestTriclinicDistances(object): """Unit tests for the Triclinic PBC functions. @@ -596,6 +763,20 @@ def test_issue_3725(box): np.testing.assert_allclose(self_da_serial, self_da_openmp) +def conv_dtype_if_ndarr(a, dtype): + if isinstance(a, np.ndarray): + return a.astype(dtype) + else: + return a + + +def convert_position_dtype_if_ndarray(a, b, c, d, dtype): + return (conv_dtype_if_ndarr(a, dtype), + conv_dtype_if_ndarr(b, dtype), + conv_dtype_if_ndarr(c, dtype), + conv_dtype_if_ndarr(d, dtype)) + + @pytest.mark.parametrize('backend', ['serial', 'openmp']) class TestCythonFunctions(object): # Unit tests for calc_bonds calc_angles and calc_dihedrals in lib.distances @@ -612,7 +793,7 @@ def box(): @pytest.fixture() def triclinic_box(): box_vecs = np.array([[10., 0., 0.], [1., 10., 0., ], [1., 0., 10.]], - dtype=np.float32) + dtype=np.float32) return mdamath.triclinic_box(box_vecs[0], box_vecs[1], box_vecs[2]) @staticmethod @@ -626,8 +807,15 @@ def positions(): return a, b, c, d @staticmethod - def convert_position_dtype(a, b, c, d, dtype): - return a.astype(dtype), b.astype(dtype), c.astype(dtype), d.astype(dtype) + @pytest.fixture() + def positions_atomgroups(positions): + a, b, c, d = positions + arrs = [a, b, c, d] + universes = [MDAnalysis.Universe.empty(arr.shape[0], + trajectory=True) for arr in arrs] + for u, a in zip(universes, arrs): + u.atoms.positions = a + return tuple([u.atoms for u in universes]) @staticmethod @pytest.fixture() @@ -643,8 +831,10 @@ def wronglength(): ((4, 3, -2), (-2, 2, 2), (-5, 2, 2), (0, 2, 2))] # multiple boxlengths @pytest.mark.parametrize('dtype', (np.float32, np.float64)) - def test_bonds(self, positions, box, backend, dtype): - a, b, c, d = self.convert_position_dtype(*positions, dtype=dtype) + @pytest.mark.parametrize('pos', ['positions', 'positions_atomgroups']) + def test_bonds(self, box, backend, dtype, pos, request): + a, b, c, d = request.getfixturevalue(pos) + a, b, c, d = convert_position_dtype_if_ndarray(a, b, c, d, dtype) dists = distances.calc_bonds(a, b, backend=backend) assert_equal(len(dists), 4, err_msg="calc_bonds results have wrong length") dists_pbc = distances.calc_bonds(a, b, box=box, backend=backend) @@ -679,8 +869,12 @@ def test_bonds_badresult(self, positions, backend): with pytest.raises(ValueError): distances.calc_bonds(a, b, result=badresult, backend=backend) - def test_bonds_triclinic(self, positions, triclinic_box, backend): - a, b, c, d = positions + @pytest.mark.parametrize('dtype', (np.float32, np.float64)) + @pytest.mark.parametrize('pos', ['positions', 'positions_atomgroups']) + def test_bonds_triclinic(self, triclinic_box, backend, dtype, pos, + request): + a, b, c, d = request.getfixturevalue(pos) + a, b, c, d = convert_position_dtype_if_ndarray(a, b, c, d, dtype) dists = distances.calc_bonds(a, b, box=triclinic_box, backend=backend) reference = np.array([0.0, 1.7320508, 1.4142136, 2.82842712]) assert_almost_equal(dists, reference, self.prec, err_msg="calc_bonds with triclinic box failed") @@ -705,8 +899,10 @@ def test_bonds_single_coords(self, shift, periodic, backend): assert_almost_equal(result, reference, decimal=self.prec) @pytest.mark.parametrize('dtype', (np.float32, np.float64)) - def test_angles(self, positions, backend, dtype): - a, b, c, d = self.convert_position_dtype(*positions, dtype=dtype) + @pytest.mark.parametrize('pos', ['positions', 'positions_atomgroups']) + def test_angles(self, backend, dtype, pos, request): + a, b, c, d = request.getfixturevalue(pos) + a, b, c, d = convert_position_dtype_if_ndarray(a, b, c, d, dtype) angles = distances.calc_angles(a, b, c, backend=backend) # Check calculated values assert_equal(len(angles), 4, err_msg="calc_angles results have wrong length") @@ -751,8 +947,10 @@ def manual_angle(x, y, z): assert_almost_equal(result, reference, decimal=4) @pytest.mark.parametrize('dtype', (np.float32, np.float64)) - def test_dihedrals(self, positions, backend, dtype): - a, b, c, d = self.convert_position_dtype(*positions, dtype=dtype) + @pytest.mark.parametrize('pos', ['positions', 'positions_atomgroups']) + def test_dihedrals(self, backend, dtype, pos, request): + a, b, c, d = request.getfixturevalue(pos) + a, b, c, d = convert_position_dtype_if_ndarray(a, b, c, d, dtype) dihedrals = distances.calc_dihedrals(a, b, c, d, backend=backend) # Check calculated values assert_equal(len(dihedrals), 4, err_msg="calc_dihedrals results have wrong length") @@ -842,24 +1040,60 @@ def test_numpy_compliance(self, positions, backend): class Test_apply_PBC(object): prec = 6 - def test_ortho_PBC(self, backend): - U = MDAnalysis.Universe(PSF, DCD) + @pytest.fixture() + def DCD_universe_pos(self, DCD_Universe): + U, _ = DCD_Universe + return U.atoms.positions + + @pytest.fixture() + def DCD_universe_ag(self, DCD_Universe): + U, _ = DCD_Universe + return U.atoms + + @pytest.fixture() + def Triclinic_universe_pos_box(self, Triclinic_Universe): + U, _ = Triclinic_Universe + atoms = U.atoms.positions + box = U.dimensions + return atoms, box + + @pytest.fixture() + def Triclinic_universe_pos_box(self, Triclinic_Universe): + U, _ = Triclinic_Universe atoms = U.atoms.positions + box = U.dimensions + return atoms, box + + @pytest.fixture() + def Triclinic_universe_ag_box(self, Triclinic_Universe): + U, _ = Triclinic_Universe + atoms = U.atoms + box = U.dimensions + return atoms, box + + @pytest.mark.parametrize('pos', ['DCD_universe_pos', 'DCD_universe_ag']) + def test_ortho_PBC(self, backend, pos, request, DCD_universe_pos): + positions = request.getfixturevalue(pos) box = np.array([2.5, 2.5, 3.5, 90., 90., 90.], dtype=np.float32) with pytest.raises(ValueError): - cyth1 = distances.apply_PBC(atoms, box[:3], backend=backend) - cyth2 = distances.apply_PBC(atoms, box, backend=backend) - reference = atoms - np.floor(atoms / box[:3]) * box[:3] + cyth1 = distances.apply_PBC(positions, box[:3], backend=backend) + cyth2 = distances.apply_PBC(positions, box, backend=backend) + reference = (DCD_universe_pos - + np.floor(DCD_universe_pos / box[:3]) * box[:3]) assert_almost_equal(cyth2, reference, self.prec, err_msg="Ortho apply_PBC #2 failed comparison with np") - def test_tric_PBC(self, backend): - U = MDAnalysis.Universe(TRIC) - atoms = U.atoms.positions - box = U.dimensions - + @pytest.mark.parametrize('pos', ['Triclinic_universe_pos_box', + 'Triclinic_universe_ag_box']) + def test_tric_PBC(self, backend, pos, request): + positions, box = request.getfixturevalue(pos) def numpy_PBC(coords, box): + # need this to allow both AtomGroup and array + if isinstance(coords, MDAnalysis.core.AtomGroup): + coords = coords.positions + else: + pass # move to fractional coordinates fractional = distances.transform_RtoS(coords, box) # move fractional coordinates to central cell @@ -867,8 +1101,9 @@ def numpy_PBC(coords, box): # move back to real coordinates return distances.transform_StoR(fractional, box) - cyth1 = distances.apply_PBC(atoms, box, backend=backend) - reference = numpy_PBC(atoms, box) + cyth1 = distances.apply_PBC(positions, box, backend=backend) + + reference = numpy_PBC(positions, box) assert_almost_equal(cyth1, reference, decimal=4, err_msg="Triclinic apply_PBC failed comparison with np") @@ -1009,6 +1244,15 @@ def coords(): np.array([[0.1, 1.9, 1.9], [-0.9, 0.9, 0.9]], dtype=np.float32), np.array([[0.1, 1.9, 0.1], [-0.9, 0.9, -0.9]], dtype=np.float32)] + @staticmethod + @pytest.fixture() + def coords_atomgroups(coords): + universes = [MDAnalysis.Universe.empty(arr.shape[0], trajectory=True) + for arr in coords] + for u, a in zip(universes, coords): + u.atoms.positions = a + return [u.atoms for u in universes] + @pytest.mark.parametrize('box', boxes) @pytest.mark.parametrize('backend', ['serial', 'openmp']) def test_input_unchanged_distance_array(self, coords, box, backend): @@ -1018,6 +1262,16 @@ def test_input_unchanged_distance_array(self, coords, box, backend): backend=backend) assert_equal(crds, refs) + @pytest.mark.parametrize('box', boxes) + @pytest.mark.parametrize('backend', ['serial', 'openmp']) + def test_input_unchanged_distance_array_atomgroup(self, coords_atomgroups, + box, backend): + crds = coords_atomgroups[:2] + refs = [crd.positions.copy() for crd in crds] + res = distances.distance_array(crds[0], crds[1], box=box, + backend=backend) + assert_equal([crd.positions for crd in crds], refs) + @pytest.mark.parametrize('box', boxes) @pytest.mark.parametrize('backend', ['serial', 'openmp']) def test_input_unchanged_self_distance_array(self, coords, box, backend): @@ -1026,6 +1280,16 @@ def test_input_unchanged_self_distance_array(self, coords, box, backend): res = distances.self_distance_array(crd, box=box, backend=backend) assert_equal(crd, ref) + @pytest.mark.parametrize('box', boxes) + @pytest.mark.parametrize('backend', ['serial', 'openmp']) + def test_input_unchanged_self_distance_array_atomgroup(self, + coords_atomgroups, + box, backend): + crd = coords_atomgroups[0] + ref = crd.positions.copy() + res = distances.self_distance_array(crd, box=box, backend=backend) + assert_equal(crd.positions, ref) + @pytest.mark.parametrize('box', boxes) @pytest.mark.parametrize('met', ["bruteforce", "pkdtree", "nsgrid", None]) def test_input_unchanged_capped_distance(self, coords, box, met): @@ -1065,6 +1329,15 @@ def test_input_unchanged_calc_bonds(self, coords, box, backend): res = distances.calc_bonds(crds[0], crds[1], box=box, backend=backend) assert_equal(crds, refs) + @pytest.mark.parametrize('box', boxes) + @pytest.mark.parametrize('backend', ['serial', 'openmp']) + def test_input_unchanged_calc_bonds_atomgroup(self, coords_atomgroups, + box, backend): + crds = coords_atomgroups[:2] + refs = [crd.positions.copy() for crd in crds] + res = distances.calc_bonds(crds[0], crds[1], box=box, backend=backend) + assert_equal([crd.positions for crd in crds], refs) + @pytest.mark.parametrize('box', boxes) @pytest.mark.parametrize('backend', ['serial', 'openmp']) def test_input_unchanged_calc_angles(self, coords, box, backend): @@ -1074,6 +1347,16 @@ def test_input_unchanged_calc_angles(self, coords, box, backend): backend=backend) assert_equal(crds, refs) + @pytest.mark.parametrize('box', boxes) + @pytest.mark.parametrize('backend', ['serial', 'openmp']) + def test_input_unchanged_calc_angles_atomgroup(self, coords_atomgroups, + box, backend): + crds = coords_atomgroups[:3] + refs = [crd.positions.copy() for crd in crds] + res = distances.calc_angles(crds[0], crds[1], crds[2], box=box, + backend=backend) + assert_equal([crd.positions for crd in crds], refs) + @pytest.mark.parametrize('box', boxes) @pytest.mark.parametrize('backend', ['serial', 'openmp']) def test_input_unchanged_calc_dihedrals(self, coords, box, backend): @@ -1083,6 +1366,16 @@ def test_input_unchanged_calc_dihedrals(self, coords, box, backend): box=box, backend=backend) assert_equal(crds, refs) + @pytest.mark.parametrize('box', boxes) + @pytest.mark.parametrize('backend', ['serial', 'openmp']) + def test_input_unchanged_calc_dihedrals_atomgroup(self, coords_atomgroups, + box, backend): + crds = coords_atomgroups + refs = [crd.positions.copy() for crd in crds] + res = distances.calc_dihedrals(crds[0], crds[1], crds[2], crds[3], + box=box, backend=backend) + assert_equal([crd.positions for crd in crds], refs) + @pytest.mark.parametrize('box', boxes[:2]) @pytest.mark.parametrize('backend', ['serial', 'openmp']) def test_input_unchanged_apply_PBC(self, coords, box, backend): @@ -1091,6 +1384,14 @@ def test_input_unchanged_apply_PBC(self, coords, box, backend): res = distances.apply_PBC(crd, box, backend=backend) assert_equal(crd, ref) + @pytest.mark.parametrize('box', boxes[:2]) + @pytest.mark.parametrize('backend', ['serial', 'openmp']) + def test_input_unchanged_apply_PBC_atomgroup(self, coords_atomgroups, box, + backend): + crd = coords_atomgroups[0] + ref = crd.positions.copy() + res = distances.apply_PBC(crd, box, backend=backend) + assert_equal(crd.positions, ref) class TestEmptyInputCoordinates(object): """Tests ensuring that the following functions in MDAnalysis.lib.distances diff --git a/testsuite/MDAnalysisTests/lib/test_util.py b/testsuite/MDAnalysisTests/lib/test_util.py index 72ad883240a..d95a6008745 100644 --- a/testsuite/MDAnalysisTests/lib/test_util.py +++ b/testsuite/MDAnalysisTests/lib/test_util.py @@ -43,9 +43,8 @@ check_coords, store_init_arguments,) from MDAnalysis.core.topologyattrs import Bonds from MDAnalysis.exceptions import NoDataError, DuplicateWarning - - -from MDAnalysisTests.datafiles import ( +from MDAnalysis.core.groups import AtomGroup +from MDAnalysisTests.datafiles import (PSF, DCD, Make_Whole, TPR, GRO, fullerene, two_water_gro, ) @@ -1738,6 +1737,59 @@ def func(a, b): with pytest.raises(ValueError): res = func(a_in, b_in2) + @pytest.fixture() + def atomgroup(self): + u = mda.Universe(PSF, DCD) + return u.atoms + + # check atomgroup handling with every option except allow_atomgroup + @pytest.mark.parametrize('enforce_copy', [True, False]) + @pytest.mark.parametrize('enforce_dtype', [True, False]) + @pytest.mark.parametrize('allow_single', [True, False]) + @pytest.mark.parametrize('convert_single', [True, False]) + @pytest.mark.parametrize('reduce_result_if_single', [True, False]) + @pytest.mark.parametrize('check_lengths_match', [True, False]) + def test_atomgroup(self, atomgroup, enforce_copy, enforce_dtype, + allow_single, convert_single, reduce_result_if_single, + check_lengths_match): + ag1 = atomgroup + ag2 = atomgroup + + @check_coords('ag1', 'ag2', enforce_copy=enforce_copy, + enforce_dtype=enforce_dtype, allow_single=allow_single, + convert_single=convert_single, + reduce_result_if_single=reduce_result_if_single, + check_lengths_match=check_lengths_match, + allow_atomgroup=True) + def func(ag1, ag2): + assert_allclose(ag1, ag2) + assert isinstance(ag1, np.ndarray) + assert isinstance(ag2, np.ndarray) + assert ag1.dtype == ag2.dtype == np.float32 + return ag1 + ag2 + + res = func(ag1, ag2) + + assert_allclose(res, atomgroup.positions*2) + + def test_atomgroup_not_allowed(self, atomgroup): + + @check_coords('ag1', allow_atomgroup=False) + def func(ag1): + return ag1 + + with pytest.raises(TypeError, match="allow_atomgroup is False"): + _ = func(atomgroup) + + def test_atomgroup_not_allowed_default(self, atomgroup): + + @check_coords('ag1') + def func_default(ag1): + return ag1 + + with pytest.raises(TypeError, match="allow_atomgroup is False"): + _ = func_default(atomgroup) + def test_enforce_copy(self): a_2d = np.ones((1, 3), dtype=np.float32) @@ -1836,6 +1888,21 @@ def func(a, b): assert res_a is a_2d assert res_b is b_2d + def test_atomgroup_mismatched_lengths(self): + u = mda.Universe(PSF, DCD) + ag1 = u.select_atoms("index 0 to 10") + ag2 = u.atoms + + @check_coords('ag1', 'ag2', check_lengths_match=True, + allow_atomgroup=True) + def func(ag1, ag2): + + return ag1, ag2 + + with pytest.raises(ValueError, match="must contain the same number of " + "coordinates"): + _, _ = func(ag1, ag2) + def test_invalid_input(self): a_inv_dtype = np.array([['hello', 'world', '!']]) From 50b8fdb3b026fd98fe4f588028e24dce0364b013 Mon Sep 17 00:00:00 2001 From: Raymond Zhao <7199958+rzhao271@users.noreply.github.com> Date: Fri, 22 Jul 2022 01:38:53 -0700 Subject: [PATCH 2/2] Modernize tests under converters and visualization (#3757) * Improve float testing for converters and visualization --- package/AUTHORS | 1 + package/CHANGELOG | 2 +- .../MDAnalysisTests/converters/test_openmm.py | 48 ++++++++++--------- .../MDAnalysisTests/converters/test_parmed.py | 32 +++++++------ .../MDAnalysisTests/converters/test_rdkit.py | 13 +++-- .../visualization/test_streamlines.py | 8 ++-- 6 files changed, 57 insertions(+), 47 deletions(-) diff --git a/package/AUTHORS b/package/AUTHORS index b80e0f8e67f..f0e6dddbf46 100644 --- a/package/AUTHORS +++ b/package/AUTHORS @@ -196,6 +196,7 @@ Chronological list of authors - Ricky Sexton - Rafael R. Pappalardo - Tengyu Xie + - Raymond Zhao External code diff --git a/package/CHANGELOG b/package/CHANGELOG index 3770922ee6c..e43b8c85766 100644 --- a/package/CHANGELOG +++ b/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 diff --git a/testsuite/MDAnalysisTests/converters/test_openmm.py b/testsuite/MDAnalysisTests/converters/test_openmm.py index 8155a4ea8e5..7f695b37992 100644 --- a/testsuite/MDAnalysisTests/converters/test_openmm.py +++ b/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) diff --git a/testsuite/MDAnalysisTests/converters/test_parmed.py b/testsuite/MDAnalysisTests/converters/test_parmed.py index ac7a889e57d..38ef20caaab 100644 --- a/testsuite/MDAnalysisTests/converters/test_parmed.py +++ b/testsuite/MDAnalysisTests/converters/test_parmed.py @@ -23,9 +23,7 @@ import pytest import MDAnalysis as mda -from numpy.testing import (assert_equal, - assert_almost_equal, - assert_almost_equal) +from numpy.testing import (assert_allclose, assert_equal) from MDAnalysisTests.coordinates.base import _SingleFrameReader from MDAnalysisTests.coordinates.reference import RefAdKSmall @@ -51,38 +49,40 @@ class TestParmEdReaderGRO: universe = mda.Universe(pmd.load_file(GRO)) ref = mda.Universe(GRO) - prec = 3 def test_dimensions(self): - assert_almost_equal( + assert_allclose( self.universe.trajectory.ts.dimensions, self.ref.trajectory.ts.dimensions, - self.prec, - "ParmEdReader failed to get unitcell dimensions from ParmEd") + rtol=0, + atol=1e-3, + err_msg=("ParmEdReader failed to get unitcell dimensions " + "from ParmEd")) 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 BaseTestParmEdReader(_SingleFrameReader): def setUp(self): self.universe = mda.Universe(pmd.load_file(self.ref_filename)) self.ref = mda.Universe(self.ref_filename) - self.prec = 3 def test_dimensions(self): - assert_almost_equal( + assert_allclose( self.universe.trajectory.ts.dimensions, self.ref.trajectory.ts.dimensions, - self.prec, - "ParmEdReader failed to get unitcell dimensions from ParmEd") + rtol=0, + atol=1e-3, + err_msg=("ParmEdReader failed to get unitcell dimensions " + "from ParmEd")) 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 TestParmEdReaderPDB(BaseTestParmEdReader): @@ -187,7 +187,9 @@ def test_equivalent_atoms(self, ref, output): for attr in self.almost_equal_atom_attrs: ra = getattr(r, attr) oa = getattr(o, attr) - assert_almost_equal(ra, oa, decimal=2, err_msg='atom {} not almost equal for atoms {} and {}'.format(attr, r, o)) + assert_allclose(ra, oa, rtol=0, atol=1e-2, + err_msg=(f'atom {attr} not almost equal for atoms ' + f'{r} and {o}')) @pytest.mark.parametrize('attr', ('bonds', 'angles', 'impropers', 'cmaps')) @@ -292,7 +294,7 @@ class TestParmEdConverterPDB(BaseTestParmEdConverter): almost_equal_atom_attrs = ('charge', 'occupancy') def test_equivalent_coordinates(self, ref, output): - assert_almost_equal(ref.coordinates, output.coordinates, decimal=3) + assert_allclose(ref.coordinates, output.coordinates, rtol=0, atol=1e-3) def test_incorrect_object_passed_typeerror(): diff --git a/testsuite/MDAnalysisTests/converters/test_rdkit.py b/testsuite/MDAnalysisTests/converters/test_rdkit.py index 24a6066b4eb..5ff1ab903a9 100644 --- a/testsuite/MDAnalysisTests/converters/test_rdkit.py +++ b/testsuite/MDAnalysisTests/converters/test_rdkit.py @@ -27,8 +27,7 @@ import MDAnalysis as mda from MDAnalysis.topology.guessers import guess_atom_element import numpy as np -from numpy.testing import (assert_equal, - assert_almost_equal) +from numpy.testing import (assert_allclose, assert_equal) from MDAnalysisTests.datafiles import mol2_molecule, PDB_full, GRO, PDB_helix from MDAnalysisTests.util import import_not_available @@ -216,7 +215,11 @@ def test_identical_topology(self, rdmol): assert_equal(u.atoms.bonds, u2.atoms.bonds) assert_equal(u.atoms.elements, u2.atoms.elements) assert_equal(u.atoms.names, u2.atoms.names) - assert_almost_equal(u.atoms.positions, u2.atoms.positions, decimal=7) + assert_allclose( + u.atoms.positions, + u2.atoms.positions, + rtol=0, + atol=1e-7) def test_raise_requires_elements(self): u = mda.Universe(mol2_molecule) @@ -304,7 +307,7 @@ def test_index_property(self, pdb, sel_str): def test_assign_coordinates(self, pdb): mol = pdb.atoms.convert_to.rdkit(NoImplicit=False) positions = mol.GetConformer().GetPositions() - assert_almost_equal(positions, pdb.atoms.positions) + assert_allclose(positions, pdb.atoms.positions) def test_assign_stereochemistry(self, mol2): umol = mol2.atoms.convert_to("RDKIT") @@ -317,7 +320,7 @@ def test_trajectory_coords(self): for ts in u.trajectory: mol = u.atoms.convert_to("RDKIT") positions = mol.GetConformer().GetPositions() - assert_almost_equal(positions, ts.positions) + assert_allclose(positions, ts.positions) def test_nan_coords(self): u = mda.Universe.from_smiles("CCO") diff --git a/testsuite/MDAnalysisTests/visualization/test_streamlines.py b/testsuite/MDAnalysisTests/visualization/test_streamlines.py index 3c8a1b23585..8569e1498cb 100644 --- a/testsuite/MDAnalysisTests/visualization/test_streamlines.py +++ b/testsuite/MDAnalysisTests/visualization/test_streamlines.py @@ -21,13 +21,13 @@ # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787 # import numpy as np -from numpy.testing import assert_almost_equal import MDAnalysis from MDAnalysis.visualization import (streamlines, streamlines_3D) from MDAnalysis.coordinates.XTC import XTCWriter from MDAnalysisTests.datafiles import Martini_membrane_gro import pytest +from pytest import approx import matplotlib.pyplot as plt import os @@ -105,6 +105,6 @@ def test_streamplot_3D(membrane_xtc, univ, tmpdir): assert dx.shape == (5, 5, 2) assert dy.shape == (5, 5, 2) assert dz.shape == (5, 5, 2) - assert_almost_equal(dx[4, 4, 0], 0.700004, decimal=5) - assert_almost_equal(dy[0, 0, 0], 0.460000, decimal=5) - assert_almost_equal(dz[2, 2, 0], 0.240005, decimal=5) + assert dx[4, 4, 0] == approx(0.700004, abs=1e-5) + assert dy[0, 0, 0] == approx(0.460000, abs=1e-5) + assert dz[2, 2, 0] == approx(0.240005, abs=1e-5)