Skip to content

Commit

Permalink
Added optional forces and velocities to MemoryReader (#2080)
Browse files Browse the repository at this point in the history
Universe.empty now creates a MemoryReader trajectory

Removed DummyReader (it was buggy)

added varying box size to MemoryReader

Universe.transfer_to_memory now also transfer velocities and forces

MemoryReader now always has array of dimensions

fixes Issue #2081 #2076 #2077 #1041
  • Loading branch information
richardjgowers authored Oct 9, 2018
1 parent d37f87b commit 97db6d1
Show file tree
Hide file tree
Showing 9 changed files with 356 additions and 106 deletions.
7 changes: 7 additions & 0 deletions package/CHANGELOG
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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.
Expand Down
1 change: 0 additions & 1 deletion package/MDAnalysis/coordinates/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -748,4 +748,3 @@ class can choose an appropriate reader automatically.
from . import MMTF
from . import GSD
from . import null
from . import dummy
61 changes: 0 additions & 61 deletions package/MDAnalysis/coordinates/dummy.py

This file was deleted.

103 changes: 93 additions & 10 deletions package/MDAnalysis/coordinates/memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand All @@ -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
----------
Expand All @@ -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
------
Expand All @@ -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__()
Expand All @@ -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 "
Expand All @@ -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):
Expand All @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down
Loading

0 comments on commit 97db6d1

Please sign in to comment.