65 lines
1.6 KiB
Python
65 lines
1.6 KiB
Python
"""
|
|
Geometric transforms
|
|
"""
|
|
from collections.abc import Sequence
|
|
from functools import lru_cache
|
|
|
|
import numpy
|
|
from numpy.typing import NDArray
|
|
from numpy import pi
|
|
|
|
|
|
@lru_cache
|
|
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
|
|
"""
|
|
arr = numpy.array([[numpy.cos(theta), -numpy.sin(theta)],
|
|
[numpy.sin(theta), +numpy.cos(theta)]])
|
|
|
|
# If this was a manhattan rotation, round to remove some inacuraccies in sin & cos
|
|
if numpy.isclose(theta % (pi / 2), 0):
|
|
arr = numpy.round(arr)
|
|
|
|
arr.flags.writeable = False
|
|
return arr
|
|
|
|
|
|
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
|
|
"""
|
|
|
|
mirrored_x, mirrored_y = 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.
|
|
"""
|
|
offsets -= pivot
|
|
offsets[:] = (rotation_matrix_2d(angle) @ offsets.T).T
|
|
offsets += pivot
|
|
return offsets
|