From 05f2dddd2b8c40633650d6a5dac1c8fc2cdcd5d0 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 12:07:16 +0200 Subject: [PATCH 01/26] refactor transformations into class --- package/MDAnalysis/transformations/fit.py | 102 +++++++++--------- .../transformations/positionaveraging.py | 22 ++-- package/MDAnalysis/transformations/rotate.py | 79 +++++++------- .../MDAnalysis/transformations/translate.py | 72 +++++++------ package/MDAnalysis/transformations/wrap.py | 38 +++---- 5 files changed, 166 insertions(+), 147 deletions(-) diff --git a/package/MDAnalysis/transformations/fit.py b/package/MDAnalysis/transformations/fit.py index 084868a0e55..c5301e23dcf 100644 --- a/package/MDAnalysis/transformations/fit.py +++ b/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 diff --git a/package/MDAnalysis/transformations/positionaveraging.py b/package/MDAnalysis/transformations/positionaveraging.py index 179b98163f1..a09a60f15c7 100644 --- a/package/MDAnalysis/transformations/positionaveraging.py +++ b/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 - diff --git a/package/MDAnalysis/transformations/rotate.py b/package/MDAnalysis/transformations/rotate.py index 4e27f81012d..94e6e8d0ce6 100644 --- a/package/MDAnalysis/transformations/rotate.py +++ b/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 - diff --git a/package/MDAnalysis/transformations/translate.py b/package/MDAnalysis/transformations/translate.py index 6a4d28113be..b0ad43d2963 100644 --- a/package/MDAnalysis/transformations/translate.py +++ b/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 - diff --git a/package/MDAnalysis/transformations/wrap.py b/package/MDAnalysis/transformations/wrap.py index 99b60f943a1..e32428884a4 100644 --- a/package/MDAnalysis/transformations/wrap.py +++ b/package/MDAnalysis/transformations/wrap.py @@ -37,7 +37,7 @@ from ..lib._cutil import make_whole -def wrap(ag, compound='atoms'): +class wrap(object): """ Shift the contents of a given AtomGroup back into the unit cell. :: @@ -81,16 +81,15 @@ def wrap(ag, compound='atoms'): MDAnalysis.coordinates.base.Timestep """ - - def wrapped(ts): - ag.wrap(compound=compound) - - return ts - - return wrapped + def __init__(self, ag, compound='atoms'): + self.ag = ag + self.compound = compound + def __call__(self, ts): + self.ag.wrap(compound=self.compound) + return ts -def unwrap(ag): +class unwrap(object): """ Move all atoms in an AtomGroup so that bonds don't split over images @@ -131,17 +130,14 @@ def unwrap(ag): MDAnalysis.coordinates.base.Timestep """ - - try: - ag.fragments - except AttributeError: - raise AttributeError("{} has no fragments".format(ag)) - - def wrapped(ts): - for frag in ag.fragments: + def __init__(self, ag): + self.ag = ag + + def __call__(self, ts): + try: + self.ag.fragments + except AttributeError: + raise AttributeError("{} has no fragments".format(self.ag)) + for frag in self.ag.fragments: make_whole(frag) - return ts - - return wrapped - From 92a9653abb5bdf4cc258793f9803232965b1a3a5 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 12:32:43 +0200 Subject: [PATCH 02/26] pep8 and test fix --- package/MDAnalysis/transformations/fit.py | 67 ++++++++++++------- .../transformations/test_fit.py | 4 +- 2 files changed, 47 insertions(+), 24 deletions(-) diff --git a/package/MDAnalysis/transformations/fit.py b/package/MDAnalysis/transformations/fit.py index c5301e23dcf..3b383d84723 100644 --- a/package/MDAnalysis/transformations/fit.py +++ b/package/MDAnalysis/transformations/fit.py @@ -33,10 +33,8 @@ """ import numpy as np -from functools import partial from ..analysis import align -from ..lib.util import get_weights from ..lib.transformations import euler_from_matrix, euler_matrix @@ -48,8 +46,8 @@ class fit_translation(object): Example ------- - Removing the translations of a given AtomGroup `ag` on the XY plane by fitting - its center of mass to the center of mass of a reference `ref`: + Removing the translations of a given AtomGroup `ag` on the XY plane by + fitting its center of mass to the center of mass of a reference `ref`: .. code-block:: python @@ -66,11 +64,12 @@ class fit_translation(object): :class:`~MDAnalysis.core.groups.AtomGroup` or a whole :class:`~MDAnalysis.core.universe.Universe` reference : Universe or AtomGroup - reference structure, a :class:`~MDAnalysis.core.groups.AtomGroup` or a whole - :class:`~MDAnalysis.core.universe.Universe` + reference structure, a :class:`~MDAnalysis.core.groups.AtomGroup` or a + whole :class:`~MDAnalysis.core.universe.Universe` plane: str, optional - used to define the plane on which the translations will be removed. Defined as a - string of the plane. Suported planes are yz, xz and xy planes. + used to define the plane on which the translations will be removed. + Defined as a string of the plane. + Suported planes are yz, xz and xy planes. weights : {"mass", ``None``} or array_like, optional choose weights. With ``"mass"`` uses masses as weights; with ``None`` weigh each atom equally. If a float array of the same length as @@ -87,22 +86,30 @@ def __init__(self, ag, reference, plane=None, weights=None): self.plane = plane self.weights = weights - def __call__(self, ts): if self.plane is not None: - axes = {'yz' : 0, 'xz' : 1, 'xy' : 2} + 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 + raise ValueError(f'{self.plane} is not a valid plane') \ + from None try: if self.ag.atoms.n_residues != self.reference.atoms.n_residues: - errmsg = f"{self.ag} and {self.reference} have mismatched number of residues" + errmsg = ( + f"{self.ag} and {self.reference} have mismatched" + f"number of residues" + ) + raise ValueError(errmsg) except AttributeError: - errmsg = f"{self.ag} or {self.reference} is not valid Universe/AtomGroup" + errmsg = ( + f"{self.ag} or {self.reference} is not valid" + f"Universe/AtomGroup" + ) raise AttributeError(errmsg) from None - ref, mobile = align.get_matching_atoms(self.reference.atoms, self.ag.atoms) + 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 @@ -168,36 +175,50 @@ def __init__(self, ag, reference, plane=None, weights=None): 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} + 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: if self.ag.atoms.n_residues != self.reference.atoms.n_residues: - errmsg = f"{self.ag} and {self.reference} have mismatched number of residues" + errmsg = ( + f"{self.ag} and {self.reference} have mismatched " + f"number of residues" + ) raise ValueError(errmsg) except AttributeError: - errmsg = f"{self.ag} or {self.reference} is not valid Universe/AtomGroup" + errmsg = ( + f"{self.ag} or {self.reference} is not valid " + f"Universe/AtomGroup" + ) raise AttributeError(errmsg) from None - ref, mobile = align.get_matching_atoms(self.reference.atoms, self.ag.atoms) + 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=self.weights) + rotation, dump = align.rotation_matrix(mobile_coordinates, + ref_coordinates, + weights=weights) vector = ref_com if self.plane is not None: - matrix = np.r_[rotation, np.zeros(3).reshape(1,3)] + 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) + euler_angs = np.asarray(euler_from_matrix(matrix, axes='sxyz'), + np.float32) for i in range(0, euler_angs.size): - euler_angs[i] = ( euler_angs[plane] if i == plane else 0) - rotation = euler_matrix(euler_angs[0], euler_angs[1], euler_angs[2], axes='sxyz')[:3, :3] + euler_angs[i] = (euler_angs[plane] if i == plane else 0) + rotation = euler_matrix(euler_angs[0], + euler_angs[1], + euler_angs[2], + axes='sxyz')[:3, :3] vector[plane] = mobile_com[plane] ts.positions = ts.positions - mobile_com ts.positions = np.dot(ts.positions, rotation.T) diff --git a/testsuite/MDAnalysisTests/transformations/test_fit.py b/testsuite/MDAnalysisTests/transformations/test_fit.py index 36469ec6d7f..1b7fab35220 100644 --- a/testsuite/MDAnalysisTests/transformations/test_fit.py +++ b/testsuite/MDAnalysisTests/transformations/test_fit.py @@ -183,7 +183,9 @@ def test_fit_rot_trans_shorter_universe(fit_universe): bad_u =fit_universe[0].atoms[0:5] test_u= bad_u with pytest.raises(ValueError): - fit_rot_trans(test_u, ref_u)(test_u.trajectory.ts) + # To make sure we catch the right error + # test_u is changed to ref_u here. + fit_rot_trans(test_u, ref_u)(ref_u.trajectory.ts) @pytest.mark.parametrize('weights', ( From 9fc50712d02cfb2096064431825499260641a7f6 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 12:48:50 +0200 Subject: [PATCH 03/26] move fit prep toinit --- package/MDAnalysis/transformations/fit.py | 48 +++++++++---------- .../transformations/test_fit.py | 4 +- 2 files changed, 25 insertions(+), 27 deletions(-) diff --git a/package/MDAnalysis/transformations/fit.py b/package/MDAnalysis/transformations/fit.py index 3b383d84723..e2fba978c56 100644 --- a/package/MDAnalysis/transformations/fit.py +++ b/package/MDAnalysis/transformations/fit.py @@ -86,11 +86,10 @@ def __init__(self, ag, reference, plane=None, weights=None): 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] + self.plane = axes[self.plane] except (TypeError, KeyError): raise ValueError(f'{self.plane} is not a valid plane') \ from None @@ -108,16 +107,17 @@ def __call__(self, ts): f"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 + self.ref, self.mobile = align.get_matching_atoms(self.reference.atoms, + self.ag.atoms) + self.weights = align.get_weights(self.ref.atoms, weights=self.weights) + self.ref_com = self.ref.center(self.weights) - mobile_com = np.asarray(mobile.atoms.center(weights), np.float32) - vector = ref_com - mobile_com + def __call__(self, ts): + mobile_com = np.asarray(self.mobile.atoms.center(self.weights), + np.float32) + vector = self.ref_com - mobile_com if self.plane is not None: - vector[plane] = 0 + vector[self.plane] = 0 ts.positions += vector return ts @@ -176,11 +176,10 @@ def __init__(self, ag, reference, plane=None, weights=None): 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] + self.plane = axes[self.plane] except (TypeError, KeyError): raise ValueError(f'{self.plane} is not a valid plane') from None try: @@ -196,32 +195,33 @@ def __call__(self, ts): f"Universe/AtomGroup" ) raise AttributeError(errmsg) from None - ref, mobile = align.get_matching_atoms(self.reference.atoms, + self.ref, self.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 + self.weights = align.get_weights(self.ref.atoms, weights=self.weights) + self.ref_com = self.ref.center(self.weights) + self.ref_coordinates = self.ref.atoms.positions - self.ref_com - mobile_com = mobile.atoms.center(self.weights) - mobile_coordinates = mobile.atoms.positions - mobile_com + def __call__(self, ts): + mobile_com = self.mobile.atoms.center(self.weights) + mobile_coordinates = self.mobile.atoms.positions - mobile_com rotation, dump = align.rotation_matrix(mobile_coordinates, - ref_coordinates, - weights=weights) - vector = ref_com + self.ref_coordinates, + weights=self.weights) + vector = self.ref_com 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) for i in range(0, euler_angs.size): - euler_angs[i] = (euler_angs[plane] if i == plane else 0) + euler_angs[i] = (euler_angs[self.plane] if i == self.plane + else 0) rotation = euler_matrix(euler_angs[0], euler_angs[1], euler_angs[2], axes='sxyz')[:3, :3] - vector[plane] = mobile_com[plane] + vector[self.plane] = mobile_com[self.plane] ts.positions = ts.positions - mobile_com ts.positions = np.dot(ts.positions, rotation.T) ts.positions = ts.positions + vector - return ts diff --git a/testsuite/MDAnalysisTests/transformations/test_fit.py b/testsuite/MDAnalysisTests/transformations/test_fit.py index 1b7fab35220..36469ec6d7f 100644 --- a/testsuite/MDAnalysisTests/transformations/test_fit.py +++ b/testsuite/MDAnalysisTests/transformations/test_fit.py @@ -183,9 +183,7 @@ def test_fit_rot_trans_shorter_universe(fit_universe): bad_u =fit_universe[0].atoms[0:5] test_u= bad_u with pytest.raises(ValueError): - # To make sure we catch the right error - # test_u is changed to ref_u here. - fit_rot_trans(test_u, ref_u)(ref_u.trajectory.ts) + fit_rot_trans(test_u, ref_u)(test_u.trajectory.ts) @pytest.mark.parametrize('weights', ( From 6b7c45bb337abb9f6290dbf42ed3b3ef699654b8 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 12:56:22 +0200 Subject: [PATCH 04/26] move rotate prep to init --- package/MDAnalysis/transformations/rotate.py | 48 ++++++++++++-------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/package/MDAnalysis/transformations/rotate.py b/package/MDAnalysis/transformations/rotate.py index 94e6e8d0ce6..577baefac34 100644 --- a/package/MDAnalysis/transformations/rotate.py +++ b/package/MDAnalysis/transformations/rotate.py @@ -31,7 +31,6 @@ .. autofunction:: rotateby """ -import math import numpy as np from functools import partial @@ -105,7 +104,13 @@ class rotateby(object): after rotating the trajectory. ''' - def __init__(self, angle, direction, point=None, ag=None, weights=None, wrap=False): + def __init__(self, + angle, + direction, + point=None, + ag=None, + weights=None, + wrap=False): self.angle = angle self.direction = direction self.point = point @@ -113,45 +118,48 @@ def __init__(self, angle, direction, point=None, ag=None, weights=None, wrap=Fal self.weights = weights self.wrap = wrap - def __call__(self, ts): - angle = np.deg2rad(self.angle) + self.angle = np.deg2rad(self.angle) try: - direction = np.asarray(self.direction, np.float32) - if direction.shape != (3, ) and direction.shape != (1, 3): + self.direction = np.asarray(self.direction, np.float32) + if self.direction.shape != (3, ) and self.direction.shape != (1, 3): raise ValueError('{} is not a valid direction' - .format(direction)) - direction = direction.reshape(3, ) + .format(self.direction)) + self.direction = self.direction.reshape(3, ) except ValueError: - raise ValueError(f'{self.direction} is not a valid direction') from None + 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, ) + self.point = np.asarray(self.point, np.float32) + if self.point.shape != (3, ) and self.point.shape != (1, 3): + raise ValueError('{} is not a valid point'.format(self.point)) + self.point = self.point.reshape(3, ) elif self.ag: try: - atoms = self.ag.atoms + self.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) + self.weights = get_weights(self.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) + self.center_method = partial(self.atoms.center, + self.weights, + pbc=self.wrap) else: raise ValueError('A point or an AtomGroup must be specified') + def __call__(self, ts): if self.point is None: - position = center_method() + position = self.center_method() else: - position = point - matrix = rotation_matrix(angle, direction, position) + position = self.point + matrix = rotation_matrix(self.angle, self.direction, position) rotation = matrix[:3, :3].T translation = matrix[:3, 3] - ts.positions= np.dot(ts.positions, rotation) + ts.positions = np.dot(ts.positions, rotation) ts.positions += translation return ts From ae51c5bc6c34cc4869d0587261d5ca2c31bac5e1 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 13:01:09 +0200 Subject: [PATCH 05/26] move translate prep to init --- .../MDAnalysis/transformations/translate.py | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/package/MDAnalysis/transformations/translate.py b/package/MDAnalysis/transformations/translate.py index b0ad43d2963..097c37a2dc6 100644 --- a/package/MDAnalysis/transformations/translate.py +++ b/package/MDAnalysis/transformations/translate.py @@ -39,7 +39,6 @@ import numpy as np from functools import partial -from ..lib.mdamath import triclinic_vectors class translate(object): """ @@ -63,16 +62,16 @@ class translate(object): def __init__(self, vector): self.vector = vector - def __call__(self, ts): - if len(self.vector)>2: - vector = np.float32(self.vector) + if len(self.vector) > 2: + self.vector = np.float32(self.vector) else: raise ValueError("{} vector is too short".format(self.vector)) - ts.positions += vector - + def __call__(self, ts): + ts.positions += self.vector return ts + class center_in_box(object): """ Translates the coordinates of a given :class:`~MDAnalysis.coordinates.base.Timestep` @@ -115,32 +114,35 @@ def __init__(self, ag, center='geometry', point=None, wrap=False): 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)) + self.point = np.asarray(self.point, np.float32) + if self.point.shape != (3, ) and self.point.shape != (1, 3): + raise ValueError('{} is not a valid point'.format(self.point)) try: if self.center == 'geometry': - center_method = partial(self.ag.center_of_geometry, pbc=pbc_arg) + self.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) + self.center_method = partial(self.ag.center_of_mass, + pbc=pbc_arg) else: - raise ValueError('{} is not a valid argument for center'.format(self.center)) + raise ValueError(f'{self.center} is valid for 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 + raise ValueError(f'{self.ag} is not an AtomGroup object') \ + from None + def __call__(self, ts): if self.point is None: boxcenter = np.sum(ts.triclinic_dimensions, axis=0) / 2 else: - boxcenter = point + boxcenter = self.point - ag_center = center_method() + ag_center = self.center_method() vector = boxcenter - ag_center ts.positions += vector From a542792c44ac8521af2b01b285319b6c8f05e790 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 13:02:22 +0200 Subject: [PATCH 06/26] move wrap prep to init --- package/MDAnalysis/transformations/wrap.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/package/MDAnalysis/transformations/wrap.py b/package/MDAnalysis/transformations/wrap.py index e32428884a4..e4cbb0fc392 100644 --- a/package/MDAnalysis/transformations/wrap.py +++ b/package/MDAnalysis/transformations/wrap.py @@ -133,11 +133,12 @@ class unwrap(object): def __init__(self, ag): self.ag = ag - def __call__(self, ts): try: self.ag.fragments except AttributeError: raise AttributeError("{} has no fragments".format(self.ag)) + + def __call__(self, ts): for frag in self.ag.fragments: make_whole(frag) return ts From 1729ee6129eea87cce32ec95f97b5d843e0a5a88 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 13:12:26 +0200 Subject: [PATCH 07/26] pep8 --- package/MDAnalysis/transformations/fit.py | 7 ++-- .../transformations/positionaveraging.py | 40 +++++++++---------- package/MDAnalysis/transformations/rotate.py | 10 +++-- package/MDAnalysis/transformations/wrap.py | 1 + 4 files changed, 30 insertions(+), 28 deletions(-) diff --git a/package/MDAnalysis/transformations/fit.py b/package/MDAnalysis/transformations/fit.py index e2fba978c56..17079549f04 100644 --- a/package/MDAnalysis/transformations/fit.py +++ b/package/MDAnalysis/transformations/fit.py @@ -181,7 +181,8 @@ def __init__(self, ag, reference, plane=None, weights=None): try: self.plane = axes[self.plane] except (TypeError, KeyError): - raise ValueError(f'{self.plane} is not a valid plane') from None + raise ValueError(f'{self.plane} is not a valid plane') \ + from None try: if self.ag.atoms.n_residues != self.reference.atoms.n_residues: errmsg = ( @@ -196,7 +197,7 @@ def __init__(self, ag, reference, plane=None, weights=None): ) raise AttributeError(errmsg) from None self.ref, self.mobile = align.get_matching_atoms(self.reference.atoms, - self.ag.atoms) + self.ag.atoms) self.weights = align.get_weights(self.ref.atoms, weights=self.weights) self.ref_com = self.ref.center(self.weights) self.ref_coordinates = self.ref.atoms.positions - self.ref_com @@ -215,7 +216,7 @@ def __call__(self, ts): np.float32) for i in range(0, euler_angs.size): euler_angs[i] = (euler_angs[self.plane] if i == self.plane - else 0) + else 0) rotation = euler_matrix(euler_angs[0], euler_angs[1], euler_angs[2], diff --git a/package/MDAnalysis/transformations/positionaveraging.py b/package/MDAnalysis/transformations/positionaveraging.py index a09a60f15c7..45a89090cee 100644 --- a/package/MDAnalysis/transformations/positionaveraging.py +++ b/package/MDAnalysis/transformations/positionaveraging.py @@ -36,15 +36,14 @@ import warnings -class PositionAverager(object): +class PositionAverager(object): """ - Averages the coordinates of a given timestep so that the coordinates of the AtomGroup correspond to the average positions of the N previous frames. For frames < N, the average of the frames iterated up to that point will be returned. - + Example ------- @@ -55,13 +54,13 @@ class PositionAverager(object): complete, or if the frames iterated are not sequential. .. code-block:: python - + N=3 transformation = PositionAverager(N, check_reset=True) u.trajectory.add_transformations(transformation) for ts in u.trajectory: print(ts.positions) - + In this case, ``ts.positions`` will return the average coordinates of the last N iterated frames. @@ -136,43 +135,42 @@ class PositionAverager(object): """ - + def __init__(self, avg_frames, check_reset=True): self.avg_frames = avg_frames 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) self.idx_array[:] = np.nan - - def rollidx(self,ts): - self.idx_array = np.roll(self.idx_array, 1) + + def rollidx(self, ts): + self.idx_array = np.roll(self.idx_array, 1) self.idx_array[0] = ts.frame - - def rollposx(self,ts): + + def rollposx(self, ts): try: self.coord_array.size except AttributeError: size = (ts.positions.shape[0], ts.positions.shape[1], self.avg_frames) self.coord_array = np.empty(size) - + self.coord_array = np.roll(self.coord_array, 1, axis=2) - self.coord_array[...,0] = ts.positions.copy() - - + self.coord_array[..., 0] = ts.positions.copy() + def __call__(self, ts): # calling the same timestep will not add new data to coord_array - # This can prevent from getting different values when + # 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()): + and not np.isnan(self.idx_array).all()): test = ~np.isnan(self.idx_array) - ts.positions = np.mean(self.coord_array[...,test], axis=2) + ts.positions = np.mean(self.coord_array[..., test], axis=2) return ts else: self.current_frame = ts.frame @@ -182,7 +180,7 @@ def __call__(self, ts): self.current_avg = sum(test) if self.check_reset: sign = np.sign(np.diff(self.idx_array[test])) - if not (np.all(sign == 1) or np.all(sign==-1)): + 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) @@ -190,6 +188,6 @@ def __call__(self, ts): return self(ts) self.rollposx(ts) - ts.positions = np.mean(self.coord_array[...,test], axis=2) + ts.positions = np.mean(self.coord_array[..., test], axis=2) return ts diff --git a/package/MDAnalysis/transformations/rotate.py b/package/MDAnalysis/transformations/rotate.py index 577baefac34..f9f9b2d6b87 100644 --- a/package/MDAnalysis/transformations/rotate.py +++ b/package/MDAnalysis/transformations/rotate.py @@ -121,7 +121,8 @@ def __init__(self, self.angle = np.deg2rad(self.angle) try: self.direction = np.asarray(self.direction, np.float32) - if self.direction.shape != (3, ) and self.direction.shape != (1, 3): + if self.direction.shape != (3, ) and \ + self.direction.shape != (1, 3): raise ValueError('{} is not a valid direction' .format(self.direction)) self.direction = self.direction.reshape(3, ) @@ -141,10 +142,11 @@ def __init__(self, from None else: try: - self.weights = get_weights(self.atoms, weights=self.weights) + self.weights = get_weights(self.atoms, + weights=self.weights) except (ValueError, TypeError): - errmsg = ("weights must be {'mass', None} or an iterable of" - "the same size as the atomgroup.") + errmsg = ("weights must be {'mass', None} or an iterable " + "of the same size as the atomgroup.") raise TypeError(errmsg) from None self.center_method = partial(self.atoms.center, self.weights, diff --git a/package/MDAnalysis/transformations/wrap.py b/package/MDAnalysis/transformations/wrap.py index e4cbb0fc392..e6499c13ab6 100644 --- a/package/MDAnalysis/transformations/wrap.py +++ b/package/MDAnalysis/transformations/wrap.py @@ -89,6 +89,7 @@ def __call__(self, ts): self.ag.wrap(compound=self.compound) return ts + class unwrap(object): """ Move all atoms in an AtomGroup so that bonds don't split over images From 26950b28af21eefc07b0a09a36fd9b2a2cde05f9 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 13:36:57 +0200 Subject: [PATCH 08/26] pep --- package/MDAnalysis/transformations/positionaveraging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package/MDAnalysis/transformations/positionaveraging.py b/package/MDAnalysis/transformations/positionaveraging.py index 45a89090cee..930a77774b3 100644 --- a/package/MDAnalysis/transformations/positionaveraging.py +++ b/package/MDAnalysis/transformations/positionaveraging.py @@ -148,7 +148,7 @@ def resetarrays(self): self.idx_array[:] = np.nan def rollidx(self, ts): - self.idx_array = np.roll(self.idx_array, 1) + self.idx_array = np.roll(self.idx_array, 1) self.idx_array[0] = ts.frame def rollposx(self, ts): From 5aa8d2bed94e335d6575324410ae519216bf4137 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 13:54:40 +0200 Subject: [PATCH 09/26] test pickle --- .../parallelism/test_pickle_transformation.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py new file mode 100644 index 00000000000..90d5fee8d79 --- /dev/null +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -0,0 +1,76 @@ +# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*- +# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8 +# +# 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. +# doi: 10.25080/majora-629e541a-00e +# +# 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 +# +import multiprocessing + +import numpy as np +import pytest +import pickle +from numpy.testing import assert_equal + +import MDAnalysis as mda + +from MDAnalysis.transformations.fit import fit_translation, fit_rot_trans +from MDAnalysis.transformations.positionaveraging import PositionAverager +from MDAnalysis.transformations.rotate import rotateby +from MDAnalysis.transformations.translate import translate, center_in_box +from MDAnalysis.transformations.wrap import wrap, unwrap + +from MDAnalysisTests.datafiles import TPR, XTC + +@pytest.fixture() +def test_u(): + u = mda.Universe(TPR, XTC) + return u + + +u = mda.Universe(TPR, XTC) +ag = u.atoms[0:10] + +@pytest.fixture(params=[ + (fit_translation, [ag, ag]), + (fit_rot_trans, [ag, ag]), + (PositionAverager, [3]), + (rotateby, [90, [0,0,1], [1,2,3]]), + (translate,[[1,2,3]]), + (center_in_box,[ag]), + (wrap,[u.atoms]), + (unwrap,[u.atoms]) +]) +def transformation(request): + transform = request.param[0](*request.param[1]) + return transform + +def test_transformation_pickle(transformation, test_u): + ref_result = transformation(test_u.trajectory[5]).positions + transformation_p = pickle.loads(pickle.dumps(transformation)) + result = transformation_p(test_u.trajectory[5]).positions + assert_equal(ref_result, result) + +@pytest.mark.skip(reason="`Universe` cannot be picked now") +def test_add_transformation_pickle(transformation, test_u): + test_u.trajectory.add_transformations(transformation) + test_u_p = pickle.loads(pickle.dumps(test_u)) + test_u.trajectory[0] + for u_ts, u_p_ts in zip(test_u.trajectory[:5], test_u_p.trajectory[:5]): + assert_equal(u_ts.positions, u_p_ts.positions) + From 1c5bbdbdd353fbaf7fa380aada9ea07e3e8d35ee Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 20 Jul 2020 13:57:33 +0200 Subject: [PATCH 10/26] pep --- .../parallelism/test_pickle_transformation.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py index 90d5fee8d79..ccb22e022df 100644 --- a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -20,9 +20,6 @@ # MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations. # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787 # -import multiprocessing - -import numpy as np import pytest import pickle from numpy.testing import assert_equal @@ -37,6 +34,7 @@ from MDAnalysisTests.datafiles import TPR, XTC + @pytest.fixture() def test_u(): u = mda.Universe(TPR, XTC) @@ -46,26 +44,29 @@ def test_u(): u = mda.Universe(TPR, XTC) ag = u.atoms[0:10] + @pytest.fixture(params=[ (fit_translation, [ag, ag]), (fit_rot_trans, [ag, ag]), (PositionAverager, [3]), - (rotateby, [90, [0,0,1], [1,2,3]]), - (translate,[[1,2,3]]), - (center_in_box,[ag]), - (wrap,[u.atoms]), - (unwrap,[u.atoms]) + (rotateby, [90, [0, 0, 1], [1, 2, 3]]), + (translate, [[1, 2, 3]]), + (center_in_box, [ag]), + (wrap, [u.atoms]), + (unwrap, [u.atoms]) ]) def transformation(request): transform = request.param[0](*request.param[1]) return transform + def test_transformation_pickle(transformation, test_u): ref_result = transformation(test_u.trajectory[5]).positions transformation_p = pickle.loads(pickle.dumps(transformation)) result = transformation_p(test_u.trajectory[5]).positions assert_equal(ref_result, result) + @pytest.mark.skip(reason="`Universe` cannot be picked now") def test_add_transformation_pickle(transformation, test_u): test_u.trajectory.add_transformations(transformation) @@ -73,4 +74,3 @@ def test_add_transformation_pickle(transformation, test_u): test_u.trajectory[0] for u_ts, u_p_ts in zip(test_u.trajectory[:5], test_u_p.trajectory[:5]): assert_equal(u_ts.positions, u_p_ts.positions) - From b09bef6379a6a6c71d90011356b9c9e5218ad8c3 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 27 Jul 2020 14:04:46 +0200 Subject: [PATCH 11/26] doc for each module --- .../MDAnalysis/transformations/__init__.py | 23 +++++++++++++++++++ package/MDAnalysis/transformations/fit.py | 8 +++++-- package/MDAnalysis/transformations/rotate.py | 5 +++- .../MDAnalysis/transformations/translate.py | 7 ++++-- package/MDAnalysis/transformations/wrap.py | 7 ++++-- 5 files changed, 43 insertions(+), 7 deletions(-) diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index 13ce68dfbd4..0b810bac7d7 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -53,7 +53,30 @@ def wrapped(ts): See `MDAnalysis.transformations.translate` for a simple example. +To meet the need of serialization of universe, transformations are converted +into classes. They retain the similar API and functionality by implementing the +aforementioned `wrapped(ts)` as a special method `__call__`. For example: +.. code-blocks:: python + + class transfomrations(object): + def __init__(self, *args, **kwargs): + # do some things + # save needed args as attributes. + self.needed_var = args[0] + + def __call__(self, ts): + # apply changes to the Timestep object + return ts + + +Note it does not mean that the old closure/wrapped function implementation will fail. +They can still be used if one does not need serialization, which is a prerequisite for parallel +analysis. + + +.. versionchanged:: 2.0.0 + All transformations are classes now. """ from .translate import translate, center_in_box diff --git a/package/MDAnalysis/transformations/fit.py b/package/MDAnalysis/transformations/fit.py index 17079549f04..5de4af56168 100644 --- a/package/MDAnalysis/transformations/fit.py +++ b/package/MDAnalysis/transformations/fit.py @@ -27,9 +27,9 @@ Translate and/or rotates the coordinates of a given trajectory to align a given AtomGroup to a reference structure. -.. autofunction:: fit_translation +.. autoclass:: fit_translation -.. autofunction:: fit_rot_trans +.. autoclass:: fit_rot_trans """ import numpy as np @@ -79,6 +79,10 @@ class fit_translation(object): Returns ------- MDAnalysis.coordinates.base.Timestep + + + .. versionchanged:: 2.0.0 + Now it is a class with `__call__`. """ def __init__(self, ag, reference, plane=None, weights=None): self.ag = ag diff --git a/package/MDAnalysis/transformations/rotate.py b/package/MDAnalysis/transformations/rotate.py index f9f9b2d6b87..e7aafb56975 100644 --- a/package/MDAnalysis/transformations/rotate.py +++ b/package/MDAnalysis/transformations/rotate.py @@ -28,7 +28,7 @@ Rotates the coordinates by a given angle arround an axis formed by a direction and a point. -.. autofunction:: rotateby +.. autoclass:: rotateby """ import numpy as np @@ -103,6 +103,9 @@ class rotateby(object): Wrapping/unwrapping the trajectory or performing PBC corrections may not be possible after rotating the trajectory. + + .. versionchanged:: 2.0.0 + Now it is a class with `__call__`. ''' def __init__(self, angle, diff --git a/package/MDAnalysis/transformations/translate.py b/package/MDAnalysis/transformations/translate.py index 097c37a2dc6..599e1bc1e72 100644 --- a/package/MDAnalysis/transformations/translate.py +++ b/package/MDAnalysis/transformations/translate.py @@ -30,9 +30,9 @@ or defined by centering an AtomGroup in the unit cell using the function :func:`center_in_box` -.. autofunction:: translate +.. autoclass:: translate -.. autofunction:: center_in_box +.. autoclass:: center_in_box """ @@ -107,6 +107,9 @@ class center_in_box(object): ------- :class:`~MDAnalysis.coordinates.base.Timestep` object + + .. versionchanged:: 2.0.0 + Now it is a class with `__call__`. """ def __init__(self, ag, center='geometry', point=None, wrap=False): self.ag = ag diff --git a/package/MDAnalysis/transformations/wrap.py b/package/MDAnalysis/transformations/wrap.py index e6499c13ab6..b9e368f8d7d 100644 --- a/package/MDAnalysis/transformations/wrap.py +++ b/package/MDAnalysis/transformations/wrap.py @@ -28,9 +28,9 @@ translates the atoms back in the unit cell. :func:`unwrap` translates the atoms of each molecule so that bons don't split over images. -.. autofunction:: wrap +.. autoclass:: wrap -.. autofunction:: unwrap +.. autoclass:: unwrap """ @@ -80,6 +80,9 @@ class wrap(object): ------- MDAnalysis.coordinates.base.Timestep + + .. versionchanged:: 2.0.0 + Now it is a class with `__call__`. """ def __init__(self, ag, compound='atoms'): self.ag = ag From 235e8a918a613061ac9461195aad24a41d9e3e63 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Thu, 30 Jul 2020 11:13:20 +0200 Subject: [PATCH 12/26] read_offset issue? --- .../parallelism/test_pickle_transformation.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py index ccb22e022df..8aa6ac5c058 100644 --- a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -41,8 +41,8 @@ def test_u(): return u -u = mda.Universe(TPR, XTC) -ag = u.atoms[0:10] +uni = mda.Universe(TPR, XTC) +ag = uni.atoms[0:10] @pytest.fixture(params=[ @@ -52,8 +52,8 @@ def test_u(): (rotateby, [90, [0, 0, 1], [1, 2, 3]]), (translate, [[1, 2, 3]]), (center_in_box, [ag]), - (wrap, [u.atoms]), - (unwrap, [u.atoms]) + (wrap, [uni.atoms]), + (unwrap, [uni.atoms]) ]) def transformation(request): transform = request.param[0](*request.param[1]) From 5ef7d3f0a2c7bc51c8a1c95ca8f26fc114af2a78 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Thu, 30 Jul 2020 11:55:48 +0200 Subject: [PATCH 13/26] change to dcd --- .../parallelism/test_pickle_transformation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py index 8aa6ac5c058..b25f187cffb 100644 --- a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -32,16 +32,16 @@ from MDAnalysis.transformations.translate import translate, center_in_box from MDAnalysis.transformations.wrap import wrap, unwrap -from MDAnalysisTests.datafiles import TPR, XTC +from MDAnalysisTests.datafiles import PSF_TRICLINIC, DCD_TRICLINIC @pytest.fixture() def test_u(): - u = mda.Universe(TPR, XTC) + u = mda.Universe(PSF_TRICLINIC, DCD_TRICLINIC) return u -uni = mda.Universe(TPR, XTC) +uni = mda.Universe(PSF_TRICLINIC, DCD_TRICLINIC) ag = uni.atoms[0:10] From f6331d3221f73a9495d3f0c2864535135bc2b3db Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Fri, 31 Jul 2020 16:18:30 +0200 Subject: [PATCH 14/26] doc transformation --- .../MDAnalysis/transformations/__init__.py | 81 +++++++++++-------- package/MDAnalysis/transformations/fit.py | 3 +- package/MDAnalysis/transformations/rotate.py | 3 +- .../MDAnalysis/transformations/translate.py | 3 +- package/MDAnalysis/transformations/wrap.py | 11 ++- 5 files changed, 61 insertions(+), 40 deletions(-) diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index 0b810bac7d7..c994e23ac99 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -26,40 +26,17 @@ Trajectory transformations --- :mod:`MDAnalysis.transformations` ================================================================ -The transformations submodule contains a collection of functions to modify the -trajectory. Coordinate transformations, such as PBC corrections and molecule fitting -are often required for some analyses and visualization, and the functions in this -module allow transformations to be applied on-the-fly. -These transformation functions can be called by the user for any given -timestep of the trajectory, added as a workflow using :meth:`add_transformations` -of the :mod:`~MDAnalysis.coordinates.base` module, or upon Universe creation using -the keyword argument `transformations`. Note that in the two latter cases, the -workflow cannot be changed after being defined. - -In addition to the specific arguments that each transformation can take, they also -contain a wrapped function that takes a `Timestep` object as argument. -So, a transformation can be roughly defined as follows: - -.. code-block:: python - - def transformations(*args,**kwargs): - # do some things - def wrapped(ts): - # apply changes to the Timestep object - return ts - - return wrapped - +The transformations submodule contains a collection of function-like classes to +modify the trajectory. +Coordinate transformations, such as PBC corrections and molecule fitting +are often required for some analyses and visualization, and the functions in +this module allow transformations to be applied on-the-fly. -See `MDAnalysis.transformations.translate` for a simple example. - -To meet the need of serialization of universe, transformations are converted -into classes. They retain the similar API and functionality by implementing the -aforementioned `wrapped(ts)` as a special method `__call__`. For example: +A typical transformation class looks like this: .. code-blocks:: python - class transfomrations(object): + class transfomration(object): def __init__(self, *args, **kwargs): # do some things # save needed args as attributes. @@ -69,14 +46,50 @@ def __call__(self, ts): # apply changes to the Timestep object return ts +See `MDAnalysis.transformations.translate` for a simple example. + +These transformation functions can be called by the user for any given timestep +of the trajectory, added as a workflow using :meth:`add_transformations` +of the :mod:`~MDAnalysis.coordinates.base`, or upon Universe creation using +the keyword argument `transformations`. Note that in the two latter cases, the +workflow cannot be changed after being defined. for example: + +.. code-block:: python + + u = mda.Universe(GRO, XTC) + ts = u.trajectory[0] + trans = transformation(args) + ts = trans(ts) + + # or add as a workflow + u.trajectory.add_transformations(trans) + +Transformations can also be created as a closure/nested function. +In addition to the specific arguments that each transformation can take, they +also contain a wrapped function that takes a `Timestep` object as argument. +So, a closure-style transformation can be roughly defined as follows: + +.. code-block:: python + + def transformation(*args,**kwargs): + # do some things + def wrapped(ts): + # apply changes to the Timestep object + return ts + + return wrapped -Note it does not mean that the old closure/wrapped function implementation will fail. -They can still be used if one does not need serialization, which is a prerequisite for parallel -analysis. +Note, to meet the need of serialization of universe, only transformation class +are used after MDAnlaysis 2.0.0. One can still write functions (closures) as in +MDA 1.x, but that these cannot be serialized and thus will not work with all +forms of parallel analysis. For detailed descriptions about how to write a +closure-style transformation, read the code in MDA 1.x as a reference +or read MDAnalysis UserGuide. .. versionchanged:: 2.0.0 - All transformations are classes now. + Transformations should now be created as classes with a :meth:`__call__` + method instead of being written as a function/closure. """ from .translate import translate, center_in_box diff --git a/package/MDAnalysis/transformations/fit.py b/package/MDAnalysis/transformations/fit.py index 5de4af56168..85a2c93e615 100644 --- a/package/MDAnalysis/transformations/fit.py +++ b/package/MDAnalysis/transformations/fit.py @@ -82,7 +82,8 @@ class fit_translation(object): .. versionchanged:: 2.0.0 - Now it is a class with `__call__`. + The transformation was changed from a function/closure to a class + with ``__call__``. """ def __init__(self, ag, reference, plane=None, weights=None): self.ag = ag diff --git a/package/MDAnalysis/transformations/rotate.py b/package/MDAnalysis/transformations/rotate.py index e7aafb56975..3c4d3612984 100644 --- a/package/MDAnalysis/transformations/rotate.py +++ b/package/MDAnalysis/transformations/rotate.py @@ -105,7 +105,8 @@ class rotateby(object): .. versionchanged:: 2.0.0 - Now it is a class with `__call__`. + The transformation was changed from a function/closure to a class + with ``__call__``. ''' def __init__(self, angle, diff --git a/package/MDAnalysis/transformations/translate.py b/package/MDAnalysis/transformations/translate.py index 599e1bc1e72..1b1be2d92ed 100644 --- a/package/MDAnalysis/transformations/translate.py +++ b/package/MDAnalysis/transformations/translate.py @@ -109,7 +109,8 @@ class center_in_box(object): .. versionchanged:: 2.0.0 - Now it is a class with `__call__`. + The transformation was changed from a function/closure to a class + with ``__call__``. """ def __init__(self, ag, center='geometry', point=None, wrap=False): self.ag = ag diff --git a/package/MDAnalysis/transformations/wrap.py b/package/MDAnalysis/transformations/wrap.py index b9e368f8d7d..10e565fc072 100644 --- a/package/MDAnalysis/transformations/wrap.py +++ b/package/MDAnalysis/transformations/wrap.py @@ -79,10 +79,11 @@ class wrap(object): Returns ------- MDAnalysis.coordinates.base.Timestep - + .. versionchanged:: 2.0.0 - Now it is a class with `__call__`. + The transformation was changed from a function/closure to a class + with ``__call__``. """ def __init__(self, ag, compound='atoms'): self.ag = ag @@ -132,7 +133,11 @@ class unwrap(object): Returns ------- MDAnalysis.coordinates.base.Timestep - + + + .. versionchanged:: 2.0.0 + The transformation was changed from a function/closure to a class + with ``__call__``. """ def __init__(self, ag): self.ag = ag From 1ebcc334c61bb8cf5d07620707224c439579b096 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Sat, 1 Aug 2020 15:12:19 +0200 Subject: [PATCH 15/26] note revise --- package/MDAnalysis/transformations/__init__.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index c994e23ac99..4db168ff40a 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -79,12 +79,14 @@ def wrapped(ts): return wrapped -Note, to meet the need of serialization of universe, only transformation class -are used after MDAnlaysis 2.0.0. One can still write functions (closures) as in -MDA 1.x, but that these cannot be serialized and thus will not work with all -forms of parallel analysis. For detailed descriptions about how to write a -closure-style transformation, read the code in MDA 1.x as a reference -or read MDAnalysis UserGuide. +.. Note:: + Although functions (closures) work as transformations, they are not used in + in MDAnalysis from release 2.0.0 onwards because they cannot be reliably + serialized and thus a :class:`Universe` with such transformations cannot be + used with common parallelization schemes (e.g., ones based on + :mod:`multiprocessing`). + For detailed descriptions about how to write a closure-style transformation, + please refer to MDAnalysis 1.x documentation. .. versionchanged:: 2.0.0 From c4121044aebec6f3ac5dfed1fefd4c6e83303441 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Mon, 3 Aug 2020 12:05:24 +0200 Subject: [PATCH 16/26] doc for transformation --- .../trajectory_transformations.rst | 94 ++++++++++++------- 1 file changed, 62 insertions(+), 32 deletions(-) diff --git a/package/doc/sphinx/source/documentation_pages/trajectory_transformations.rst b/package/doc/sphinx/source/documentation_pages/trajectory_transformations.rst index 8b91466e856..9662f687442 100644 --- a/package/doc/sphinx/source/documentation_pages/trajectory_transformations.rst +++ b/package/doc/sphinx/source/documentation_pages/trajectory_transformations.rst @@ -7,8 +7,8 @@ Trajectory transformations ("on-the-fly" transformations) .. module:: MDAnalysis.transformations -In MDAnalysis, a *transformation* is a function that modifies the data -for the current :class:`Timestep` and returns the +In MDAnalysis, a *transformation* is a function/function-like class +that modifies the data for the current :class:`Timestep` and returns the :class:`Timestep`. For instance, coordinate transformations, such as PBC corrections and molecule fitting are often required for some analyses and visualization. Transformation functions @@ -34,7 +34,7 @@ trajectory is **transformed on-the-fly**, i.e., the data read from the trajectory file will be changed before it is made available in, say, the :attr:`AtomGroup.positions` attribute. -The submodule :mod:`MDAnalysis.transformations` contains a +The submodule :mod:`MDAnalysis.transformations` contains a collection of transformations (see :ref:`transformations-module`) that can be immediately used but one can always write custom transformations (see :ref:`custom-transformations`). @@ -90,40 +90,67 @@ being added. Creating transformations ------------------------ -A *transformation* is a function that takes a +A simple *transformation* can also be a function that takes a :class:`~MDAnalysis.coordinates.base.Timestep` as input, modifies it, and -returns it. - -A simple transformation that takes no other arguments but a :class:`Timestep` +returns it. If it takes no other arguments but a :class:`Timestep` can be defined as the following example: .. code-block:: python def up_by_2(ts): - """ - Translate all coordinates by 2 angstroms up along the Z dimension. - """ - ts.positions = ts.positions + np.array([0, 0, 2], dtype=np.float32) - return ts - + """ + Translate all coordinates by 2 angstroms up along the Z dimension. + """ + ts.positions = ts.positions + np.array([0, 0, 2], dtype=np.float32) + return ts If the transformation requires other arguments besides the :class:`Timestep`, -the transformation takes these arguments, while a wrapped function takes the -:class:`Timestep` object as argument. So, a transformation can be roughly -defined as follows: +the following two methods can be used to create such transformation: + + +.. _custom-transformations-class: + +Creating complex transformation classes +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +It is implemented by defining :func:`__call__` for the transformation class +and can be applied directly to a :class:`Timestep`. +So, a transformation class can be roughly defined as follows: .. code-block:: python - def up_by_x(distance): - """ - Creates a transformation that will translate all coordinates by a given amount along the Z dimension. - """ - def wrapped(ts): - ts.positions = ts.positions + np.array([0, 0, distance], dtype=np.float32) - return ts - return wrapped - - + class up_by_x_class(object): + def __init__(self, distance): + self.distance = distance + + def __call__(self, ts): + ts.positions = ts.positions + np.array([0, 0, self.distance], dtype=np.float32) + return ts + +It is the default construction method in :mod:`MDAnalysis.transformations` +from release 2.0.0 onwards because it can be reliably serialized. +See :class:`MDAnalysis.transformations.translate` for a simple example. + + +.. _custom-transformations-closure: + +Creating complex transformation closure functions +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Transformation can also be a wrapped function takes the :class:`Timestep` object as argument. +So in this case, a transformation function (closure) can be roughly defined as follows: + +.. code-block:: python + + def up_by_x_func(distance): + """ + Creates a transformation that will translate all coordinates by a given amount along the Z dimension. + """ + def wrapped(ts): + ts.positions = ts.positions + np.array([0, 0, distance], dtype=np.float32) + return ts + return wrapped + An alternative to using a wrapped function is using partials from :mod:`functools`. The above function can be written as: @@ -132,15 +159,18 @@ above function can be written as: import functools def up_by_x(ts, distance): - ts.positions = ts.positions + np.array([0, 0, distance], dtype=np.float32) - return ts + ts.positions = ts.positions + np.array([0, 0, distance], dtype=np.float32) + return ts up_by_2 = functools.partial(up_by_x, distance=2) - -See :func:`MDAnalysis.transformations.translate` for a simple -example of such a type of function. - +Although functions (closures) work as transformations, they are not used in +in MDAnalysis from release 2.0.0 onwards because they cannot be reliably +serialized and thus a :class:`Universe` with such transformations cannot be +used with common parallelization schemes (e.g., ones based on +:mod:`multiprocessing`). +For detailed descriptions about how to write a closure-style transformation, +please refer to MDAnalysis 1.x documentation. .. _transformations-module: From 3b1c470e25baadde9ffe44218f3996f5f0a9582d Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Tue, 4 Aug 2020 09:35:26 +0200 Subject: [PATCH 17/26] changelog --- package/CHANGELOG | 2 ++ package/MDAnalysis/transformations/__init__.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/package/CHANGELOG b/package/CHANGELOG index 6440e9cef5a..e3833cb46af 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -82,6 +82,8 @@ Changes * Changes the minimal NumPy version to 1.16.0 (Issue #2827, PR #2831) * Sets the minimal RDKit version for CI to 2020.03.1 (Issue #2827, PR #2831) * Removes deprecated waterdynamics.HydrogenBondLifetimes (PR #2842) + * The transformation was changed from a function/closure to a class with + `__call__` (Issue #2860, PR #2859) Deprecations diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index 4db168ff40a..cfc724aa247 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -36,7 +36,7 @@ .. code-blocks:: python - class transfomration(object): + class transformation(object): def __init__(self, *args, **kwargs): # do some things # save needed args as attributes. From 11f36441dd8b30154fdc060a1501ed020c68c35d Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Sun, 9 Aug 2020 16:45:13 +0200 Subject: [PATCH 18/26] enable universe pickle --- .../MDAnalysisTests/parallelism/test_pickle_transformation.py | 1 - 1 file changed, 1 deletion(-) diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py index b25f187cffb..d576ed0493d 100644 --- a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -67,7 +67,6 @@ def test_transformation_pickle(transformation, test_u): assert_equal(ref_result, result) -@pytest.mark.skip(reason="`Universe` cannot be picked now") def test_add_transformation_pickle(transformation, test_u): test_u.trajectory.add_transformations(transformation) test_u_p = pickle.loads(pickle.dumps(test_u)) From 58f6e4e94e0f77b17de2234454da8901f869ce2f Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Wed, 12 Aug 2020 18:09:34 +0200 Subject: [PATCH 19/26] transformation doc --- package/MDAnalysis/transformations/__init__.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index cfc724aa247..a0dcf3fda74 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -43,7 +43,9 @@ def __init__(self, *args, **kwargs): self.needed_var = args[0] def __call__(self, ts): - # apply changes to the Timestep object + # apply changes to the Timestep, + # or modify an AtomGroup and return Timestep + return ts See `MDAnalysis.transformations.translate` for a simple example. @@ -57,13 +59,13 @@ def __call__(self, ts): .. code-block:: python u = mda.Universe(GRO, XTC) - ts = u.trajectory[0] trans = transformation(args) - ts = trans(ts) - - # or add as a workflow u.trajectory.add_transformations(trans) + # it is equivalent to applying this transforamtion to each Timestep by + ts = u.trajectory[0] + ts_trans = trans(ts) + Transformations can also be created as a closure/nested function. In addition to the specific arguments that each transformation can take, they also contain a wrapped function that takes a `Timestep` object as argument. @@ -73,8 +75,11 @@ def __call__(self, ts): def transformation(*args,**kwargs): # do some things + def wrapped(ts): - # apply changes to the Timestep object + # apply changes to the Timestep, + # or modify an AtomGroup and return Timestep + return ts return wrapped From 3f083ae5bb3f7dbf7618534c28bc3b0c88c1d064 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Thu, 27 Aug 2020 18:59:01 +0200 Subject: [PATCH 20/26] change universe pickle to reduce --- package/MDAnalysis/core/universe.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/package/MDAnalysis/core/universe.py b/package/MDAnalysis/core/universe.py index cbf8041359a..0a94739491f 100644 --- a/package/MDAnalysis/core/universe.py +++ b/package/MDAnalysis/core/universe.py @@ -687,16 +687,23 @@ def __repr__(self): return "".format( n_atoms=len(self.atoms)) - def __getstate__(self): - # Universe's two "legs" of topology and traj both serialise themselves - return self._topology, self._trajectory - def __setstate__(self, args): - self._topology = args[0] - _generate_from_topology(self) + @classmethod + def _unpickle_U(cls, top, traj): + """Special method used by __reduce__ to deserialise a Universe""" + # top is a Topology object at this point, but Universe can handle that + u = cls(top) + # maybe this is None, but that's still cool + u.trajectory = traj - self._trajectory = args[1] + return u + def __reduce__(self): + # __setstate__/__getstate__ will raise an error when Universe has a + # transformation (that has AtomGroup inside). Use __reduce__ instead. + # Universe's two "legs" of topology and traj both serialise themselves. + return (self._unpickle_U, (self._topology, + self._trajectory)) # Properties @property def dimensions(self): From 9ae7f83490591409958be470758c5874bcb49998 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Thu, 27 Aug 2020 18:59:44 +0200 Subject: [PATCH 21/26] remove pure pickle/unpickle tests --- .../parallelism/test_pickle_transformation.py | 135 ++++++++++++++---- 1 file changed, 107 insertions(+), 28 deletions(-) diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py index d576ed0493d..fe0fa85ac3c 100644 --- a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -35,41 +35,120 @@ from MDAnalysisTests.datafiles import PSF_TRICLINIC, DCD_TRICLINIC +@pytest.fixture(params=[ + (PSF_TRICLINIC, DCD_TRICLINIC), +]) +def u(request): + top, traj = request.param + return mda.Universe(top, traj) + + @pytest.fixture() -def test_u(): - u = mda.Universe(PSF_TRICLINIC, DCD_TRICLINIC) - return u +def fit_translation_transformation(u): + ag = u.atoms[0:10] + return fit_translation(ag, ag) -uni = mda.Universe(PSF_TRICLINIC, DCD_TRICLINIC) -ag = uni.atoms[0:10] +@pytest.fixture() +def fit_rot_trans_transformation(u): + ag = u.atoms[0:10] + return fit_rot_trans(ag, ag) -@pytest.fixture(params=[ - (fit_translation, [ag, ag]), - (fit_rot_trans, [ag, ag]), - (PositionAverager, [3]), - (rotateby, [90, [0, 0, 1], [1, 2, 3]]), - (translate, [[1, 2, 3]]), - (center_in_box, [ag]), - (wrap, [uni.atoms]), - (unwrap, [uni.atoms]) -]) -def transformation(request): - transform = request.param[0](*request.param[1]) - return transform +@pytest.fixture() +def PositionAverager_transformation(u): + return PositionAverager(3) + + +@pytest.fixture() +def rotateby_transformation(u): + ag = u.atoms[0:10] + return rotateby(90, [0, 0, 1], [1, 2, 3]) + + +@pytest.fixture() +def translate_transformation(u): + return translate([1, 2, 3]) + + +@pytest.fixture() +def center_in_box_transformation(u): + ag = u.atoms[0:10] + return center_in_box(ag) + + +@pytest.fixture() +def wrap_transformation(u): + ag = u.atoms + return wrap(ag) + + +@pytest.fixture() +def unwrap_transformation(u): + ag = u.atoms + return unwrap(ag) + + +def test_add_fit_translation_pickle(fit_translation_transformation, u): + u.trajectory.add_transformations(fit_translation_transformation) + u_p = pickle.loads(pickle.dumps(u)) + u.trajectory[0] + for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): + assert_equal(u_ts.positions, u_p_ts.positions) + + +def test_add_fit_rot_trans_pickle(fit_rot_trans_transformation, + u): + u.trajectory.add_transformations(fit_rot_trans_transformation) + u_p = pickle.loads(pickle.dumps(u)) + u.trajectory[0] + for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): + assert_equal(u_ts.positions, u_p_ts.positions) + + +def test_add_PositionAverager_pickle(PositionAverager_transformation, u): + u.trajectory.add_transformations(PositionAverager_transformation) + u_p = pickle.loads(pickle.dumps(u)) + u.trajectory[0] + for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): + assert_equal(u_ts.positions, u_p_ts.positions) -def test_transformation_pickle(transformation, test_u): - ref_result = transformation(test_u.trajectory[5]).positions - transformation_p = pickle.loads(pickle.dumps(transformation)) - result = transformation_p(test_u.trajectory[5]).positions - assert_equal(ref_result, result) +def test_add_rotateby_pickle(rotateby_transformation, u): + u.trajectory.add_transformations(rotateby_transformation) + u_p = pickle.loads(pickle.dumps(u)) + u.trajectory[0] + for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): + assert_equal(u_ts.positions, u_p_ts.positions) + + +def test_add_translate_pickle(translate_transformation, u): + u.trajectory.add_transformations(translate_transformation) + u_p = pickle.loads(pickle.dumps(u)) + u.trajectory[0] + for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): + assert_equal(u_ts.positions, u_p_ts.positions) + + +def test_add_center_in_box_pickle(center_in_box_transformation, u): + u.trajectory.add_transformations(center_in_box_transformation) + u_p = pickle.loads(pickle.dumps(u)) + u.trajectory[0] + for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): + assert_equal(u_ts.positions, u_p_ts.positions) + + +def test_add_wrap_pickle(wrap_transformation, u): + u.trajectory.add_transformations(wrap_transformation) + u_p = pickle.loads(pickle.dumps(u)) + u.trajectory[0] + for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): + assert_equal(u_ts.positions, u_p_ts.positions) -def test_add_transformation_pickle(transformation, test_u): - test_u.trajectory.add_transformations(transformation) - test_u_p = pickle.loads(pickle.dumps(test_u)) - test_u.trajectory[0] - for u_ts, u_p_ts in zip(test_u.trajectory[:5], test_u_p.trajectory[:5]): +def test_add_unwrap_pickle(unwrap_transformation, u): + u.trajectory.add_transformations(unwrap_transformation) + u_p = pickle.loads(pickle.dumps(u)) + u.trajectory[0] + for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): assert_equal(u_ts.positions, u_p_ts.positions) From ae8a2a2bcd7aabd78743a050aa123d379a639f58 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Thu, 27 Aug 2020 19:43:07 +0200 Subject: [PATCH 22/26] pep --- package/MDAnalysis/core/universe.py | 7 +++---- .../parallelism/test_pickle_transformation.py | 3 +-- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/package/MDAnalysis/core/universe.py b/package/MDAnalysis/core/universe.py index 0a94739491f..8e11e25f8c3 100644 --- a/package/MDAnalysis/core/universe.py +++ b/package/MDAnalysis/core/universe.py @@ -687,13 +687,11 @@ def __repr__(self): return "".format( n_atoms=len(self.atoms)) - @classmethod def _unpickle_U(cls, top, traj): """Special method used by __reduce__ to deserialise a Universe""" - # top is a Topology object at this point, but Universe can handle that + # top is a Topology obj at this point, but Universe can handle that. u = cls(top) - # maybe this is None, but that's still cool u.trajectory = traj return u @@ -701,9 +699,10 @@ def _unpickle_U(cls, top, traj): def __reduce__(self): # __setstate__/__getstate__ will raise an error when Universe has a # transformation (that has AtomGroup inside). Use __reduce__ instead. - # Universe's two "legs" of topology and traj both serialise themselves. + # Universe's two "legs" of top and traj both serialise themselves. return (self._unpickle_U, (self._topology, self._trajectory)) + # Properties @property def dimensions(self): diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py index fe0fa85ac3c..f6c4d506ede 100644 --- a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -62,7 +62,6 @@ def PositionAverager_transformation(u): @pytest.fixture() def rotateby_transformation(u): - ag = u.atoms[0:10] return rotateby(90, [0, 0, 1], [1, 2, 3]) @@ -98,7 +97,7 @@ def test_add_fit_translation_pickle(fit_translation_transformation, u): def test_add_fit_rot_trans_pickle(fit_rot_trans_transformation, - u): + u): u.trajectory.add_transformations(fit_rot_trans_transformation) u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] From 499efc10a1314aa1549c6b06fbee86a12250770b Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Fri, 28 Aug 2020 10:33:13 +0200 Subject: [PATCH 23/26] example --- .../MDAnalysis/transformations/__init__.py | 31 +++++++++++++++++-- .../parallelism/test_pickle_transformation.py | 18 +++++------ 2 files changed, 38 insertions(+), 11 deletions(-) diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index a0dcf3fda74..b99696625f4 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -32,7 +32,9 @@ are often required for some analyses and visualization, and the functions in this module allow transformations to be applied on-the-fly. -A typical transformation class looks like this: +A typical transformation class looks like this (note that we keep its name +lowercase because we will treat it as a function, thanks to the ``__call__`` +method): .. code-blocks:: python @@ -48,7 +50,32 @@ def __call__(self, ts): return ts -See `MDAnalysis.transformations.translate` for a simple example. +As a concrete example we will write a transformation that rotates a group of +atoms around the z-axis through the center of geometry by a fixed increment +for every time step. We will use +:meth:`MDAnalysis.core.groups.AtomGroup.rotateby` +and simply increment the rotation angle every time the +transformation is called:: + + class spin_atoms(object): + def __init__(self, atoms, dphi): + """Rotate atoms by dphi degrees for every ts (around the z axis)""" + self.atoms = atoms + self.dphi = dphi + self.axis = np.array([0, 0, 1]) + self.phi = 0 + + def __call__(self, ts): + self.atoms.rotateby(self.phi, self.axis) + self.phi += self.dphi + return ts + +This transformation can be used as :: + + u = mda.Universe(PSF, DCD) + u.trajectory.add_transformations(spin_atoms(u.select_atoms("protein"), 1.0)) + +Alos see :mod:`MDAnalysis.transformations.translate` for a simple example. These transformation functions can be called by the user for any given timestep of the trajectory, added as a workflow using :meth:`add_transformations` diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py index f6c4d506ede..d663c75cff5 100644 --- a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -22,7 +22,7 @@ # import pytest import pickle -from numpy.testing import assert_equal +from numpy.testing import assert_almost_equal import MDAnalysis as mda @@ -93,7 +93,7 @@ def test_add_fit_translation_pickle(fit_translation_transformation, u): u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): - assert_equal(u_ts.positions, u_p_ts.positions) + assert_almost_equal(u_ts.positions, u_p_ts.positions) def test_add_fit_rot_trans_pickle(fit_rot_trans_transformation, @@ -102,7 +102,7 @@ def test_add_fit_rot_trans_pickle(fit_rot_trans_transformation, u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): - assert_equal(u_ts.positions, u_p_ts.positions) + assert_almost_equal(u_ts.positions, u_p_ts.positions) def test_add_PositionAverager_pickle(PositionAverager_transformation, u): @@ -110,7 +110,7 @@ def test_add_PositionAverager_pickle(PositionAverager_transformation, u): u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): - assert_equal(u_ts.positions, u_p_ts.positions) + assert_almost_equal(u_ts.positions, u_p_ts.positions) def test_add_rotateby_pickle(rotateby_transformation, u): @@ -118,7 +118,7 @@ def test_add_rotateby_pickle(rotateby_transformation, u): u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): - assert_equal(u_ts.positions, u_p_ts.positions) + assert_almost_equal(u_ts.positions, u_p_ts.positions) def test_add_translate_pickle(translate_transformation, u): @@ -126,7 +126,7 @@ def test_add_translate_pickle(translate_transformation, u): u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): - assert_equal(u_ts.positions, u_p_ts.positions) + assert_almost_equal(u_ts.positions, u_p_ts.positions) def test_add_center_in_box_pickle(center_in_box_transformation, u): @@ -134,7 +134,7 @@ def test_add_center_in_box_pickle(center_in_box_transformation, u): u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): - assert_equal(u_ts.positions, u_p_ts.positions) + assert_almost_equal(u_ts.positions, u_p_ts.positions) def test_add_wrap_pickle(wrap_transformation, u): @@ -142,7 +142,7 @@ def test_add_wrap_pickle(wrap_transformation, u): u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): - assert_equal(u_ts.positions, u_p_ts.positions) + assert_almost_equal(u_ts.positions, u_p_ts.positions) def test_add_unwrap_pickle(unwrap_transformation, u): @@ -150,4 +150,4 @@ def test_add_unwrap_pickle(unwrap_transformation, u): u_p = pickle.loads(pickle.dumps(u)) u.trajectory[0] for u_ts, u_p_ts in zip(u.trajectory[:5], u_p.trajectory[:5]): - assert_equal(u_ts.positions, u_p_ts.positions) + assert_almost_equal(u_ts.positions, u_p_ts.positions) From 2793ba25295d31b04f3f9a46c118da906d774e10 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Fri, 28 Aug 2020 11:15:17 +0200 Subject: [PATCH 24/26] doc --- package/MDAnalysis/transformations/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index b99696625f4..55e74b3941f 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -22,7 +22,7 @@ # -"""\ +""" Trajectory transformations --- :mod:`MDAnalysis.transformations` ================================================================ @@ -55,11 +55,11 @@ def __call__(self, ts): for every time step. We will use :meth:`MDAnalysis.core.groups.AtomGroup.rotateby` and simply increment the rotation angle every time the -transformation is called:: +transformation is called :: class spin_atoms(object): def __init__(self, atoms, dphi): - """Rotate atoms by dphi degrees for every ts (around the z axis)""" + # Rotate atoms by dphi degrees for every ts around the z axis self.atoms = atoms self.dphi = dphi self.axis = np.array([0, 0, 1]) From 1080678a25b8f0a11f3972b0a1eaa8a200d03b1a Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Fri, 4 Sep 2020 09:24:55 +0200 Subject: [PATCH 25/26] Update package/MDAnalysis/transformations/__init__.py Co-authored-by: Oliver Beckstein --- package/MDAnalysis/transformations/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index 55e74b3941f..45a0488baa9 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -75,7 +75,7 @@ def __call__(self, ts): u = mda.Universe(PSF, DCD) u.trajectory.add_transformations(spin_atoms(u.select_atoms("protein"), 1.0)) -Alos see :mod:`MDAnalysis.transformations.translate` for a simple example. +Also see :mod:`MDAnalysis.transformations.translate` for a simple example. These transformation functions can be called by the user for any given timestep of the trajectory, added as a workflow using :meth:`add_transformations` From 06552750027c0c23e6c393c52137ceec8d29eea1 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhuang Date: Sun, 6 Sep 2020 19:15:12 +0200 Subject: [PATCH 26/26] change snippet --- package/MDAnalysis/transformations/__init__.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index 45a0488baa9..91083fff18b 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -63,11 +63,10 @@ def __init__(self, atoms, dphi): self.atoms = atoms self.dphi = dphi self.axis = np.array([0, 0, 1]) - self.phi = 0 def __call__(self, ts): - self.atoms.rotateby(self.phi, self.axis) - self.phi += self.dphi + phi = self.dphi * ts.frame + self.atoms.rotateby(phi, self.axis) return ts This transformation can be used as ::