corrected rotateby transformation calculation

package/MDAnalysis/transformations/rotate.py

@@ -21,12 +21,14 @@
 #
 
 """\
-Rotate trajectory --- :mod:`MDAnalysis.transformations.translate`
-=================================================================
+Trajectory rotation --- :mod:`MDAnalysis.transformations.rotate`
+================================================================
 
 Rotates the coordinates by a given angle arround an axis formed by a direction and a
 point 
-
+
+.. autofunction:: rotateby
+
 """
 from __future__ import absolute_import
 
@@ -35,21 +37,28 @@
 from functools import partial
 
 from ..lib.transformations import rotation_matrix
+from ..lib.util import get_weights
 
-def rotateby(angle, direction, point=None, center="geometry", wrap=False, ag=None):
+def rotateby(angle, direction, point=None, ag=None, weights=None, wrap=False):
     '''
     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:
+    of geometry or the center of mass of a user defined AtomGroup, or an array defining 
+    custom coordinates.
+
+    Examples
+    --------
+
+    e.g. rotate the coordinates by 90 degrees on a axis formed by the [0,0,1] vector and 
+    the center of geometry of a given AtomGroup:
 
     .. code-block:: python
 
         ts = u.trajectory.ts
         angle = 90
         ag = u.atoms()
-        rotated = MDAnalysis.transformations.rotate(angle, ag)(ts)
+        d = [0,0,1]
+        rotated = MDAnalysis.transformations.rotate(angle, direction=d, ag=ag)(ts)
 
     e.g. rotate the coordinates by a custom axis:
 
@@ -59,58 +68,68 @@ def rotateby(angle, direction, point=None, center="geometry", wrap=False, ag=Non
         angle = 90
         p = [1,2,3]
         d = [0,0,1]
-        rotated = MDAnalysis.transformations.rotate(angle, point=point, direction=d)(ts) 
+        rotated = MDAnalysis.transformations.rotate(angle, direction=d, point=point)(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.
+        provided point. Expected shapes are (3, ) or (1, 3). 
     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'
+        use the weighted center of an AtomGroup as the point from where the rotation axis
+        will be defined. If no AtomGroup is given, the `point` argument becomes mandatory
+    point: array-like, optional
+        list of the coordinates of the point from where a custom axis of rotation will
+        be defined. Expected shapes are (3, ) or (1, 3). If no point is given, the
+        `ag` argument becomes mandatory.
+    weights: {"mass", ``None``} or array_like, optional
+        define the weights of the atoms when calculating the center of the AtomGroup.
+        With ``"mass"`` uses masses as weights; with ``None`` weigh each atom equally.
+        If a float array of the same length as `ag` is provided, use each element of
+        the `array_like` as a weight for the corresponding atom in `ag`. Default is 
+        None.
     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
+    MDAnalysis.coordinates.base.Timestep
 
     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:
+    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('{} is not a valid direction'.format(direction))
+    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:
         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))
+            atoms = ag.atoms
         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))
+            raise ValueError('{} is not an AtomGroup object'.format(ag))
+        else:
+            try:
+                weights = get_weights(atoms, weights=weights)
+            except (ValueError, TypeError):
+                raise TypeError("weights must be {'mass', None} or an iterable of the "
+                                "same size as the atomgroup.")
+        center_method = partial(atoms.center, weights, pbc=wrap)
     else:
         raise ValueError('A point or an AtomGroup must be specified')
 
@@ -119,8 +138,11 @@ def wrapped(ts):
             position = center_method()
         else:
             position = point
-        rotation = rotation_matrix(angle, direction, position)[:3, :3]
+        matrix = rotation_matrix(angle, direction, position)
+        rotation = matrix[:3, :3].T
+        translation = matrix[:3, 3]
         ts.positions= np.dot(ts.positions, rotation)
+        ts.positions += translation
 
         return ts
 

package/doc/sphinx/source/documentation_pages/trajectory_transformations.rst

@@ -96,3 +96,4 @@ e.g. giving a workflow as a keyword argument when defining the universe:
 .. toctree::
 
    ./transformations/translate
+   ./transformations/rotate

package/doc/sphinx/source/documentation_pages/transformations/rotate.rst

@@ -0,0 +1 @@
+.. automodule:: MDAnalysis.transformations.rotate