masque/masque/utils/transform.py

"""
Geometric transforms
"""
from collections.abc import Sequence
from functools import lru_cache
from math import acos

import numpy
from numpy.typing import NDArray, ArrayLike
from numpy import pi


# Constants for shorthand rotations
R90 = pi / 2
R180 = pi

# Preserve the effective tolerance of the historical
# ``isclose(cos(4 * theta), 1, atol=1e-12)`` Manhattan check. Expressing it as
# an angular distance lets us canonicalize before consulting the matrix cache.
_MANHATTAN_SNAP_ATOL = acos(1 - (1e-5 + 1e-12)) / 4
_CARDINAL_ANGLES = (0.0, R90, R180, 3 * R90)


@lru_cache
def _rotation_matrix_2d(theta: float) -> NDArray[numpy.float64]:
    """Build and cache an immutable matrix for a canonicalized angle."""
    arr = numpy.array([[numpy.cos(theta), -numpy.sin(theta)],
                       [numpy.sin(theta), +numpy.cos(theta)]])

    if theta in _CARDINAL_ANGLES:
        arr = numpy.round(arr)

    arr.flags.writeable = False
    return arr


def rotation_matrix_2d(theta: float) -> NDArray[numpy.float64]:
    """
    2D rotation matrix for rotating counterclockwise around the origin.

    Args:
        theta: Angle to rotate, in radians

    Returns:
        rotation matrix
    """
    theta = float(theta)
    quarter_turn = round(theta / R90)
    if abs(theta - quarter_turn * R90) <= _MANHATTAN_SNAP_ATOL:
        theta = _CARDINAL_ANGLES[quarter_turn % 4]
    return _rotation_matrix_2d(theta)


def normalize_mirror(mirrored: Sequence[bool]) -> tuple[bool, float]:
    """
    Converts 0-2 mirror operations `(mirror_across_x_axis, mirror_across_y_axis)`
    into 0-1 mirror operations and a rotation

    Args:
        mirrored: `(mirror_across_x_axis, mirror_across_y_axis)`

    Returns:
        `mirror_across_x_axis` (bool) and
        `angle_to_rotate` in radians
    """

    if len(mirrored) != 2:
        raise ValueError(f'mirrored must be a 2-item sequence, got length {len(mirrored)}')

    mirrored_x, mirrored_y = (bool(value) for value in mirrored)
    mirror_x = (mirrored_x != mirrored_y)  # XOR
    angle = numpy.pi if mirrored_y else 0
    return mirror_x, angle


def rotate_offsets_around(
        offsets: NDArray[numpy.float64],
        pivot: NDArray[numpy.float64],
        angle: float,
        ) -> NDArray[numpy.float64]:
    """
    Rotates offsets around a pivot point.

    Args:
        offsets: Nx2 array, rows are (x, y) offsets
        pivot: (x, y) location to rotate around
        angle: rotation angle in radians

    Returns:
        Nx2 ndarray of (x, y) position after the rotation is applied.
    """
    offsets -= pivot
    offsets[:] = (rotation_matrix_2d(angle) @ offsets.T).T
    offsets += pivot
    return offsets


def apply_transforms(
        outer: ArrayLike,
        inner: ArrayLike,
        tensor: bool = False,
        ) -> NDArray[numpy.float64]:
    """
    Apply a set of transforms (`outer`) to a second set (`inner`).
    This is used to find the "absolute" transform for nested `Ref`s.

    The two transforms should be of shape Ox5 and Ix5.
    Rows should be of the form `(x_offset, y_offset, rotation_ccw_rad, mirror_across_x, scale)`.
    The output will be of the form (O*I)x5 (if `tensor=False`) or OxIx5 (`tensor=True`).

    Args:
        outer: Transforms for the container refs. Shape Ox5.
        inner: Transforms for the contained refs. Shape Ix5.
        tensor: If `True`, an OxIx5 array is returned, with `result[oo, ii, :]` corresponding
            to the `oo`th `outer` transform applied to the `ii`th inner transform.
            If `False` (default), this is concatenated into `(O*I)x5` to allow simple
            chaining into additional `apply_transforms()` calls.

    Returns:
        OxIx5 or (O*I)x5 array. Final dimension is
            `(total_x, total_y, total_rotation_ccw_rad, net_mirrored_x, total_scale)`.
    """
    outer = numpy.atleast_2d(outer).astype(float, copy=False)
    inner = numpy.atleast_2d(inner).astype(float, copy=False)

    if outer.shape[1] == 4:
        outer = numpy.pad(outer, ((0, 0), (0, 1)), constant_values=1.0)
    if inner.shape[1] == 4:
        inner = numpy.pad(inner, ((0, 0), (0, 1)), constant_values=1.0)

    if outer.shape[0] == 0 or inner.shape[0] == 0:
        if tensor:
            return numpy.empty((outer.shape[0], inner.shape[0], 5))
        return numpy.empty((0, 5))

    # If mirrored, flip y's
    xy_mir = numpy.tile(inner[:, :2], (outer.shape[0], 1, 1))   # dims are outer, inner, xyrm
    xy_mir[outer[:, 3].astype(bool), :, 1] *= -1

    # Apply outer scale to inner offset
    xy_mir *= outer[:, None, 4, None]

    rot_mats = [rotation_matrix_2d(angle) for angle in outer[:, 2]]
    xy = numpy.einsum('ort,oit->oir', rot_mats, xy_mir)

    tot = numpy.empty((outer.shape[0], inner.shape[0], 5))
    tot[:, :, :2] = outer[:, None, :2] + xy

    # If mirrored, flip inner rotation
    mirrored_outer = outer[:, None, 3].astype(bool)
    rotations = outer[:, None, 2] + numpy.where(mirrored_outer, -inner[None, :, 2], inner[None, :, 2])

    tot[:, :, 2] = rotations % (2 * pi)
    tot[:, :, 3] = (outer[:, None, 3] + inner[None, :, 3]) % 2             # net mirrored
    tot[:, :, 4] = outer[:, None, 4] * inner[None, :, 4]                   # net scale

    if tensor:
        return tot
    return numpy.concatenate(tot)