From 615212bcfdf9ac5293f70b7af402fc0f0215d09b Mon Sep 17 00:00:00 2001 From: Allen Goodman Date: Fri, 5 Apr 2024 16:37:56 -0400 Subject: [PATCH] rotations --- .github/workflows/push.yml | 23 +- .pre-commit-config.yaml | 13 + pyproject.toml | 7 + src/beignet/__init__.py | 91 +++++++ src/beignet/_apply_euler_angle.py | 63 +++++ src/beignet/_apply_rotation_matrix.py | 45 ++++ src/beignet/_apply_rotation_quaternion.py | 63 +++++ src/beignet/_apply_rotation_vector.py | 58 +++++ src/beignet/_compose_rotation_quaternion.py | 91 +++++++ src/beignet/_euler_angle_identity.py | 70 +++++ src/beignet/_euler_angle_magnitude.py | 41 +++ .../_euler_angle_to_rotation_matrix.py | 140 ++++++++++ .../_euler_angle_to_rotation_quaternion.py | 170 ++++++++++++ .../_euler_angle_to_rotation_vector.py | 241 ++++++++++++++++++ src/beignet/_invert_euler_angle.py | 48 ++++ src/beignet/_invert_rotation_matrix.py | 19 ++ src/beignet/_invert_rotation_quaternion.py | 31 +++ src/beignet/_invert_rotation_vector.py | 20 ++ src/beignet/_mean_rotation_quaternion.py | 41 +++ src/beignet/_random_euler_angle.py | 91 +++++++ src/beignet/_random_rotation_matrix.py | 69 +++++ src/beignet/_random_rotation_quaternion.py | 98 +++++++ src/beignet/_random_rotation_vector.py | 73 ++++++ src/beignet/_rotation_matrix_identity.py | 58 +++++ src/beignet/_rotation_matrix_magnitude.py | 27 ++ .../_rotation_matrix_to_euler_angle.py | 51 ++++ ..._rotation_matrix_to_rotation_quaternion.py | 88 +++++++ .../_rotation_matrix_to_rotation_vector.py | 37 +++ src/beignet/_rotation_quaternion_identity.py | 64 +++++ src/beignet/_rotation_quaternion_magnitude.py | 44 ++++ .../_rotation_quaternion_to_euler_angle.py | 143 +++++++++++ ..._rotation_quaternion_to_rotation_matrix.py | 43 ++++ ..._rotation_quaternion_to_rotation_vector.py | 57 +++++ src/beignet/_rotation_vector_identity.py | 64 +++++ src/beignet/_rotation_vector_magnitude.py | 34 +++ .../_rotation_vector_to_euler_angle.py | 48 ++++ .../_rotation_vector_to_rotation_matrix.py | 37 +++ ..._rotation_vector_to_rotation_quaternion.py | 71 ++++++ src/beignet/_slerp.py | 176 +++++++++++++ src/beignet/_translation_vector_identity.py | 51 ++++ tests/beignet/test__apply_euler_angle.py | 73 ++++++ tests/beignet/test__apply_rotation_matrix.py | 52 ++++ .../test__apply_rotation_quaternion.py | 52 ++++ tests/beignet/test__apply_rotation_vector.py | 53 ++++ .../test__compose_rotation_quaternion.py | 58 +++++ tests/beignet/test__euler_angle_identity.py | 11 + tests/beignet/test__euler_angle_magnitude.py | 56 ++++ .../test__euler_angle_to_rotation_matrix.py | 65 +++++ ...est__euler_angle_to_rotation_quaternion.py | 75 ++++++ .../test__euler_angle_to_rotation_vector.py | 63 +++++ tests/beignet/test__invert_euler_angle.py | 56 ++++ tests/beignet/test__invert_rotation_matrix.py | 33 +++ .../test__invert_rotation_quaternion.py | 36 +++ tests/beignet/test__invert_rotation_vector.py | 35 +++ .../beignet/test__mean_rotation_quaternion.py | 52 ++++ tests/beignet/test__random_euler_angle.py | 51 ++++ tests/beignet/test__random_rotation_matrix.py | 28 ++ .../test__random_rotation_quaternion.py | 31 +++ tests/beignet/test__random_rotation_vector.py | 31 +++ .../beignet/test__rotation_matrix_identity.py | 34 +++ .../test__rotation_matrix_magnitude.py | 33 +++ .../test__rotation_matrix_to_euler_angle.py | 65 +++++ ..._rotation_matrix_to_rotation_quaternion.py | 49 ++++ ...est__rotation_matrix_to_rotation_vector.py | 45 ++++ .../test__rotation_quaternion_identity.py | 37 +++ .../test__rotation_quaternion_magnitude.py | 36 +++ ...est__rotation_quaternion_to_euler_angle.py | 67 +++++ ..._rotation_quaternion_to_rotation_matrix.py | 41 +++ ..._rotation_quaternion_to_rotation_vector.py | 46 ++++ .../beignet/test__rotation_vector_identity.py | 36 +++ .../test__rotation_vector_magnitude.py | 36 +++ .../test__rotation_vector_to_euler_angle.py | 67 +++++ ...est__rotation_vector_to_rotation_matrix.py | 42 +++ ..._rotation_vector_to_rotation_quaternion.py | 54 ++++ tests/beignet/test__slerp.py | 149 +++++++++++ 75 files changed, 4446 insertions(+), 1 deletion(-) create mode 100644 .pre-commit-config.yaml create mode 100644 src/beignet/_apply_euler_angle.py create mode 100644 src/beignet/_apply_rotation_matrix.py create mode 100644 src/beignet/_apply_rotation_quaternion.py create mode 100644 src/beignet/_apply_rotation_vector.py create mode 100644 src/beignet/_compose_rotation_quaternion.py create mode 100644 src/beignet/_euler_angle_identity.py create mode 100644 src/beignet/_euler_angle_magnitude.py create mode 100644 src/beignet/_euler_angle_to_rotation_matrix.py create mode 100644 src/beignet/_euler_angle_to_rotation_quaternion.py create mode 100644 src/beignet/_euler_angle_to_rotation_vector.py create mode 100644 src/beignet/_invert_euler_angle.py create mode 100644 src/beignet/_invert_rotation_matrix.py create mode 100644 src/beignet/_invert_rotation_quaternion.py create mode 100644 src/beignet/_invert_rotation_vector.py create mode 100644 src/beignet/_mean_rotation_quaternion.py create mode 100644 src/beignet/_random_euler_angle.py create mode 100644 src/beignet/_random_rotation_matrix.py create mode 100644 src/beignet/_random_rotation_quaternion.py create mode 100644 src/beignet/_random_rotation_vector.py create mode 100644 src/beignet/_rotation_matrix_identity.py create mode 100644 src/beignet/_rotation_matrix_magnitude.py create mode 100644 src/beignet/_rotation_matrix_to_euler_angle.py create mode 100644 src/beignet/_rotation_matrix_to_rotation_quaternion.py create mode 100644 src/beignet/_rotation_matrix_to_rotation_vector.py create mode 100644 src/beignet/_rotation_quaternion_identity.py create mode 100644 src/beignet/_rotation_quaternion_magnitude.py create mode 100644 src/beignet/_rotation_quaternion_to_euler_angle.py create mode 100644 src/beignet/_rotation_quaternion_to_rotation_matrix.py create mode 100644 src/beignet/_rotation_quaternion_to_rotation_vector.py create mode 100644 src/beignet/_rotation_vector_identity.py create mode 100644 src/beignet/_rotation_vector_magnitude.py create mode 100644 src/beignet/_rotation_vector_to_euler_angle.py create mode 100644 src/beignet/_rotation_vector_to_rotation_matrix.py create mode 100644 src/beignet/_rotation_vector_to_rotation_quaternion.py create mode 100644 src/beignet/_slerp.py create mode 100644 src/beignet/_translation_vector_identity.py create mode 100644 tests/beignet/test__apply_euler_angle.py create mode 100644 tests/beignet/test__apply_rotation_matrix.py create mode 100644 tests/beignet/test__apply_rotation_quaternion.py create mode 100644 tests/beignet/test__apply_rotation_vector.py create mode 100644 tests/beignet/test__compose_rotation_quaternion.py create mode 100644 tests/beignet/test__euler_angle_identity.py create mode 100644 tests/beignet/test__euler_angle_magnitude.py create mode 100644 tests/beignet/test__euler_angle_to_rotation_matrix.py create mode 100644 tests/beignet/test__euler_angle_to_rotation_quaternion.py create mode 100644 tests/beignet/test__euler_angle_to_rotation_vector.py create mode 100644 tests/beignet/test__invert_euler_angle.py create mode 100644 tests/beignet/test__invert_rotation_matrix.py create mode 100644 tests/beignet/test__invert_rotation_quaternion.py create mode 100644 tests/beignet/test__invert_rotation_vector.py create mode 100644 tests/beignet/test__mean_rotation_quaternion.py create mode 100644 tests/beignet/test__random_euler_angle.py create mode 100644 tests/beignet/test__random_rotation_matrix.py create mode 100644 tests/beignet/test__random_rotation_quaternion.py create mode 100644 tests/beignet/test__random_rotation_vector.py create mode 100644 tests/beignet/test__rotation_matrix_identity.py create mode 100644 tests/beignet/test__rotation_matrix_magnitude.py create mode 100644 tests/beignet/test__rotation_matrix_to_euler_angle.py create mode 100644 tests/beignet/test__rotation_matrix_to_rotation_quaternion.py create mode 100644 tests/beignet/test__rotation_matrix_to_rotation_vector.py create mode 100644 tests/beignet/test__rotation_quaternion_identity.py create mode 100644 tests/beignet/test__rotation_quaternion_magnitude.py create mode 100644 tests/beignet/test__rotation_quaternion_to_euler_angle.py create mode 100644 tests/beignet/test__rotation_quaternion_to_rotation_matrix.py create mode 100644 tests/beignet/test__rotation_quaternion_to_rotation_vector.py create mode 100644 tests/beignet/test__rotation_vector_identity.py create mode 100644 tests/beignet/test__rotation_vector_magnitude.py create mode 100644 tests/beignet/test__rotation_vector_to_euler_angle.py create mode 100644 tests/beignet/test__rotation_vector_to_rotation_matrix.py create mode 100644 tests/beignet/test__rotation_vector_to_rotation_quaternion.py create mode 100644 tests/beignet/test__slerp.py diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 241d2fee55..07fbaddeaa 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -14,6 +14,27 @@ jobs: with: name: "python-package-distributions" path: "dist/" + pytest: + strategy: + matrix: + platform: + - "macos-latest" + - "ubuntu-latest" + - "windows-latest" + python: + - "3.10" + - "3.11" + runs-on: ${{ matrix.platform }} + steps: + - uses: "actions/checkout@v4" + - uses: "actions/setup-python@v5" + with: + python-version: ${{ matrix.python }} + - run: "python -m pip install --editable '.[test]'" + - run: "python -m pytest" + - env: + CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} + uses: "codecov/codecov-action@v3" pypi: environment: name: "pypi.org" @@ -58,7 +79,7 @@ jobs: - uses: "chartboost/ruff-action@v1" with: args: "format --check" - test-pypi: + testpypi: environment: name: "test.pypi.org" url: "https://test.pypi.org/project/beignet" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000000..32e8d38b71 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,13 @@ +repos: + - hooks: + - id: "check-toml" + - id: "check-yaml" + repo: "https://github.com/pre-commit/pre-commit-hooks" + rev: "v4.5.0" + - hooks: + - args: + - "--fix" + id: "ruff" + - id: "ruff-format" + repo: "https://github.com/astral-sh/ruff-pre-commit" + rev: "v0.3.5" diff --git a/pyproject.toml b/pyproject.toml index 7f9d47afbf..4b12a3b1b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,13 @@ name = "beignet" readme = "README.md" requires-python = ">=3.10" +[project.optional-dependencies] +test = [ + "hypothesis", + "pytest", + "scipy", +] + [tool.ruff.format] docstring-code-format = true diff --git a/src/beignet/__init__.py b/src/beignet/__init__.py index 8e489b9a12..cfc5dfdcc0 100644 --- a/src/beignet/__init__.py +++ b/src/beignet/__init__.py @@ -1,5 +1,96 @@ import importlib.metadata from importlib.metadata import PackageNotFoundError +from ._apply_euler_angle import apply_euler_angle +from ._apply_rotation_matrix import apply_rotation_matrix +from ._apply_rotation_quaternion import ( + apply_rotation_quaternion, +) +from ._apply_rotation_vector import apply_rotation_vector +from ._compose_rotation_quaternion import compose_rotation_quaternion +from ._euler_angle_identity import euler_angle_identity +from ._euler_angle_magnitude import euler_angle_magnitude +from ._euler_angle_to_rotation_matrix import euler_angle_to_rotation_matrix +from ._euler_angle_to_rotation_quaternion import ( + euler_angle_to_rotation_quaternion, +) +from ._euler_angle_to_rotation_vector import euler_angle_to_rotation_vector +from ._invert_euler_angle import invert_euler_angle +from ._invert_rotation_matrix import invert_rotation_matrix +from ._invert_rotation_quaternion import invert_rotation_quaternion +from ._invert_rotation_vector import invert_rotation_vector +from ._mean_rotation_quaternion import mean_rotation_quaternion +from ._random_euler_angle import random_euler_angle +from ._random_rotation_matrix import random_rotation_matrix +from ._random_rotation_quaternion import random_rotation_quaternion +from ._random_rotation_vector import random_rotation_vector +from ._rotation_matrix_identity import rotation_matrix_identity +from ._rotation_matrix_magnitude import rotation_matrix_magnitude +from ._rotation_matrix_to_euler_angle import rotation_matrix_to_euler_angle +from ._rotation_matrix_to_rotation_quaternion import ( + rotation_matrix_to_rotation_quaternion, +) +from ._rotation_matrix_to_rotation_vector import ( + rotation_matrix_to_rotation_vector, +) +from ._rotation_quaternion_identity import rotation_quaternion_identity +from ._rotation_quaternion_magnitude import rotation_quaternion_magnitude +from ._rotation_quaternion_to_euler_angle import ( + rotation_quaternion_to_euler_angle, +) +from ._rotation_quaternion_to_rotation_matrix import ( + rotation_quaternion_to_rotation_matrix, +) +from ._rotation_quaternion_to_rotation_vector import ( + rotation_quaternion_to_rotation_vector, +) +from ._rotation_vector_identity import rotation_vector_identity +from ._rotation_vector_magnitude import rotation_vector_magnitude +from ._rotation_vector_to_euler_angle import rotation_vector_to_euler_angle +from ._rotation_vector_to_rotation_matrix import ( + rotation_vector_to_rotation_matrix, +) +from ._rotation_vector_to_rotation_quaternion import ( + rotation_vector_to_rotation_quaternion, +) +from ._slerp import slerp + +__all__ = [ + "apply_euler_angle", + "apply_rotation_matrix", + "apply_rotation_quaternion", + "apply_rotation_vector", + "compose_rotation_quaternion", + "euler_angle_identity", + "euler_angle_magnitude", + "euler_angle_to_rotation_matrix", + "euler_angle_to_rotation_quaternion", + "euler_angle_to_rotation_vector", + "invert_euler_angle", + "invert_rotation_matrix", + "invert_rotation_quaternion", + "invert_rotation_vector", + "mean_rotation_quaternion", + "random_euler_angle", + "random_rotation_matrix", + "random_rotation_quaternion", + "random_rotation_vector", + "rotation_matrix_identity", + "rotation_matrix_magnitude", + "rotation_matrix_to_euler_angle", + "rotation_matrix_to_rotation_quaternion", + "rotation_matrix_to_rotation_vector", + "rotation_quaternion_identity", + "rotation_quaternion_magnitude", + "rotation_quaternion_to_euler_angle", + "rotation_quaternion_to_rotation_matrix", + "rotation_quaternion_to_rotation_vector", + "rotation_vector_identity", + "rotation_vector_magnitude", + "rotation_vector_to_euler_angle", + "rotation_vector_to_rotation_matrix", + "rotation_vector_to_rotation_quaternion", + "slerp", +] try: __version__ = importlib.metadata.version("beignet") diff --git a/src/beignet/_apply_euler_angle.py b/src/beignet/_apply_euler_angle.py new file mode 100644 index 0000000000..8d77b2b787 --- /dev/null +++ b/src/beignet/_apply_euler_angle.py @@ -0,0 +1,63 @@ +from torch import Tensor + +from ._apply_rotation_matrix import apply_rotation_matrix +from ._euler_angle_to_rotation_matrix import euler_angle_to_rotation_matrix + + +def apply_euler_angle( + input: Tensor, + rotation: Tensor, + axes: str, + degrees: bool = False, + inverse: bool = False, +) -> Tensor: + r""" + Rotates vectors in three-dimensional space using Euler angles. + + Note + ---- + This function interprets the rotation of the original frame to the final + frame as either a projection, where it maps the components of vectors from + the final frame to the original frame, or as a physical rotation, + integrating the vectors into the original frame during the rotation + process. Consequently, the vector components are maintained in the original + frame’s perspective both before and after the rotation. + + Parameters + ---------- + input : Tensor + Vectors in three-dimensional space with the shape $(\ldots \times 3)$. + Euler angles and vectors must conform to PyTorch broadcasting rules. + + rotation : Tensor + Euler angles with the shape $(\ldots \times 3)$, specifying the + rotation in three-dimensional space. + + axes : str + Specifies the sequence of axes for the rotations, using one to three + characters from the set ${X, Y, Z}$ for intrinsic rotations, or + ${x, y, z}$ for extrinsic rotations. Mixing extrinsic and intrinsic + rotations raises a `ValueError`. + + degrees : bool, optional + Indicates whether the Euler angles are provided in degrees. If `False`, + angles are assumed to be in radians. Default, `False`. + + inverse : bool, optional + If `True`, applies the inverse rotation using the Euler angles to the + input vectors. Default, `False`. + + Returns + ------- + rotated_vectors : Tensor + A tensor of the same shape as `input`, containing the rotated vectors. + """ + return apply_rotation_matrix( + input, + euler_angle_to_rotation_matrix( + rotation, + axes, + degrees, + ), + inverse, + ) diff --git a/src/beignet/_apply_rotation_matrix.py b/src/beignet/_apply_rotation_matrix.py new file mode 100644 index 0000000000..ced6bb16ad --- /dev/null +++ b/src/beignet/_apply_rotation_matrix.py @@ -0,0 +1,45 @@ +import torch +from torch import Tensor + + +def apply_rotation_matrix( + input: Tensor, + rotation: Tensor, + inverse: bool | None = False, +) -> Tensor: + r""" + Rotates vectors in three-dimensional space using rotation matrices. + + Note + ---- + This function interprets the rotation of the original frame to the final + frame as either a projection, where it maps the components of vectors from + the final frame to the original frame, or as a physical rotation, + integrating the vectors into the original frame during the rotation + process. Consequently, the vector components are maintained in the original + frame’s perspective both before and after the rotation. + + Parameters + ---------- + input : Tensor, shape (..., 3) + Each vector represents a vector in three-dimensional space. The number + of rotation matrices and number of vectors must follow standard + broadcasting rules: either one of them equals unity or they both equal + each other. + + rotation : Tensor, shape (..., 3, 3) + Rotation matrices. + + inverse : bool, optional + If `True` the inverse of the rotation matrices are applied to the input + vectors. Default, `False`. + + Returns + ------- + rotated_vectors : Tensor, shape (..., 3) + Rotated vectors. + """ + if inverse: + return torch.einsum("ikj, ik -> ij", rotation, input) + + return torch.einsum("ijk, ik -> ij", rotation, input) diff --git a/src/beignet/_apply_rotation_quaternion.py b/src/beignet/_apply_rotation_quaternion.py new file mode 100644 index 0000000000..dd7548aaee --- /dev/null +++ b/src/beignet/_apply_rotation_quaternion.py @@ -0,0 +1,63 @@ +import torch +from torch import Tensor + +from ._rotation_quaternion_to_rotation_matrix import ( + rotation_quaternion_to_rotation_matrix, +) + + +def apply_rotation_quaternion( + input: Tensor, + rotation: Tensor, + inverse: bool | None = False, +) -> Tensor: + r""" + Rotates vectors in three-dimensional space using rotation quaternions. + + Note + ---- + This function interprets the rotation of the original frame to the final + frame as either a projection, where it maps the components of vectors from + the final frame to the original frame, or as a physical rotation, + integrating the vectors into the original frame during the rotation + process. Consequently, the vector components are maintained in the original + frame’s perspective both before and after the rotation. + + Parameters + ---------- + input : Tensor, shape (..., 3) + Each vector represents a vector in three-dimensional space. The number + of rotation quaternions and number of vectors must follow standard + broadcasting rules: either one of them equals unity or they both equal + each other. + + rotation : Tensor, shape (..., 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + inverse : bool, optional + If `True` the inverse of the rotation quaternions are applied to the + input vectors. Default, `False`. + + Returns + ------- + output : Tensor, shape (..., 3) + Rotated vectors. + """ + if inverse: + output = torch.einsum( + "ikj, ik -> ij", + rotation_quaternion_to_rotation_matrix( + rotation, + ), + input, + ) + else: + output = torch.einsum( + "ijk, ik -> ij", + rotation_quaternion_to_rotation_matrix( + rotation, + ), + input, + ) + + return output diff --git a/src/beignet/_apply_rotation_vector.py b/src/beignet/_apply_rotation_vector.py new file mode 100644 index 0000000000..15c648fab0 --- /dev/null +++ b/src/beignet/_apply_rotation_vector.py @@ -0,0 +1,58 @@ +from torch import Tensor + +from ._apply_rotation_matrix import apply_rotation_matrix +from ._rotation_vector_to_rotation_matrix import ( + rotation_vector_to_rotation_matrix, +) + + +def apply_rotation_vector( + input: Tensor, + rotation: Tensor, + degrees: bool | None = False, + inverse: bool | None = False, +) -> Tensor: + r""" + Rotates vectors in three-dimensional space using rotation vectors. + + Note + ---- + This function interprets the rotation of the original frame to the final + frame as either a projection, where it maps the components of vectors from + the final frame to the original frame, or as a physical rotation, + integrating the vectors into the original frame during the rotation + process. Consequently, the vector components are maintained in the original + frame’s perspective both before and after the rotation. + + Parameters + ---------- + input : Tensor, shape (..., 3) + Each vector represents a vector in three-dimensional space. The number + of rotation vectors and number of vectors must follow standard + broadcasting rules: either one of them equals unity or they both equal + each other. + + rotation : Tensor, shape (..., 4) + Rotation vectors. + + degrees : bool, optional + If `True`, rotation vector magnitudes are assumed to be in degrees. + Default, `False`. + + inverse : bool, optional + If `True` the inverse of the rotation vectors are applied to the input + vectors. Default, `False`. + + Returns + ------- + rotated_vectors : Tensor, shape (..., 3) + Rotated vectors. + """ + return apply_rotation_matrix( + input, + rotation_vector_to_rotation_matrix( + rotation, + degrees, + ), + inverse, + ) diff --git a/src/beignet/_compose_rotation_quaternion.py b/src/beignet/_compose_rotation_quaternion.py new file mode 100644 index 0000000000..47c1cd8a79 --- /dev/null +++ b/src/beignet/_compose_rotation_quaternion.py @@ -0,0 +1,91 @@ +import torch +from torch import Tensor + + +def compose_rotation_quaternion( + input: Tensor, + other: Tensor, + canonical: bool = False, +) -> Tensor: + """ + Mean rotation quaternions. + + Parameters + ---------- + input : Tensor, shape=(..., 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + other : Tensor, shape=(..., 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + canonical : bool, optional + Whether to map the redundant double cover of rotation space to a unique + canonical single cover. If `True`, then the rotation quaternion is + chosen from :math:`{q, -q}` such that the :math:`w` term is positive. + If the :math:`w` term is :math:`0`, then the rotation quaternion is + chosen such that the first non-zero term of the :math:`x`, :math:`y`, + and :math:`z` terms is positive. + + Returns + ------- + output : Tensor, shape=(..., 4) + Mean rotation quaternions. + """ + output = torch.empty( + [max(input.shape[0], other.shape[0]), 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + cross = torch.empty( + [3], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j in range(max(input.shape[0], other.shape[0])): + p = input[j, 0] + q = input[j, 1] + r = input[j, 2] + s = input[j, 3] + + t = other[j, 0] + u = other[j, 1] + v = other[j, 2] + w = other[j, 3] + + cross[0] = q * v - r * u + cross[1] = r * t - p * v + cross[2] = p * u - q * t + + output[j, 0] = s * t + w * p + cross[0] + output[j, 1] = s * u + w * q + cross[1] + output[j, 2] = s * v + w * r + cross[2] + output[j, 3] = s * w - p * t - q * u - r * v + + for j in range(input.shape[0]): + p = output[j, 0] + q = output[j, 1] + v = output[j, 2] + r = output[j, 3] + + norm = torch.sqrt(p**2.0 + q**2.0 + v**2.0 + r**2.0) + + if norm == 0.0: + output[j] = torch.nan + + output[j] = output[j] / norm + + if canonical: + for j in range(output.shape[0]): + a = output[j, 0] + b = output[j, 1] + c = output[j, 2] + d = output[j, 3] + + if d == 0 and (a == 0 & (b == 0 & c < 0 | b < 0) | a < 0) | d < 0: + output[j] = -output[j] + + return output diff --git a/src/beignet/_euler_angle_identity.py b/src/beignet/_euler_angle_identity.py new file mode 100644 index 0000000000..37337806dc --- /dev/null +++ b/src/beignet/_euler_angle_identity.py @@ -0,0 +1,70 @@ +import torch +from torch import Tensor + +from ._rotation_quaternion_identity import rotation_quaternion_identity +from ._rotation_quaternion_to_euler_angle import ( + rotation_quaternion_to_euler_angle, +) + + +def euler_angle_identity( + size: int, + axes: str, + degrees: bool | None = False, + *, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, +) -> Tensor: + """ + Identity Euler angles. + + Parameters + ---------- + size : int + Output size. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees. Default, `False`. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + Returns + ------- + identity_euler_angles : Tensor, shape (size, 3) + Identity Euler angles. + """ + return rotation_quaternion_to_euler_angle( + rotation_quaternion_identity( + size, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ), + axes, + degrees, + ) diff --git a/src/beignet/_euler_angle_magnitude.py b/src/beignet/_euler_angle_magnitude.py new file mode 100644 index 0000000000..901f677222 --- /dev/null +++ b/src/beignet/_euler_angle_magnitude.py @@ -0,0 +1,41 @@ +from torch import Tensor + +from ._euler_angle_to_rotation_quaternion import ( + euler_angle_to_rotation_quaternion, +) +from ._rotation_quaternion_magnitude import rotation_quaternion_magnitude + + +def euler_angle_magnitude( + input: Tensor, + axes: str, + degrees: bool | None = False, +) -> Tensor: + r""" + Euler angle magnitudes. + + Parameters + ---------- + input : Tensor, shape (..., 3) + Euler angles. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees. Default, `False`. + + Returns + ------- + euler_angle_magnitudes: Tensor, shape (...) + Angles in radians. Magnitudes will be in the range :math:`[0, \pi]`. + """ + return rotation_quaternion_magnitude( + euler_angle_to_rotation_quaternion( + input, + axes, + degrees, + ), + ) diff --git a/src/beignet/_euler_angle_to_rotation_matrix.py b/src/beignet/_euler_angle_to_rotation_matrix.py new file mode 100644 index 0000000000..02dcf74792 --- /dev/null +++ b/src/beignet/_euler_angle_to_rotation_matrix.py @@ -0,0 +1,140 @@ +import torch +from torch import Tensor + + +def euler_angle_to_rotation_matrix( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + """ + Convert Euler angles to rotation matrices. + + Parameters + ---------- + input : Tensor, shape (..., 3) + Euler angles. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees. Default, `False`. + + Returns + ------- + rotation_matrices : Tensor, shape (..., 3, 3) + Rotation matrices. + """ + if degrees: + input = torch.deg2rad(input) + + for index, axis in enumerate(axes): + cos = torch.cos(input[..., index]) + sin = torch.sin(input[..., index]) + + match axis.lower(): + case "x": + matrix = torch.stack( + [ + torch.stack( + [ + torch.full_like(cos, 1.0), + torch.full_like(cos, 0.0), + torch.full_like(cos, 0.0), + ], + dim=-1, + ), + torch.stack( + [ + torch.full_like(cos, 0.0), + cos, + torch.negative(sin), + ], + dim=-1, + ), + torch.stack( + [ + torch.full_like(cos, 0.0), + sin, + cos, + ], + dim=-1, + ), + ], + dim=-2, + ) + case "y": + matrix = torch.stack( + [ + torch.stack( + [ + cos, + torch.full_like(cos, 0.0), + sin, + ], + dim=-1, + ), + torch.stack( + [ + torch.full_like(cos, 0.0), + torch.full_like(cos, 1.0), + torch.full_like(cos, 0.0), + ], + dim=-1, + ), + torch.stack( + [ + torch.negative(sin), + torch.full_like(cos, 0.0), + cos, + ], + dim=-1, + ), + ], + dim=-2, + ) + case "z": + matrix = torch.stack( + [ + torch.stack( + [ + cos, + torch.negative(sin), + torch.full_like(cos, 0.0), + ], + dim=-1, + ), + torch.stack( + [ + sin, + cos, + torch.full_like(cos, 0.0), + ], + dim=-1, + ), + torch.stack( + [ + torch.full_like(cos, 0.0), + torch.full_like(cos, 0.0), + torch.full_like(cos, 1.0), + ], + dim=-1, + ), + ], + dim=-2, + ) + case _: + raise ValueError + + if index == 0: + rotation_matrix = matrix + else: + if axes.islower(): + rotation_matrix = matrix @ rotation_matrix + else: + rotation_matrix = rotation_matrix @ matrix + + return rotation_matrix diff --git a/src/beignet/_euler_angle_to_rotation_quaternion.py b/src/beignet/_euler_angle_to_rotation_quaternion.py new file mode 100644 index 0000000000..7f932825ba --- /dev/null +++ b/src/beignet/_euler_angle_to_rotation_quaternion.py @@ -0,0 +1,170 @@ +import re + +import torch +from torch import Tensor + + +def euler_angle_to_rotation_quaternion( + input: Tensor, + axes: str, + degrees: bool = False, + canonical: bool | None = False, +) -> Tensor: + """ + Convert Euler angles to rotation quaternions. + + Parameters + ---------- + input : Tensor + Euler angles. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees. Default, `False`. + + canonical : bool, optional + Whether to map the redundant double cover of rotation space to a unique + canonical single cover. If `True`, then the rotation quaternion is + chosen from :math:`{q, -q}` such that the :math:`w` term is positive. + If the :math:`w` term is :math:`0`, then the rotation quaternion is + chosen such that the first non-zero term of the :math:`x`, :math:`y`, + and :math:`z` terms is positive. + + Returns + ------- + rotation_quaternions : Tensor, shape (..., 4) + Rotation quaternions. + """ + intrinsic = re.match(r"^[XYZ]{1,3}$", axes) is not None + + if degrees: + input = torch.deg2rad(input) + + if len(axes) == 1: + if input.ndim == 0: + input = input.reshape((1, 1)) + elif input.ndim == 1: + input = input[:, None] + elif input.ndim == 2 and input.shape[-1] != 1: + raise ValueError + elif input.ndim > 2: + raise ValueError + else: + if input.ndim not in [1, 2] or input.shape[-1] != len(axes): + raise ValueError + + if input.ndim == 1: + input = input[None, :] + + if input.ndim != 2 or input.shape[-1] != len(axes): + raise ValueError + + output = torch.zeros( + [input.shape[0], 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + match axes.lower()[0]: + case "x": + k = 0 + case "y": + k = 1 + case "z": + k = 2 + case _: + raise ValueError + + for j in range(input[:, 0].shape[0]): + output[j, 3] = torch.cos(input[:, 0][j] / 2) + output[j, k] = torch.sin(input[:, 0][j] / 2) + + z = output + + c = torch.empty( + [3], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j in range(1, len(axes.lower())): + y = torch.zeros( + [input.shape[0], 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + r = torch.empty( + [max(y.shape[0], z.shape[0]), 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + match axes.lower()[j]: + case "x": + p = 0 + case "y": + p = 1 + case "z": + p = 2 + case _: + raise ValueError + + for k in range(input[:, j].shape[0]): + y[k, 3] = torch.cos(input[:, j][k] / 2) + y[k, p] = torch.sin(input[:, j][k] / 2) + + if intrinsic: + for k in range(max(y.shape[0], z.shape[0])): + c[0] = z[k][:3][1] * y[k][:3][2] - z[k][:3][2] * y[k][:3][1] + c[1] = z[k][:3][2] * y[k][:3][0] - z[k][:3][0] * y[k][:3][2] + c[2] = z[k][:3][0] * y[k][:3][1] - z[k][:3][1] * y[k][:3][0] + + t = z[k, 0] + u = z[k, 1] + v = z[k, 2] + w = z[k, 3] + + r[k, 0] = w * y[k, 0] + y[k, 3] * t + c[0] + r[k, 1] = w * y[k, 1] + y[k, 3] * u + c[1] + r[k, 2] = w * y[k, 2] + y[k, 3] * v + c[2] + r[k, 3] = w * y[k, 3] - t * y[k, 0] - u * y[k, 1] - v * y[k, 2] + + z = r + else: + for k in range(max(y.shape[0], z.shape[0])): + c[0] = y[k][:3][1] * z[k][:3][2] - y[k][:3][2] * z[k][:3][1] + c[1] = y[k][:3][2] * z[k][:3][0] - y[k][:3][0] * z[k][:3][2] + c[2] = y[k][:3][0] * z[k][:3][1] - y[k][:3][1] * z[k][:3][0] + + t = z[k, 0] + u = z[k, 1] + v = z[k, 2] + w = z[k, 3] + + r[k, 0] = y[k, 3] * t + w * y[k, 0] + c[0] + r[k, 1] = y[k, 3] * u + w * y[k, 1] + c[1] + r[k, 2] = y[k, 3] * v + w * y[k, 2] + c[2] + r[k, 3] = y[k, 3] * w - y[k, 0] * t - y[k, 1] * u - y[k, 2] * v + + z = r + + if canonical: + for j in range(z.shape[0]): + a = z[j, 0] + b = z[j, 1] + c = z[j, 2] + d = z[j, 3] + + if d == 0 and (a == 0 & (b == 0 & c < 0 | b < 0) | a < 0) | d < 0: + z[j] = -z[j] + + return z diff --git a/src/beignet/_euler_angle_to_rotation_vector.py b/src/beignet/_euler_angle_to_rotation_vector.py new file mode 100644 index 0000000000..5a3e0831c6 --- /dev/null +++ b/src/beignet/_euler_angle_to_rotation_vector.py @@ -0,0 +1,241 @@ +import re + +import torch +from torch import Tensor + + +def _quat_canonical_single(q): + if ( + (q[3] < 0) + or (q[3] == 0 and q[0] < 0) + or (q[3] == 0 and q[0] == 0 and q[1] < 0) + or (q[3] == 0 and q[0] == 0 and q[1] == 0 and q[2] < 0) + ): + q[0] *= -1.0 + q[1] *= -1.0 + q[2] *= -1.0 + q[3] *= -1.0 + + return q + + +def _quat_canonical(q): + n = q.shape[0] + for ind in range(n): + q[ind] = _quat_canonical_single(q[ind]) + + return q + + +def _cross3(a, b): + cross = torch.empty([3], dtype=a.dtype, device=a.device) + + cross[0] = a[1] * b[2] - a[2] * b[1] + cross[1] = a[2] * b[0] - a[0] * b[2] + cross[2] = a[0] * b[1] - a[1] * b[0] + + return cross + + +def _compose_quat_single(p, q, r): + a = p[:3] + b = q[:3] + + cross = torch.empty([3], dtype=a.dtype, device=a.device) + + cross[0] = a[1] * b[2] - a[2] * b[1] + cross[1] = a[2] * b[0] - a[0] * b[2] + cross[2] = a[0] * b[1] - a[1] * b[0] + + r[0] = p[3] * q[0] + q[3] * p[0] + cross[0] + r[1] = p[3] * q[1] + q[3] * p[1] + cross[1] + r[2] = p[3] * q[2] + q[3] * p[2] + cross[2] + r[3] = p[3] * q[3] - p[0] * q[0] - p[1] * q[1] - p[2] * q[2] + + return r + + +def _compose_quat(p, q): + n = max(p.shape[0], q.shape[0]) + + r = torch.empty([n, 4], dtype=p.dtype, device=p.device) + + if p.shape[0] == 1: + for j in range(n): + r[j] = _compose_quat_single(p[0], q[j], r[j]) + elif q.shape[0] == 1: + for j in range(n): + r[j] = _compose_quat_single(p[j], q[0], r[j]) + else: + for j in range(n): + r[j] = _compose_quat_single(p[j], q[j], r[j]) + + return r + + +def _make_elementary_quat(input: Tensor, axis: str) -> Tensor: + quaternion = torch.zeros( + [input.shape[0], 4], + dtype=input.dtype, + device=input.device, + ) + + match axis: + case "x": + k = 0 + case "y": + k = 1 + case "z": + k = 2 + case _: + raise ValueError + + for j in range(input.shape[0]): + quaternion[j, 3] = torch.cos(input[j] / 2) + quaternion[j, k] = torch.sin(input[j] / 2) + + return quaternion + + +def _elementary_quat_compose(seq, angles, intrinsic=False): + result = _make_elementary_quat(angles[:, 0], seq[0]) + + seq_len = len(seq) + + for idx in range(1, seq_len): + if intrinsic: + result = _compose_quat( + result, _make_elementary_quat(angles[:, idx], seq[idx]) + ) + else: + result = _compose_quat( + _make_elementary_quat(angles[:, idx], seq[idx]), result + ) + + return result + + +def _format_angles(input, degrees, num_axes): + if degrees: + input = torch.deg2rad(input) + + if num_axes == 1: + if input.ndim == 0: + input = input.reshape((1, 1)) + elif input.ndim == 1: + input = input[:, None] + elif input.ndim == 2 and input.shape[-1] != 1: + raise ValueError + elif input.ndim > 2: + raise ValueError + else: + if input.ndim not in [1, 2] or input.shape[-1] != num_axes: + raise ValueError + + if input.ndim == 1: + input = input[None, :] + + if input.ndim != 2 or input.shape[-1] != num_axes: + raise ValueError + + return input + + +def _quat_canonical_single(q): + if ( + (q[3] < 0) + or (q[3] == 0 and q[0] < 0) + or (q[3] == 0 and q[0] == 0 and q[1] < 0) + or (q[3] == 0 and q[0] == 0 and q[1] == 0 and q[2] < 0) + ): + q[0] *= -1.0 + q[1] *= -1.0 + q[2] *= -1.0 + q[3] *= -1.0 + + return q + + +def _dot3(a, b): + return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] + + +def _norm3(elems): + return torch.sqrt(_dot3(elems, elems)) + + +def euler_angle_to_rotation_vector( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + """ + Convert Euler angles to rotation vectors. + + Parameters + ---------- + input : Tensor + Euler angles. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees and returned + rotation vector magnitudes are in degrees. Default, `False`. + + Returns + ------- + rotation_vector : Tensor, shape=(..., 3) + Rotation vectors. + """ + num_axes = len(axes) + + if num_axes < 1 or num_axes > 3: + raise ValueError + + intrinsic = re.match(r"^[XYZ]{1,3}$", axes) is not None + extrinsic = re.match(r"^[xyz]{1,3}$", axes) is not None + + if not (intrinsic or extrinsic): + raise ValueError + + if any(axes[i] == axes[i + 1] for i in range(num_axes - 1)): + raise ValueError + + axes = axes.lower() + + input = _format_angles(input, degrees, num_axes) + + quaternion = _elementary_quat_compose(axes, input, intrinsic) + + output = torch.empty( + [quaternion.shape[0], 3], + dtype=input.dtype, + device=input.device, + ) + + for j in range(quaternion.shape[0]): + rotation_quaternion = quaternion[j, :].clone() + + _quat_canonical_single(rotation_quaternion) + + angle = 2.0 * torch.atan2(_norm3(rotation_quaternion), rotation_quaternion[3]) + + if angle <= 0.001: + angle2 = torch.square(angle) + + scale = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880 + else: + scale = angle / torch.sin(angle / 2.0) + + output[j, 0] = scale * rotation_quaternion[0] + output[j, 1] = scale * rotation_quaternion[1] + output[j, 2] = scale * rotation_quaternion[2] + + if degrees: + output = torch.rad2deg(output) + + return output diff --git a/src/beignet/_invert_euler_angle.py b/src/beignet/_invert_euler_angle.py new file mode 100644 index 0000000000..f19e0779e8 --- /dev/null +++ b/src/beignet/_invert_euler_angle.py @@ -0,0 +1,48 @@ +from torch import Tensor + +from ._euler_angle_to_rotation_quaternion import ( + euler_angle_to_rotation_quaternion, +) +from ._invert_rotation_quaternion import invert_rotation_quaternion +from ._rotation_quaternion_to_euler_angle import ( + rotation_quaternion_to_euler_angle, +) + + +def invert_euler_angle( + input: Tensor, + axes: str, + degrees: bool | None = False, +) -> Tensor: + """ + Invert Euler angles. + + Parameters + ---------- + input : Tensor, shape (..., 3) + Euler angles. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees. Default, `False`. + + Returns + ------- + inverted_euler_angles : Tensor, shape (..., 3) + Inverted Euler angles. + """ + return rotation_quaternion_to_euler_angle( + invert_rotation_quaternion( + euler_angle_to_rotation_quaternion( + input, + axes, + degrees, + ), + ), + axes, + degrees, + ) diff --git a/src/beignet/_invert_rotation_matrix.py b/src/beignet/_invert_rotation_matrix.py new file mode 100644 index 0000000000..b6aa259bed --- /dev/null +++ b/src/beignet/_invert_rotation_matrix.py @@ -0,0 +1,19 @@ +import torch +from torch import Tensor + + +def invert_rotation_matrix(input: Tensor) -> Tensor: + """ + Invert rotation matrices. + + Parameters + ---------- + input : Tensor, shape (..., 3, 3) + Rotation matrices. + + Returns + ------- + inverted_rotation_matrices : Tensor, shape (..., 3, 3) + Inverted rotation matrices. + """ + return torch.transpose(input, -2, -1) diff --git a/src/beignet/_invert_rotation_quaternion.py b/src/beignet/_invert_rotation_quaternion.py new file mode 100644 index 0000000000..856a6b226f --- /dev/null +++ b/src/beignet/_invert_rotation_quaternion.py @@ -0,0 +1,31 @@ +from torch import Tensor + + +def invert_rotation_quaternion( + input: Tensor, + canonical: bool = False, +) -> Tensor: + """ + Invert rotation quaternions. + + Parameters + ---------- + input : Tensor, shape (..., 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + canonical : bool, optional + Whether to map the redundant double cover of rotation space to a unique + canonical single cover. If `True`, then the rotation quaternion is + chosen from :math:`{q, -q}` such that the :math:`w` term is positive. + If the :math:`w` term is :math:`0`, then the rotation quaternion is + chosen such that the first non-zero term of the :math:`x`, :math:`y`, + and :math:`z` terms is positive. + + Returns + ------- + inverted_rotation_quaternions : Tensor, shape (..., 4) + Inverted rotation quaternions. + """ + input[:, :3] = -input[:, :3] + + return input diff --git a/src/beignet/_invert_rotation_vector.py b/src/beignet/_invert_rotation_vector.py new file mode 100644 index 0000000000..4d886aad6f --- /dev/null +++ b/src/beignet/_invert_rotation_vector.py @@ -0,0 +1,20 @@ +from torch import Tensor + + +def invert_rotation_vector( + input: Tensor, +) -> Tensor: + """ + Invert rotation vectors. + + Parameters + ---------- + input : Tensor, shape (..., 3) + Rotation vectors. + + Returns + ------- + inverted_rotation_vectors : Tensor, shape (..., 3) + Inverted rotation vectors. + """ + return -input diff --git a/src/beignet/_mean_rotation_quaternion.py b/src/beignet/_mean_rotation_quaternion.py new file mode 100644 index 0000000000..b3284c85a9 --- /dev/null +++ b/src/beignet/_mean_rotation_quaternion.py @@ -0,0 +1,41 @@ +import torch +from torch import Tensor + + +def mean_rotation_quaternion( + input: Tensor, + weight: Tensor | None = None, + canonical: bool = False, +) -> Tensor: + """ + Mean rotation quaternions. + + Parameters + ---------- + input : Tensor, shape=(..., 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + weight : Tensor, shape=(..., 4), optional + Relative importance of rotation quaternions. + + canonical : bool, optional + Whether to map the redundant double cover of rotation space to a unique + canonical single cover. If `True`, then the rotation quaternion is + chosen from :math:`{q, -q}` such that the :math:`w` term is positive. + If the :math:`w` term is :math:`0`, then the rotation quaternion is + chosen such that the first non-zero term of the :math:`x`, :math:`y`, + and :math:`z` terms is positive. + + Returns + ------- + output : Tensor, shape=(..., 4) + Mean rotation quaternions. + """ + if weight is None: + weight = torch.ones(input.shape[0]) + + _, output = torch.linalg.eigh((input.T * weight) @ input) + + output = output[:, -1] + + return output diff --git a/src/beignet/_random_euler_angle.py b/src/beignet/_random_euler_angle.py new file mode 100644 index 0000000000..9477148520 --- /dev/null +++ b/src/beignet/_random_euler_angle.py @@ -0,0 +1,91 @@ +import torch +from torch import Generator, Tensor + +from ._random_rotation_quaternion import random_rotation_quaternion +from ._rotation_quaternion_to_euler_angle import ( + rotation_quaternion_to_euler_angle, +) + + +def random_euler_angle( + size: int, + axes: str, + degrees: bool | None = False, + *, + generator: Generator | None = None, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, + pin_memory: bool | None = False, +) -> Tensor: + """ + Generate random Euler angles. + + Parameters + ---------- + size : int + Output size. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees. Default, `False`. + + generator : torch.Generator, optional + Psuedo-random number generator. Default, `None`. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + pin_memory : bool, optional + If `True`, returned tensor is allocated in pinned memory. Default, + `False`. + + Returns + ------- + random_euler_angles : Tensor, shape (..., 3) + Random Euler angles. + + The returned Euler angles are in the range: + + * First angle: :math:`(-180, 180]` degrees (inclusive) + * Second angle: + * :math:`[-90, 90]` degrees if all axes are different + (e.g., :math:`xyz`) + * :math:`[0, 180]` degrees if first and third axes are + the same (e.g., :math:`zxz`) + * Third angle: :math:`[-180, 180]` degrees (inclusive) + """ + return rotation_quaternion_to_euler_angle( + random_rotation_quaternion( + size, + generator=generator, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + pin_memory=pin_memory, + ), + axes, + degrees, + ) diff --git a/src/beignet/_random_rotation_matrix.py b/src/beignet/_random_rotation_matrix.py new file mode 100644 index 0000000000..0bd654529a --- /dev/null +++ b/src/beignet/_random_rotation_matrix.py @@ -0,0 +1,69 @@ +import torch +from torch import Generator, Tensor + +from ._random_rotation_quaternion import random_rotation_quaternion +from ._rotation_quaternion_to_rotation_matrix import ( + rotation_quaternion_to_rotation_matrix, +) + + +def random_rotation_matrix( + size: int, + *, + generator: Generator | None = None, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, + pin_memory: bool | None = False, +) -> Tensor: + """ + Generate random rotation matrices. + + Parameters + ---------- + size : int + Output size. + + generator : torch.Generator, optional + Psuedo-random number generator. Default, `None`. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + pin_memory : bool, optional + If `True`, returned tensor is allocated in pinned memory. Default, + `False`. + + Returns + ------- + random_rotation_matrices : Tensor, shape (..., 3, 3) + Random rotation matrices. + """ + return rotation_quaternion_to_rotation_matrix( + random_rotation_quaternion( + size, + generator=generator, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + pin_memory=pin_memory, + ), + ) diff --git a/src/beignet/_random_rotation_quaternion.py b/src/beignet/_random_rotation_quaternion.py new file mode 100644 index 0000000000..4d9683d55e --- /dev/null +++ b/src/beignet/_random_rotation_quaternion.py @@ -0,0 +1,98 @@ +import torch +from torch import Generator, Tensor + + +def random_rotation_quaternion( + size: int, + canonical: bool = False, + *, + generator: Generator | None = None, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, + pin_memory: bool | None = False, +) -> Tensor: + """ + Generate random rotation quaternions. + + Parameters + ---------- + size : int + Output size. + + canonical : bool, optional + Whether to map the redundant double cover of rotation space to a unique + canonical single cover. If `True`, then the rotation quaternion is + chosen from :math:`{q, -q}` such that the :math:`w` term is positive. + If the :math:`w` term is :math:`0`, then the rotation quaternion is + chosen such that the first non-zero term of the :math:`x`, :math:`y`, + and :math:`z` terms is positive. + + generator : torch.Generator, optional + Psuedo-random number generator. Default, `None`. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + pin_memory : bool, optional + If `True`, returned tensor is allocated in pinned memory. Default, + `False`. + + Returns + ------- + random_rotation_quaternions : Tensor, shape (..., 4) + Random rotation quaternions. + """ + rotation_quaternions = torch.rand( + [size, 4], + generator=generator, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + pin_memory=pin_memory, + ) + + if canonical: + for index in range(rotation_quaternions.size(0)): + if ( + (rotation_quaternions[index][3] < 0) + or ( + rotation_quaternions[index][3] == 0 + and rotation_quaternions[index][0] < 0 + ) + or ( + rotation_quaternions[index][3] == 0 + and rotation_quaternions[index][0] == 0 + and rotation_quaternions[index][1] < 0 + ) + or ( + rotation_quaternions[index][3] == 0 + and rotation_quaternions[index][0] == 0 + and rotation_quaternions[index][1] == 0 + and rotation_quaternions[index][2] < 0 + ) + ): + rotation_quaternions[index][0] *= -1.0 + rotation_quaternions[index][1] *= -1.0 + rotation_quaternions[index][2] *= -1.0 + rotation_quaternions[index][3] *= -1.0 + + return rotation_quaternions diff --git a/src/beignet/_random_rotation_vector.py b/src/beignet/_random_rotation_vector.py new file mode 100644 index 0000000000..dce8f276c4 --- /dev/null +++ b/src/beignet/_random_rotation_vector.py @@ -0,0 +1,73 @@ +import torch +from torch import Generator, Tensor + +from ._random_rotation_quaternion import random_rotation_quaternion +from ._rotation_quaternion_to_rotation_vector import ( + rotation_quaternion_to_rotation_vector, +) + + +def random_rotation_vector( + size: int, + degrees: bool = False, + *, + generator: Generator | None = None, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, + pin_memory: bool | None = False, +) -> Tensor: + """ + Generate random rotation vectors. + + Parameters + ---------- + size : int + Output size. + + degrees + + generator : torch.Generator, optional + Psuedo-random number generator. Default, `None`. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + pin_memory : bool, optional + If `True`, returned tensor is allocated in pinned memory. Default, + `False`. + + Returns + ------- + random_rotation_vectors : Tensor, shape (..., 3) + Random rotation vectors. + """ + return rotation_quaternion_to_rotation_vector( + random_rotation_quaternion( + size, + generator=generator, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + pin_memory=pin_memory, + ), + degrees, + ) diff --git a/src/beignet/_rotation_matrix_identity.py b/src/beignet/_rotation_matrix_identity.py new file mode 100644 index 0000000000..0e958fd03a --- /dev/null +++ b/src/beignet/_rotation_matrix_identity.py @@ -0,0 +1,58 @@ +import torch +from torch import Tensor + +from ._rotation_quaternion_identity import rotation_quaternion_identity +from ._rotation_quaternion_to_rotation_matrix import ( + rotation_quaternion_to_rotation_matrix, +) + + +def rotation_matrix_identity( + size: int, + *, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, +) -> Tensor: + """ + Identity rotation matrices. + + Parameters + ---------- + size : int + Output size. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + Returns + ------- + identity_rotation_matrices : Tensor, shape (size, 3, 3) + Identity rotation matrices. + """ + return rotation_quaternion_to_rotation_matrix( + rotation_quaternion_identity( + size, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ), + ) diff --git a/src/beignet/_rotation_matrix_magnitude.py b/src/beignet/_rotation_matrix_magnitude.py new file mode 100644 index 0000000000..92ddb91ee0 --- /dev/null +++ b/src/beignet/_rotation_matrix_magnitude.py @@ -0,0 +1,27 @@ +from torch import Tensor + +from ._rotation_matrix_to_rotation_quaternion import ( + rotation_matrix_to_rotation_quaternion, +) +from ._rotation_quaternion_magnitude import rotation_quaternion_magnitude + + +def rotation_matrix_magnitude(input: Tensor) -> Tensor: + r""" + Rotation matrix magnitudes. + + Parameters + ---------- + input : Tensor, shape (..., 3, 3) + Rotation matrices. + + Returns + ------- + rotation_matrix_magnitudes: Tensor, shape (...) + Angles in radians. Magnitudes will be in the range :math:`[0, \pi]`. + """ + return rotation_quaternion_magnitude( + rotation_matrix_to_rotation_quaternion( + input, + ), + ) diff --git a/src/beignet/_rotation_matrix_to_euler_angle.py b/src/beignet/_rotation_matrix_to_euler_angle.py new file mode 100644 index 0000000000..f125a7e5dc --- /dev/null +++ b/src/beignet/_rotation_matrix_to_euler_angle.py @@ -0,0 +1,51 @@ +from torch import Tensor + +from ._rotation_matrix_to_rotation_quaternion import ( + rotation_matrix_to_rotation_quaternion, +) +from ._rotation_quaternion_to_euler_angle import ( + rotation_quaternion_to_euler_angle, +) + + +def rotation_matrix_to_euler_angle( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + """ + Convert rotation matrices to Euler angles. + + Parameters + ---------- + input : Tensor, shape (..., 3, 3) + Rotation matrices. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees. Default, `False`. + + Returns + ------- + euler_angles : Tensor, shape (..., 3) + Euler angles. The returned Euler angles are in the range: + + * First angle: :math:`(-180, 180]` degrees (inclusive) + * Second angle: + * :math:`[-90, 90]` degrees if all axes are different + (e.g., :math:`xyz`) + * :math:`[0, 180]` degrees if first and third axes are the same + (e.g., :math:`zxz`) + * Third angle: :math:`[-180, 180]` degrees (inclusive) + """ + return rotation_quaternion_to_euler_angle( + rotation_matrix_to_rotation_quaternion( + input, + ), + axes, + degrees, + ) diff --git a/src/beignet/_rotation_matrix_to_rotation_quaternion.py b/src/beignet/_rotation_matrix_to_rotation_quaternion.py new file mode 100644 index 0000000000..829c7535bf --- /dev/null +++ b/src/beignet/_rotation_matrix_to_rotation_quaternion.py @@ -0,0 +1,88 @@ +import torch +from torch import Tensor + + +def rotation_matrix_to_rotation_quaternion( + input: Tensor, + canonical: bool | None = False, +) -> Tensor: + """ + Convert rotation matrices to rotation quaternions. + + Parameters + ---------- + input : Tensor, shape=(..., 3, 3) + Rotation matrices. + + canonical : bool, optional + Whether to map the redundant double cover of rotation space to a unique + canonical single cover. If `True`, then the rotation quaternion is + chosen from :math:`{q, -q}` such that the :math:`w` term is positive. + If the :math:`w` term is :math:`0`, then the rotation quaternion is + chosen such that the first non-zero term of the :math:`x`, :math:`y`, + and :math:`z` terms is positive. + + Returns + ------- + output : Tensor, shape=(..., 4) + Rotation quaternion. + """ + indexes = torch.empty( + [4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + output = torch.empty( + [input.shape[0], 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j in range(input.shape[0]): + indexes[0] = input[j, 0, 0] + indexes[1] = input[j, 1, 1] + indexes[2] = input[j, 2, 2] + indexes[3] = input[j, 0, 0] + input[j, 1, 1] + input[j, 2, 2] + + index, maximum = 0, indexes[0] + + for k in range(1, 4): + if indexes[k] > maximum: + index, maximum = k, indexes[k] + + if index == 3: + output[j, 0] = input[j, 2, 1] - input[j, 1, 2] + output[j, 1] = input[j, 0, 2] - input[j, 2, 0] + output[j, 2] = input[j, 1, 0] - input[j, 0, 1] + output[j, 3] = 1.0 + indexes[3] + else: + t = index + u = (t + 1) % 3 + v = (u + 1) % 3 + + output[j, t] = 1.0 - indexes[3] + 2.0 * input[j, t, t] + output[j, u] = input[j, u, t] + input[j, t, u] + output[j, v] = input[j, v, t] + input[j, t, v] + output[j, 3] = input[j, v, u] - input[j, u, v] + + a = output[j, 0] ** 2.0 + b = output[j, 1] ** 2.0 + c = output[j, 2] ** 2.0 + d = output[j, 3] ** 2.0 + + output[j] = output[j] / torch.sqrt(a + b + c + d) + + if canonical: + for j in range(output.shape[0]): + a = output[j, 0] + b = output[j, 1] + c = output[j, 2] + d = output[j, 3] + + if d == 0 and (a == 0 and (b == 0 and c < 0 or b < 0) or a < 0) or d < 0: + output[j] = -output[j] + + return output diff --git a/src/beignet/_rotation_matrix_to_rotation_vector.py b/src/beignet/_rotation_matrix_to_rotation_vector.py new file mode 100644 index 0000000000..56b4383046 --- /dev/null +++ b/src/beignet/_rotation_matrix_to_rotation_vector.py @@ -0,0 +1,37 @@ +from torch import Tensor + +from ._rotation_matrix_to_rotation_quaternion import ( + rotation_matrix_to_rotation_quaternion, +) +from ._rotation_quaternion_to_rotation_vector import ( + rotation_quaternion_to_rotation_vector, +) + + +def rotation_matrix_to_rotation_vector( + input: Tensor, + degrees: bool = False, +) -> Tensor: + """ + Convert rotation matrices to rotation vectors. + + Parameters + ---------- + input : Tensor, shape=(..., 3, 3) + Rotation matrices. + + degrees : bool + If `True`, rotation vector magnitudes are assumed to be in degrees. + Default, `False`. + + Returns + ------- + output : Tensor, shape=(..., 3) + Rotation vectors. + """ + return rotation_quaternion_to_rotation_vector( + rotation_matrix_to_rotation_quaternion( + input, + ), + degrees, + ) diff --git a/src/beignet/_rotation_quaternion_identity.py b/src/beignet/_rotation_quaternion_identity.py new file mode 100644 index 0000000000..31c5616837 --- /dev/null +++ b/src/beignet/_rotation_quaternion_identity.py @@ -0,0 +1,64 @@ +import torch +from torch import Tensor + + +def rotation_quaternion_identity( + size: int, + canonical: bool | None = False, + *, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, +) -> Tensor: + """ + Identity rotation quaternions. + + Parameters + ---------- + size : int + Output size. + + canonical : bool, optional + Whether to map the redundant double cover of rotation space to a unique + canonical single cover. If `True`, then the rotation quaternion is + chosen from :math:`{q, -q}` such that the :math:`w` term is positive. + If the :math:`w` term is :math:`0`, then the rotation quaternion is + chosen such that the first non-zero term of the :math:`x`, :math:`y`, + and :math:`z` terms is positive. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + Returns + ------- + identity_rotation_quaternions : Tensor, shape (size, 4) + Identity rotation quaternions. + """ + rotation = torch.zeros( + [size, 4], + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + rotation[:, 3] = 1.0 + + return rotation diff --git a/src/beignet/_rotation_quaternion_magnitude.py b/src/beignet/_rotation_quaternion_magnitude.py new file mode 100644 index 0000000000..360c196602 --- /dev/null +++ b/src/beignet/_rotation_quaternion_magnitude.py @@ -0,0 +1,44 @@ +import torch +from torch import Tensor + + +def rotation_quaternion_magnitude(input: Tensor, canonical=False) -> Tensor: + r""" + Rotation quaternion magnitudes. + + Parameters + ---------- + input : Tensor, shape (..., 4) + Rotation quaternions. + + Returns + ------- + rotation_quaternion_magnitudes: Tensor, shape (...) + Angles in radians. Magnitudes will be in the range :math:`[0, \pi]`. + """ + angles = torch.empty( + input.shape[0], + dtype=input.dtype, + layout=input.layout, + device=input.device, + requires_grad=input.requires_grad, + ) + + for index in range(input.shape[0]): + angles[index] = torch.multiply( + torch.atan2( + torch.sqrt( + torch.add( + torch.add( + torch.square(input[index, 0]), + torch.square(input[index, 1]), + ), + torch.square(input[index, 2]), + ), + ), + torch.abs(input[index, 3]), + ), + 2.0, + ) + + return angles diff --git a/src/beignet/_rotation_quaternion_to_euler_angle.py b/src/beignet/_rotation_quaternion_to_euler_angle.py new file mode 100644 index 0000000000..6e12fa665c --- /dev/null +++ b/src/beignet/_rotation_quaternion_to_euler_angle.py @@ -0,0 +1,143 @@ +import math +import re + +import torch +from torch import Tensor + + +def rotation_quaternion_to_euler_angle( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + """ + Convert rotation quaternions to Euler angles. + + Parameters + ---------- + input : Tensor, shape=(..., 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + axes : str + Axes. 1-3 characters belonging to the set {‘X’, ‘Y’, ‘Z’} for intrinsic + rotations, or {‘x’, ‘y’, ‘z’} for extrinsic rotations. Extrinsic and + intrinsic rotations cannot be mixed. + + degrees : bool, optional + If `True`, Euler angles are assumed to be in degrees. Default, `False`. + + Returns + ------- + output : Tensor, shape=(..., 3) + Euler angles. The returned Euler angles are in the range: + + * First angle: :math:`(-180, 180]` degrees (inclusive) + * Second angle: + * :math:`[-90, 90]` degrees if all axes are different + (e.g., :math:`xyz`) + * :math:`[0, 180]` degrees if first and third axes are the same + (e.g., :math:`zxz`) + * Third angle: :math:`[-180, 180]` degrees (inclusive) + """ + epsilon = torch.finfo(input.dtype).eps + + output = torch.empty( + [input.shape[0], 3], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + extrinsic = re.match(r"^[xyz]{1,3}$", axes) is not None + + axes = axes.lower() + + if not extrinsic: + axes = axes[::-1] + + match axes[0]: + case "x": + p = 0 + case "y": + p = 1 + case "z": + p = 2 + case _: + raise ValueError + + match axes[1]: + case "x": + q = 0 + case "y": + q = 1 + case "z": + q = 2 + case _: + raise ValueError + + match axes[2]: + case "x": + r = 0 + case "y": + r = 1 + case "z": + r = 2 + case _: + raise ValueError + + if p == r: + r = 3 - p - q + + s = (p - q) * (q - r) * (r - p) // 2 + + for j in range(input.shape[0]): + if p == r: + t = input[j, 3] + u = input[j, p] + v = input[j, q] + w = input[j, r] * s + else: + t = input[j, 3] - input[j, q] + u = input[j, p] + input[j, r] * s + v = input[j, q] + input[j, 3] + w = input[j, r] * s - input[j, p] + + if extrinsic: + a = 0 + c = 2 + else: + a = 2 + c = 0 + + output[j, 1] = 2.0 * torch.atan2(torch.hypot(v, w), torch.hypot(t, u)) + + match output[j, 1]: + case _ if abs(output[j, 1]) < epsilon: + output[j, 0] = 2.0 * torch.atan2(u, t) + output[j, 2] = 0.0 + case _ if abs(output[j, 1] - math.pi) < epsilon: + if extrinsic: + output[j, 0] = 2.0 * -torch.atan2(w, v) + else: + output[j, 0] = 2.0 * +torch.atan2(w, v) + + output[j, 2] = 0.0 + case _: + output[j, a] = torch.atan2(u, t) - torch.atan2(w, v) + output[j, c] = torch.atan2(u, t) + torch.atan2(w, v) + + if not p == r: + output[j, 1] = output[j, 1] - math.pi / 2.0 + output[j, c] = output[j, c] * s + + for k in range(3): + if output[j, k] <= -math.pi: + output[j, k] = output[j, k] + math.tau + + if output[j, k] >= +math.pi: + output[j, k] = output[j, k] - math.tau + + if degrees: + output = torch.rad2deg(output) + + return output diff --git a/src/beignet/_rotation_quaternion_to_rotation_matrix.py b/src/beignet/_rotation_quaternion_to_rotation_matrix.py new file mode 100644 index 0000000000..d571fd5db7 --- /dev/null +++ b/src/beignet/_rotation_quaternion_to_rotation_matrix.py @@ -0,0 +1,43 @@ +import torch +from torch import Tensor + + +def rotation_quaternion_to_rotation_matrix(input: Tensor) -> Tensor: + """ + Convert rotation quaternions to rotation matrices. + + Parameters + ---------- + input : Tensor, shape=(..., 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + Returns + ------- + output : Tensor, shape=(..., 3, 3) + Rotation matrices. + """ + output = torch.empty( + [input.shape[0], 3, 3], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j in range(input.shape[0]): + a = input[j, 0] + b = input[j, 1] + c = input[j, 2] + d = input[j, 3] + + output[j, 0, 0] = +(a**2.0) - b**2.0 - c**2.0 + d**2.0 + output[j, 1, 1] = -(a**2.0) + b**2.0 - c**2.0 + d**2.0 + output[j, 2, 2] = -(a**2.0) - b**2.0 + c**2.0 + d**2.0 + + output[j, 0, 1] = 2.0 * (a * b) - 2.0 * (c * d) + output[j, 0, 2] = 2.0 * (a * c) + 2.0 * (b * d) + output[j, 1, 0] = 2.0 * (a * b) + 2.0 * (c * d) + output[j, 1, 2] = 2.0 * (b * c) - 2.0 * (a * d) + output[j, 2, 0] = 2.0 * (a * c) - 2.0 * (b * d) + output[j, 2, 1] = 2.0 * (b * c) + 2.0 * (a * d) + + return output diff --git a/src/beignet/_rotation_quaternion_to_rotation_vector.py b/src/beignet/_rotation_quaternion_to_rotation_vector.py new file mode 100644 index 0000000000..b1eb8feb69 --- /dev/null +++ b/src/beignet/_rotation_quaternion_to_rotation_vector.py @@ -0,0 +1,57 @@ +import torch +from torch import Tensor + + +def rotation_quaternion_to_rotation_vector( + input: Tensor, + degrees: bool | None = False, +) -> Tensor: + """ + Convert rotation quaternions to rotation vectors. + + Parameters + ---------- + input : Tensor, shape=(..., 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + degrees : bool, optional + + Returns + ------- + output : Tensor, shape=(..., 3) + Rotation vectors. + """ + output = torch.empty( + [input.shape[0], 3], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j in range(input.shape[0]): + a = input[j, 0] + b = input[j, 1] + c = input[j, 2] + d = input[j, 3] + + if d == 0 and (a == 0 and (b == 0 and c < 0 or b < 0) or a < 0) or d < 0: + input[j] = -input[j] + + t = input[j, 0] ** 2.0 + u = input[j, 1] ** 2.0 + v = input[j, 2] ** 2.0 + w = input[j, 3] ** 1.0 + + y = 2.0 * torch.atan2(torch.sqrt(t + u + v), w) + + if y < 0.001: + y = 0.00243056 * y**4.0 + 0.08333333 * y**2.0 + 2.0 + else: + y = y / torch.sin(y / 2.0) + + output[j] = input[j, :-1] * y + + if degrees: + output = torch.rad2deg(output) + + return output diff --git a/src/beignet/_rotation_vector_identity.py b/src/beignet/_rotation_vector_identity.py new file mode 100644 index 0000000000..df308fd550 --- /dev/null +++ b/src/beignet/_rotation_vector_identity.py @@ -0,0 +1,64 @@ +import torch +from torch import Tensor + +from ._rotation_quaternion_identity import rotation_quaternion_identity +from ._rotation_quaternion_to_rotation_vector import ( + rotation_quaternion_to_rotation_vector, +) + + +def rotation_vector_identity( + size: int, + degrees: bool = False, + *, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, +) -> Tensor: + """ + Identity rotation vectors. + + Parameters + ---------- + size : int + Output size. + + degrees : bool + If `True`, rotation vector magnitudes are assumed to be in degrees. + Default, `False`. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + Returns + ------- + identity_rotation_vectors : Tensor, shape (size, 3) + Identity rotation vectors. + """ + return rotation_quaternion_to_rotation_vector( + rotation_quaternion_identity( + size, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ), + degrees, + ) diff --git a/src/beignet/_rotation_vector_magnitude.py b/src/beignet/_rotation_vector_magnitude.py new file mode 100644 index 0000000000..34cff10570 --- /dev/null +++ b/src/beignet/_rotation_vector_magnitude.py @@ -0,0 +1,34 @@ +from torch import Tensor + +from ._rotation_quaternion_magnitude import rotation_quaternion_magnitude +from ._rotation_vector_to_rotation_quaternion import ( + rotation_vector_to_rotation_quaternion, +) + + +def rotation_vector_magnitude( + input: Tensor, + degrees: bool | None = False, +) -> Tensor: + r""" + Rotation vector magnitudes. + + Parameters + ---------- + input : Tensor, shape (..., 3) + Rotation vectors. + + degrees : bool, optional + If `True`, magnitudes are assumed to be in degrees. Default, `False`. + + Returns + ------- + rotation_vector_magnitudes : Tensor, shape (...) + Angles in radians. Magnitudes will be in the range :math:`[0, \pi]`. + """ + return rotation_quaternion_magnitude( + rotation_vector_to_rotation_quaternion( + input, + degrees, + ), + ) diff --git a/src/beignet/_rotation_vector_to_euler_angle.py b/src/beignet/_rotation_vector_to_euler_angle.py new file mode 100644 index 0000000000..8378aa6795 --- /dev/null +++ b/src/beignet/_rotation_vector_to_euler_angle.py @@ -0,0 +1,48 @@ +from torch import Tensor + +from ._rotation_quaternion_to_euler_angle import ( + rotation_quaternion_to_euler_angle, +) +from ._rotation_vector_to_rotation_quaternion import ( + rotation_vector_to_rotation_quaternion, +) + + +def rotation_vector_to_euler_angle( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + """ + Convert rotation vectors to Euler angles. + + Parameters + ---------- + input : Tensor, shape=(..., 3) + Rotation vectors. + + degrees : bool, optional + If `True`, rotation vector magnitudes are assumed to be in degrees. + Default, `False`. + + Returns + ------- + output : Tensor, shape=(..., 3) + Euler angles. The returned Euler angles are in the range: + + * First angle: :math:`(-180, 180]` degrees (inclusive) + * Second angle: + * :math:`[-90, 90]` degrees if all axes are different + (e.g., :math:`xyz`) + * :math:`[0, 180]` degrees if first and third axes are the same + (e.g., :math:`zxz`) + * Third angle: :math:`[-180, 180]` degrees (inclusive) + """ + return rotation_quaternion_to_euler_angle( + rotation_vector_to_rotation_quaternion( + input, + degrees, + ), + axes, + degrees, + ) diff --git a/src/beignet/_rotation_vector_to_rotation_matrix.py b/src/beignet/_rotation_vector_to_rotation_matrix.py new file mode 100644 index 0000000000..96b84773ba --- /dev/null +++ b/src/beignet/_rotation_vector_to_rotation_matrix.py @@ -0,0 +1,37 @@ +from torch import Tensor + +from ._rotation_quaternion_to_rotation_matrix import ( + rotation_quaternion_to_rotation_matrix, +) +from ._rotation_vector_to_rotation_quaternion import ( + rotation_vector_to_rotation_quaternion, +) + + +def rotation_vector_to_rotation_matrix( + input: Tensor, + degrees: bool | None = False, +) -> Tensor: + """ + Convert rotation vectors to rotation matrices. + + Parameters + ---------- + input : Tensor, shape=(..., 3) + Rotation vectors. + + degrees : bool, optional + If `True`, rotation vector magnitudes are assumed to be in degrees. + Default, `False`. + + Returns + ------- + output : Tensor, shape=(..., 3, 3) + Rotation matrices. + """ + return rotation_quaternion_to_rotation_matrix( + rotation_vector_to_rotation_quaternion( + input, + degrees, + ), + ) diff --git a/src/beignet/_rotation_vector_to_rotation_quaternion.py b/src/beignet/_rotation_vector_to_rotation_quaternion.py new file mode 100644 index 0000000000..23fb7b8a6b --- /dev/null +++ b/src/beignet/_rotation_vector_to_rotation_quaternion.py @@ -0,0 +1,71 @@ +import torch +from torch import Tensor + + +def rotation_vector_to_rotation_quaternion( + input: Tensor, + degrees: bool | None = False, + canonical: bool | None = False, +) -> Tensor: + r""" + Convert rotation vector to rotation quaternion. + + Parameters + ---------- + input : Tensor, shape=(..., 3) + Rotation vector. + + degrees : bool, optional + If `True`, rotation vector magnitudes are assumed to be in degrees. + Default, `False`. + + canonical : bool, optional + Whether to map the redundant double cover of rotation space to a unique + canonical single cover. If `True`, then the rotation quaternion is + chosen from :math:`{q, -q}` such that the :math:`w` term is positive. + If the :math:`w` term is :math:`0`, then the rotation quaternion is + chosen such that the first non-zero term of the :math:`x`, :math:`y`, + and :math:`z` terms is positive. + + Returns + ------- + output : Tensor, shape=(..., 4) + Rotation quaternion. + """ + if degrees: + input = torch.deg2rad(input) + + output = torch.empty( + [input.shape[0], 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j in range(input.shape[0]): + t = input[j, 0] ** 2.0 + u = input[j, 1] ** 2.0 + v = input[j, 2] ** 2.0 + + y = torch.sqrt(t + u + v) + + if y < 0.001: + scale = 0.5 - y**2.0 / 48.0 + y**2.0 * y**2.0 / 3840.0 + else: + scale = torch.sin(y / 2) / y + + output[j, :-1] = input[j] * scale + + output[j, 3] = torch.cos(y / 2) + + if canonical: + for j in range(output.shape[0]): + a = output[j, 0] + b = output[j, 1] + c = output[j, 2] + d = output[j, 3] + + if d == 0 and (a == 0 and (b == 0 and c < 0 or b < 0) or a < 0) or d < 0: + output[j] = -output[j] + + return output diff --git a/src/beignet/_slerp.py b/src/beignet/_slerp.py new file mode 100644 index 0000000000..2ce6a612ef --- /dev/null +++ b/src/beignet/_slerp.py @@ -0,0 +1,176 @@ +import torch +import torch.testing +from torch import Tensor + + +def slerp( + input: Tensor, + time: Tensor, + rotation: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Interpolate between two or more points on a sphere. + + Unlike linear interpolation, which can result in changes in speed when + interpolating between orientations or positions on a sphere, spherical + linear interpolation ensures that the interpolation occurs at a constant + rate and follows the shortest path on the surface of the sphere. + The process is useful for rotations and orientation interpolation in + three-dimensional spaces, smoothly transitioning between different + orientations. + + Mathematically, spherical linear interpolation interpolates between two + points on a sphere using a parameter $t$, where $t = 0$ represents the + start point and $t = n$ represents the end point. For two rotation + quaternions $q_{1}$ and $q_{2}$ representing the start and end + orientations: + + $$\text{slerp}(q_{1}, q_{2}; t) = q_{1}\frac{\sin((1 - t)\theta)}{\sin(\theta)} + q_{2}\frac{\sin(t\theta)}{\sin(\theta)}$$ + + where $\theta$ is the angle between $q_{1}$ and $q_{2}$, and is computed + using the dot product of $q_{1}$ and $q_{2}$. This formula ensures that the + interpolation moves along the shortest path on the four-dimensional sphere + of rotation quaternions, resulting in a smooth and constant-speed rotation. + + Parameters + ---------- + input : Tensor, shape (..., N) + Times. + + time : Tensor, shape (..., N) + Times of the known rotations. At least 2 times must be specified. + + rotation : Tensor, shape (..., N, 4) + Rotation quaternions. Rotation quaternions are normalized to unit norm. + + out : Tensor + Output. + """ + if time.shape[-1] != rotation.shape[-2]: + raise ValueError("`times` and `rotations` must match in size.") + + interpolated_rotation_quaternions = torch.empty( + [*input.shape, 4], + dtype=input.dtype, + ) + + for index, t in enumerate(input): + # FIND KEYFRAMES IN INTERVAL, ASSUME `times` IS SORTED: + b = torch.min(torch.nonzero(torch.greater_equal(time, t))) + + if b > 0: + a = b - 1 + else: + a = b + + # IF `times` MATCHES A KEYFRAME: + if time[b] == t or b == a: + interpolated_rotation_quaternions[index] = rotation[b] + + continue + + time_0, time_1 = time[a], time[b] + + relative_time = torch.divide( + torch.subtract( + t, + time_0, + ), + torch.subtract( + time_1, + time_0, + ), + ) + + quaternion_0, quaternion_1 = rotation[a], rotation[b] + + cos_theta = torch.dot(quaternion_0, quaternion_1) + + # NOTE: IF `cos_theta` IS NEGATIVE, `rotations` HAS OPPOSITE HANDEDNESS + # SO `slerp` WON’T TAKE THE SHORTER PATH. INVERT `quaternion_1`. + # @0X00B1, FEBRUARY, 26, 2024 + if cos_theta < 0.0: + quaternion_1 = torch.negative(quaternion_1) + + cos_theta = torch.negative(cos_theta) + + # USE LINEAR INTERPOLATION IF QUATERNIONS ARE CLOSE: + if cos_theta > 0.9995: + interpolated_rotation_quaternion = torch.add( + torch.multiply( + torch.subtract( + torch.tensor(1.0), + relative_time, + ), + quaternion_0, + ), + torch.multiply( + relative_time, + quaternion_1, + ), + ) + else: + interpolated_rotation_quaternion = torch.add( + torch.multiply( + quaternion_0, + torch.divide( + torch.sin( + torch.multiply( + torch.subtract( + torch.tensor(1.0), + relative_time, + ), + torch.atan2( + torch.sqrt( + torch.subtract( + torch.tensor(1.0), + torch.square(cos_theta), + ) + ), + cos_theta, + ), + ), + ), + torch.sqrt( + torch.subtract( + torch.tensor(1.0), + torch.square(cos_theta), + ) + ), + ), + ), + torch.multiply( + quaternion_1, + torch.divide( + torch.sin( + torch.multiply( + relative_time, + torch.atan2( + torch.sqrt( + torch.subtract( + torch.tensor(1.0), + torch.square(cos_theta), + ) + ), + cos_theta, + ), + ), + ), + torch.sqrt( + torch.subtract( + torch.tensor(1.0), + torch.square(cos_theta), + ), + ), + ), + ), + ) + + interpolated_rotation_quaternions[index] = torch.divide( + interpolated_rotation_quaternion, + torch.norm(interpolated_rotation_quaternion), + ) + + return interpolated_rotation_quaternions diff --git a/src/beignet/_translation_vector_identity.py b/src/beignet/_translation_vector_identity.py new file mode 100644 index 0000000000..2847b041f5 --- /dev/null +++ b/src/beignet/_translation_vector_identity.py @@ -0,0 +1,51 @@ +import torch +from torch import Tensor + + +def translation_vector_identity( + size: int, + *, + out: Tensor | None = None, + dtype: torch.dtype | None = None, + layout: torch.layout | None = torch.strided, + device: torch.device | None = None, + requires_grad: bool | None = False, +) -> Tensor: + """ + Identity translation vectors. + + Parameters + ---------- + size : int + Output size. + + out : Tensor, optional + Output tensor. Default, `None`. + + dtype : torch.dtype, optional + Type of the returned tensor. Default, global default. + + layout : torch.layout, optional + Layout of the returned tensor. Default, `torch.strided`. + + device : torch.device, optional + Device of the returned tensor. Default, current device for the default + tensor type. + + requires_grad : bool, optional + Whether autograd records operations on the returned tensor. Default, + `False`. + + Returns + ------- + output : Tensor, shape=(size, 3) + Identity rotation quaternions. + """ + return torch.zeros( + [size, 3], + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) diff --git a/tests/beignet/test__apply_euler_angle.py b/tests/beignet/test__apply_euler_angle.py new file mode 100644 index 0000000000..2943616868 --- /dev/null +++ b/tests/beignet/test__apply_euler_angle.py @@ -0,0 +1,73 @@ +import hypothesis.extra.numpy +import hypothesis.strategies +import numpy +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + batch_size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = function( + hypothesis.extra.numpy.arrays( + numpy.float64, + (batch_size, 3), + elements={ + "allow_infinity": False, + "allow_nan": False, + }, + ), + ) + + rotation = Rotation.random(batch_size) + + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + inverse = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy(input), + "rotation": torch.from_numpy(rotation.as_euler(axes, degrees)), + "axes": axes, + "degrees": degrees, + "inverse": inverse, + }, + torch.from_numpy(rotation.apply(input, inverse)), + ) + + +@hypothesis.given(_strategy()) +def test_apply_rotation_matrix(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.apply_euler_angle(**parameters), + expected, + ) diff --git a/tests/beignet/test__apply_rotation_matrix.py b/tests/beignet/test__apply_rotation_matrix.py new file mode 100644 index 0000000000..9336bd6c8f --- /dev/null +++ b/tests/beignet/test__apply_rotation_matrix.py @@ -0,0 +1,52 @@ +import hypothesis.extra.numpy +import hypothesis.strategies +import numpy +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + batch_size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = function( + hypothesis.extra.numpy.arrays( + numpy.float64, + (batch_size, 3), + elements={ + "allow_infinity": False, + "allow_nan": False, + "min_value": -1e3, + "max_value": +1e3, + }, + ), + ) + + rotation = Rotation.random(batch_size) + + inverse = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy(input), + "rotation": torch.from_numpy(rotation.as_matrix()), + "inverse": inverse, + }, + torch.from_numpy(rotation.apply(input, inverse)), + ) + + +@hypothesis.given(_strategy()) +def test_apply_rotation_matrix(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.apply_rotation_matrix(**parameters), + expected, + ) diff --git a/tests/beignet/test__apply_rotation_quaternion.py b/tests/beignet/test__apply_rotation_quaternion.py new file mode 100644 index 0000000000..82ae8e8c22 --- /dev/null +++ b/tests/beignet/test__apply_rotation_quaternion.py @@ -0,0 +1,52 @@ +import hypothesis.extra.numpy +import hypothesis.strategies +import numpy +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + batch_size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = function( + hypothesis.extra.numpy.arrays( + numpy.float64, + (batch_size, 3), + elements={ + "allow_infinity": False, + "allow_nan": False, + }, + ), + ) + + rotation = Rotation.random(batch_size) + + canonical = function(hypothesis.strategies.booleans()) + + inverse = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy(input), + "rotation": torch.from_numpy(rotation.as_quat(canonical)), + "inverse": inverse, + }, + torch.from_numpy(rotation.apply(input, inverse)), + ) + + +@hypothesis.given(_strategy()) +def test_apply_rotation_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.apply_rotation_quaternion(**parameters), + expected, + ) diff --git a/tests/beignet/test__apply_rotation_vector.py b/tests/beignet/test__apply_rotation_vector.py new file mode 100644 index 0000000000..78cd4b4526 --- /dev/null +++ b/tests/beignet/test__apply_rotation_vector.py @@ -0,0 +1,53 @@ +import hypothesis.extra.numpy +import hypothesis.strategies +import numpy +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + batch_size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = function( + hypothesis.extra.numpy.arrays( + numpy.float64, + (batch_size, 3), + elements={ + "allow_infinity": False, + "allow_nan": False, + }, + ), + ) + + rotation = Rotation.random(batch_size) + + degrees = function(hypothesis.strategies.booleans()) + + inverse = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy(input), + "rotation": torch.from_numpy(rotation.as_rotvec(degrees)), + "degrees": degrees, + "inverse": inverse, + }, + torch.from_numpy(rotation.apply(input, inverse)), + ) + + +@hypothesis.given(_strategy()) +def test_apply_rotation_vector(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.apply_rotation_vector(**parameters), + expected, + ) diff --git a/tests/beignet/test__compose_rotation_quaternion.py b/tests/beignet/test__compose_rotation_quaternion.py new file mode 100644 index 0000000000..54e3ece053 --- /dev/null +++ b/tests/beignet/test__compose_rotation_quaternion.py @@ -0,0 +1,58 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = Rotation.random(size) + other = Rotation.random(size) + + canonical = function( + hypothesis.strategies.booleans(), + ) + + return ( + { + "input": torch.from_numpy( + input.as_quat( + canonical, + ), + ), + "other": torch.from_numpy( + other.as_quat( + canonical, + ), + ), + "canonical": canonical, + }, + torch.abs( + torch.from_numpy( + (input * other).as_quat( + canonical, + ), + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_compose_rotation_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.compose_rotation_quaternion( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/test__euler_angle_identity.py b/tests/beignet/test__euler_angle_identity.py new file mode 100644 index 0000000000..f9706d7bd8 --- /dev/null +++ b/tests/beignet/test__euler_angle_identity.py @@ -0,0 +1,11 @@ +import hypothesis.strategies + + +@hypothesis.strategies.composite +def strategy(f): + return + + +@hypothesis.given(strategy()) +def test_euler_angle_identity(data): + assert True diff --git a/tests/beignet/test__euler_angle_magnitude.py b/tests/beignet/test__euler_angle_magnitude.py new file mode 100644 index 0000000000..3122372bbe --- /dev/null +++ b/tests/beignet/test__euler_angle_magnitude.py @@ -0,0 +1,56 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + rotations = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + return ( + { + "input": torch.from_numpy(rotations.as_euler(axes, degrees)), + "axes": axes, + "degrees": degrees, + }, + torch.from_numpy(rotations.magnitude()), + ) + + +@hypothesis.given(_strategy()) +def test_euler_angle_magnitude(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.euler_angle_magnitude(**parameters), + expected, + ) diff --git a/tests/beignet/test__euler_angle_to_rotation_matrix.py b/tests/beignet/test__euler_angle_to_rotation_matrix.py new file mode 100644 index 0000000000..50ee344460 --- /dev/null +++ b/tests/beignet/test__euler_angle_to_rotation_matrix.py @@ -0,0 +1,65 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_euler( + axes, + degrees, + ), + ), + "axes": axes, + "degrees": degrees, + }, + torch.from_numpy( + rotation.as_matrix(), + ), + ) + + +@hypothesis.given(_strategy()) +def test_euler_angle_to_rotation_matrix(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.euler_angle_to_rotation_matrix( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__euler_angle_to_rotation_quaternion.py b/tests/beignet/test__euler_angle_to_rotation_quaternion.py new file mode 100644 index 0000000000..7e78964dea --- /dev/null +++ b/tests/beignet/test__euler_angle_to_rotation_quaternion.py @@ -0,0 +1,75 @@ +import hypothesis.extra.numpy +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + canonical = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_euler( + axes, + degrees, + ), + ), + "axes": axes, + "degrees": degrees, + "canonical": canonical, + }, + torch.abs( + torch.from_numpy( + rotation.as_quat( + canonical, + ), + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_euler_angle_to_rotation_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.euler_angle_to_rotation_quaternion( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/test__euler_angle_to_rotation_vector.py b/tests/beignet/test__euler_angle_to_rotation_vector.py new file mode 100644 index 0000000000..dc349fc5e2 --- /dev/null +++ b/tests/beignet/test__euler_angle_to_rotation_vector.py @@ -0,0 +1,63 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_euler( + axes, + degrees, + ), + ), + "axes": axes, + "degrees": degrees, + }, + torch.from_numpy(rotation.as_rotvec(degrees)), + ) + + +@hypothesis.given(_strategy()) +def test_euler_angle_to_rotation_vector(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.euler_angle_to_rotation_vector( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__invert_euler_angle.py b/tests/beignet/test__invert_euler_angle.py new file mode 100644 index 0000000000..975f502282 --- /dev/null +++ b/tests/beignet/test__invert_euler_angle.py @@ -0,0 +1,56 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + rotations = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + return ( + { + "input": torch.from_numpy(rotations.as_euler(axes, degrees)), + "axes": axes, + "degrees": degrees, + }, + torch.from_numpy(rotations.inv().as_euler(axes, degrees)), + ) + + +@hypothesis.given(_strategy()) +def test_invert_euler_angle(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.invert_euler_angle(**parameters), + expected, + ) diff --git a/tests/beignet/test__invert_rotation_matrix.py b/tests/beignet/test__invert_rotation_matrix.py new file mode 100644 index 0000000000..824cef407e --- /dev/null +++ b/tests/beignet/test__invert_rotation_matrix.py @@ -0,0 +1,33 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotations = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + return ( + { + "input": torch.from_numpy(rotations.as_matrix()), + }, + torch.from_numpy(rotations.inv().as_matrix()), + ) + + +@hypothesis.given(_strategy()) +def test_invert_rotation_matrix(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.invert_rotation_matrix(**parameters), + expected, + ) diff --git a/tests/beignet/test__invert_rotation_quaternion.py b/tests/beignet/test__invert_rotation_quaternion.py new file mode 100644 index 0000000000..17041f41a3 --- /dev/null +++ b/tests/beignet/test__invert_rotation_quaternion.py @@ -0,0 +1,36 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + canonical = function(hypothesis.strategies.booleans()) + + rotations = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + return ( + { + "input": torch.from_numpy(rotations.as_quat(canonical)), + "canonical": canonical, + }, + torch.from_numpy(rotations.inv().as_quat(canonical)), + ) + + +@hypothesis.given(_strategy()) +def test_invert_rotation_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.invert_rotation_quaternion(**parameters), + expected, + ) diff --git a/tests/beignet/test__invert_rotation_vector.py b/tests/beignet/test__invert_rotation_vector.py new file mode 100644 index 0000000000..9b58aceeb4 --- /dev/null +++ b/tests/beignet/test__invert_rotation_vector.py @@ -0,0 +1,35 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + degrees = function(hypothesis.strategies.booleans()) + + rotations = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + return ( + { + "input": torch.from_numpy(rotations.as_rotvec(degrees)), + }, + torch.from_numpy(rotations.inv().as_rotvec(degrees)), + ) + + +@hypothesis.given(_strategy()) +def test_invert_rotation_vector(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.invert_rotation_vector(**parameters), + expected, + ) diff --git a/tests/beignet/test__mean_rotation_quaternion.py b/tests/beignet/test__mean_rotation_quaternion.py new file mode 100644 index 0000000000..1c464653af --- /dev/null +++ b/tests/beignet/test__mean_rotation_quaternion.py @@ -0,0 +1,52 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + canonical = function( + hypothesis.strategies.booleans(), + ) + + return ( + { + "input": torch.from_numpy( + rotation.as_quat( + canonical, + ), + ), + "canonical": canonical, + }, + torch.abs( + torch.from_numpy( + rotation.mean().as_quat( + canonical, + ), + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_mean_rotation_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.mean_rotation_quaternion( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/test__random_euler_angle.py b/tests/beignet/test__random_euler_angle.py new file mode 100644 index 0000000000..52f2a321dd --- /dev/null +++ b/tests/beignet/test__random_euler_angle.py @@ -0,0 +1,51 @@ +import hypothesis.strategies +import beignet + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "size": size, + "axes": axes, + "degrees": degrees, + }, + None, + ) + + +@hypothesis.given(_strategy()) +def test_random_euler_angle(data): + parameters, _ = data + + assert beignet.random_euler_angle( + **parameters, + ).shape == (parameters["size"], 3) diff --git a/tests/beignet/test__random_rotation_matrix.py b/tests/beignet/test__random_rotation_matrix.py new file mode 100644 index 0000000000..c5e635d0b1 --- /dev/null +++ b/tests/beignet/test__random_rotation_matrix.py @@ -0,0 +1,28 @@ +import hypothesis.strategies +import beignet + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + return ( + { + "size": size, + }, + None, + ) + + +@hypothesis.given(_strategy()) +def test_random_rotation_matrix(data): + parameters, _ = data + + assert beignet.random_rotation_matrix( + **parameters, + ).shape == (parameters["size"], 3, 3) diff --git a/tests/beignet/test__random_rotation_quaternion.py b/tests/beignet/test__random_rotation_quaternion.py new file mode 100644 index 0000000000..8b46956f6c --- /dev/null +++ b/tests/beignet/test__random_rotation_quaternion.py @@ -0,0 +1,31 @@ +import hypothesis.strategies +import beignet + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + canonical = function(hypothesis.strategies.booleans()) + + return ( + { + "size": size, + "canonical": canonical, + }, + None, + ) + + +@hypothesis.given(_strategy()) +def test_random_rotation_quaternion(data): + parameters, _ = data + + assert beignet.random_rotation_quaternion( + **parameters, + ).shape == (parameters["size"], 4) diff --git a/tests/beignet/test__random_rotation_vector.py b/tests/beignet/test__random_rotation_vector.py new file mode 100644 index 0000000000..ad444c88c6 --- /dev/null +++ b/tests/beignet/test__random_rotation_vector.py @@ -0,0 +1,31 @@ +import hypothesis.strategies +import beignet + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "size": size, + "degrees": degrees, + }, + None, + ) + + +@hypothesis.given(_strategy()) +def test_random_rotation_vector(data): + parameters, _ = data + + assert beignet.random_rotation_vector( + **parameters, + ).shape == (parameters["size"], 3) diff --git a/tests/beignet/test__rotation_matrix_identity.py b/tests/beignet/test__rotation_matrix_identity.py new file mode 100644 index 0000000000..69c823dfef --- /dev/null +++ b/tests/beignet/test__rotation_matrix_identity.py @@ -0,0 +1,34 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + rotation = Rotation.identity(size) + + return ( + { + "size": size, + "dtype": torch.float64, + }, + torch.from_numpy(rotation.as_matrix()), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_matrix_identity(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_matrix_identity(**parameters), + expected, + ) diff --git a/tests/beignet/test__rotation_matrix_magnitude.py b/tests/beignet/test__rotation_matrix_magnitude.py new file mode 100644 index 0000000000..2880e292df --- /dev/null +++ b/tests/beignet/test__rotation_matrix_magnitude.py @@ -0,0 +1,33 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotations = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + return ( + { + "input": torch.from_numpy(rotations.as_matrix()), + }, + torch.from_numpy(rotations.magnitude()), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_matrix_magnitude(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_matrix_magnitude(**parameters), + expected, + ) diff --git a/tests/beignet/test__rotation_matrix_to_euler_angle.py b/tests/beignet/test__rotation_matrix_to_euler_angle.py new file mode 100644 index 0000000000..0a0dff2a46 --- /dev/null +++ b/tests/beignet/test__rotation_matrix_to_euler_angle.py @@ -0,0 +1,65 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_matrix(), + ), + "axes": axes, + "degrees": degrees, + }, + torch.from_numpy( + rotation.as_euler( + axes, + degrees, + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_matrix_to_euler_angle(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_matrix_to_euler_angle( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__rotation_matrix_to_rotation_quaternion.py b/tests/beignet/test__rotation_matrix_to_rotation_quaternion.py new file mode 100644 index 0000000000..85e009c2ab --- /dev/null +++ b/tests/beignet/test__rotation_matrix_to_rotation_quaternion.py @@ -0,0 +1,49 @@ +import hypothesis.extra.numpy +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + canonical = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_matrix(), + ), + "canonical": canonical, + }, + torch.abs( + torch.from_numpy( + rotation.as_quat( + canonical, + ), + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_matrix_to_rotation_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.rotation_matrix_to_rotation_quaternion( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/test__rotation_matrix_to_rotation_vector.py b/tests/beignet/test__rotation_matrix_to_rotation_vector.py new file mode 100644 index 0000000000..cf838a8ca4 --- /dev/null +++ b/tests/beignet/test__rotation_matrix_to_rotation_vector.py @@ -0,0 +1,45 @@ +import hypothesis.extra.numpy +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_matrix(), + ), + "degrees": degrees, + }, + torch.from_numpy( + rotation.as_rotvec( + degrees, + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_matrix_to_rotation_vector(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_matrix_to_rotation_vector( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__rotation_quaternion_identity.py b/tests/beignet/test__rotation_quaternion_identity.py new file mode 100644 index 0000000000..8c759de741 --- /dev/null +++ b/tests/beignet/test__rotation_quaternion_identity.py @@ -0,0 +1,37 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + rotation = Rotation.identity(size) + + canonical = function(hypothesis.strategies.booleans()) + + return ( + { + "size": size, + "canonical": canonical, + "dtype": torch.float64, + }, + torch.from_numpy(rotation.as_quat(canonical)), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_quaternion_identity(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_quaternion_identity(**parameters), + expected, + ) diff --git a/tests/beignet/test__rotation_quaternion_magnitude.py b/tests/beignet/test__rotation_quaternion_magnitude.py new file mode 100644 index 0000000000..07cac4c4b5 --- /dev/null +++ b/tests/beignet/test__rotation_quaternion_magnitude.py @@ -0,0 +1,36 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + canonical = function(hypothesis.strategies.booleans()) + + rotations = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + return ( + { + "input": torch.from_numpy(rotations.as_quat(canonical)), + "canonical": canonical, + }, + torch.from_numpy(rotations.magnitude()), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_quaternion_magnitude(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_quaternion_magnitude(**parameters), + expected, + ) diff --git a/tests/beignet/test__rotation_quaternion_to_euler_angle.py b/tests/beignet/test__rotation_quaternion_to_euler_angle.py new file mode 100644 index 0000000000..e1fdadb0f4 --- /dev/null +++ b/tests/beignet/test__rotation_quaternion_to_euler_angle.py @@ -0,0 +1,67 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_quat( + canonical=False, + ), + ), + "axes": axes, + "degrees": degrees, + }, + torch.from_numpy( + rotation.as_euler( + axes, + degrees, + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_quaternion_to_euler_angle(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_quaternion_to_euler_angle( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__rotation_quaternion_to_rotation_matrix.py b/tests/beignet/test__rotation_quaternion_to_rotation_matrix.py new file mode 100644 index 0000000000..f0d8228bfb --- /dev/null +++ b/tests/beignet/test__rotation_quaternion_to_rotation_matrix.py @@ -0,0 +1,41 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + return ( + { + "input": torch.from_numpy( + rotation.as_quat( + canonical=False, + ), + ), + }, + torch.from_numpy( + rotation.as_matrix(), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_quaternion_to_rotation_matrix(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_quaternion_to_rotation_matrix( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__rotation_quaternion_to_rotation_vector.py b/tests/beignet/test__rotation_quaternion_to_rotation_vector.py new file mode 100644 index 0000000000..8c9f3c0389 --- /dev/null +++ b/tests/beignet/test__rotation_quaternion_to_rotation_vector.py @@ -0,0 +1,46 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_quat( + canonical=False, + ), + ), + "degrees": degrees, + }, + torch.from_numpy( + rotation.as_rotvec( + degrees, + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_quaternion_to_rotation_vector(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_quaternion_to_rotation_vector( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__rotation_vector_identity.py b/tests/beignet/test__rotation_vector_identity.py new file mode 100644 index 0000000000..c5e24d7115 --- /dev/null +++ b/tests/beignet/test__rotation_vector_identity.py @@ -0,0 +1,36 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + rotation = Rotation.identity(size) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "size": size, + "dtype": torch.float64, + }, + torch.from_numpy(rotation.as_rotvec(degrees)), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_vector_identity(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_vector_identity(**parameters), + expected, + ) diff --git a/tests/beignet/test__rotation_vector_magnitude.py b/tests/beignet/test__rotation_vector_magnitude.py new file mode 100644 index 0000000000..66dfc05f26 --- /dev/null +++ b/tests/beignet/test__rotation_vector_magnitude.py @@ -0,0 +1,36 @@ +import hypothesis.strategies +import torch.testing +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + degrees = function(hypothesis.strategies.booleans()) + + rotations = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ), + ) + + return ( + { + "input": torch.from_numpy(rotations.as_rotvec(degrees)), + "degrees": degrees, + }, + torch.from_numpy(rotations.magnitude()), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_vector_magnitude(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_vector_magnitude(**parameters), + expected, + ) diff --git a/tests/beignet/test__rotation_vector_to_euler_angle.py b/tests/beignet/test__rotation_vector_to_euler_angle.py new file mode 100644 index 0000000000..b0c178d7bf --- /dev/null +++ b/tests/beignet/test__rotation_vector_to_euler_angle.py @@ -0,0 +1,67 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + axes = function( + hypothesis.strategies.sampled_from( + [ + "xyz", + "xzy", + "yxz", + "yzx", + "zxy", + "zyx", + "XYZ", + "XZY", + "YXZ", + "YZX", + "ZXY", + "ZYX", + ] + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_rotvec( + degrees, + ), + ), + "axes": axes, + "degrees": degrees, + }, + torch.from_numpy( + rotation.as_euler( + axes, + degrees, + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_vector_to_euler_angle(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_vector_to_euler_angle( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__rotation_vector_to_rotation_matrix.py b/tests/beignet/test__rotation_vector_to_rotation_matrix.py new file mode 100644 index 0000000000..66a132e41a --- /dev/null +++ b/tests/beignet/test__rotation_vector_to_rotation_matrix.py @@ -0,0 +1,42 @@ +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_rotvec(degrees), + ), + "degrees": degrees, + }, + torch.from_numpy( + rotation.as_matrix(), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_vector_to_rotation_matrix(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.rotation_vector_to_rotation_matrix( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/test__rotation_vector_to_rotation_quaternion.py b/tests/beignet/test__rotation_vector_to_rotation_quaternion.py new file mode 100644 index 0000000000..34730c206f --- /dev/null +++ b/tests/beignet/test__rotation_vector_to_rotation_quaternion.py @@ -0,0 +1,54 @@ +import hypothesis.extra.numpy +import hypothesis.strategies +import torch +import beignet +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + rotation = Rotation.random( + function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ), + ) + + degrees = function(hypothesis.strategies.booleans()) + + canonical = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy( + rotation.as_rotvec( + degrees, + ), + ), + "degrees": degrees, + "canonical": canonical, + }, + torch.abs( + torch.from_numpy( + rotation.as_quat( + canonical, + ), + ), + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_vector_to_rotation_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.rotation_vector_to_rotation_quaternion( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/test__slerp.py b/tests/beignet/test__slerp.py new file mode 100644 index 0000000000..c742f6ed01 --- /dev/null +++ b/tests/beignet/test__slerp.py @@ -0,0 +1,149 @@ +import hypothesis.strategies +import numpy +import torch +import beignet +from scipy.spatial.transform import Rotation, Slerp + + +def test_slerp(): + # t = 0 + torch.testing.assert_close( + beignet.slerp( + torch.tensor([+0.00000]), + torch.tensor([+0.00000, +1.00000]), + torch.tensor( + [ + [+1.00000, +0.00000, +0.00000, +0.00000], + [+0.00000, +1.00000, +0.00000, +0.00000], + ] + ), + ), + torch.tensor([[+1.00000, +0.00000, +0.00000, +0.00000]]), + ) + + # t = 1 + torch.testing.assert_close( + beignet.slerp( + torch.tensor([+1.00000]), + torch.tensor([+0.00000, +1.00000]), + torch.tensor( + [ + [+1.00000, +0.00000, +0.00000, +0.00000], + [+0.00000, +1.00000, +0.00000, +0.00000], + ] + ), + ), + torch.tensor([[+0.00000, +1.00000, +0.00000, +0.00000]]), + ) + + # SMALL (ACUTE) ANGLE BETWEEN QUATERNIONS + torch.testing.assert_close( + beignet.slerp( + torch.tensor([+0.50000]), + torch.tensor([+0.00000, +1.00000]), + torch.tensor( + [ + [+1.00000, +0.00000, +0.00000, +0.00000], + [+0.70710, +0.70710, +0.00000, +0.00000], + ], + ), + ), + torch.reshape( + torch.tensor([+0.92388, +0.38268, +0.00000, +0.00000]), + [1, -1], + ), + ) + + # LARGE (OBTUSE) ANGLE BETWEEN QUATERNIONS + torch.testing.assert_close( + beignet.slerp( + torch.tensor([+0.50000]), + torch.tensor([+0.00000, +1.00000]), + torch.tensor( + [ + [+1.00000, +0.00000, +0.00000, +0.00000], + [-1.00000, +0.00000, +0.00000, +0.00000], + ] + ), + ), + torch.reshape( + torch.tensor([+1.00000, +0.00000, +0.00000, +0.00000]), + [1, -1], + ), + ) + + +@hypothesis.strategies.composite +def slerp_parameters(f): + n = f( + hypothesis.strategies.integers( + min_value=2, + max_value=8, + ), + ) + + times = numpy.sort( + f( + hypothesis.strategies.lists( + hypothesis.strategies.floats( + allow_infinity=False, + allow_nan=False, + ), + min_size=n, + max_size=n, + unique=True, + ), + ), + ) + + min_value = numpy.min(times) + max_value = numpy.max(times) + + input = numpy.sort( + f( + hypothesis.strategies.lists( + hypothesis.strategies.floats( + min_value=min_value, + max_value=max_value, + ), + min_size=1, + max_size=8, + unique=True, + ), + ), + ) + + rotations = f( + hypothesis.strategies.lists( + hypothesis.strategies.lists( + hypothesis.strategies.floats( + numpy.finfo(numpy.float32).eps, + 1.0, + ), + min_size=4, + max_size=4, + ), + min_size=n, + max_size=n, + ), + ) + + rotations = Rotation.from_quat(rotations) + + return [ + [ + torch.from_numpy(input), + torch.from_numpy(times), + torch.from_numpy(rotations.as_quat(canonical=True)), + ], + torch.from_numpy( + Slerp(times, rotations)(input).as_quat(canonical=True), + ), + ] + + +@hypothesis.given(slerp_parameters()) +def test_slerp_properties(data): + parameters, expected_rotations = data + + torch.testing.assert_close(beignet.slerp(*parameters), expected_rotations)