Merge pull request #992 from duncandc/rotations_devel

Adding vector rotation utilities to halotools.utils

halotools/utils/__init__.py

@@ -14,3 +14,5 @@
 from .distribution_matching import *
 from .probabilistic_binning import fuzzy_digitize
 from .conditional_percentile import sliding_conditional_percentile
+from .vector_utilities import *
+from .rotate_vector_collection import rotate_vector_collection

halotools/utils/mcrotations.py

@@ -0,0 +1,195 @@
+"""
+A set of rotation utilites that apply monte carlo
+roations to collections of 2- and 3-dimensional vectors
+"""
+
+
+from __future__ import (division, print_function, absolute_import,
+                        unicode_literals)
+import numpy as np
+from astropy.utils.misc import NumpyRNGContext
+from .vector_utilities import *
+from .rotate_vector_collection import rotate_vector_collection
+from .rotations2d import rotation_matrices_from_angles as rotation_matrices_from_angles_2d
+from .rotations3d import rotation_matrices_from_angles as rotation_matrices_from_angles_3d
+
+
+__all__=['random_rotation_3d',
+         'random_rotation_2d',
+         'random_perpendicular_directions',
+         'random_unit_vectors_3d',
+         'random_unit_vectors_2d']
+__author__ = ['Duncan Campbell']
+
+
+def random_rotation_3d(vectors, seed=None):
+    r"""
+    Apply a random rotation to a set of 3d vectors.
+
+    Parameters
+    ----------
+    vectors : ndarray
+        Numpy array of shape (npts, 3) storing a collection of 3d vectors
+
+    seed : int, optional
+        Random number seed
+
+    Returns
+    -------
+    rotated_vectors : ndarray
+        Numpy array of shape (npts, 3) storing a collection of 3d vectors
+
+    Example
+    -------
+    Create a random set of 3D unit vectors.
+
+    >>> npts = 1000
+    >>> x1 = random_unit_vectors_3d(npts)
+
+    Randomly rotate these vectors.
+
+    >>> x2 = random_rotation_3d(x1)
+    """
+
+    with NumpyRNGContext(seed):
+        ran_direction = random_unit_vectors_3d(1)[0]
+        ran_angle = np.random.random(size=1)*(np.pi)
+
+    ran_rot = rotation_matrices_from_angles_3d(ran_angle, ran_direction)
+
+    return rotate_vector_collection(ran_rot, vectors)
+
+
+def random_rotation_2d(vectors, seed=None):
+    r"""
+    Apply a random rotation to a set of 2d vectors.
+
+    Parameters
+    ----------
+    vectors : ndarray
+        Numpy array of shape (npts, 2) storing a collection of 2d vectors
+
+    seed : int, optional
+        Random number seed
+
+    Returns
+    -------
+    rotated_vectors : ndarray
+        Numpy array of shape (npts, 2) storing a collection of 2d vectors
+
+    Example
+    -------
+    Create a random set of 2D unit vectors.
+
+    >>> npts = 1000
+    >>> x1 = random_unit_vectors_2d(npts)
+
+    Randomly rotate these vectors.
+
+    >>> x2 = random_rotation_2d(x1)
+    """
+
+    with NumpyRNGContext(seed):
+        ran_angle = np.random.random(size=1)*(np.pi)
+
+    ran_rot = rotation_matrices_from_angles_2d(ran_angle)
+
+    return rotate_vector_collection(ran_rot, vectors)
+
+
+def random_perpendicular_directions(v, seed=None):
+    r"""
+    Given an input list of 3D vectors, v, return a list of 3D vectors
+    such that each returned vector has unit-length and is
+    orthogonal to the corresponding vector in v.
+
+    Parameters
+    ----------
+    v : ndarray
+        Numpy array of shape (npts, 3) storing a collection of 3d vectors
+
+    seed : int, optional
+        Random number seed used to choose a random orthogonal direction
+
+    Returns
+    -------
+    result : ndarray
+        Numpy array of shape (npts, 3)
+
+    Example
+    -------
+    Create a random set of 3D unit vectors.
+
+    >>> npts = 1000
+    >>> x1 = random_unit_vectors_3d(npts)
+
+    For each vector in x1, create a perpendicular vector
+
+    >>> x2 = random_perpendicular_directions(x1)
+    """
+
+    v = np.atleast_2d(v)
+    npts = v.shape[0]
+
+    with NumpyRNGContext(seed):
+        w = random_unit_vectors_3d(npts)
+
+    vnorms = elementwise_norm(v).reshape((npts, 1))
+    wnorms = elementwise_norm(w).reshape((npts, 1))
+
+    e_v = v/vnorms
+    e_w = w/wnorms
+
+    v_dot_w = elementwise_dot(e_v, e_w).reshape((npts, 1))
+
+    e_v_perp = e_w - v_dot_w*e_v
+    e_v_perp_norm = elementwise_norm(e_v_perp).reshape((npts, 1))
+
+    return e_v_perp/e_v_perp_norm
+
+
+def random_unit_vectors_3d(npts):
+    r"""
+    Generate random 3D unit vectors.
+
+    Parameters
+    ----------
+    npts : int
+        number of vectors
+
+    Returns
+    -------
+    result : numpy.array
+        Numpy array of shape (npts, 3) containing random unit vectors
+    """
+
+    ndim = 3
+    x = np.random.normal(size=(npts,ndim), scale=1.0)
+    r = np.sqrt(np.sum((x)**2, axis=-1))
+
+    return (1.0/r[:,np.newaxis])*x
+
+
+def random_unit_vectors_2d(npts):
+    r"""
+    Generate random 2D unit vectors.
+
+    Parameters
+    ----------
+    npts : int
+        number of vectors
+
+    Returns
+    -------
+    result : numpy.array
+        Numpy array of shape (npts, 2) containing random unit vectors
+    """
+
+    r = 1.0
+    phi = np.random.uniform(0.0, 2.0*np.pi, size=(npts,))
+    x = r*np.cos(phi)
+    y = r*np.sin(phi)
+
+    return np.vstack((x,y)).T
+
+

