Module molgri.rotations

Expand source code
import numpy as np
from numpy.typing import ArrayLike
from scipy.spatial.transform import Rotation


def quaternion2grid(array_quaternions: np.ndarray) -> np.ndarray:
    """
    Take an array where each row is a quaternion and convert each rotation into a point on a unit sphere.

    Args:
        array_quaternions: an array (N, 4) where each row is a quaternion needed for a rotation

    Returns:
        an array (N, 3) where each row is a coordinate on a 3D sphere
    """
    base_vector = np.array([0, 0, 1])
    points = np.zeros(array_quaternions[:, 1:].shape)
    for i, quat in enumerate(array_quaternions):
        my_rotation = Rotation.from_quat(quat)
        points[i] = my_rotation.apply(base_vector)
    return points


def grid2quaternion(grid_3D):
    """
    Take an array where each row is a point on a unit sphere and convert each rotation into a set of euler_123 angles.

    Args:
        grid_3D: an array (N, 3) where each row is a coordinate on a 3D sphere

    Returns:
        an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation
    """

    assert np.allclose(np.linalg.norm(grid_3D, axis=1), 1), "Points on a grid must be unit vectors!"
    base_vector = np.array([0, 0, 1])
    points = np.zeros((grid_3D.shape[0], 4))
    for i, point in enumerate(grid_3D):
        my_matrix = two_vectors2rot(base_vector, point)
        my_rotation = Rotation.from_matrix(my_matrix)
        points[i] = my_rotation.as_quat()
    return points


def euler2grid(array_euler_angles: np.ndarray) -> np.ndarray:
    """
    Take an array where each row is a set of euler_123 angles and convert each rotation into a point on a unit sphere.

    Args:
        array_euler_angles: an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation

    Returns:
        an array (N, 3) where each row is a coordinate on a 3D sphere
    """
    base_vector = np.array([0, 0, 1])
    points = np.zeros(array_euler_angles.shape)
    for i, euler in enumerate(array_euler_angles):
        my_rotation = Rotation.from_euler("ZYX", euler)
        points[i] = my_rotation.apply(base_vector)
    assert np.allclose(np.linalg.norm(points, axis=1), 1), "Points on a grid must be unit vectors!"
    return points


def grid2euler(grid_3D: np.ndarray) -> np.ndarray:
    """
    Take an array where each row is a point on a unit sphere and convert each rotation into a set of euler_123 angles.

    Args:
        grid_3D: an array (N, 3) where each row is a coordinate on a 3D sphere

    Returns:
        an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation
    """
    assert np.allclose(np.linalg.norm(grid_3D, axis=1), 1), "Points on a grid must be unit vectors!"
    base_vector = np.array([0, 0, 1])
    points = np.zeros(grid_3D.shape)
    for i, point in enumerate(grid_3D):
        my_matrix = two_vectors2rot(base_vector, point)
        my_rotation = Rotation.from_matrix(my_matrix)
        points[i] = my_rotation.as_euler("ZYX")
    np.nan_to_num(points, copy=False)
    return points


# ########################## HELPER FUNCTIONS ################################

def skew(x: np.ndarray) -> np.ndarray:
    """
    Take a matrix and return its skew matrix.

    Args:
        x: a matrix

    Returns:
        skew matrix, see structure below
    """
    return np.array([[0, -x[2], x[1]],
                     [x[2], 0, -x[0]],
                     [-x[1], x[0], 0]])


