rodrigues.py

# Copyright 2025 Toon Verstraelen
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the “Software”),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#

"""Forward and inverse Rodrigues rotation formula.

See:
- https://en.wikipedia.org/wiki/Rodrigues'_rotation_formula
- https://mathworld.wolfram.com/RodriguesRotationFormula.html
"""

import pytest
import numpy as np


def rodrigues(angle, axis):
    """Compute a rotation matrix about the given axis over the given angle.

    Parameters
    ----------
    angle
        A scalar or a numpy array with angles.
        `shape = (*angle_shape, )
    axis
        A rotation axis (3D vector) or a NumPy array with 3D vectors.
        `shape = (*angle_shape, 3)`

    Returns
    -------
    mat
        The rotation matrix or matrices. `shape = (*angle_shape, 3, 3)`
    """
    angle = np.asarray(angle)
    axis = np.asarray(axis)
    if axis.shape != (*angle.shape, 3):
        raise TypeError("The shape of the angle and axis arrays does not match.")
    c = np.cos(angle)
    s = np.sin(angle)
    cm = 1 - c
    x = axis[..., 0]
    y = axis[..., 1]
    z = axis[..., 2]
    n = np.sqrt(x * x + y * y + z * z)
    x = x / n
    y = y / n
    z = z / n
    mat = np.zeros((*angle.shape, 3, 3))
    mat[..., 0, 0] = c + x * x * cm
    mat[..., 1, 1] = c + y * y * cm
    mat[..., 2, 2] = c + z * z * cm
    mat[..., 0, 1] = x * y * cm - z * s
    mat[..., 1, 2] = y * z * cm - x * s
    mat[..., 2, 0] = z * x * cm - y * s
    mat[..., 1, 0] = x * y * cm + z * s
    mat[..., 2, 1] = y * z * cm + x * s
    mat[..., 0, 2] = z * x * cm + y * s
    return mat


def inverse_rodrigues(mat):
    """Inverrsion of Rodrigues' formula.

    Parameters
    ----------
    mat
        The rotation matrix or matrices. `shape = (*angle_shape, 3, 3)`

    Returns
    -------
    angle
        A scalar or a numpy array with angles.
        `shape = (*angle_shape, )
    axis
        A rotation axis (3D vector) or a NumPy array with 3D vectors.
        `shape = (*angle_shape, 3)`
    """
    if mat.ndim < 2:
        raise TypeError("The parameter mat must be at least a 2D array.")
    if mat.shape[-2:] != (3, 3):
        raise TypeError(
            "The last two dimensions of the mat array must have shape (3, 3)."
        )
    angle_shape = mat.shape[:-2]
    xs = (mat[..., 2, 1] - mat[..., 1, 2]) / 2
    ys = (mat[..., 0, 2] - mat[..., 2, 0]) / 2
    zs = (mat[..., 1, 0] - mat[..., 0, 1]) / 2
    s = np.sqrt(xs * xs + ys * ys + zs * zs)
    x = xs / s
    y = ys / s
    z = zs / s
    cm = (
        mat[..., 2, 1]
        + mat[..., 1, 2]
        + mat[..., 0, 2]
        + mat[..., 2, 0]
        + mat[..., 1, 0]
        + mat[..., 0, 1]
    ) / (2 * (x * y + y * z + z * x))
    c = 1 - cm
    angle = np.arctan2(s, c)
    axis = np.zeros((*angle_shape, 3))
    axis[..., 0] = x
    axis[..., 1] = y
    axis[..., 2] = z
    return angle, axis


#
# Unit tests
#


def test_rodrigues_simple():
    mat = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
    assert rodrigues(np.pi, [1.25, 0, 0]) == pytest.approx(mat)


def test_vectorize():
    rng = np.random.default_rng()
    angle1 = rng.uniform(-np.pi, np.pi, (7, 5))
    axis1 = rng.normal(0, 1, (7, 5, 3))
    mat = rodrigues(angle1, axis1)
    assert mat.shape == (7, 5, 3, 3)
    for i in range(7):
        for j in range(5):
            assert rodrigues(angle1[i, j], axis1[i, j]) == pytest.approx(mat[i, j])
    norms = np.sqrt(np.einsum("ijk,ijk->ij", axis1, axis1))
    axis1 /= norms.reshape((7, 5, 1))
    angle2, axis2 = inverse_rodrigues(mat)
    assert abs(angle1) == pytest.approx(angle2)
    assert axis2 == pytest.approx(axis1 * (2 * (angle1 > 0) - 1).reshape((7, 5, 1)))


@pytest.mark.parametrize("seed", range(10))
def test_consisteny(seed: int):
    rng = np.random.default_rng()
    angle1 = rng.uniform(-np.pi, np.pi)
    axis1 = rng.normal(0, 1, (3,))
    axis1 /= np.linalg.norm(axis1)
    mat = rodrigues(angle1, axis1)
    angle2, axis2 = inverse_rodrigues(mat)
    assert abs(angle1) == pytest.approx(angle2)
    assert abs(np.dot(axis2, axis2)) == pytest.approx(1)
    assert abs(np.dot(axis1, axis2)) == pytest.approx(1)
    if angle1 < 0:
        assert axis1 == pytest.approx(-axis2)
    else:
        assert axis1 == pytest.approx(axis2)