halotools/utils/rotate_vector_collection.py

@@ -0,0 +1,88 @@
+"""
+A function to rotate collections of n-dimensional vectors
+"""
+
+from __future__ import (division, print_function, absolute_import,
+                        unicode_literals)
+import numpy as np
+from .vector_utilities import (elementwise_dot, elementwise_norm,
+                               normalized_vectors, angles_between_list_of_vectors)
+
+
+__all__=['rotate_vector_collection',]
+__author__ = ['Duncan Campbell', 'Andrew Hearin']
+
+def rotate_vector_collection(rotation_matrices, vectors, optimize=False):
+    r"""
+    Given a collection of rotation matrices and a collection of n-dimensional vectors,
+    apply an asscoiated matrix to rotate corresponding vector(s).
+
+    Parameters
+    ----------
+    rotation_matrices : ndarray
+        The options are:
+        1.) array of shape (npts, ndim, ndim) storing a collection of rotation matrices.
+        2.) array of shape (ndim, ndim) storing a single rotation matrix
+        
+    vectors : ndarray
+        The corresponding options for above are:
+        1.) array of shape (npts, ndim) storing a collection of ndim-dimensional vectors
+        2.) array of shape (npts, ndim) storing a collection of ndim-dimensional vectors
+
+    Returns
+    -------
+    rotated_vectors : ndarray
+        Numpy array of shape (npts, ndim) storing a collection of ndim-dimensional vectors
+
+    Notes
+    -----
+    This function is set up to preform either rotation operations on a single collection \
+    of vectors, either applying a single rotation matrix to all vectors in the collection,
+    or applying a unique rotation matrix to each vector in the set.
+
+    The behavior of the function is determined by the arguments supplied by the user.
+
+    Examples
+    --------
+    In this example, we'll randomly generate two sets of unit-vectors, `v0` and `v1`.
+    We'll use the `rotation_matrices_from_vectors` function to generate the
+    rotation matrices that rotate each `v0` into the corresponding `v1`.
+    Then we'll use the `rotate_vector_collection` function to apply each
+    rotation, and verify that we recover each of the `v1`.
+
+    >>> from halotools.utils.rotations3d import rotation_matrices_from_vectors
+    >>> from halotools.utils import normalized_vectors
+    >>> npts, ndim = int(1e4), 3
+    >>> v0 = normalized_vectors(np.random.random((npts, ndim)))
+    >>> v1 = normalized_vectors(np.random.random((npts, ndim)))
+    >>> rotation_matrices = rotation_matrices_from_vectors(v0, v1)
+    >>> v2 = rotate_vector_collection(rotation_matrices, v0)
+    >>> assert np.allclose(v1, v2)
+    """
+
+    ndim_rotm = np.shape(rotation_matrices)[-1]
+    ndim_vec = np.shape(vectors)[-1]
+    assert ndim_rotm==ndim_vec
+
+    if len(np.shape(vectors))==2:
+        ntps, ndim = np.shape(vectors)
+        nsets = 0
+    elif len(np.shape(vectors))==3:
+        nsets, ntps, ndim = np.shape(vectors)
+
+    # apply same rotation matrix to all vectors
+    if (len(np.shape(rotation_matrices)) == 2):
+        if nsets == 1:
+            vectors = vectors[0]
+        return np.dot(rotation_matrices, vectors.T).T
+    # rotate each vector by associated rotation matrix
+    else:
+        ein_string = 'ijk,ik->ij'
+        n1, ndim = np.shape(vectors)
+
+        try:
+            return np.einsum(ein_string, rotation_matrices, vectors, optimize=optimize)
+        except TypeError:
+            return np.einsum(ein_string, rotation_matrices, vectors)
+
+