def two_vectors2rot(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """
    Take vectors x and y and return the rotational matrix that transforms x into y.

    Args:
        x: an array of shape (3,), first vector
        y: an array of shape (3,),  second vector

    Returns:
        a 3x3 rotational matrix
    """
    x = x[:, np.newaxis]
    y = y[:, np.newaxis]
    assert y.shape == x.shape == (3, 1)
    assert np.isclose(np.linalg.norm(x), 1) and np.isclose(np.linalg.norm(y), 1)
    v = np.cross(x.T, y.T)
    s = np.linalg.norm(v)
    c = np.dot(x.T, y)[0, 0]
    if s != 0:
        my_matrix = np.eye(3) + skew(v[0]) + skew(v[0]).dot(skew(v[0])) * (1 - c) / s ** 2
    else:
        # if sin = 0, meaning that there is no rotation (or half circle)
        my_matrix = np.eye(3)
    return my_matrix


class Rotation2D:

    def __init__(self, alpha: float or ArrayLike):
        """
        Initializes 2D rotation matrix with an angle.

        Args:
            alpha: angle in radians
        """
        rot_matrix = np.array([[np.cos(alpha), np.sin(alpha)],
                               [-np.sin(alpha), np.cos(alpha)]])
        self.rot_matrix = rot_matrix

    def apply(self, vector_set: ArrayLike, inverse: bool = False) -> ArrayLike:
        """
        Applies 2D rotational matrix to a set of vectors of shape (N, 2)

        Args:
            vector_set: array (each column a vector) that should be rotated
            inverse: True if the rotation should be inverted

        Returns:
            rotated vector of shape (N, 2)
        """
        if inverse:
            inverted_mat = self.rot_matrix.T
            result = vector_set.dot(inverted_mat)
        else:
            result = vector_set.dot(self.rot_matrix)
        result = result.squeeze()
        return result

Functions

def euler2grid(array_euler_angles: numpy.ndarray) ‑> numpy.ndarray

Take an array where each row is a set of euler_123 angles and convert each rotation into a point on a unit sphere.

Args

array_euler_angles
an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation

Returns

an array (N, 3) where each row is a coordinate on a 3D sphere

Expand source code
def euler2grid(array_euler_angles: np.ndarray) -> np.ndarray:
    """
    Take an array where each row is a set of euler_123 angles and convert each rotation into a point on a unit sphere.

    Args:
        array_euler_angles: an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation

    Returns:
        an array (N, 3) where each row is a coordinate on a 3D sphere
    """
    base_vector = np.array([0, 0, 1])
    points = np.zeros(array_euler_angles.shape)
    for i, euler in enumerate(array_euler_angles):
        my_rotation = Rotation.from_euler("ZYX", euler)
        points[i] = my_rotation.apply(base_vector)
    assert np.allclose(np.linalg.norm(points, axis=1), 1), "Points on a grid must be unit vectors!"
    return points
def grid2euler(grid_3D: numpy.ndarray) ‑> numpy.ndarray

Take an array where each row is a point on a unit sphere and convert each rotation into a set of euler_123 angles.

Args

grid_3D
an array (N, 3) where each row is a coordinate on a 3D sphere

Returns

an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation

Expand source code
def grid2euler(grid_3D: np.ndarray) -> np.ndarray:
    """
    Take an array where each row is a point on a unit sphere and convert each rotation into a set of euler_123 angles.

    Args:
        grid_3D: an array (N, 3) where each row is a coordinate on a 3D sphere

    Returns:
        an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation
    """
    assert np.allclose(np.linalg.norm(grid_3D, axis=1), 1), "Points on a grid must be unit vectors!"
    base_vector = np.array([0, 0, 1])
    points = np.zeros(grid_3D.shape)
    for i, point in enumerate(grid_3D):
        my_matrix = two_vectors2rot(base_vector, point)
        my_rotation = Rotation.from_matrix(my_matrix)
        points[i] = my_rotation.as_euler("ZYX")
    np.nan_to_num(points, copy=False)
    return points
def grid2quaternion(grid_3D)

Take an array where each row is a point on a unit sphere and convert each rotation into a set of euler_123 angles.

Args

grid_3D
an array (N, 3) where each row is a coordinate on a 3D sphere

Returns

an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation

Expand source code
def grid2quaternion(grid_3D):
    """
    Take an array where each row is a point on a unit sphere and convert each rotation into a set of euler_123 angles.

    Args:
        grid_3D: an array (N, 3) where each row is a coordinate on a 3D sphere

    Returns:
        an array (N, 3) where each row is a set of the 3 euler angles needed for a rotation
    """

    assert np.allclose(np.linalg.norm(grid_3D, axis=1), 1), "Points on a grid must be unit vectors!"
    base_vector = np.array([0, 0, 1])
    points = np.zeros((grid_3D.shape[0], 4))
    for i, point in enumerate(grid_3D):
        my_matrix = two_vectors2rot(base_vector, point)
        my_rotation = Rotation.from_matrix(my_matrix)
        points[i] = my_rotation.as_quat()
    return points
def quaternion2grid(array_quaternions: numpy.ndarray) ‑> numpy.ndarray

Take an array where each row is a quaternion and convert each rotation into a point on a unit sphere.

Args

array_quaternions
an array (N, 4) where each row is a quaternion needed for a rotation

Returns

an array (N, 3) where each row is a coordinate on a 3D sphere

Expand source code
def quaternion2grid(array_quaternions: np.ndarray) -> np.ndarray:
    """
    Take an array where each row is a quaternion and convert each rotation into a point on a unit sphere.

    Args:
        array_quaternions: an array (N, 4) where each row is a quaternion needed for a rotation

    Returns:
        an array (N, 3) where each row is a coordinate on a 3D sphere
    """
    base_vector = np.array([0, 0, 1])
    points = np.zeros(array_quaternions[:, 1:].shape)
    for i, quat in enumerate(array_quaternions):
        my_rotation = Rotation.from_quat(quat)
        points[i] = my_rotation.apply(base_vector)
    return points
def skew(x: numpy.ndarray) ‑> numpy.ndarray

Take a matrix and return its skew matrix.

Args

x
a matrix

Returns

skew matrix, see structure below

Expand source code
def skew(x: np.ndarray) -> np.ndarray:
    """
    Take a matrix and return its skew matrix.

    Args:
        x: a matrix

    Returns:
        skew matrix, see structure below
    """
    return np.array([[0, -x[2], x[1]],
                     [x[2], 0, -x[0]],
                     [-x[1], x[0], 0]])
def two_vectors2rot(x: numpy.ndarray, y: numpy.ndarray) ‑> numpy.ndarray

Take vectors x and y and return the rotational matrix that transforms x into y.

Args

x
an array of shape (3,), first vector
y
an array of shape (3,), second vector

Returns

a 3x3 rotational matrix

Expand source code
def two_vectors2rot(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """
    Take vectors x and y and return the rotational matrix that transforms x into y.

    Args:
        x: an array of shape (3,), first vector
        y: an array of shape (3,),  second vector

    Returns:
        a 3x3 rotational matrix
    """
    x = x[:, np.newaxis]
    y = y[:, np.newaxis]
    assert y.shape == x.shape == (3, 1)
    assert np.isclose(np.linalg.norm(x), 1) and np.isclose(np.linalg.norm(y), 1)
    v = np.cross(x.T, y.T)
    s = np.linalg.norm(v)
    c = np.dot(x.T, y)[0, 0]
    if s != 0:
        my_matrix = np.eye(3) + skew(v[0]) + skew(v[0]).dot(skew(v[0])) * (1 - c) / s ** 2
    else:
        # if sin = 0, meaning that there is no rotation (or half circle)
        my_matrix = np.eye(3)
    return my_matrix

Classes

class Rotation2D (alpha: float)

Initializes 2D rotation matrix with an angle.

Args

alpha
angle in radians
Expand source code
class Rotation2D:

    def __init__(self, alpha: float or ArrayLike):
        """
        Initializes 2D rotation matrix with an angle.

        Args:
            alpha: angle in radians
        """
        rot_matrix = np.array([[np.cos(alpha), np.sin(alpha)],
                               [-np.sin(alpha), np.cos(alpha)]])
        self.rot_matrix = rot_matrix

    def apply(self, vector_set: ArrayLike, inverse: bool = False) -> ArrayLike:
        """
        Applies 2D rotational matrix to a set of vectors of shape (N, 2)

        Args:
            vector_set: array (each column a vector) that should be rotated
            inverse: True if the rotation should be inverted

        Returns:
            rotated vector of shape (N, 2)
        """
        if inverse:
            inverted_mat = self.rot_matrix.T
            result = vector_set.dot(inverted_mat)
        else:
            result = vector_set.dot(self.rot_matrix)
        result = result.squeeze()
        return result

Methods

def apply(self, vector_set: Union[numpy.__array_like._SupportsArray[numpy.dtype], numpy.__nested_sequence._NestedSequence[numpy.__array_like._SupportsArray[numpy.dtype]], bool, int, float, complex, str, bytes, numpy.__nested_sequence._NestedSequence[Union[bool, int, float, complex, str, bytes]]], inverse: bool = False) ‑> Union[numpy.__array_like._SupportsArray[numpy.dtype], numpy.__nested_sequence._NestedSequence[numpy.__array_like._SupportsArray[numpy.dtype]], bool, int, float, complex, str, bytes, numpy.__nested_sequence._NestedSequence[Union[bool, int, float, complex, str, bytes]]]

Applies 2D rotational matrix to a set of vectors of shape (N, 2)

Args

vector_set
array (each column a vector) that should be rotated
inverse
True if the rotation should be inverted

Returns

rotated vector of shape (N, 2)

Expand source code
def apply(self, vector_set: ArrayLike, inverse: bool = False) -> ArrayLike:
    """
    Applies 2D rotational matrix to a set of vectors of shape (N, 2)

    Args:
        vector_set: array (each column a vector) that should be rotated
        inverse: True if the rotation should be inverted

    Returns:
        rotated vector of shape (N, 2)
    """
    if inverse:
        inverted_mat = self.rot_matrix.T
        result = vector_set.dot(inverted_mat)
    else:
        result = vector_set.dot(self.rot_matrix)
    result = result.squeeze()
    return result