MemoryReader facelift

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.

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

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):