Rotating points with SO3 using the matrix exponential for so3

so3_mat_exp.py

"""
Code from the video:

https://youtu.be/QcfYTagj3pk

For two gifs made from it, see
https://thumbs.gfycat.com/BlissfulMenacingCanine-mobile.mp4
https://thumbs.gfycat.com/MinorHealthyCub-mobile.mp4

Note that this code can't be executed on it's own - it's just all the functions about geometry used in the animation.
For the animation code, see the "Spherical Pendulum" animation video/code.
"""

import numpy as np

import random

from utils.utils import naturals, _inner, norm2, normalize
from utils.algebra import transpose
from utils.animation import Animation

from pendulum_animation import _set_up_plot, _plot_bowl

def matrix_shape(mat):
    """
    :param mat: Any matrix.
    Returns the pair (n, m) where n, m are the dimensions of mat.
    """
    num_rows = len(mat)
    if num_rows == 0:
        return 0, 0

    num_cols = len(mat[0])
    assert all(len(col) == num_cols for col in mat)

    return num_rows, num_cols


def transpose(points) -> list:
    _num, dim = matrix_shape(points)
    return [[point[i] for point in points] for i in range(dim)]


def matrix_exp(n: int, mat):
    """
    :param n: Accuracy parameter.
    :param mat: d x d matrix.
    Return (1 + mat / k)^k with k=2^n
    """
    num_rows, num_cols = matrix_shape(mat)
    assert num_rows == num_cols, f"Matrix is not square. Dimensions: {(num_rows, num_cols)}"

    id = np.identity(num_rows)

    s = 2 ** -n  # s = 1 / k
    p = id + s * mat
    for _k in range(n):
        p = np.matmul(p, p)  # Square

    return p


DIM = 3
ORIGIN_IN_WORLD = np.zeros(DIM)
#ID_ROT = Rotation.from_rotvec(ORIGIN_IN_WORLD)

ID_3 = np.identity(DIM)
EX, EY, EZ = ID_3

LX = np.array([[ 0, 0, 0], [ 0, 0,-1], [ 0, 1, 0]])
LY = np.array([[ 0, 0, 1], [ 0, 0, 0], [-1, 0, 0]])
LZ = np.array([[ 0,-1, 0], [ 1, 0, 0], [ 0, 0, 0]])
BASE_so3 = (LX, LY, LZ)


def to_so3(vec):
    assert len(vec) == DIM, vec

    components = (v * b for v, b in zip(vec, BASE_so3))

    mat = np.zeros((DIM, DIM))  # Init as zero matrix
    for c in components:
        mat = mat + c

    return mat


def rot_mat_from_rotvec(n: int, angle_axis_vector):
    """
    :param n: Accuracy parameter.
    """
    lie_algebra_element = to_so3(angle_axis_vector)

    rot_mat = matrix_exp(n, lie_algebra_element)

    return rot_mat


def _ramp(start: int, num_steps: int, idx: int) -> float:
    """
    Ramp of length 'steps' over the naturals, rising from 0 to 1.
    """
    assert num_steps > 0

    idx_rel = idx - start

    h: float = idx_rel / num_steps

    return max(0, min(h, 1))


class Config:
    EXP_ACCURACY = 10

    RADIUS = 2  # Irrelevant
    BOX_SIZE = RADIUS * 1.5

    _IDX_SIM_START = 20
    _ITER_STEPS = 50

    RAMP = (
        _ramp(Config._IDX_SIM_START, Config._ITER_STEPS, n)
        for n in naturals()
    )


def _make_random_points_3D(dim: int, num_points: int, interval: float) -> list:
    def make() -> float:
        return random.uniform(-interval, interval)

    points = [
        np.array(list(make() for _comp in range(dim)))
        for _i in range(num_points)
    ]

    return points


def _final_points_in_world():
    DIM = 3  # Necessary
    NUM_FIXED_POINTS = 3

    # Compute fixed vectors and angles
    rand_points: list = _make_random_points_3D(DIM, NUM_FIXED_POINTS, Config.BOX_SIZE)
    fixed_points, hovering_extra_point = rand_points[:-1], rand_points[-1]

    normalized_fixed_points: list = [normalize(pt) for pt in fixed_points]
    fixed_points_on_sphere: list = [Config.RADIUS * pt for pt in normalized_fixed_points]
    points_to_rotate: list = [fixed_points_on_sphere[0], hovering_extra_point]

    normalized_angle_axis_vector = normalize(np.cross(*fixed_points))

    def angle_to_rotated_points(angle: float, points):
        angle_axis_vector = angle * normalized_angle_axis_vector
        rot_mat = rot_mat_from_rotvec(Config.EXP_ACCURACY, angle_axis_vector)
        rotated_points = [np.matmul(rot_mat, pt) for pt in points]
        return rotated_points

    full_angle: float = np.arccos(_inner(*normalized_fixed_points))

    # Make generator of moving points and return generator for all points
    angles = (full_angle * fraction for fraction in Config.RAMP)
    moving_points = (
        angle_to_rotated_points(a, points_to_rotate)
        for a in angles
    )
    stream = (
        dict(
            points=pts + fixed_points + fixed_points_on_sphere,
            normalized_angle_axis_vector=normalized_angle_axis_vector,
        )
        for pts in moving_points
    )

    return stream


def _plot_point_rotation(ax, data: list) -> None:
    n = len(data) - 1
    print(f"#{n}", end=" " if n % 10 else "\n")

    _set_up_plot(ax, data, Config.BOX_SIZE)

    FULL_SPHERE = 1
    _plot_bowl(ax, Config.RADIUS, FULL_SPHERE)

    current_datum = data[-1]
    current_points = current_datum["points"]

    ra = 5 * Config.BOX_SIZE * current_datum["normalized_angle_axis_vector"]
    ax.plot3D(*transpose([-ra, ra]), "gray", linewidth=1.0)
    FIXED_POINTS_INDEX = 2
    ax.plot3D(*transpose([ORIGIN_IN_WORLD, current_points[FIXED_POINTS_INDEX + 0]]), "blue", linewidth=2.0)
    ax.plot3D(*transpose([ORIGIN_IN_WORLD, current_points[FIXED_POINTS_INDEX + 1]]), "green", linewidth=2.0)


if __name__ == "__main__":
    a = Animation(_final_points_in_world, _plot_point_rotation, fps=3)
    a.run()