Refactor Transformations from closure into class

package/MDAnalysis/transformations/fit.py

@@ -40,8 +40,7 @@
 from ..lib.transformations import euler_from_matrix, euler_matrix
 
 
-def fit_translation(ag, reference, plane=None, weights=None):
-
+class fit_translation(object):
     """Translates a given AtomGroup so that its center of geometry/mass matches
     the respective center of the given reference. A plane can be given by the
     user using the option `plane`, and will result in the removal of
@@ -82,38 +81,42 @@ def fit_translation(ag, reference, plane=None, weights=None):
     -------
     MDAnalysis.coordinates.base.Timestep
     """
-
-    if plane is not None:
-        axes = {'yz' : 0, 'xz' : 1, 'xy' : 2}
+    def __init__(self, ag, reference, plane=None, weights=None):
+        self.ag = ag
+        self.reference = reference
+        self.plane = plane
+        self.weights = weights
+
+
+    def __call__(self, ts):
+        if self.plane is not None:
+            axes = {'yz' : 0, 'xz' : 1, 'xy' : 2}
+            try:
+                plane = axes[self.plane]
+            except (TypeError, KeyError):
+                raise ValueError(f'{self.plane} is not a valid plane') from None
         try:
-            plane = axes[plane]
-        except (TypeError, KeyError):
-            raise ValueError(f'{plane} is not a valid plane') from None
-    try:
-        if ag.atoms.n_residues != reference.atoms.n_residues:
-            errmsg = f"{ag} and {reference} have mismatched number of residues"
-            raise ValueError(errmsg)
-    except AttributeError:
-        errmsg = f"{ag} or {reference} is not valid Universe/AtomGroup"
-        raise AttributeError(errmsg) from None
-    ref, mobile = align.get_matching_atoms(reference.atoms, ag.atoms)
-    weights = align.get_weights(ref.atoms, weights=weights)
-    ref_com = ref.center(weights)
-    ref_coordinates = ref.atoms.positions - ref_com
-
-    def wrapped(ts):
+            if self.ag.atoms.n_residues != self.reference.atoms.n_residues:
+                errmsg = f"{self.ag} and {self.reference} have mismatched number of residues"
+                raise ValueError(errmsg)
+        except AttributeError:
+            errmsg = f"{self.ag} or {self.reference} is not valid Universe/AtomGroup"
+            raise AttributeError(errmsg) from None
+        ref, mobile = align.get_matching_atoms(self.reference.atoms, self.ag.atoms)
+        weights = align.get_weights(ref.atoms, weights=self.weights)
+        ref_com = ref.center(weights)
+        ref_coordinates = ref.atoms.positions - ref_com
+
         mobile_com = np.asarray(mobile.atoms.center(weights), np.float32)
         vector = ref_com - mobile_com
-        if plane is not None:
+        if self.plane is not None:
             vector[plane] = 0
         ts.positions += vector
 
         return ts
 
-    return wrapped
 
-
-def fit_rot_trans(ag, reference, plane=None, weights=None):
+class fit_rot_trans(object):
     """Perform a spatial superposition by minimizing the RMSD.
 
     Spatially align the group of atoms `ag` to `reference` by doing a RMSD
@@ -160,30 +163,35 @@ def fit_rot_trans(ag, reference, plane=None, weights=None):
     -------
     MDAnalysis.coordinates.base.Timestep
     """
-    if plane is not None:
-        axes = {'yz' : 0, 'xz' : 1, 'xy' : 2}
+    def __init__(self, ag, reference, plane=None, weights=None):
+        self.ag = ag
+        self.reference = reference
+        self.plane = plane
+        self.weights = weights
+    def __call__(self, ts):
+        if self.plane is not None:
+            axes = {'yz' : 0, 'xz' : 1, 'xy' : 2}
+            try:
+                plane = axes[self.plane]
+            except (TypeError, KeyError):
+                raise ValueError(f'{self.plane} is not a valid plane') from None
         try:
