barebones working prototype of trajectory transformations

package/CHANGELOG

@@ -14,18 +14,21 @@ The rules for this file:
 
 ------------------------------------------------------------------------------
 ??/??/18 tylerjereddy, richardjgowers, palnabarun, orbeckst, kain88-de, zemanj,
-         VOD555
+         VOD555, davidercruz
 
   * 0.18.1
 
 Enhancements
+  * Added a on-the-fly trajectory transformations API and a coordinate translation
+  	function (Issue #786)
   * Added various duecredit stubs
   * Import time reduced by a factor two (PR #1881)
   * added compound kwarg to center, centroid, center_of_geometry, center_of_mass (PR #1903)
   * Added rdf.InterRDF_s to calculate site-specific pair distribution
     functions (PR #1815)
 
 Fixes
+  * rewind in the SingleFrameReader now reads the frame from the file (Issue #1929)
   * Fixed order of indices in Angle/Dihedral/Improper repr
   * coordinates.memory.MemoryReader now takes np.ndarray only (Issue #1685)
   * Silenced warning when duecredit is not installed (Issue #1872)

package/MDAnalysis/coordinates/base.py

@@ -278,6 +278,7 @@ def __init__(self, n_atoms, **kwargs):
 
         # set up aux namespace for adding auxiliary data
         self.aux = Namespace()
+
 
     @classmethod
     def from_timestep(cls, other, **kwargs):
@@ -1170,6 +1171,7 @@ def __init__(self):
         # initialise list to store added auxiliary readers in
         # subclasses should now call super
         self._auxs = {}
+        self._transformations=[]
 
     def __len__(self):
         return self.n_frames
@@ -1201,6 +1203,9 @@ def next(self):
         else:
             for auxname in self.aux_list:
                 ts = self._auxs[auxname].update_ts(ts)
+
+            ts = self._apply_transformations(ts)
+
         return ts
 
     def __next__(self):
@@ -1356,10 +1361,13 @@ def _read_frame(self, frame):
         # return ts
 
     def _read_frame_with_aux(self, frame):
-        """Move to *frame*, updating ts with trajectory and auxiliary data."""
+        """Move to *frame*, updating ts with trajectory, transformations and auxiliary data."""
         ts = self._read_frame(frame)
         for aux in self.aux_list:
             ts = self._auxs[aux].update_ts(ts)
+
+        ts = self._apply_transformations(ts)
+
         return ts
 
     def _sliced_iter(self, start, stop, step):
@@ -1741,7 +1749,67 @@ def get_aux_descriptions(self, auxnames=None):
             auxnames = self.aux_list
         descriptions = [self._auxs[aux].get_description() for aux in auxnames]
         return descriptions
-
+
+    @property
+    def transformations(self):
+        """ Returns the list of transformations"""
+        return self._transformations
+
+    @transformations.setter
+    def transformations(self, transformations):
+        if not self._transformations:
+            self._transformations = transformations
+        else:
+            raise ValueError("Transformations are already set")
+
+    def add_transformations(self, *transformations):
+        """ Add all transformations to be applied to the trajectory.
+
+        This function take as list of transformations as an argument. These
+        transformations are functions that will be called by the Reader and given
+        a :class:`Timestep` object as argument, which will be transformed and returned
+        to the Reader.
+        The transformations can be part of the :mod:`~MDAnalysis.transformations` 
+        module, or created by the user, and are stored as a list `transformations`. 
+        This list can only be modified once, and further calls of this function will
+        raise an exception.
+
+        .. code-block:: python
+
+          u = MDAnalysis.Universe(topology, coordinates)
+          workflow = [some_transform, another_transform, this_transform]
+          u.trajectory.add_transformations(*workflow)
+
+        Parameters
+        ----------
+        transform_list : list
+            list of all the transformations that will be applied to the coordinates
+
+        See Also
+        --------
+        :mod:`MDAnalysis.transformations`
+        """
+
+        try:
+            self.transformations = transformations
+        except ValueError:
+            raise ValueError("Can't add transformations again. Please create new Universe object")  
+        else:
+            self.ts = self._apply_transformations(self.ts)
+
+
+        # call reader here to apply the newly added transformation on the
+        # current loaded frame?
+
+    def _apply_transformations(self, ts):
+        """Applies all the transformations given by the user """
+
+        for transform in self.transformations:
+            ts = transform(ts)
+
+        return ts
+
+
 
 class ReaderBase(ProtoReader):
     """Base class for trajectory readers that extends :class:`ProtoReader` with a
@@ -1802,6 +1870,8 @@ def copy(self):
         """
         new = self.__class__(self.filename,
                              n_atoms=self.n_atoms)
+        if self.transformations:
+            new.add_transformations(*self.transformations)
         # seek the new reader to the same frame we started with
         new[self.ts.frame]
         # then copy over the current Timestep in case it has
@@ -1976,14 +2046,21 @@ def copy(self):
         new.ts = self.ts.copy()
         for auxname, auxread in self._auxs.items():
             new.add_auxiliary(auxname, auxread.copy())
+        # since the transformations have already been applied to the frame
+        # simply copy the property
+        new.transformations = self.transformations
+
         return new
 
     def _read_first_frame(self):  # pragma: no cover
         # Override this in subclasses to create and fill a Timestep
         pass
 
     def rewind(self):
-        pass
+        self._read_first_frame()
+        for auxname, auxread in self._auxs.items():
+            self.ts = auxread.update_ts(self.ts)
+        super(SingleFrameReaderBase, self)._apply_transformations(self.ts)
 
     def _reopen(self):
         pass
@@ -2006,3 +2083,47 @@ def close(self):
         # self.filename. Explicitly setting it to the null action in case
         # the IOBase.close method is ever changed from that.
         pass
+
+    def add_transformations(self, *transformations):
+        """ Add all transformations to be applied to the trajectory.
+
+        This function take as list of transformations as an argument. These
+        transformations are functions that will be called by the Reader and given
+        a :class:`Timestep` object as argument, which will be transformed and returned
+        to the Reader.
+        The transformations can be part of the :mod:`~MDAnalysis.transformations` 
+        module, or created by the user, and are stored as a list `transformations`. 
+        This list can only be modified once, and further calls of this function will
+        raise an exception.
+
+        .. code-block:: python
+
+          u = MDAnalysis.Universe(topology, coordinates)
+          workflow = [some_transform, another_transform, this_transform]
+          u.trajectory.add_transformations(*workflow)
+
+        Parameters
+        ----------
+        transform_list : list
+            list of all the transformations that will be applied to the coordinates
+
+        See Also
+        --------
+        :mod:`MDAnalysis.transformations`
+        """
+        #Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations`
+        #to avoid unintended behaviour where the coordinates of each frame are transformed
+        #multiple times when iterating over the trajectory.
+        #In this method, the trajectory is modified all at once and once only.
+
+        super(SingleFrameReaderBase, self).add_transformations(*transformations)
+        for transform in self.transformations:
+            self.ts = transform(self.ts)
+
+    def _apply_transformations(self, ts):
+        """ Applies the transformations to the timestep."""
+        # Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations`
+        # to avoid applying the same transformations multiple times on each frame
+
+        return ts
+

package/MDAnalysis/coordinates/chain.py

@@ -397,4 +397,45 @@ def __repr__(self):
                     nframes=self.n_frames,
                     natoms=self.n_atoms))
 
-
+    def add_transformations(self, *transformations):
+        """ Add all transformations to be applied to the trajectory.
+
+        This function take as list of transformations as an argument. These
+        transformations are functions that will be called by the Reader and given
+        a :class:`Timestep` object as argument, which will be transformed and returned
+        to the Reader.
+        The transformations can be part of the :mod:`~MDAnalysis.transformations` 
+        module, or created by the user, and are stored as a list `transformations`. 
+        This list can only be modified once, and further calls of this function will
+        raise an exception.
+
+        .. code-block:: python
+
+          u = MDAnalysis.Universe(topology, coordinates)
+          workflow = [some_transform, another_transform, this_transform]
+          u.trajectory.add_transformations(*workflow)
+
+        Parameters
+        ----------
+        transform_list : list
+            list of all the transformations that will be applied to the coordinates
+
+        See Also
+        --------
+        :mod:`MDAnalysis.transformations`
+        """
+        #Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations`
+        #to avoid unintended behaviour where the coordinates of each frame are transformed
+        #multiple times when iterating over the trajectory.
+        #In this method, the trajectory is modified all at once and once only.
+
+        super(ChainReader, self).add_transformations(*transformations)
+        for r in self.readers:
+            r.add_transformations(*transformations)
+
+    def _apply_transformations(self, ts):
+        """ Applies the transformations to the timestep."""
+        # Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations`
+        # to avoid applying the same transformations multiple times on each frame
+
+        return ts

package/MDAnalysis/coordinates/memory.py

@@ -339,6 +339,9 @@ def copy(self):
         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
+        # simply copy the property
+        new.transformations = self.transformations
 
         return new
 
@@ -480,3 +483,47 @@ def __repr__(self):
                     nframes=self.n_frames,
                     natoms=self.n_atoms
                 ))
+
+    def add_transformations(self, *transformations):
+        """ Add all transformations to be applied to the trajectory.
+
+        This function take as list of transformations as an argument. These
+        transformations are functions that will be called by the Reader and given
+        a :class:`Timestep` object as argument, which will be transformed and returned
+        to the Reader.
+        The transformations can be part of the :mod:`~MDAnalysis.transformations` 
+        module, or created by the user, and are stored as a list `transformations`. 
+        This list can only be modified once, and further calls of this function will
+        raise an exception.
+
+        .. code-block:: python
+
+          u = MDAnalysis.Universe(topology, coordinates)
+          workflow = [some_transform, another_transform, this_transform]
+          u.trajectory.add_transformations(*workflow)
+
+        Parameters
+        ----------
+        transform_list : list
+            list of all the transformations that will be applied to the coordinates
+
+        See Also
+        --------
+        :mod:`MDAnalysis.transformations`
+        """
+        #Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations`
+        #to avoid unintended behaviour where the coordinates of each frame are transformed
+        #multiple times when iterating over the trajectory.
+        #In this method, the trajectory is modified all at once and once only.
+
+        super(MemoryReader, self).add_transformations(*transformations)
+        for i, ts in enumerate(self):
+            for transform in self.transformations:
+                ts = transform(ts)
+
+    def _apply_transformations(self, ts):
+        """ Applies the transformations to the timestep."""
+        # Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations`
+        # to avoid applying the same transformations multiple times on each frame
+
+        return ts

package/MDAnalysis/core/universe.py

@@ -213,6 +213,10 @@ class Universe(object):
         Universe to only unpickle if a compatible Universe with matching
         *anchor_name* is found. Even if *anchor_name* is set *is_anchor* will
         still be honored when unpickling.
+    transformations: function or list, optional
+        Provide a list of transformations that you wish to apply to the 
+        trajectory upon reading. Transformations can be found in 
+        :mod:`MDAnalysis.transformations`, or can be user-created.
     in_memory
         After reading in the trajectory, transfer it to an in-memory
         representations, which allow for manipulation of coordinates.
@@ -321,7 +325,13 @@ def __init__(self, *args, **kwargs):
 
             if not coordinatefile:
                 coordinatefile = None
+
             self.load_new(coordinatefile, **kwargs)
+            # parse transformations
+            trans_arg = kwargs.pop('transformations', None)
+            if trans_arg:
+                transforms =[trans_arg] if callable(trans_arg) else trans_arg
+                self.trajectory.add_transformations(*transforms)
 
         # Check for guess_bonds
         if kwargs.pop('guess_bonds', False):
@@ -333,6 +343,7 @@ def __init__(self, *args, **kwargs):
         self._anchor_name = kwargs.get('anchor_name', None)
         # Universes are anchors by default
         self.is_anchor = kwargs.get('is_anchor', True)
+
 
     def copy(self):
         """Return an independent copy of this Universe"""