diff --git a/package/CHANGELOG b/package/CHANGELOG index 83a6dbe1228..b25b1c3cbc6 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -61,6 +61,7 @@ Enhancements * added unwrap keyword to center_of_mass (PR #2286) * added unwrap keyword to moment_of_inertia (PR #2287) * added unwrap keyword to asphericity (PR #2290) + * added unwrap keyword to radius_of_gyration (PR #2294) Changes * added official support for Python 3.7 (PR #1963) diff --git a/package/MDAnalysis/core/topologyattrs.py b/package/MDAnalysis/core/topologyattrs.py index 22a71302b30..2ba01f3bccf 100644 --- a/package/MDAnalysis/core/topologyattrs.py +++ b/package/MDAnalysis/core/topologyattrs.py @@ -937,6 +937,7 @@ def moment_of_inertia(group, **kwargs): ('moment_of_inertia', moment_of_inertia)) @warn_if_not_unique + @check_pbc_and_unwrap def radius_of_gyration(group, **kwargs): """Radius of gyration. @@ -945,6 +946,10 @@ def radius_of_gyration(group, **kwargs): pbc : bool, optional If ``True``, move all atoms within the primary unit cell before calculation. [``False``] + unwrap : bool, optional + If ``True``, compounds will be unwrapped before computing their centers. + compound : {'group', 'segments', 'residues', 'molecules', 'fragments'}, optional + Which type of component to keep together during unwrapping. Note ---- @@ -953,15 +958,22 @@ def radius_of_gyration(group, **kwargs): .. versionchanged:: 0.8 Added *pbc* keyword - + .. versionchanged:: 0.20.0 Added *unwrap* and *compound* parameter """ atomgroup = group.atoms pbc = kwargs.pop('pbc', flags['use_pbc']) masses = atomgroup.masses + unwrap = kwargs.pop('unwrap', False) + compound = kwargs.pop('compound', 'group') + + com = atomgroup.center_of_mass(pbc=pbc, unwrap=unwrap, compound=compound) + if compound is not 'group': + com = (com * group.masses[:, None]).sum(axis=0) / group.masses.sum() - com = atomgroup.center_of_mass(pbc=pbc) if pbc: recenteredpos = atomgroup.pack_into_box(inplace=False) - com + elif unwrap: + recenteredpos = atomgroup.unwrap(compound=compound) - com else: recenteredpos = atomgroup.positions - com @@ -984,6 +996,10 @@ def shape_parameter(group, **kwargs): pbc : bool, optional If ``True``, move all atoms within the primary unit cell before calculation. [``False``] + unwrap : bool, optional + If ``True``, compounds will be unwrapped before computing their centers. + compound : {'group', 'segments', 'residues', 'molecules', 'fragments'}, optional + Which type of component to keep together during unwrapping. Note ---- diff --git a/package/MDAnalysis/transformations/__init__.py b/package/MDAnalysis/transformations/__init__.py index b1776527c59..31c7241889c 100644 --- a/package/MDAnalysis/transformations/__init__.py +++ b/package/MDAnalysis/transformations/__init__.py @@ -96,3 +96,4 @@ def wrapped(ts): from .translate import translate, center_in_box from .rotate import rotateby from .positionaveraging import PositionAverager +from .fit import fit_translation, fit_rot_trans diff --git a/package/MDAnalysis/transformations/fit.py b/package/MDAnalysis/transformations/fit.py new file mode 100644 index 00000000000..63ee6213544 --- /dev/null +++ b/package/MDAnalysis/transformations/fit.py @@ -0,0 +1,206 @@ +# -*- 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 +# + +"""\ +Fitting transformations --- :mod:`MDAnalysis.transformations.fit` +================================================================= + +Translate and/or rotates the coordinates of a given trajectory to align +a given AtomGroup to a reference structure. + +.. autofunction:: fit_translation + +.. autofunction:: fit_rot_trans + +""" +from __future__ import absolute_import + +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): + + """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 + the translation motions of the AtomGroup over that particular plane. + + 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`: + + .. code-block:: python + + ag = u.select_atoms("protein") + ref = mda.Universe("reference.pdb") + transform = mda.transformations.fit_translation(ag, ref, plane="xy", + weights="mass") + u.trajectory.add_transformations(transform) + + Parameters + ---------- + ag : Universe or AtomGroup + structure to translate, a + :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` + 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. + 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 + `ag` is provided, use each element of the `array_like` as a + weight for the corresponding atom in `ag`. + + Returns + ------- + MDAnalysis.coordinates.base.Timestep + """ + + if plane is not None: + axes = {'yz' : 0, 'xz' : 1, 'xy' : 2} + try: + plane = axes[plane] + except (TypeError, KeyError): + raise ValueError('{} is not a valid plane'.format(plane)) + try: + if ag.atoms.n_residues != reference.atoms.n_residues: + raise ValueError("{} and {} have mismatched number of residues".format(ag,reference)) + except AttributeError: + raise AttributeError("{} or {} is not valid Universe/AtomGroup".format(ag,reference)) + ref, mobile = align.get_matching_atoms(reference.atoms, ag.atoms) + try: + weights = align.get_weights(ref.atoms, weights=weights) + except (ValueError, TypeError): + raise ValueError("weights must be {'mass', None} or an iterable of the " + "same size as the atomgroup.") + + ref_com = np.asarray(ref.center(weights), np.float32) + + 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 + ts.positions += vector + + return ts + + return wrapped + + +def fit_rot_trans(ag, reference, plane=None, weights=None): + """Perform a spatial superposition by minimizing the RMSD. + + Spatially align the group of atoms `ag` to `reference` by doing a RMSD + fit. + + This fit works as a way to remove translations and rotations of a given + AtomGroup in a trajectory. A plane can be given using the flag `plane` + so that only translations and rotations in that particular plane are + removed. This is useful for protein-membrane systems to where the membrane + must remain in the same orientation. + + Example + ------- + Removing the translations and rotations of a given AtomGroup `ag` on the XY plane + by fitting it to a reference `ref`, using the masses as weights for the RMSD fit: + + .. code-block:: python + + ag = u.select_atoms("protein") + ref = mda.Universe("reference.pdb") + transform = mda.transformations.fit_rot_trans(ag, ref, plane="xy", + weights="mass") + u.trajectory.add_transformations(transform) + + Parameters + ---------- + ag : Universe or AtomGroup + structure to translate and rotate, a + :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` + plane: str, optional + used to define the plane on which the rotations and translations will be removed. + Defined as a string of the plane. Supported 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 + `ag` is provided, use each element of the `array_like` as a + weight for the corresponding atom in `ag`. + + Returns + ------- + MDAnalysis.coordinates.base.Timestep + """ + if plane is not None: + axes = {'yz' : 0, 'xz' : 1, 'xy' : 2} + try: + plane = axes[plane] + except (TypeError, KeyError): + raise ValueError('{} is not a valid plane'.format(plane)) + try: + if ag.atoms.n_residues != reference.atoms.n_residues: + raise ValueError("{} and {} have mismatched number of residues".format(ag,reference)) + except AttributeError: + raise AttributeError("{} or {} is not valid Universe/AtomGroup".format(ag,reference)) + ref, mobile = align.get_matching_atoms(reference.atoms, ag.atoms) + try: + weights = align.get_weights(ref.atoms, weights=weights) + except (ValueError, TypeError): + raise ValueError("weights must be {'mass', None} or an iterable of the " + "same size as the atomgroup.") + 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)] + 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) + 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) + ts.positions = ts.positions + vector + + 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 a5f29f4a371..051f7b726fa 100644 --- a/package/doc/sphinx/source/documentation_pages/trajectory_transformations.rst +++ b/package/doc/sphinx/source/documentation_pages/trajectory_transformations.rst @@ -126,3 +126,4 @@ e.g. giving a workflow as a keyword argument when defining the universe: ./transformations/translate ./transformations/rotate ./transformations/positionaveraging + ./transformations/fit diff --git a/package/doc/sphinx/source/documentation_pages/transformations/fit.rst b/package/doc/sphinx/source/documentation_pages/transformations/fit.rst new file mode 100644 index 00000000000..a5fe6293a83 --- /dev/null +++ b/package/doc/sphinx/source/documentation_pages/transformations/fit.rst @@ -0,0 +1 @@ +.. automodule:: MDAnalysis.transformations.fit \ No newline at end of file diff --git a/testsuite/MDAnalysisTests/core/test_atomgroup.py b/testsuite/MDAnalysisTests/core/test_atomgroup.py index c17c3013233..f7f4412f903 100644 --- a/testsuite/MDAnalysisTests/core/test_atomgroup.py +++ b/testsuite/MDAnalysisTests/core/test_atomgroup.py @@ -539,7 +539,6 @@ def test_center_unwrap(self, level, compound, is_triclinic): # get the expected results center = group.center(weights=None, pbc=False, compound=compound, unwrap=True) - ref_center = u.center(compound=compound) assert_almost_equal(ref_center, center, decimal=4) @@ -839,6 +838,7 @@ class TestUnwrapFlag(object): prec = 3 @pytest.fixture() +<<<<<<< HEAD def ag(self): universe = mda.Universe(TRZ_psf, TRZ) group = universe.residues[0:3] @@ -859,6 +859,7 @@ def ref_noUnwrap_residues(self): [-211.8997285, 7059.07470427, -91.32156884], [-721.50785456, -91.32156884, 6509.31735029]]), 'Asph': 0.02060121, + 'ROG': 27.713009, } @pytest.fixture() @@ -874,6 +875,7 @@ def ref_Unwrap_residues(self): [-1330.489, 19315.253, 3306.212], [ 2938.243, 3306.212, 8990.481]]), 'Asph': 0.2969491080, + 'ROG': 40.686541, } @pytest.fixture() @@ -886,6 +888,7 @@ def ref_noUnwrap(self): [0.0, 98.6542, 0.0], [0.0, 0.0, 98.65421327]]), 'Asph': 1.0, + 'ROG': 0.9602093, } @pytest.fixture() @@ -898,6 +901,7 @@ def ref_Unwrap(self): [0.0, 132.673, 0.0], [0.0, 0.0, 132.673]]), 'Asph': 1.0, + 'ROG': 1.1135230, } def test_default_residues(self, ag, ref_noUnwrap_residues): @@ -905,12 +909,14 @@ def test_default_residues(self, ag, ref_noUnwrap_residues): assert_almost_equal(ag.center_of_mass(compound='residues'), ref_noUnwrap_residues['COM'], self.prec) assert_almost_equal(ag.moment_of_inertia(compound='residues'), ref_noUnwrap_residues['MOI'], self.prec) assert_almost_equal(ag.asphericity(compound='residues'), ref_noUnwrap_residues['Asph'], self.prec) + assert_almost_equal(ag.radius_of_gyration(compound='residues'), ref_noUnwrap_residues['ROG'], self.prec) def test_UnWrapFlag_residues(self, ag, ref_Unwrap_residues): assert_almost_equal(ag.center_of_geometry(unwrap=True, compound='residues'), ref_Unwrap_residues['COG'], self.prec) assert_almost_equal(ag.center_of_mass(unwrap=True, compound='residues'), ref_Unwrap_residues['COM'], self.prec) assert_almost_equal(ag.moment_of_inertia(unwrap=True, compound='residues'), ref_Unwrap_residues['MOI'], self.prec) assert_almost_equal(ag.asphericity(unwrap=True, compound='residues'), ref_Unwrap_residues['Asph'], self.prec) + assert_almost_equal(ag.radius_of_gyration(unwrap=True, compound='residues'), ref_Unwrap_residues['ROG'], self.prec) def test_default(self, ref_noUnwrap): u = UnWrapUniverse(is_triclinic=False) @@ -922,16 +928,19 @@ def test_default(self, ref_noUnwrap): assert_almost_equal(group.center_of_mass(), ref_noUnwrap['COM'], self.prec) assert_almost_equal(group.moment_of_inertia(), ref_noUnwrap['MOI'], self.prec) assert_almost_equal(group.asphericity(), ref_noUnwrap['Asph'], self.prec) + assert_almost_equal(group.radius_of_gyration(), ref_noUnwrap['ROG'], self.prec) def test_UnWrapFlag(self, ref_Unwrap): u = UnWrapUniverse(is_triclinic=False) group = u.atoms[31:39] # molecules 11 + # Changing masses for center_of_mass group.masses = [100.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] assert_almost_equal(group.center_of_geometry(unwrap=True), ref_Unwrap['COG'], self.prec) assert_almost_equal(group.center_of_mass(unwrap=True), ref_Unwrap['COM'], self.prec) assert_almost_equal(group.moment_of_inertia(unwrap=True), ref_Unwrap['MOI'], self.prec) assert_almost_equal(group.asphericity(unwrap=True), ref_Unwrap['Asph'], self.prec) + assert_almost_equal(group.radius_of_gyration(unwrap=True), ref_Unwrap['ROG'], self.prec) class TestPBCFlag(object): diff --git a/testsuite/MDAnalysisTests/core/util.py b/testsuite/MDAnalysisTests/core/util.py index 653c743d19c..56b4c958250 100644 --- a/testsuite/MDAnalysisTests/core/util.py +++ b/testsuite/MDAnalysisTests/core/util.py @@ -601,8 +601,11 @@ def center(self, compound): for base in range(23, 47, 8): loc_center = np.mean(relpos[base:base + 8, :], axis=0) center_pos[pos,:] = loc_center + loc_centre = np.mean(relpos[base:base + 4, :], axis=0) + center_pos[pos,:] = loc_centre pos+=1 + if compound == "group": center_pos = center_pos[11] elif compound == "segments": diff --git a/testsuite/MDAnalysisTests/transformations/test_fit.py b/testsuite/MDAnalysisTests/transformations/test_fit.py new file mode 100644 index 00000000000..84b7f71b954 --- /dev/null +++ b/testsuite/MDAnalysisTests/transformations/test_fit.py @@ -0,0 +1,271 @@ +# -*- 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 division, print_function, absolute_import + +import numpy as np +import pytest +from numpy.testing import assert_array_almost_equal + +from MDAnalysisTests import make_Universe +from MDAnalysis.transformations.fit import fit_translation, fit_rot_trans +from MDAnalysis.lib.transformations import rotation_matrix + + +@pytest.fixture() +def fit_universe(): + # make a test universe + test = make_Universe(('masses', ), trajectory=True) + ref = make_Universe(('masses', ), trajectory=True) + ref.atoms.positions += np.asarray([10, 10, 10], np.float32) + return test, ref + + +@pytest.mark.parametrize('universe', ( + [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]]), + 'thisisnotanag', + 1) +) +def test_fit_translation_bad_ag(fit_universe, universe): + ts = fit_universe[0].trajectory.ts + test_u = fit_universe[0] + ref_u = fit_universe[1] + # what happens if something other than an AtomGroup or Universe is given? + with pytest.raises(AttributeError): + fit_translation(universe, ref_u)(ts) + + +@pytest.mark.parametrize('weights', ( + " ", + "totallynotmasses", + 123456789, + [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]])) +) +def test_fit_translation_bad_weights(fit_universe, weights): + ts = fit_universe[0].trajectory.ts + test_u = fit_universe[0] + ref_u = fit_universe[1] + # what happens if a bad string for center is given? + with pytest.raises(ValueError): + fit_translation(test_u, ref_u, weights=weights)(ts) + + +@pytest.mark.parametrize('plane', ( + 1, + [0, 1], + [0, 1, 2, 3, 4], + np.array([0, 1]), + "xyz", + "notaplane") +) +def test_fit_translation_bad_plane(fit_universe, plane): + ts = fit_universe[0].trajectory.ts + test_u = fit_universe[0] + ref_u = fit_universe[1] + # what happens if a bad string for center is given? + with pytest.raises(ValueError): + fit_translation(test_u, ref_u, plane=plane)(ts) + + +def test_fit_translation_no_masses(fit_universe): + ts = fit_universe[0].trajectory.ts + test_u = fit_universe[0] + # create a universe without masses + ref_u = make_Universe() + # what happens Universe without masses is given? + with pytest.raises(AttributeError): + fit_translation(test_u, ref_u, weights="mass")(ts) + + +def test_fit_translation_no_options(fit_universe): + test_u = fit_universe[0] + ref_u = fit_universe[1] + fit_translation(test_u, ref_u)(test_u.trajectory.ts) + # what happens when no options are passed? + assert_array_almost_equal(test_u.trajectory.ts.positions, ref_u.trajectory.ts.positions, decimal=6) + +def test_fit_translation_residue_mismatch(fit_universe): + test_u = fit_universe[0] + ref_u = fit_universe[1].residues[:-1].atoms + with pytest.raises(ValueError, match='number of residues'): + fit_translation(test_u, ref_u)(test_u.trajectory.ts) + +def test_fit_translation_com(fit_universe): + test_u = fit_universe[0] + ref_u = fit_universe[1] + fit_translation(test_u, ref_u, weights="mass")(test_u.trajectory.ts) + # what happens when the center o mass is used? + assert_array_almost_equal(test_u.trajectory.ts.positions, ref_u.trajectory.ts.positions, decimal=6) + + +@pytest.mark.parametrize('plane', ( + "yz", + "xz", + "xy") +) +def test_fit_translation_plane(fit_universe, plane): + test_u = fit_universe[0] + ref_u = fit_universe[1] + axes = {'yz' : 0, 'xz' : 1, 'xy' : 2} + idx = axes[plane] + # translate the test universe on the plane coordinates only + fit_translation(test_u, ref_u, plane=plane)(test_u.trajectory.ts) + # the reference is 10 angstrom in all coordinates above the test universe + shiftz = np.asanyarray([0, 0, 0], np.float32) + shiftz[idx] = -10 + ref_coordinates = ref_u.trajectory.ts.positions + shiftz + assert_array_almost_equal(test_u.trajectory.ts.positions, ref_coordinates, decimal=6) + + +def test_fit_translation_all_options(fit_universe): + test_u = fit_universe[0] + ref_u = fit_universe[1] + # translate the test universe on the x and y coordinates only + fit_translation(test_u, ref_u, plane="xy", weights="mass")(test_u.trajectory.ts) + # the reference is 10 angstrom in the z coordinate above the test universe + shiftz = np.asanyarray([0, 0, -10], np.float32) + ref_coordinates = ref_u.trajectory.ts.positions + shiftz + assert_array_almost_equal(test_u.trajectory.ts.positions, ref_coordinates, decimal=6) + + +def test_fit_translation_transformations_api(fit_universe): + test_u = fit_universe[0] + ref_u = fit_universe[1] + transform = fit_translation(test_u, ref_u) + test_u.trajectory.add_transformations(transform) + assert_array_almost_equal(test_u.trajectory.ts.positions, ref_u.trajectory.ts.positions, decimal=6) + + +@pytest.mark.parametrize('universe', ( + [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]]), + 'thisisnotanag', + 1) +) +def test_fit_rot_trans_bad_universe(fit_universe, universe): + test_u = fit_universe[0] + ref_u= universe + with pytest.raises(AttributeError): + fit_rot_trans(test_u, ref_u)(test_u.trajectory.ts) + + +def test_fit_rot_trans_shorter_universe(fit_universe): + ref_u = fit_universe[1] + 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) + + +@pytest.mark.parametrize('weights', ( + " ", + "totallynotmasses", + 123456789, + [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]])) +) +def test_fit_rot_trans_bad_weights(fit_universe, weights): + test_u = fit_universe[0] + ref_u = fit_universe[1] + bad_weights = weights + with pytest.raises(ValueError): + fit_rot_trans(test_u, ref_u, weights=bad_weights)(test_u.trajectory.ts) + + +@pytest.mark.parametrize('plane', ( + " ", + "totallynotaplane", + "xyz", + 123456789, + [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]])) +) +def test_fit_rot_trans_bad_plane(fit_universe, plane): + test_u = fit_universe[0] + ref_u = fit_universe[1] + with pytest.raises(ValueError): + fit_rot_trans(test_u, ref_u, plane=plane)(test_u.trajectory.ts) + + +def test_fit_rot_trans_no_options(fit_universe): + test_u = fit_universe[0] + ref_u = fit_universe[1] + ref_com = ref_u.atoms.center(None) + ref_u.trajectory.ts.positions -= ref_com + R = rotation_matrix(np.pi/3, [1,0,0])[:3,:3] + ref_u.trajectory.ts.positions = np.dot(ref_u.trajectory.ts.positions, R) + ref_u.trajectory.ts.positions += ref_com + fit_rot_trans(test_u, ref_u)(test_u.trajectory.ts) + assert_array_almost_equal(test_u.trajectory.ts.positions, ref_u.trajectory.ts.positions, decimal=3) + + +@pytest.mark.parametrize('plane', ( + "yz", + "xz", + "xy") +) +def test_fit_rot_trans_plane(fit_universe, plane): + # the reference is rotated in the x axis so removing the translations and rotations + # in the yz plane should return the same as the fitting without specifying a plane + test_u = fit_universe[0] + ref_u = fit_universe[1] + ref_com = ref_u.atoms.center(None) + mobile_com = test_u.atoms.center(None) + axes = {'yz' : 0, 'xz' : 1, 'xy' : 2} + idx = axes[plane] + rotaxis = np.asarray([0,0,0]) + rotaxis[idx]=1 + ref_u.trajectory.ts.positions -= ref_com + R = rotation_matrix(np.pi/3, rotaxis)[:3,:3] + ref_u.trajectory.ts.positions = np.dot(ref_u.trajectory.ts.positions, R) + ref_com[idx] = mobile_com[idx] + ref_u.trajectory.ts.positions += ref_com + fit_rot_trans(test_u, ref_u, plane=plane)(test_u.trajectory.ts) + assert_array_almost_equal(test_u.trajectory.ts.positions[:,idx], ref_u.trajectory.ts.positions[:,idx], decimal=3) + + +def test_fit_rot_trans_transformations_api(fit_universe): + test_u = fit_universe[0] + ref_u = fit_universe[1] + ref_com = ref_u.atoms.center(None) + ref_u.trajectory.ts.positions -= ref_com + R = rotation_matrix(np.pi/3, [1,0,0])[:3,:3] + ref_u.trajectory.ts.positions = np.dot(ref_u.trajectory.ts.positions, R) + ref_u.trajectory.ts.positions += ref_com + transform = fit_rot_trans(test_u, ref_u) + test_u.trajectory.add_transformations(transform) + assert_array_almost_equal(test_u.trajectory.ts.positions, ref_u.trajectory.ts.positions, decimal=3)