-            plane = axes[plane]
-        except (TypeError, KeyError):
-            raise ValueError(f'{plane} is not a valid plane') from None
-    try:
-        if ag.atoms.n_residues != reference.atoms.n_residues:
-            errmsg = f"{ag} and {reference} have mismatched number of residues"
-            raise ValueError(errmsg)
-    except AttributeError:
-        errmsg = f"{ag} or {reference} is not valid Universe/AtomGroup"
-        raise AttributeError(errmsg) from None
-    ref, mobile = align.get_matching_atoms(reference.atoms, ag.atoms)
-    weights = align.get_weights(ref.atoms, weights=weights)
-    ref_com = ref.center(weights)
-    ref_coordinates = ref.atoms.positions - ref_com
-
-    def wrapped(ts):
-        mobile_com = mobile.atoms.center(weights)
+            if self.ag.atoms.n_residues != self.reference.atoms.n_residues:
+                errmsg = f"{self.ag} and {self.reference} have mismatched number of residues"
+                raise ValueError(errmsg)
+        except AttributeError:
+            errmsg = f"{self.ag} or {self.reference} is not valid Universe/AtomGroup"
+            raise AttributeError(errmsg) from None
+        ref, mobile = align.get_matching_atoms(self.reference.atoms, self.ag.atoms)
+        weights = align.get_weights(ref.atoms, weights=self.weights)
+        ref_com = ref.center(self.weights)
+        ref_coordinates = ref.atoms.positions - ref_com
+
+        mobile_com = mobile.atoms.center(self.weights)
         mobile_coordinates = mobile.atoms.positions - mobile_com
-        rotation, dump = align.rotation_matrix(mobile_coordinates, ref_coordinates, weights=weights)
+        rotation, dump = align.rotation_matrix(mobile_coordinates, ref_coordinates, weights=self.weights)
         vector = ref_com
-        if plane is not None:
+        if self.plane is not None:
             matrix = np.r_[rotation, np.zeros(3).reshape(1,3)]
             matrix = np.c_[matrix, np.zeros(4)]
             euler_angs = np.asarray(euler_from_matrix(matrix, axes='sxyz'), np.float32)
@@ -196,5 +204,3 @@ def wrapped(ts):
         ts.positions = ts.positions + vector
 
         return ts
-
-    return wrapped

package/MDAnalysis/transformations/positionaveraging.py

@@ -142,6 +142,7 @@ def __init__(self, avg_frames, check_reset=True):
         self.check_reset = check_reset
         self.current_avg = 0
         self.resetarrays()
+        self.current_frame = 0
 
     def resetarrays(self):
         self.idx_array = np.empty(self.avg_frames)
@@ -164,24 +165,31 @@ def rollposx(self,ts):
 
 
     def __call__(self, ts):
+        #  calling the same timestep will not add new data to coord_array
+        #  This can prevent from getting different values when 
+        #  call `u.trajectory[i]` multiple times.
+        if (ts.frame == self.current_frame
+                and hasattr(self, 'coord_array')
+                    and not np.isnan(self.idx_array).all()):
+            test = ~np.isnan(self.idx_array)
+            ts.positions = np.mean(self.coord_array[...,test], axis=2)
+            return ts
+        else:
+            self.current_frame = ts.frame
+
         self.rollidx(ts)
         test = ~np.isnan(self.idx_array)
         self.current_avg = sum(test)
-        if self.current_avg == 1:
-            return ts
-
         if self.check_reset:
             sign = np.sign(np.diff(self.idx_array[test]))
-
             if not (np.all(sign == 1) or np.all(sign==-1)):
                 warnings.warn('Cannot average position for non sequential'
                               'iterations. Averager will be reset.',
                               Warning)
                 self.resetarrays()
                 return self(ts)
