diff --git a/package/CHANGELOG b/package/CHANGELOG index acb0c32cbb5..70691ee318b 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -85,6 +85,9 @@ Enhancements * Performance improvements to make_whole (PR #1965) * Performance improvements to fragment finding (PR #2028) * Added user-defined boxes in density code (PR #2005) + * MemoryReader now can accept velocities and forces (PR #2080) + * Universe.transfer_to_memory now copies dimensions, velocities and forces + (where possible) (Issue #1041 PR #2080) Fixes * Rewind in the SingleFrameReader now reads the frame from the file (Issue #1929) @@ -111,6 +114,9 @@ Fixes center of the reference group (Issue #1795) * PCA analysis now uses start frame as reference frame rather than 0th frame (PR #2055) + * Fixed trajectory iteration from a MDAnalysis.Universe.empty (#2076) + * Fixed copies MemoryReader not linking to the underlying coordinate array + on initial Timestep (Issue #2081 PR #2080) Changes * TopologyAttrs are now statically typed (Issue #1876) @@ -136,6 +142,7 @@ Changes otherwise a ValueError is raised (PR #2055) * The quiet keyword has been removed and replaced with verbose (Issue #1975 PR #2055) + * MDAnalysis.Universe.empty now creates a MemoryReader trajectory (#2076 #2077) Deprecations * start/stop/step are deprecated in the initialization of Analysis classes. diff --git a/package/MDAnalysis/coordinates/__init__.py b/package/MDAnalysis/coordinates/__init__.py index 2706f492151..8b58f434bd3 100644 --- a/package/MDAnalysis/coordinates/__init__.py +++ b/package/MDAnalysis/coordinates/__init__.py @@ -748,4 +748,3 @@ class can choose an appropriate reader automatically. from . import MMTF from . import GSD from . import null -from . import dummy diff --git a/package/MDAnalysis/coordinates/dummy.py b/package/MDAnalysis/coordinates/dummy.py deleted file mode 100644 index d4ee963be3e..00000000000 --- a/package/MDAnalysis/coordinates/dummy.py +++ /dev/null @@ -1,61 +0,0 @@ -# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding: utf-8 -*- -# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 -# -# MDAnalysis --- https://www.mdanalysis.org -# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors -# (see the file AUTHORS for the full list of names) -# -# Released under the GNU Public Licence, v2 or any higher version -# -# Please cite your use of MDAnalysis in published work: -# -# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler, -# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein. -# MDAnalysis: A Python package for the rapid analysis of molecular dynamics -# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th -# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy. -# -# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein. -# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations. -# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787 -# - -""" -Dummy coordinate reader -======================= - - -Classes -------- - -.. autoclass:: DummyReader - :members: - - -""" -from __future__ import absolute_import - - -import numpy as np - -from .base import SingleFrameReaderBase - - -class DummyReader(SingleFrameReaderBase): - """Basic Reader which does not read from any file - - .. versionadded:: 0.17.0 - """ - format = 'dummy' - - def __init__(self, n_atoms=None, velocities=False, forces=False): - self.n_atoms = n_atoms - self.filename = 'DummyReader' - self.n_frames = 1 - self._read_first_frame(velocities, forces) - self._transformations = [] - - def _read_first_frame(self, velocities=False, forces=False): - ts = self.ts = self._Timestep(self.n_atoms, positions=True, - velocities=velocities, forces=forces) - return ts diff --git a/package/MDAnalysis/coordinates/memory.py b/package/MDAnalysis/coordinates/memory.py index 617d4095f40..c03421a43cb 100644 --- a/package/MDAnalysis/coordinates/memory.py +++ b/package/MDAnalysis/coordinates/memory.py @@ -226,6 +226,17 @@ def _replace_positions_array(self, new): self.has_positions = True self._pos = new + def _replace_velocities_array(self, new): + self.has_velocities = True + self._velocities = new + + def _replace_forces_array(self, new): + self.has_forces = True + self._forces = new + + def _replace_dimensions(self, new): + self._unitcell = new + class MemoryReader(base.ProtoReader): """ @@ -243,7 +254,9 @@ class MemoryReader(base.ProtoReader): _Timestep = Timestep def __init__(self, coordinate_array, order='fac', - dimensions=None, dt=1, filename=None, **kwargs): + dimensions=None, dt=1, filename=None, + velocities=None, forces=None, + **kwargs): """ Parameters ---------- @@ -259,13 +272,20 @@ def __init__(self, coordinate_array, order='fac', dimensions: [A, B, C, alpha, beta, gamma] (optional) unitcell dimensions (*A*, *B*, *C*, *alpha*, *beta*, *gamma*) lengths *A*, *B*, *C* are in the MDAnalysis length unit (Å), and - angles are in degrees. + angles are in degrees. An array of dimensions can be given, + which must then be shape (nframes, 6) dt: float (optional) The time difference between frames (ps). If :attr:`time` is set, then `dt` will be ignored. filename: string (optional) The name of the file from which this instance is created. Set to ``None`` when created from an array + velocities : numpy.ndarray (optional) + Atom velocities. Must match shape of coordinate_array. Will share order + with coordinates. + forces : numpy.ndarray (optional) + Atom forces. Must match shape of coordinate_array Will share order + with coordinates Raises ------ @@ -281,7 +301,7 @@ def __init__(self, coordinate_array, order='fac', .. versionchanged:: 0.18.1 The input to the MemoryReader now must be a np.ndarray - + Added optional velocities and forces """ super(MemoryReader, self).__init__() @@ -292,8 +312,8 @@ def __init__(self, coordinate_array, order='fac', # passed is of shape (N, 3) and if it is, the coordiante array is # reshaped to (1, N, 3) try: - if len(coordinate_array.shape) == 2 and coordinate_array.shape[1] == 3: - coordinate_array = coordinate_array[np.newaxis, :, :] + if coordinate_array.ndim == 2 and coordinate_array.shape[1] == 3: + coordinate_array = coordinate_array[np.newaxis, :, :] except AttributeError as e: raise TypeError("The input has to be a numpy.ndarray that " "corresponds to the layout specified by the " @@ -305,6 +325,41 @@ def __init__(self, coordinate_array, order='fac', self.n_atoms = \ self.coordinate_array.shape[self.stored_order.find('a')] + if velocities is not None: + try: + velocities = np.asarray(velocities, dtype=np.float32) + except ValueError: + raise TypeError("'velocities' must be array-like got {}" + "".format(type(velocities))) + # if single frame, make into array of 1 frame + if velocities.ndim == 2: + velocities = velocities[np.newaxis, :, :] + if not velocities.shape == self.coordinate_array.shape: + raise ValueError('Velocities has wrong shape {} ' + 'to match coordinates {}' + ''.format(velocities.shape, + self.coordinate_array.shape)) + self.velocity_array = velocities.astype(np.float32, copy=False) + else: + self.velocity_array = None + + if forces is not None: + try: + forces = np.asarray(forces, dtype=np.float32) + except ValueError: + raise TypeError("'forces' must be array like got {}" + "".format(type(forces))) + if forces.ndim == 2: + forces = forces[np.newaxis, :, :] + if not forces.shape == self.coordinate_array.shape: + raise ValueError('Forces has wrong shape {} ' + 'to match coordinates {}' + ''.format(forces.shape, + self.coordinate_array.shape)) + self.force_array = forces.astype(np.float32, copy=False) + else: + self.force_array = None + provided_n_atoms = kwargs.pop("n_atoms", None) if (provided_n_atoms is not None and provided_n_atoms != self.n_atoms): @@ -315,8 +370,23 @@ def __init__(self, coordinate_array, order='fac', self.ts = self._Timestep(self.n_atoms, **kwargs) self.ts.dt = dt - if dimensions is not None: - self.ts.dimensions = dimensions + if dimensions is None: + dimensions = np.zeros((self.n_frames, 6), dtype=np.float64) + else: + try: + dimensions = np.asarray(dimensions, dtype=np.float64) + except ValueError: + raise TypeError("'dimensions' must be array-like got {}" + "".format(type(dimensions))) + if dimensions.shape == (6,): + # single box, tile this to trajectory length + # allows modifying the box of some frames + dimensions = np.tile(dimensions, (self.n_frames, 1)) + elif dimensions.shape != (self.n_frames, 6): + raise ValueError("Provided dimensions array has shape {}. " + "This must be a array of shape (6,) or " + "(n_frames, 6)".format(dimensions.shape)) + self.dimensions_array = dimensions self.ts.frame = -1 self.ts.time = -1 self._read_next_timestep() @@ -346,15 +416,23 @@ def parse_n_atoms(filename, order='fac', **kwargs): def copy(self): """Return a copy of this Memory Reader""" + vels = (self.velocity_array.copy() + if self.velocity_array is not None else None) + fors = (self.force_array.copy() + if self.force_array is not None else None) + dims = self.dimensions_array.copy() + new = self.__class__( self.coordinate_array.copy(), order=self.stored_order, - dimensions=self.ts.dimensions, + dimensions=dims, + velocities=vels, + forces=fors, dt=self.ts.dt, filename=self.filename, ) new[self.ts.frame] - new.ts = self.ts.copy() + for auxname, auxread in self._auxs.items(): new.add_auxiliary(auxname, auxread.copy()) # since transformations are already applied to the whole trajectory @@ -483,8 +561,13 @@ def _read_next_timestep(self, ts=None): [self.ts.frame] + [slice(None)]*(2-f_index)) ts._replace_positions_array(self.coordinate_array[basic_slice]) + ts._replace_dimensions(self.dimensions_array[self.ts.frame]) + if self.velocity_array is not None: + ts._replace_velocities_array(self.velocity_array[basic_slice]) + if self.force_array is not None: + ts._replace_forces_array(self.force_array[basic_slice]) - ts.time = self.ts.frame*self.dt + ts.time = self.ts.frame * self.dt return ts def _read_frame(self, i): diff --git a/package/MDAnalysis/core/universe.py b/package/MDAnalysis/core/universe.py index 139b0a02fd0..522f2e1c971 100644 --- a/package/MDAnalysis/core/universe.py +++ b/package/MDAnalysis/core/universe.py @@ -417,6 +417,10 @@ def empty(cls, n_atoms, n_residues=None, n_segments=None, Useful for building a Universe without requiring existing files, for example for system building. + If `trajectory` is set to True, a + :class:`MDAnalysis.coordinates.memory.MemoryReader` will be + attached to the Universe. + Parameters ---------- n_atoms : int @@ -425,17 +429,19 @@ def empty(cls, n_atoms, n_residues=None, n_segments=None, number of Residues in the Universe, defaults to 1 n_segments : int, optional number of Segments in the Universe, defaults to 1 - atom_resindex : numpy.array, optional - mapping of atoms to residues - residue_segindex : numpy.array, optional + atom_resindex : array like, optional + mapping of atoms to residues, e.g. with 6 atoms, + `atom_resindex=[0, 0, 1, 1, 2, 2]` would put 2 atoms + into each of 3 residues. + residue_segindex : array like, optional mapping of residues to segments trajectory : bool, optional - if True, attaches a dummy reader to the Universe, therefore + if True, attaches a :class:`MDAnalysis.coordinates.memory.MemoryReader` allowing coordinates to be set and written. Default is False velocities : bool, optional - include velocities in the dummy Reader + include velocities in the :class:`MDAnalysis.coordinates.memory.MemoryReader` forces : bool, optional - include forces in the dummy Reader + include forces in the :class:`MDAnalysis.coordinates.memory.MemoryReader` Returns ------- @@ -453,6 +459,8 @@ def empty(cls, n_atoms, n_residues=None, n_segments=None, >>> u.add_TopologyAttr('masses') .. versionadded:: 0.17.0 + .. versionchanged:: 0.19.0 + The attached Reader when trajectory=True is now a MemoryReader """ if n_residues is None: n_residues = 1 @@ -477,9 +485,15 @@ def empty(cls, n_atoms, n_residues=None, n_segments=None, u = cls(top) if trajectory: - u.trajectory = get_reader_for('', format='dummy')( - n_atoms=n_atoms, - velocities=velocities, forces=forces) + coords = np.zeros((1, n_atoms, 3), dtype=np.float32) + dims = np.zeros(6, dtype=np.float64) + vels = np.zeros_like(coords) if velocities else None + forces = np.zeros_like(coords) if forces else None + + # grab and attach a MemoryReader + u.trajectory = get_reader_for(coords)( + coords, order='fac', n_atoms=n_atoms, + dimensions=dims, velocities=vels, forces=forces) return u @@ -627,26 +641,33 @@ def transfer_to_memory(self, start=None, stop=None, step=None, from ..coordinates.memory import MemoryReader if not isinstance(self.trajectory, MemoryReader): - # Try to extract coordinates using Timeseries object - # This is significantly faster, but only implemented for certain - # trajectory file formats - try: - coordinates = self.trajectory.timeseries( - self.atoms, start=start, stop=stop, step=step, order='fac') - # if the Timeseries extraction fails, - # fall back to a slower approach - except AttributeError: - n_frames = len(range( - *self.trajectory.check_slice_indices(start, stop, step) - )) - pm_format = '{step}/{numsteps} frames copied to memory (frame {frame})' - pm = ProgressMeter(n_frames, interval=1, - verbose=verbose, format=pm_format) - coordinates = [] # TODO: use pre-allocated array - for i, ts in enumerate(self.trajectory[start:stop:step]): - coordinates.append(np.copy(ts.positions)) - pm.echo(i, frame=ts.frame) - coordinates = np.array(coordinates) + n_frames = len(range( + *self.trajectory.check_slice_indices(start, stop, step) + )) + n_atoms = len(self.atoms) + pm_format = '{step}/{numsteps} frames copied to memory (frame {frame})' + pm = ProgressMeter(n_frames, interval=1, + verbose=verbose, format=pm_format) + coordinates = np.zeros((n_frames, n_atoms, 3), dtype=np.float32) + ts = self.trajectory.ts + has_vels = ts.has_velocities + has_fors = ts.has_forces + has_dims = ts.dimensions is not None + + velocities = np.zeros_like(coordinates) if has_vels else None + forces = np.zeros_like(coordinates) if has_fors else None + dimensions = (np.zeros((n_frames, 6), dtype=np.float64) + if has_dims else None) + + for i, ts in enumerate(self.trajectory[start:stop:step]): + np.copyto(coordinates[i], ts.positions) + if has_vels: + np.copyto(velocities[i], ts.velocities) + if has_fors: + np.copyto(forces[i], ts.forces) + if has_dims: + np.copyto(dimensions[i], ts.dimensions) + pm.echo(i, frame=ts.frame) # Overwrite trajectory in universe with an MemoryReader # object, to provide fast access and allow coordinates @@ -655,9 +676,12 @@ def transfer_to_memory(self, start=None, stop=None, step=None, step = 1 self.trajectory = MemoryReader( coordinates, - dimensions=self.trajectory.ts.dimensions, + dimensions=dimensions, dt=self.trajectory.ts.dt * step, - filename=self.trajectory.filename) + filename=self.trajectory.filename, + velocities=velocities, + forces=forces, + ) # python 2 doesn't allow an efficient splitting of kwargs in function # argument signatures. diff --git a/package/doc/sphinx/source/documentation_pages/coordinates/dummy.rst b/package/doc/sphinx/source/documentation_pages/coordinates/dummy.rst deleted file mode 100644 index 7f0e874a9cf..00000000000 --- a/package/doc/sphinx/source/documentation_pages/coordinates/dummy.rst +++ /dev/null @@ -1,2 +0,0 @@ -.. automodule:: MDAnalysis.coordinates.dummy - :members: diff --git a/package/doc/sphinx/source/documentation_pages/coordinates_modules.rst b/package/doc/sphinx/source/documentation_pages/coordinates_modules.rst index 359719e3330..e81313985e3 100644 --- a/package/doc/sphinx/source/documentation_pages/coordinates_modules.rst +++ b/package/doc/sphinx/source/documentation_pages/coordinates_modules.rst @@ -41,7 +41,6 @@ provide the format in the keyword argument *format* to coordinates/TRZ coordinates/memory coordinates/null - coordinates/dummy .. rubric:: Coordinate core modules diff --git a/testsuite/MDAnalysisTests/coordinates/test_memory.py b/testsuite/MDAnalysisTests/coordinates/test_memory.py index 8188c688101..0dfc0e643ef 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_memory.py +++ b/testsuite/MDAnalysisTests/coordinates/test_memory.py @@ -200,3 +200,177 @@ def test_position_assignation(self, reader): reader.ts.positions = new_positions reader[0] assert_almost_equal(reader.ts.positions, new_positions) + + +class TestMemoryReaderVelsForces(object): + @staticmethod + @pytest.fixture(params=['2d', '3d']) + def ref_pos(request): + if request.param == '2d': + return np.arange(30).reshape(10, 3) + elif request.param == '3d': + return np.arange(30).reshape(1, 10, 3) + + @staticmethod + @pytest.fixture(params=['2d', '3d']) + def ref_vels(request): + if request.param == '2d': + return np.arange(30).reshape(10, 3) + 100 + elif request.param == '3d': + return np.arange(30).reshape(1, 10, 3) + 100 + + @staticmethod + @pytest.fixture(params=['2d', '3d']) + def ref_forces(request): + if request.param == '2d': + return np.arange(30).reshape(10, 3) + 1000 + elif request.param == '3d': + return np.arange(30).reshape(1, 10, 3) + 1000 + + @staticmethod + def assert_equal_dims(arr1, arr2): + if arr2.ndim == 3: + assert_equal(arr1, arr2[0]) + elif arr2.ndim == 2: + assert_equal(arr1, arr2) + + def test_velocities(self, ref_pos, ref_vels): + mr = MemoryReader(ref_pos, + velocities=ref_vels) + + assert mr.ts.has_velocities + self.assert_equal_dims(mr.ts.velocities, ref_vels) + assert not mr.ts.has_forces + + def test_forces(self, ref_pos, ref_forces): + mr = MemoryReader(ref_pos, + forces=ref_forces) + + assert not mr.ts.has_velocities + assert mr.ts.has_forces + self.assert_equal_dims(mr.ts.forces, ref_forces) + + def test_both(self, ref_pos, ref_vels, ref_forces): + mr = MemoryReader(ref_pos, + velocities=ref_vels, + forces=ref_forces) + assert mr.ts.has_velocities + self.assert_equal_dims(mr.ts.velocities, ref_vels) + assert mr.ts.has_forces + self.assert_equal_dims(mr.ts.forces, ref_forces) + + @pytest.mark.parametrize('param', ['velocities', 'forces']) + def test_wrongshape(self, ref_pos, param): + with pytest.raises(ValueError): + mr = MemoryReader(ref_pos, **{param: np.zeros((3, 2, 1))}) + + +class TestDimensions(object): + @staticmethod + @pytest.fixture + def ref_pos(): + return np.arange(270).reshape(3, 30, 3) + + @staticmethod + @pytest.fixture + def ref_box(): + return np.arange(18).reshape(3, 6) + + def test_single_box(self, ref_pos): + box = np.array([3, 4, 5, 90, 90, 90]) + + mr = MemoryReader(ref_pos, dimensions=box) + + for ts in mr: + assert_equal(ts.dimensions, box) + + def test_varying_box(self, ref_pos, ref_box): + mr = MemoryReader(ref_pos, dimensions=ref_box) + + for i, ts in enumerate(mr): + assert_equal(ts.dimensions, ref_box[i]) + + def test_wrong_length(self, ref_pos): + bad_box = np.arange(12).reshape(2, 6) + + with pytest.raises(ValueError): + mr = MemoryReader(ref_pos, dimensions=bad_box) + + def test_wrong_shape(self, ref_pos): + bad_box = np.arange(15).reshape(3, 5) + + with pytest.raises(ValueError): + mr = MemoryReader(ref_pos, dimensions=bad_box) + + +class TestMemoryReaderModifications(object): + # check that modifying MR things behaves as expected + # in general, modifying the Timestep should be *permanent* + # this is unlike other Readers! + n_atoms = 10 + n_frames = 4 + + @pytest.fixture() + def mr_reader(self): + pos = np.arange(self.n_frames * self.n_atoms * 3).reshape( + self.n_frames, self.n_atoms, 3) + vel = np.arange(self.n_frames * self.n_atoms * 3).reshape( + self.n_frames, self.n_atoms, 3) + 200 + frc = np.arange(self.n_frames * self.n_atoms * 3).reshape( + self.n_frames, self.n_atoms, 3) + 400 + box = np.arange(self.n_frames * 6).reshape(self.n_frames, 6) + 600 + + return MemoryReader(pos, + velocities=vel, + forces=frc, + dimensions=box) + + @pytest.fixture() + def mr_universe(self, mr_reader): + u = mda.Universe.empty(self.n_atoms) + u.trajectory = mr_reader + + return u + + @pytest.mark.parametrize('attr', ['positions', 'velocities', 'forces', 'dimensions']) + def test_copying(self, mr_reader, attr): + mr2 = mr_reader.copy() + # update the attribute + ts = mr2.ts + setattr(ts, attr, 7) + # check the change worked + assert_almost_equal(getattr(ts, attr), 7) + assert ts.positions.shape == (self.n_atoms, 3) + assert ts.velocities.shape == (self.n_atoms, 3) + assert ts.forces.shape == (self.n_atoms, 3) + assert ts.dimensions.shape == (6,) + # move the Reader around, forcing updates of ts + ts = mr2[2] + ts = mr2[0] + # check our old change is still there + assert_almost_equal(getattr(ts, attr), 7) + + @pytest.mark.parametrize('attr', ['positions', 'velocities', 'forces', 'dimensions']) + def test_attr_set(self, mr_universe, attr): + # same as above, but via a Universe/AtomGroup + u = mr_universe + ts = u.trajectory[0] + + setattr(ts, attr, 7) + + assert_almost_equal(getattr(ts, attr), 7) + + ts = u.trajectory[2] + ts = u.trajectory[0] + + assert_almost_equal(getattr(ts, attr), 7) + assert u.atoms.positions.shape == (self.n_atoms, 3) + assert u.atoms.velocities.shape == (self.n_atoms, 3) + assert u.atoms.forces.shape == (self.n_atoms, 3) + assert u.atoms.dimensions.shape == (6,) + + @pytest.mark.parametrize('attr', ['velocities', 'forces', 'dimensions']) + def test_non_numpy_arr(self, attr): + with pytest.raises(TypeError): + mr = MemoryReader(np.zeros((10, 30, 3)), + **{attr: 'not an array'}) diff --git a/testsuite/MDAnalysisTests/core/test_universe.py b/testsuite/MDAnalysisTests/core/test_universe.py index ceccbd752d2..fe3672065e3 100644 --- a/testsuite/MDAnalysisTests/core/test_universe.py +++ b/testsuite/MDAnalysisTests/core/test_universe.py @@ -690,3 +690,30 @@ def test_no_segindex_warning(self): with pytest.warns(UserWarning): u = mda.Universe.empty(n_atoms=10, n_residues=2, n_segments=1, atom_resindex=res) + + def test_trajectory(self): + u = mda.Universe.empty(10, trajectory=True) + + assert len(u.atoms) == 10 + assert u.atoms.positions.shape == (10, 3) + + def test_trajectory_iteration(self): + u = mda.Universe.empty(10, trajectory=True) + + assert len(u.trajectory) == 1 + timesteps =[] + for ts in u.trajectory: + timesteps.append(ts.frame) + assert len(timesteps) == 1 + + def test_velocities(self): + u = mda.Universe.empty(10, trajectory=True, velocities=True) + + assert u.atoms.positions.shape == (10, 3) + assert u.atoms.velocities.shape == (10, 3) + + def test_forces(self): + u = mda.Universe.empty(10, trajectory=True, forces=True) + + assert u.atoms.positions.shape == (10, 3) + assert u.atoms.forces.shape == (10, 3)