Merge pull request #1937 from davidercruz/rotation-transform

Rotation transformation

package/CHANGELOG

@@ -19,6 +19,8 @@ The rules for this file:
   * 0.18.1
 
 Enhancements
+
+  * Added a rotation coordinate transformation (PR #1937)
   * Added a box centering trajectory transformation (PR #1946)
   * Added a on-the-fly trajectory transformations API and a coordinate translation
   	function (Issue #786)

package/MDAnalysis/transformations/__init__.py

@@ -56,9 +56,9 @@ def wrapped(ts):
     
     - translate: translate the coordinates of a given trajectory frame by a given vector.
     - center_in_box: translate the coordinates of a given trajectory frame so that a given
-        AtomGroup is centered in the unit cell
-
-    
+      AtomGroup is centered in the unit cell
+    - rotateby: rotates the coordinates by a given angle arround an axis formed by a direction 
+      and a point    
 
 Examples
 --------
@@ -93,6 +93,8 @@ def wrapped(ts):
 from __future__ import absolute_import
 
 from .translate import translate, center_in_box
+from .rotate import rotateby
+
 
 
 

package/MDAnalysis/transformations/rotate.py

@@ -0,0 +1,128 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
+#
+# 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.
+#
+# 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
+#
+
+"""\
+Rotate trajectory --- :mod:`MDAnalysis.transformations.translate`
+=================================================================
+
+Rotates the coordinates by a given angle arround an axis formed by a direction and a
+point 
+    
+"""
+from __future__ import absolute_import
+
+import math
+import numpy as np
+from functools import partial
+
+from ..lib.transformations import rotation_matrix
+
+def rotateby(angle, direction, point=None, center="geometry", wrap=False, ag=None):
+    '''
+    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
+    of geometry or the center of mass of a user defined AtomGroup, or a list defining custom
+    coordinates. 
+    e.g. rotate the coordinates by 90 degrees on a x axis centered on a given atom group:
+    
+    .. code-block:: python
+    
+        ts = u.trajectory.ts
+        angle = 90
+        ag = u.atoms()
+        rotated = MDAnalysis.transformations.rotate(angle, ag)(ts)
+    
+    e.g. rotate the coordinates by a custom axis:
+    
+    .. code-block:: python
+
+        ts = u.trajectory.ts
+        angle = 90
+        p = [1,2,3]
+        d = [0,0,1]
+        rotated = MDAnalysis.transformations.rotate(angle, point=point, direction=d)(ts) 
+    
+    Parameters
+    ----------
+    angle: float
+        rotation angle in degrees
+    direction: array-like
+        vector that will define the direction of a custom axis of rotation from the
+        provided point.
+    ag: AtomGroup, optional
+        use this to define the center of mass or geometry as the point from where the
+        rotation axis will be defined
+    center: str, optional
+        used to choose the method of centering on the given atom group. Can be 'geometry'
+        or 'mass'
+    wrap: bool, optional
+        If `True`, all the atoms from the given AtomGroup will be moved to the unit cell
+        before calculating the center of mass or geometry. Default is `False`, no changes
+        to the atom coordinates are done before calculating the center of the AtomGroup. 
+    point: array-like, optional
+        list of the coordinates of the point from where a custom axis of rotation will
+        be defined. 
+
+    Returns
+    -------
+    :class:`~MDAnalysis.coordinates.base.Timestep` object
+    
+    Warning
+    -------
+    Wrapping/unwrapping the trajectory or performing PBC corrections may not be possible 
+    after rotating the trajectory.
+
+    '''
+    pbc_arg = wrap
+    angle = np.deg2rad(angle)
+    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))
+    elif ag:
+        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':
+                raise AttributeError('{} is not an AtomGroup object with masses'.format(ag))
+            else:
+                raise ValueError('{} is not an AtomGroup object'.format(ag))
+    else:
+        raise ValueError('A point or an AtomGroup must be specified')
+
+    def wrapped(ts):
+        if point is None:
+            position = center_method()
+        else:
+            position = point
+        rotation = rotation_matrix(angle, direction, position)[:3, :3]
+        ts.positions= np.dot(ts.positions, rotation)
+
+        return ts
+
+    return wrapped
+

testsuite/MDAnalysisTests/transformations/test_rotate.py

@@ -0,0 +1,204 @@
+# -*- 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.
+#
+# 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
+#
+
+from __future__ import absolute_import
+
+import numpy as np
+import pytest
+from numpy.testing import assert_array_almost_equal
+
+import MDAnalysis as mda
+from MDAnalysis.transformations import rotateby
+from MDAnalysis.lib.transformations import rotation_matrix
+from MDAnalysisTests import make_Universe
+
+@pytest.fixture()
+def rotate_universes():
+    # create the Universe objects for the tests
+    reference = make_Universe(trajectory=True)
+    transformed = make_Universe(['masses'], trajectory=True)
+    transformed.trajectory.ts.dimensions = np.array([372., 373., 374., 90, 90, 90])
+    return reference, transformed
+
+def test_rotation_matrix():
+    # test if the rotation_matrix function is working properly
+    # simple angle and unit vector on origin
+    angle = 180
+    vector = [0, 0, 1]
+    pos = [0, 0, 0]
+    ref_matrix = np.asarray([[-1, 0, 0],
+                             [0, -1, 0],
+                             [0, 0, 1]], np.float64)
+    matrix = rotation_matrix(np.deg2rad(angle), vector, pos)[:3, :3]
+    assert_array_almost_equal(matrix, ref_matrix, decimal=6)
+    # another angle in a custom axis
+    angle = 60
+    vector = [1, 2, 3]
+    pos = [1, 2, 3]
+    ref_matrix = np.asarray([[ 0.53571429, -0.6229365 ,  0.57005291],
+                             [ 0.76579365,  0.64285714, -0.01716931],
+                             [-0.35576719,  0.44574074,  0.82142857]], np.float64)
+    matrix = rotation_matrix(np.deg2rad(angle), vector, pos)[:3, :3]
+    assert_array_almost_equal(matrix, ref_matrix, decimal=6)
+
+
+def test_rotateby_custom_position(rotate_universes):
+    # what happens when we use a custom point for the axis of rotation?
+    ref_u = rotate_universes[0]
+    trans_u = rotate_universes[1]
+    trans = trans_u.trajectory.ts
+    ref = ref_u.trajectory.ts
+    vector = [1,0,0]
+    pos = [0,0,0]
+    angle = 90
+    matrix = rotation_matrix(np.deg2rad(angle), vector, pos)[:3, :3]
+    ref.positions = np.dot(ref.positions, matrix)
+    transformed = rotateby(angle, vector, point=pos)(trans)
+    assert_array_almost_equal(transformed.positions, ref.positions, decimal=6)
+
+def test_rotateby_atomgroup_cog_nopbc(rotate_universes):
+    # what happens when we rotate arround the center of geometry of a residue
+    # without pbc?
+    ref_u = rotate_universes[0]
+    trans_u = rotate_universes[1]
+    trans = trans_u.trajectory.ts
+    ref = ref_u.trajectory.ts
+    center_pos = [6,7,8]
+    vector = [1,0,0]
+    angle = 90
+    matrix = rotation_matrix(np.deg2rad(angle), vector, center_pos)[:3, :3]
+    ref.positions = np.dot(ref.positions, matrix)
+    selection = trans_u.residues[0].atoms
+    transformed = rotateby(angle, vector, ag=selection, center='geometry')(trans) 
+    assert_array_almost_equal(transformed.positions, ref.positions, decimal=6)
+
+def test_rotateby_atomgroup_com_nopbc(rotate_universes):
+    # what happens when we rotate arround the center of mass of a residue
+    # without pbc?
+    ref_u = rotate_universes[0]
+    trans_u = rotate_universes[1]
+    trans = trans_u.trajectory.ts
+    ref = ref_u.trajectory.ts
+    vector = [1,0,0]
+    angle = 90
+    selection = trans_u.residues[0].atoms
+    center_pos = selection.center_of_mass()
+    matrix = rotation_matrix(np.deg2rad(angle), vector, center_pos)[:3, :3]
+    ref.positions = np.dot(ref.positions, matrix)
+    transformed = rotateby(angle, vector, ag=selection, center='mass')(trans) 
+    assert_array_almost_equal(transformed.positions, ref.positions, decimal=6)
+
+def test_rotateby_atomgroup_cog_pbc(rotate_universes):
+    # what happens when we rotate arround the center of geometry of a residue
+    # with pbc?
+    ref_u = rotate_universes[0]
+    trans_u = rotate_universes[1]
+    trans = trans_u.trajectory.ts
+    ref = ref_u.trajectory.ts
+    vector = [1,0,0]
+    angle = 90
+    selection = trans_u.residues[0].atoms
+    center_pos = selection.center_of_geometry(pbc=True)
+    matrix = rotation_matrix(np.deg2rad(angle), vector, center_pos)[:3, :3]
+    ref.positions = np.dot(ref.positions, matrix)
+    transformed = rotateby(angle, vector, ag=selection, center='geometry', wrap=True)(trans) 
+    assert_array_almost_equal(transformed.positions, ref.positions, decimal=6)
+
+def test_rotateby_atomgroup_com_pbc(rotate_universes):
+    # what happens when we rotate arround the center of mass of a residue
+    # with pbc?
+    ref_u = rotate_universes[0]
+    trans_u = rotate_universes[1]
+    trans = trans_u.trajectory.ts
+    ref = ref_u.trajectory.ts
+    vector = [1,0,0]
+    angle = 90
+    selection = trans_u.residues[0].atoms
+    center_pos = selection.center_of_mass(pbc=True)
+    matrix = rotation_matrix(np.deg2rad(angle), vector, center_pos)[:3, :3]
+    ref.positions = np.dot(ref.positions, matrix)
+    transformed = rotateby(angle, vector, ag=selection, center='mass', wrap=True)(trans) 
+    assert_array_almost_equal(transformed.positions, ref.positions, decimal=6)
+
+def test_rotateby_bad_ag(rotate_universes):
+    # this universe as a box size zero
+    ts = rotate_universes[0].trajectory.ts
+    ag = rotate_universes[0].residues[0].atoms
+    # what happens if something other than an AtomGroup is given?
+    angle = 90
+    vector = [0, 0, 1]
+    bad_ag = 1
+    with pytest.raises(ValueError): 
+        rotateby(angle, vector, ag = bad_ag)(ts)
+
+def test_rotateby_bad_position(rotate_universes):
+    # this universe as a box size zero
+    ts = rotate_universes[0].trajectory.ts
+    # what if the box is in the wrong format?
+    angle = 90
+    vector = [0, 0, 1]
+    bad_position = [1]
+    with pytest.raises(ValueError): 
+        rotateby(angle, vector, point=bad_position)(ts)
+
+def test_rotateby_bad_pbc(rotate_universes):    
+    # this universe as a box size zero
+    ts = rotate_universes[0].trajectory.ts
+    ag = rotate_universes[0].residues[0].atoms
+    # is pbc passed to the center methods?
+    # if yes it should raise an exception for boxes that are zero in size
+    angle = 90
+    vector = [0, 0, 1]
+    with pytest.raises(ValueError): 
+        rotateby(angle, vector, ag = ag, wrap=True)(ts)
+
+def test_rotateby_bad_center(rotate_universes):
+    # this universe as a box size zero
+    ts = rotate_universes[0].trajectory.ts
+    ag = rotate_universes[0].residues[0].atoms
+    # what if a wrong center type name is passed?
+    angle = 90
+    vector = [0, 0, 1]
+    bad_center = " "
+    with pytest.raises(ValueError): 
+        rotateby(angle, vector, ag = ag, center=bad_center)(ts)
+
+def test_rotateby_no_masses(rotate_universes):   
+    # this universe as a box size zero
+    ts = rotate_universes[0].trajectory.ts
+    ag = rotate_universes[0].residues[0].atoms
+    # if the universe has no masses and `mass` is passed as the center arg
+    angle = 90
+    vector = [0, 0, 1]
+    bad_center = "mass"
+    with pytest.raises(AttributeError): 
+        rotateby(angle, vector, ag = ag, center=bad_center)(ts)
+
+def test_rotateby_no_args(rotate_universes):
+    # this universe as a box size zero
+    ts = rotate_universes[0].trajectory.ts
+    angle = 90
+    vector = [0, 0, 1]
+    # if no point or AtomGroup are passed to the function
+    # it should raise a ValueError
+    with pytest.raises(ValueError): 
+        rotateby(angle, vector)(ts)