-        
+
         self.rollposx(ts)
         ts.positions = np.mean(self.coord_array[...,test], axis=2)
-            
+
         return ts
-

package/MDAnalysis/transformations/rotate.py

@@ -25,8 +25,8 @@
 Trajectory rotation --- :mod:`MDAnalysis.transformations.rotate`
 ================================================================
 
-Rotates the coordinates by a given angle arround an axis formed by a direction and a
-point 
+Rotates the coordinates by a given angle arround an axis formed by a direction
+and a point.
 
 .. autofunction:: rotateby
 
@@ -38,7 +38,8 @@
 from ..lib.transformations import rotation_matrix
 from ..lib.util import get_weights
 
-def rotateby(angle, direction, point=None, ag=None, weights=None, wrap=False):
+
+class rotateby(object):
     '''
     Rotates the trajectory by a given angle on a given axis. The axis is defined by 
     the user, combining the direction vector and a point. This point can be the center
@@ -104,37 +105,47 @@ def rotateby(angle, direction, point=None, ag=None, weights=None, wrap=False):
     after rotating the trajectory.
 
     '''
-    angle = np.deg2rad(angle)
-    try:
-        direction = np.asarray(direction, np.float32)
-        if direction.shape != (3, ) and direction.shape != (1, 3):
-            raise ValueError('{} is not a valid direction'.format(direction))
-        direction = direction.reshape(3, )
-    except ValueError:
-        raise ValueError(f'{direction} is not a valid direction') from None
-    if point is not None:
-        point = np.asarray(point, np.float32)
-        if point.shape != (3, ) and point.shape != (1, 3):
-            raise ValueError('{} is not a valid point'.format(point))
-        point = point.reshape(3, )
-    elif ag:
+    def __init__(self, angle, direction, point=None, ag=None, weights=None, wrap=False):
+        self.angle = angle
+        self.direction = direction
+        self.point = point
+        self.ag = ag
+        self.weights = weights
+        self.wrap = wrap
+
+    def __call__(self, ts):
+        angle = np.deg2rad(self.angle)
         try:
-            atoms = ag.atoms
-        except AttributeError:
-            raise ValueError(f'{ag} is not an AtomGroup object') from None
-        else:
+            direction = np.asarray(self.direction, np.float32)
+            if direction.shape != (3, ) and direction.shape != (1, 3):
+                raise ValueError('{} is not a valid direction'
+                                 .format(direction))
+            direction = direction.reshape(3, )
+        except ValueError:
+            raise ValueError(f'{self.direction} is not a valid direction') from None
+        if self.point is not None:
+            point = np.asarray(self.point, np.float32)
+            if point.shape != (3, ) and point.shape != (1, 3):
+                raise ValueError('{} is not a valid point'.format(point))
+            point = point.reshape(3, )
+        elif self.ag:
             try:
-                weights = get_weights(atoms, weights=weights)
-            except (ValueError, TypeError):
-                errmsg = ("weights must be {'mass', None} or an iterable of "
-                          "the same size as the atomgroup.")
-                raise TypeError(errmsg) from None
-        center_method = partial(atoms.center, weights, pbc=wrap)
-    else:
-        raise ValueError('A point or an AtomGroup must be specified')
-
-    def wrapped(ts):
-        if point is None:
+                atoms = self.ag.atoms
+            except AttributeError:
+                raise ValueError(f'{self.ag} is not an AtomGroup object') \
+                                from None
+            else:
+                try:
+                    weights = get_weights(atoms, weights=self.weights)
+                except (ValueError, TypeError):
+                    errmsg = ("weights must be {'mass', None} or an iterable of"
+                              "the same size as the atomgroup.")
+                    raise TypeError(errmsg) from None
+            center_method = partial(atoms.center, weights, pbc=self.wrap)
+        else:
+            raise ValueError('A point or an AtomGroup must be specified')
+
+        if self.point is None:
             position = center_method()
         else:
             position = point
