Box centering transformation

package/CHANGELOG

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

package/MDAnalysis/coordinates/dummy.py

@@ -53,6 +53,7 @@ def __init__(self, n_atoms=None, velocities=False, forces=False):
         self.filename = 'DummyReader'
         self.n_frames = 1
         self._read_first_frame(velocities, forces)
+        self._transformations = []
 
     def _read_first_frame(self, velocities=False, forces=False):
         ts = self.ts = self._Timestep(self.n_atoms, positions=True,

package/MDAnalysis/transformations/__init__.py

@@ -50,7 +50,14 @@ def wrapped(ts):
             return wrapped
 
 
-See :func:`MDAnalysis.transformations.translate` for a simple example.    
+See `MDAnalysis.transformations.translate` for a simple example.
+
+Currently implemented transformations are:
+
+    - 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
+
 
 
 Examples
@@ -85,6 +92,7 @@ def wrapped(ts):
 
 from __future__ import absolute_import
 
-from .translate import translate
+from .translate import translate, center_in_box
+
 
 

package/MDAnalysis/transformations/translate.py

@@ -21,50 +21,121 @@
 #
 
 """\
-Translations --- :mod:`MDAnalysis.transformations.translate`
-============================================================
+Trajectory translation --- :mod:`MDAnalysis.transformations.translate`
+======================================================================
 
 Translate the coordinates of a given trajectory by a given vector.
-
-.. autofunction:: translate   
-    
+The vector can either be user defined, using the function :func:`translate`
+or defined by centering an AtomGroup in the unit cell using the function
+:func:`center_in_box`
 
 """
-from __future__ import absolute_import
+from __future__ import absolute_import, division
 
 import numpy as np
+from functools import partial
+
+from ..lib.mdamath import triclinic_vectors
 
 def translate(vector):
     """
     Translates the coordinates of a given :class:`~MDAnalysis.coordinates.base.Timestep`
     instance by a given vector.
 
+    Example
+    -------
+        ts = MDAnalysis.transformations.translate([1,2,3])(ts)    
+
+    Parameters
+    ----------
+    vector: array-like
+        coordinates of the vector to which the coordinates will be translated
+
+    Returns
+    -------
+    :class:`~MDAnalysis.coordinates.base.Timestep` object
+
+    """
+    if len(vector)>2:
+        vector = np.float32(vector)
+    else:
+        raise ValueError("{} vector is too short".format(vector))
+
+    def wrapped(ts):
+        ts.positions += vector
+
+        return ts
+
+    return wrapped
+
+
+def center_in_box(ag, center='geometry', point=None, wrap=False):
+    """
+    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`
+    is centered on the unit cell. The unit cell dimensions are taken from the input Timestep object.
+    If a point is given, the center of the atomgroup will be translated to this point instead. 
+
     Example
     -------
 
     .. code-block:: python
 
-        ts = MDAnalysis.transformations.translate([1,2,3])(ts)    
+        ag = u.residues[1].atoms
+        ts = MDAnalysis.transformations.center(ag,center='mass')(ts)
 
     Parameters
     ----------
-    vector: list
-        coordinates of the vector to which the coordinates will be translated
-    ts: :class:`~MDAnalysis.coordinates.base.Timestep`
-        frame that will be transformed
-
+    ag: AtomGroup
+        atom group to be centered on the unit cell.
+    center: str, optional
+        used to choose the method of centering on the given atom group. Can be 'geometry'
+        or 'mass'
+    point: array-like, optional
+        overrides the unit cell center - the coordinates of the Timestep are translated so
+        that the center of mass/geometry of the given AtomGroup is aligned to this position
+        instead. Defined as an array of size 3.
+    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. 
+
     Returns
     -------
     :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':
+            raise AttributeError('{} is not an AtomGroup object with masses'.format(ag))
+        else:
+            raise ValueError('{} is not an AtomGroup object'.format(ag))
+
     def wrapped(ts):
-        if len(vector)>2:
-            v = np.float32(vector)
-            ts.positions += v
+        if point is None:
+            boxcenter = np.sum(ts.triclinic_dimensions, axis=0) / 2
         else:
-            raise ValueError("{} vector in too short".format(vector))
+            boxcenter = point
+
+        ag_center = center_method()
+
+        vector = boxcenter - ag_center
+        ts.positions += vector
 
         return ts
 
     return wrapped
+

testsuite/MDAnalysisTests/transformations/test_translate.py

@@ -0,0 +1,210 @@
+#-*- 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 translate, center_in_box
+from MDAnalysisTests import make_Universe
+
+
+@pytest.fixture()
+def translate_universes():
+    # create the Universe objects for the tests
+    # this universe has no masses and some tests need it as such
+    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_translate_coords(translate_universes):
+    ref_u, trans_u = translate_universes
+    ref = ref_u.trajectory.ts
+    vector = np.float32([1, 2, 3])
+    ref.positions += vector
+    trans = translate(vector)(trans_u.trajectory.ts)
+    assert_array_almost_equal(trans.positions, ref.positions, decimal=6)
+
+
+@pytest.mark.parametrize('vector', (
+    [0, 1],
+    [0, 1, 2, 3, 4],
+    np.array([0, 1]),
+    np.array([0, 1, 2, 3, 4]),
+    np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]),
+    np.array([[0], [1], [2]]))
+)
+def test_translate_vector(translate_universes, vector):
+    # what happens if the vector argument is of wrong size?
+    ts = translate_universes[0].trajectory.ts
+    with pytest.raises(ValueError):
+        translate(vector)(ts)
+
+
+def test_translate_transformations_api(translate_universes):
+    # test if the translate transformation works when using the 
+    # on-the-fly transformations API
+    ref_u, trans_u = translate_universes
+    ref = ref_u.trajectory.ts
+    vector = np.float32([1, 2, 3])
+    ref.positions += vector
+    trans_u.trajectory.add_transformations(translate(vector))
+    assert_array_almost_equal(trans_u.trajectory.ts.positions, ref.positions, decimal=6)
+
+
+def test_center_in_box_bad_ag(translate_universes):
+    # this universe has a box size zero
+    ts = translate_universes[0].trajectory.ts
+    # what happens if something other than an AtomGroup is given?
+    bad_ag = 1
+    with pytest.raises(ValueError): 
+        center_in_box(bad_ag)(ts)
+
+
+@pytest.mark.parametrize('point', (
+    [0, 1],
+    [0, 1, 2, 3, 4],
+    np.array([0, 1]),
+    np.array([0, 1, 2, 3, 4]),
+    np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]),
+    np.array([[0], [1], [2]]))
+)
+def test_center_in_box_bad_point(translate_universes, point):
+    ts = translate_universes[0].trajectory.ts
+    ag = translate_universes[0].residues[0].atoms
+    # what if the box is in the wrong format?
+    with pytest.raises(ValueError): 
+        center_in_box(ag, point=point)(ts)
+
+
+def test_center_in_box_bad_pbc(translate_universes):    
+    # this universe has a box size zero
+    ts = translate_universes[0].trajectory.ts
+    ag = translate_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
+    with pytest.raises(ValueError): 
+        center_in_box(ag, wrap=True)(ts)
+
+
+def test_center_in_box_bad_center(translate_universes):
+    # this universe has a box size zero
+    ts = translate_universes[0].trajectory.ts
+    ag = translate_universes[0].residues[0].atoms
+    # what if a wrong center type name is passed?
+    bad_center = " "
+    with pytest.raises(ValueError): 
+        center_in_box(ag, center=bad_center)(ts)
+
+
+def test_center_in_box_no_masses(translate_universes):   
+    # this universe has no masses
+    ts = translate_universes[0].trajectory.ts
+    ag = translate_universes[0].residues[0].atoms
+    # if the universe has no masses and `mass` is passed as the center arg
+    bad_center = "mass"
+    with pytest.raises(AttributeError): 
+        center_in_box(ag, center=bad_center)(ts)
+
+
+def test_center_in_box_coords_no_options(translate_universes):
+    # what happens when we center the coordinates arround the center of geometry of a residue?
+    ref_u, trans_u = translate_universes
+    ref = ref_u.trajectory.ts
+    ref_center = np.float32([6, 7, 8])
+    box_center = np.float32([186., 186.5, 187.])
+    ref.positions += box_center - ref_center
+    ag = trans_u.residues[0].atoms
+    trans = center_in_box(ag)(trans_u.trajectory.ts)
+    assert_array_almost_equal(trans.positions, ref.positions, decimal=6)
+
+
+def test_center_in_box_coords_with_pbc(translate_universes):
+    # what happens when we center the coordinates arround the center of geometry of a residue?
+    # using pbc into account for center of geometry calculation
+    ref_u, trans_u = translate_universes
+    ref = ref_u.trajectory.ts
+    trans_u.dimensions = [363., 364., 365., 90., 90., 90.]
+    ag = trans_u.residues[24].atoms
+    box_center = np.float32([181.5, 182., 182.5])
+    ref_center = np.float32([75.6, 75.8, 76.])
+    ref.positions += box_center - ref_center
+    trans = center_in_box(ag, wrap=True)(trans_u.trajectory.ts)
+    assert_array_almost_equal(trans.positions, ref.positions, decimal=6)
+
+
+def test_center_in_box_coords_with_mass(translate_universes):   
+    # using masses for calculating the center of the atomgroup
+    ref_u, trans_u = translate_universes
+    ref = ref_u.trajectory.ts
+    ag = trans_u.residues[24].atoms
+    box_center = np.float32([186., 186.5, 187.])
+    ref_center = ag.center_of_mass()
+    ref.positions += box_center - ref_center
+    trans = center_in_box(ag, center="mass")(trans_u.trajectory.ts)
+    assert_array_almost_equal(trans.positions, ref.positions, decimal=6)
+
+
+def test_center_in_box_coords_with_box(translate_universes):   
+    # using masses for calculating the center of the atomgroup
+    ref_u, trans_u = translate_universes
+    ref = ref_u.trajectory.ts
+    ag = trans_u.residues[0].atoms
+    newpoint = [1000, 1000, 1000]
+    box_center = np.float32(newpoint)
+    ref_center = np.float32([6, 7, 8])
+    ref.positions += box_center - ref_center
+    trans = center_in_box(ag, point=newpoint)(trans_u.trajectory.ts)
+    assert_array_almost_equal(trans.positions, ref.positions, decimal=6)
+
+
+def test_center_in_box_coords_all_options(translate_universes):
+    # what happens when we center the coordinates arround the center of geometry of a residue?
+    # using pbc into account for center of geometry calculation
+    ref_u, trans_u = translate_universes
+    ref = ref_u.trajectory.ts
+    ag = trans_u.residues[24].atoms
+    newpoint = [1000, 1000, 1000]
+    box_center = np.float32(newpoint)
+    ref_center = ag.center_of_mass(pbc=True)
+    ref.positions += box_center - ref_center
+    trans = center_in_box(ag, center='mass', wrap=True, point=newpoint)(trans_u.trajectory.ts)
+    assert_array_almost_equal(trans.positions, ref.positions, decimal=6)
+
+
+def test_center_transformations_api(translate_universes):
+    # test if the translate transformation works when using the 
+    # on-the-fly transformations API
+    ref_u, trans_u = translate_universes
+    ref = ref_u.trajectory.ts
+    ref_center = np.float32([6, 7, 8])
+    box_center = np.float32([186., 186.5, 187.])
+    ref.positions += box_center - ref_center
+    ag = trans_u.residues[0].atoms
+    trans_u.trajectory.add_transformations(center_in_box(ag))
+    assert_array_almost_equal(trans_u.trajectory.ts.positions, ref.positions, decimal=6)