diff --git a/package/CHANGELOG b/package/CHANGELOG index 3cf95e0f949..90d2dc6922d 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -111,7 +111,9 @@ Changes * Sets the minimal RDKit version for CI to 2020.03.1 (Issue #2827, PR #2831) * Removes deprecated waterdynamics.HydrogenBondLifetimes (PR #2842) * Make NeighborSearch return empty atomgroup, residue, segments instead of list (Issue #2892, PR #2907) - * Updated Universe creation function signatures to named arguments (Issue #2921) + * Updated Universe creation function signatures to named arguments (Issue #2921) + * The transformation was changed from a function/closure to a class with + `__call__` (Issue #2860, PR #2859) Deprecations diff --git a/package/MDAnalysis/core/universe.py b/package/MDAnalysis/core/universe.py index bbd63d29235..e4596ee6aba 100644 --- a/package/MDAnalysis/core/universe.py +++ b/package/MDAnalysis/core/universe.py @@ -666,15 +666,21 @@ 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 + @classmethod + def _unpickle_U(cls, top, traj): + """Special method used by __reduce__ to deserialise a Universe""" + # top is a Topology obj at this point, but Universe can handle that. + u = cls(top) + u.trajectory = traj - def __setstate__(self, args): - self._topology = args[0] - _generate_from_topology(self) + return u - self._trajectory = args[1] + 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 top and traj both serialise themselves. + return (self._unpickle_U, (self._topology, + self._trajectory)) # Properties @property diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index 13ce68dfbd4..91083fff18b 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -22,38 +22,107 @@ # -"""\ +""" 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 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. + +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 + + class transformation(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, + # or modify an AtomGroup and return Timestep + + return ts + +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]) + + def __call__(self, ts): + phi = self.dphi * ts.frame + self.atoms.rotateby(phi, self.axis) + 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)) + +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` +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. +workflow cannot be changed after being defined. for example: + +.. code-block:: python -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: + u = mda.Universe(GRO, XTC) + trans = transformation(args) + 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. +So, a closure-style transformation can be roughly defined as follows: .. code-block:: python - def transformations(*args,**kwargs): + 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 - -See `MDAnalysis.transformations.translate` for a simple example. +.. 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 + 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 084868a0e55..85a2c93e615 100644 --- a/package/MDAnalysis/transformations/fit.py +++ b/package/MDAnalysis/transformations/fit.py @@ -27,21 +27,18 @@ 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 -from functools import partial from ..analysis import align -from ..lib.util import get_weights 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 @@ -49,8 +46,8 @@ def fit_translation(ag, reference, plane=None, weights=None): 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 @@ -67,11 +64,12 @@ def fit_translation(ag, reference, plane=None, weights=None): :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 @@ -81,39 +79,56 @@ def fit_translation(ag, reference, plane=None, weights=None): Returns ------- MDAnalysis.coordinates.base.Timestep - """ - if plane is not None: - axes = {'yz' : 0, 'xz' : 1, 'xy' : 2} + + .. versionchanged:: 2.0.0 + 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 + self.reference = reference + self.plane = plane + self.weights = weights + + if self.plane is not None: + axes = {'yz': 0, 'xz': 1, 'xy': 2} + try: + self.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 = np.asarray(mobile.atoms.center(weights), np.float32) - vector = ref_com - mobile_com - if plane is not None: - vector[plane] = 0 + if self.ag.atoms.n_residues != self.reference.atoms.n_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" + f"Universe/AtomGroup" + ) + raise AttributeError(errmsg) from None + 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) + + 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[self.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,41 +175,59 @@ 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 + + if self.plane is not None: + axes = {'yz': 0, 'xz': 1, 'xy': 2} + try: + self.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) - mobile_coordinates = mobile.atoms.positions - mobile_com - rotation, dump = align.rotation_matrix(mobile_coordinates, ref_coordinates, weights=weights) - vector = ref_com - if plane is not None: - matrix = np.r_[rotation, np.zeros(3).reshape(1,3)] + if self.ag.atoms.n_residues != self.reference.atoms.n_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 " + f"Universe/AtomGroup" + ) + raise AttributeError(errmsg) from None + 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) + self.ref_coordinates = self.ref.atoms.positions - self.ref_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, + 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) + 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] - vector[plane] = mobile_com[plane] + 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[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 - - return wrapped diff --git a/package/MDAnalysis/transformations/positionaveraging.py b/package/MDAnalysis/transformations/positionaveraging.py index 179b98163f1..930a77774b3 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,52 +135,59 @@ 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 + # 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)): + 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) - + 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 57871292ff3..afbba495a61 100644 --- a/package/MDAnalysis/transformations/rotate.py +++ b/package/MDAnalysis/transformations/rotate.py @@ -25,20 +25,20 @@ 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 +.. autoclass:: rotateby """ -import math import numpy as np from functools import partial 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 @@ -107,48 +107,69 @@ def rotateby(angle, direction, point=None, ag=None, weights=None, wrap=False): Wrapping/unwrapping the trajectory or performing PBC corrections may not be possible after rotating the trajectory. + + .. versionchanged:: 2.0.0 + The transformation was changed from a function/closure to a class + with ``__call__``. ''' - 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 + + self.angle = np.deg2rad(self.angle) try: - atoms = ag.atoms - except AttributeError: - raise ValueError(f'{ag} is not an AtomGroup object') from None - else: + 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(self.direction)) + self.direction = self.direction.reshape(3, ) + except ValueError: + raise ValueError(f'{self.direction} is not a valid direction') \ + from None + if self.point is not None: + 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: - 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: - position = center_method() + self.atoms = self.ag.atoms + except AttributeError: + raise ValueError(f'{self.ag} is not an AtomGroup object') \ + from None + else: + try: + 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 + 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 = 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 - - return wrapped - diff --git a/package/MDAnalysis/transformations/translate.py b/package/MDAnalysis/transformations/translate.py index 6a4d28113be..1b1be2d92ed 100644 --- a/package/MDAnalysis/transformations/translate.py +++ b/package/MDAnalysis/transformations/translate.py @@ -30,18 +30,17 @@ 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 """ import numpy as np from functools import partial -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 +59,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 wrapped(ts): - ts.positions += vector + if len(self.vector) > 2: + self.vector = np.float32(self.vector) + else: + raise ValueError("{} vector is too short".format(self.vector)) + def __call__(self, ts): + ts.positions += self.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` @@ -108,39 +107,48 @@ 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: + .. versionchanged:: 2.0.0 + 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 + self.center = center + self.point = point + self.wrap = wrap + + pbc_arg = self.wrap + if self.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': + self.center_method = partial(self.ag.center_of_geometry, + pbc=pbc_arg) + elif self.center == 'mass': + self.center_method = partial(self.ag.center_of_mass, + pbc=pbc_arg) + else: + 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 + + 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 return ts - - return wrapped - diff --git a/package/MDAnalysis/transformations/wrap.py b/package/MDAnalysis/transformations/wrap.py index 99b60f943a1..10e565fc072 100644 --- a/package/MDAnalysis/transformations/wrap.py +++ b/package/MDAnalysis/transformations/wrap.py @@ -28,16 +28,16 @@ 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 """ 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. :: @@ -79,18 +79,22 @@ def wrap(ag, compound='atoms'): Returns ------- MDAnalysis.coordinates.base.Timestep - + + + .. versionchanged:: 2.0.0 + The transformation was changed from a function/closure to a class + with ``__call__``. """ - - def wrapped(ts): - ag.wrap(compound=compound) - + def __init__(self, ag, compound='atoms'): + self.ag = ag + self.compound = compound + + def __call__(self, ts): + self.ag.wrap(compound=self.compound) return ts - - return wrapped -def unwrap(ag): +class unwrap(object): """ Move all atoms in an AtomGroup so that bonds don't split over images @@ -129,19 +133,21 @@ def unwrap(ag): Returns ------- MDAnalysis.coordinates.base.Timestep - + + + .. versionchanged:: 2.0.0 + The transformation was changed from a function/closure to a class + with ``__call__``. """ - - 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 + + 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 - - return wrapped - 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: diff --git a/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py new file mode 100644 index 00000000000..d663c75cff5 --- /dev/null +++ b/testsuite/MDAnalysisTests/parallelism/test_pickle_transformation.py @@ -0,0 +1,153 @@ +# -*- 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 pytest +import pickle +from numpy.testing import assert_almost_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 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 fit_translation_transformation(u): + ag = u.atoms[0:10] + return fit_translation(ag, ag) + + +@pytest.fixture() +def fit_rot_trans_transformation(u): + ag = u.atoms[0:10] + return fit_rot_trans(ag, ag) + + +@pytest.fixture() +def PositionAverager_transformation(u): + return PositionAverager(3) + + +@pytest.fixture() +def rotateby_transformation(u): + 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_almost_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_almost_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_almost_equal(u_ts.positions, u_p_ts.positions) + + +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_almost_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_almost_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_almost_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_almost_equal(u_ts.positions, u_p_ts.positions) + + +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_almost_equal(u_ts.positions, u_p_ts.positions)