@@ -143,8 +154,4 @@ def wrapped(ts):
         translation = matrix[:3, 3]
         ts.positions= np.dot(ts.positions, rotation)
         ts.positions += translation
-
         return ts
-
-    return wrapped
-

package/MDAnalysis/transformations/translate.py

@@ -41,7 +41,7 @@
 
 from ..lib.mdamath import triclinic_vectors
 
-def translate(vector):
+class translate(object):
     """
     Translates the coordinates of a given :class:`~MDAnalysis.coordinates.base.Timestep`
     instance by a given vector.
@@ -60,20 +60,20 @@ def translate(vector):
     :class:`~MDAnalysis.coordinates.base.Timestep` object
 
     """
-    if len(vector)>2:
-        vector = np.float32(vector)
-    else:
-        raise ValueError("{} vector is too short".format(vector))
+    def __init__(self, vector):
+        self.vector = vector
+
+    def __call__(self, ts):
+        if len(self.vector)>2:
+            vector = np.float32(self.vector)
+        else:
+            raise ValueError("{} vector is too short".format(self.vector))
 
-    def wrapped(ts):
         ts.positions += vector
 
         return ts
 
-    return wrapped
-
-
-def center_in_box(ag, center='geometry', point=None, wrap=False):
+class center_in_box(object):
     """
     Translates the coordinates of a given :class:`~MDAnalysis.coordinates.base.Timestep`
     instance so that the center of geometry/mass of the given :class:`~MDAnalysis.core.groups.AtomGroup`
@@ -109,28 +109,33 @@ def center_in_box(ag, center='geometry', point=None, wrap=False):
     :class:`~MDAnalysis.coordinates.base.Timestep` object
 
     """
-
-    pbc_arg = wrap
-    if point:
-        point = np.asarray(point, np.float32)
-        if point.shape != (3, ) and point.shape != (1, 3):
-            raise ValueError('{} is not a valid point'.format(point))
-    try:
-        if center == 'geometry':
-            center_method = partial(ag.center_of_geometry, pbc=pbc_arg)
-        elif center == 'mass':
-            center_method = partial(ag.center_of_mass, pbc=pbc_arg)
-        else:
-            raise ValueError('{} is not a valid argument for center'.format(center))
-    except AttributeError:
-        if center == 'mass':
-            errmsg = f'{ag} is not an AtomGroup object with masses'
-            raise AttributeError(errmsg) from None
-        else:
-            raise ValueError(f'{ag} is not an AtomGroup object') from None
-
-    def wrapped(ts):
-        if point is None:
+    def __init__(self, ag, center='geometry', point=None, wrap=False):
+        self.ag = ag
+        self.center = center
+        self.point = point
+        self.wrap = wrap
+
+    def __call__(self, ts):
+        pbc_arg = self.wrap
+        if self.point:
+            point = np.asarray(self.point, np.float32)
+            if point.shape != (3, ) and point.shape != (1, 3):
+                raise ValueError('{} is not a valid point'.format(point))
+        try:
+            if self.center == 'geometry':
+                center_method = partial(self.ag.center_of_geometry, pbc=pbc_arg)
+            elif self.center == 'mass':
+                center_method = partial(self.ag.center_of_mass, pbc=pbc_arg)
+            else:
+                raise ValueError('{} is not a valid argument for center'.format(self.center))
+        except AttributeError:
+            if self.center == 'mass':
+                errmsg = f'{self.ag} is not an AtomGroup object with masses'
+                raise AttributeError(errmsg) from None
+            else:
+                raise ValueError(f'{self.ag} is not an AtomGroup object') from None
+
+        if self.point is None:
             boxcenter = np.sum(ts.triclinic_dimensions, axis=0) / 2
         else:
             boxcenter = point
@@ -141,6 +146,3 @@ def wrapped(ts):
         ts.positions += vector
 
         return ts
-
-    return wrapped
-