Source code for mdtraj.utils.rotation

##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2013 Stanford University and the Authors
#
# Authors: John D. Chodera and Kyle A. Beauchamp, Kim Branson,
#          Imran Haque, Michael Shirts
# Contributors: Robert McGibbon
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
#
# This code is adapted from Yank: GPU-accelerated calculation of ligand
# binding affinities in implicit and explicit solvent using alchemical
# free energy methodologies, copyright 2009-2011 University of California,
# Berkeley, Vertex Pharmaceuticals, Stanford University, University of Virginia,
# and the Authors. Yank is licensed under the GPU Lesser General Public
# License.
##############################################################################

import numpy as np

from mdtraj.utils.validation import check_random_state

__all__ = ["rotation_matrix_from_quaternion", "uniform_quaternion"]


[docs] def rotation_matrix_from_quaternion(q): """Compute a 3x3 rotation matrix from a given quaternion (4-vector). Parameters ---------- q : np.ndarray, size=(..., 4) Quaternion or array of quaternions (need not be normalized, zero norm OK) Returns ------- Rq : np.ndarray, size=(..., 3, 3) Orthogonal rotation matrices corresponding to quaternion(s) q. Given, for example, an input shape of (m, 4), Rq will have shape (m, 3, 3). Examples -------- >>> q = np.array([0.1, 0.2, 0.3, -0.4]) >>> Rq = rotation_matrix_from_quaternion(q) See Also -------- uniform_quaternion References ---------- .. [1] http://en.wikipedia.org/wiki/Rotation_matrix#Quaternion """ w = q[..., 0] x = q[..., 1] y = q[..., 2] z = q[..., 3] # w, x, y, z = q Nq = np.sum(q**2, axis=q.ndim - 1) if q.shape == (4,): s = 2.0 / Nq if Nq > 0 else 0.0 else: with np.errstate(divide="ignore"): s = 2.0 / Nq s[Nq <= 0.0] = 0.0 X = x * s Y = y * s Z = z * s wX = w * X wY = w * Y wZ = w * Z xX = x * X xY = x * Y xZ = x * Z yY = y * Y yZ = y * Z zZ = z * Z Rq = np.array( [ [1.0 - (yY + zZ), xY - wZ, xZ + wY], [xY + wZ, 1.0 - (xX + zZ), yZ - wX], [xZ - wY, yZ + wX, 1.0 - (xX + yY)], ], ) if q.shape == (4,): return Rq return np.rollaxis(Rq, Rq.ndim - 1)
[docs] def uniform_quaternion(size=None, random_state=None): """Generate uniform normalized quaternion 4-vectors Parameters ---------- size : int or tuple of ints, optional Defines the shape of the returned array of quaternions. If None (the default), returns a quaternion 4-vector. random_state : integer or numpy.RandomState, optional The generator used for random numbers. If an integer is given, it fixes the seed. Defaults to the global numpy random number generator. Returns ------- out : ndarray Array of quaternion 4-vectors. Given a ``size`` of, for example, ``(m,n)``, ``m*n`` samples are generated. Because each sample is `4-dimensional, the output shape is ``(m,n,4)``. If no shape is specified, a single (`4`-D) sample is returned. See Also -------- rotation_matrix_from_quaternion References ---------- .. [1] K. Shoemake. Uniform random rotations. In D. Kirk, editor, Graphics Gems III, pages 124-132. Academic, New York, 1992. .. [2] http://planning.cs.uiuc.edu/node198.html """ random = check_random_state(random_state) if size is None: randsize = (3, 1) elif isinstance(size, tuple): size = tuple(int(i) for i in size) randsize = (3,) + size else: randsize = (3, int(size)) u = random.random_sample(randsize) q = np.array( [ np.sqrt(1 - u[0]) * np.sin(2 * np.pi * u[1]), np.sqrt(1 - u[0]) * np.cos(2 * np.pi * u[1]), np.sqrt(u[0]) * np.sin(2 * np.pi * u[2]), np.sqrt(u[0]) * np.cos(2 * np.pi * u[2]), ], ) if size is None: return q[:, 0] return np.rollaxis(q, 0, q.ndim)