From edc79b0c8badc3ecbd59fc5b8f88c671ab637e43 Mon Sep 17 00:00:00 2001 From: Allen Goodman Date: Thu, 9 May 2024 11:42:53 -0400 Subject: [PATCH] General-purpose geometric operators (#9) general-purpose geometric operators --- .github/workflows/push.yml | 23 +- .gitignore | 3 +- .pre-commit-config.yaml | 13 + docs/beignet.ops.md | 11 + pyproject.toml | 23 +- src/beignet/ops/__init__.py | 89 ++++++ src/beignet/ops/_geometry/__init__.py | 44 +++ .../_geometry/_transformations/__init__.py | 60 +++++ .../_transformations/_apply_euler_angle.py | 63 +++++ .../_transformations/_apply_quaternion.py | 52 ++++ .../_apply_rotation_matrix.py | 47 ++++ .../_apply_rotation_vector.py | 63 +++++ .../_transformations/_compose_euler_angle.py | 53 ++++ .../_transformations/_compose_quaternion.py | 73 +++++ .../_compose_rotation_matrix.py | 33 +++ .../_compose_rotation_vector.py | 45 ++++ .../_transformations/_euler_angle_identity.py | 70 +++++ .../_euler_angle_magnitude.py | 41 +++ .../_transformations/_euler_angle_mean.py | 49 ++++ .../_euler_angle_to_quaternion.py | 170 ++++++++++++ .../_euler_angle_to_rotation_matrix.py | 80 ++++++ .../_euler_angle_to_rotation_vector.py | 253 ++++++++++++++++++ .../_transformations/_invert_euler_angle.py | 48 ++++ .../_transformations/_invert_quaternion.py | 31 +++ .../_invert_rotation_matrix.py | 19 ++ .../_invert_rotation_vector.py | 20 ++ .../_transformations/_quaternion_identity.py | 64 +++++ .../_transformations/_quaternion_magnitude.py | 37 +++ .../_transformations/_quaternion_mean.py | 34 +++ .../_transformations/_quaternion_slerp.py | 93 +++++++ .../_quaternion_to_euler_angle.py | 143 ++++++++++ .../_quaternion_to_rotation_matrix.py | 43 +++ .../_quaternion_to_rotation_vector.py | 57 ++++ .../_transformations/_random_euler_angle.py | 91 +++++++ .../_transformations/_random_quaternion.py | 95 +++++++ .../_random_rotation_matrix.py | 69 +++++ .../_random_rotation_vector.py | 73 +++++ .../_rotation_matrix_identity.py | 58 ++++ .../_rotation_matrix_magnitude.py | 27 ++ .../_transformations/_rotation_matrix_mean.py | 32 +++ .../_rotation_matrix_to_euler_angle.py | 51 ++++ .../_rotation_matrix_to_quaternion.py | 88 ++++++ .../_rotation_matrix_to_rotation_vector.py | 37 +++ .../_rotation_vector_identity.py | 64 +++++ .../_rotation_vector_magnitude.py | 34 +++ .../_transformations/_rotation_vector_mean.py | 42 +++ .../_rotation_vector_to_euler_angle.py | 48 ++++ .../_rotation_vector_to_quaternion.py | 71 +++++ .../_rotation_vector_to_rotation_matrix.py | 37 +++ .../_transformations/_translation_identity.py | 51 ++++ .../test__apply_euler_angle.py | 94 +++++++ .../test__apply_quaternion.py | 53 ++++ .../test__apply_rotation_matrix.py | 65 +++++ .../test__apply_rotation_vector.py | 54 ++++ .../test__compose_euler_angle.py | 60 +++++ .../test__compose_quaternion.py | 40 +++ .../test__compose_rotation_matrix.py | 37 +++ .../test__compose_rotation_vector.py | 40 +++ .../test__euler_angle_identity.py | 11 + .../test__euler_angle_magnitude.py | 56 ++++ .../test__euler_angle_mean.py | 61 +++++ .../test__euler_angle_to_quaternion.py | 75 ++++++ .../test__euler_angle_to_rotation_matrix.py | 65 +++++ .../test__euler_angle_to_rotation_vector.py | 63 +++++ .../test__invert_euler_angle.py | 56 ++++ .../test__invert_quaternion.py | 36 +++ .../test__invert_rotation_matrix.py | 33 +++ .../test__invert_rotation_vector.py | 35 +++ .../test__quaternion_identity.py | 37 +++ .../test__quaternion_magnitude.py | 36 +++ .../_transformations/test__quaternion_mean.py | 50 ++++ .../test__quaternion_to_euler_angle.py | 67 +++++ .../test__quaternion_to_rotation_matrix.py | 41 +++ .../test__quaternion_to_rotation_vector.py | 46 ++++ .../test__random_euler_angle.py | 51 ++++ .../test__random_quaternion.py | 31 +++ .../test__random_rotation_matrix.py | 28 ++ .../test__random_rotation_vector.py | 31 +++ .../test__rotation_matrix_identity.py | 34 +++ .../test__rotation_matrix_magnitude.py | 33 +++ .../test__rotation_matrix_mean.py | 33 +++ .../test__rotation_matrix_to_euler_angle.py | 65 +++++ .../test__rotation_matrix_to_quaternion.py | 49 ++++ ...est__rotation_matrix_to_rotation_vector.py | 45 ++++ .../test__rotation_vector_identity.py | 36 +++ .../test__rotation_vector_magnitude.py | 36 +++ .../test__rotation_vector_mean.py | 39 +++ .../test__rotation_vector_to_euler_angle.py | 67 +++++ .../test__rotation_vector_to_quaternion.py | 54 ++++ ...est__rotation_vector_to_rotation_matrix.py | 42 +++ .../_geometry/_transformations/test__slerp.py | 151 +++++++++++ 91 files changed, 4947 insertions(+), 4 deletions(-) create mode 100644 .pre-commit-config.yaml create mode 100644 docs/beignet.ops.md create mode 100644 src/beignet/ops/__init__.py create mode 100644 src/beignet/ops/_geometry/__init__.py create mode 100644 src/beignet/ops/_geometry/_transformations/__init__.py create mode 100644 src/beignet/ops/_geometry/_transformations/_apply_euler_angle.py create mode 100644 src/beignet/ops/_geometry/_transformations/_apply_quaternion.py create mode 100644 src/beignet/ops/_geometry/_transformations/_apply_rotation_matrix.py create mode 100644 src/beignet/ops/_geometry/_transformations/_apply_rotation_vector.py create mode 100644 src/beignet/ops/_geometry/_transformations/_compose_euler_angle.py create mode 100644 src/beignet/ops/_geometry/_transformations/_compose_quaternion.py create mode 100644 src/beignet/ops/_geometry/_transformations/_compose_rotation_matrix.py create mode 100644 src/beignet/ops/_geometry/_transformations/_compose_rotation_vector.py create mode 100644 src/beignet/ops/_geometry/_transformations/_euler_angle_identity.py create mode 100644 src/beignet/ops/_geometry/_transformations/_euler_angle_magnitude.py create mode 100644 src/beignet/ops/_geometry/_transformations/_euler_angle_mean.py create mode 100644 src/beignet/ops/_geometry/_transformations/_euler_angle_to_quaternion.py create mode 100644 src/beignet/ops/_geometry/_transformations/_euler_angle_to_rotation_matrix.py create mode 100644 src/beignet/ops/_geometry/_transformations/_euler_angle_to_rotation_vector.py create mode 100644 src/beignet/ops/_geometry/_transformations/_invert_euler_angle.py create mode 100644 src/beignet/ops/_geometry/_transformations/_invert_quaternion.py create mode 100644 src/beignet/ops/_geometry/_transformations/_invert_rotation_matrix.py create mode 100644 src/beignet/ops/_geometry/_transformations/_invert_rotation_vector.py create mode 100644 src/beignet/ops/_geometry/_transformations/_quaternion_identity.py create mode 100644 src/beignet/ops/_geometry/_transformations/_quaternion_magnitude.py create mode 100644 src/beignet/ops/_geometry/_transformations/_quaternion_mean.py create mode 100644 src/beignet/ops/_geometry/_transformations/_quaternion_slerp.py create mode 100644 src/beignet/ops/_geometry/_transformations/_quaternion_to_euler_angle.py create mode 100644 src/beignet/ops/_geometry/_transformations/_quaternion_to_rotation_matrix.py create mode 100644 src/beignet/ops/_geometry/_transformations/_quaternion_to_rotation_vector.py create mode 100644 src/beignet/ops/_geometry/_transformations/_random_euler_angle.py create mode 100644 src/beignet/ops/_geometry/_transformations/_random_quaternion.py create mode 100644 src/beignet/ops/_geometry/_transformations/_random_rotation_matrix.py create mode 100644 src/beignet/ops/_geometry/_transformations/_random_rotation_vector.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_matrix_identity.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_matrix_magnitude.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_matrix_mean.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_euler_angle.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_quaternion.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_rotation_vector.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_vector_identity.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_vector_magnitude.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_vector_mean.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_vector_to_euler_angle.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_vector_to_quaternion.py create mode 100644 src/beignet/ops/_geometry/_transformations/_rotation_vector_to_rotation_matrix.py create mode 100644 src/beignet/ops/_geometry/_transformations/_translation_identity.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__apply_euler_angle.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__apply_quaternion.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__apply_rotation_matrix.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__apply_rotation_vector.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__compose_euler_angle.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__compose_quaternion.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__compose_rotation_matrix.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__compose_rotation_vector.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__euler_angle_identity.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__euler_angle_magnitude.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__euler_angle_mean.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_quaternion.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_rotation_matrix.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_rotation_vector.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__invert_euler_angle.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__invert_quaternion.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__invert_rotation_matrix.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__invert_rotation_vector.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__quaternion_identity.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__quaternion_magnitude.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__quaternion_mean.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__quaternion_to_euler_angle.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__quaternion_to_rotation_matrix.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__quaternion_to_rotation_vector.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__random_euler_angle.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__random_quaternion.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__random_rotation_matrix.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__random_rotation_vector.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_identity.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_magnitude.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_mean.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_euler_angle.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_quaternion.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_rotation_vector.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_vector_identity.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_vector_magnitude.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_vector_mean.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_euler_angle.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_quaternion.py create mode 100644 tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_rotation_matrix.py create mode 100644 tests/beignet/ops/_geometry/_transformations/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/.gitignore b/.gitignore index 37d8376901..7dca454ba6 100644 --- a/.gitignore +++ b/.gitignore @@ -11,4 +11,5 @@ .ruff_cache/ __pycache__/ build/ -dist/ \ No newline at end of file +dist/ +venv/ 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/docs/beignet.ops.md b/docs/beignet.ops.md new file mode 100644 index 0000000000..f660c46747 --- /dev/null +++ b/docs/beignet.ops.md @@ -0,0 +1,11 @@ +# beignet.ops + +## Geometry + +### Transformations + +#### Rotations + +#### Translations + +## Interpolation diff --git a/pyproject.toml b/pyproject.toml index 7f9d47afbf..e0dde49845 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,14 +8,33 @@ requires = [ [project] authors = [{ email = "allen.goodman@icloud.com", name = "Allen Goodman" }] +dependencies = ["torch"] dynamic = ["version"] license = { file = "LICENSE" } name = "beignet" readme = "README.md" requires-python = ">=3.10" -[tool.ruff.format] -docstring-code-format = true +[project.optional-dependencies] +test = [ + "hypothesis", + "pytest", + "scipy", +] + +[tool.ruff] +select = [ + "B", # FLAKE8-BUGBEAR + "E", # PYCODESTYLE ERRORS + "F", # PYFLAKES + "I", # ISORT + "W", # PYCODESTYLE WARNINGS +] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = [ + "F401", # MODULE IMPORTED BUT UNUSED +] [tool.setuptools_scm] local_scheme = "no-local-version" diff --git a/src/beignet/ops/__init__.py b/src/beignet/ops/__init__.py new file mode 100644 index 0000000000..90534a8648 --- /dev/null +++ b/src/beignet/ops/__init__.py @@ -0,0 +1,89 @@ +from ._geometry import ( + apply_euler_angle, + apply_quaternion, + apply_rotation_matrix, + apply_rotation_vector, + compose_euler_angle, + compose_quaternion, + compose_rotation_matrix, + compose_rotation_vector, + euler_angle_identity, + euler_angle_magnitude, + euler_angle_mean, + euler_angle_to_quaternion, + euler_angle_to_rotation_matrix, + euler_angle_to_rotation_vector, + invert_euler_angle, + invert_quaternion, + invert_rotation_matrix, + invert_rotation_vector, + quaternion_identity, + quaternion_magnitude, + quaternion_mean, + quaternion_slerp, + quaternion_to_euler_angle, + quaternion_to_rotation_matrix, + quaternion_to_rotation_vector, + random_euler_angle, + random_quaternion, + random_rotation_matrix, + random_rotation_vector, + rotation_matrix_identity, + rotation_matrix_magnitude, + rotation_matrix_mean, + rotation_matrix_to_euler_angle, + rotation_matrix_to_quaternion, + rotation_matrix_to_rotation_vector, + rotation_vector_identity, + rotation_vector_magnitude, + rotation_vector_mean, + rotation_vector_to_euler_angle, + rotation_vector_to_quaternion, + rotation_vector_to_rotation_matrix, + translation_identity, +) + +__all__ = [ + "apply_euler_angle", + "apply_quaternion", + "apply_rotation_matrix", + "apply_rotation_vector", + "compose_euler_angle", + "compose_quaternion", + "compose_rotation_matrix", + "compose_rotation_vector", + "euler_angle_identity", + "euler_angle_magnitude", + "euler_angle_mean", + "euler_angle_to_quaternion", + "euler_angle_to_rotation_matrix", + "euler_angle_to_rotation_vector", + "invert_euler_angle", + "invert_quaternion", + "invert_rotation_matrix", + "invert_rotation_vector", + "quaternion_identity", + "quaternion_magnitude", + "quaternion_mean", + "quaternion_to_euler_angle", + "quaternion_to_rotation_matrix", + "quaternion_to_rotation_vector", + "random_euler_angle", + "random_quaternion", + "random_rotation_matrix", + "random_rotation_vector", + "rotation_matrix_identity", + "rotation_matrix_magnitude", + "rotation_matrix_mean", + "rotation_matrix_to_euler_angle", + "rotation_matrix_to_quaternion", + "rotation_matrix_to_rotation_vector", + "rotation_vector_identity", + "rotation_vector_magnitude", + "rotation_vector_mean", + "rotation_vector_to_euler_angle", + "rotation_vector_to_quaternion", + "rotation_vector_to_rotation_matrix", + "quaternion_slerp", + "translation_identity", +] diff --git a/src/beignet/ops/_geometry/__init__.py b/src/beignet/ops/_geometry/__init__.py new file mode 100644 index 0000000000..2d065e7f7f --- /dev/null +++ b/src/beignet/ops/_geometry/__init__.py @@ -0,0 +1,44 @@ +from ._transformations import ( + apply_euler_angle, + apply_quaternion, + apply_rotation_matrix, + apply_rotation_vector, + compose_euler_angle, + compose_quaternion, + compose_rotation_matrix, + compose_rotation_vector, + euler_angle_identity, + euler_angle_magnitude, + euler_angle_mean, + euler_angle_to_quaternion, + euler_angle_to_rotation_matrix, + euler_angle_to_rotation_vector, + invert_euler_angle, + invert_quaternion, + invert_rotation_matrix, + invert_rotation_vector, + quaternion_identity, + quaternion_magnitude, + quaternion_mean, + quaternion_slerp, + quaternion_to_euler_angle, + quaternion_to_rotation_matrix, + quaternion_to_rotation_vector, + random_euler_angle, + random_quaternion, + random_rotation_matrix, + random_rotation_vector, + rotation_matrix_identity, + rotation_matrix_magnitude, + rotation_matrix_mean, + rotation_matrix_to_euler_angle, + rotation_matrix_to_quaternion, + rotation_matrix_to_rotation_vector, + rotation_vector_identity, + rotation_vector_magnitude, + rotation_vector_mean, + rotation_vector_to_euler_angle, + rotation_vector_to_quaternion, + rotation_vector_to_rotation_matrix, + translation_identity, +) diff --git a/src/beignet/ops/_geometry/_transformations/__init__.py b/src/beignet/ops/_geometry/_transformations/__init__.py new file mode 100644 index 0000000000..5f64c220b5 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/__init__.py @@ -0,0 +1,60 @@ +from ._apply_euler_angle import apply_euler_angle +from ._apply_quaternion import ( + apply_quaternion, +) +from ._apply_rotation_matrix import apply_rotation_matrix +from ._apply_rotation_vector import apply_rotation_vector +from ._compose_euler_angle import compose_euler_angle +from ._compose_quaternion import compose_quaternion +from ._compose_rotation_matrix import compose_rotation_matrix +from ._compose_rotation_vector import compose_rotation_vector +from ._euler_angle_identity import euler_angle_identity +from ._euler_angle_magnitude import euler_angle_magnitude +from ._euler_angle_mean import euler_angle_mean +from ._euler_angle_to_quaternion import ( + euler_angle_to_quaternion, +) +from ._euler_angle_to_rotation_matrix import euler_angle_to_rotation_matrix +from ._euler_angle_to_rotation_vector import euler_angle_to_rotation_vector +from ._invert_euler_angle import invert_euler_angle +from ._invert_quaternion import invert_quaternion +from ._invert_rotation_matrix import invert_rotation_matrix +from ._invert_rotation_vector import invert_rotation_vector +from ._quaternion_identity import quaternion_identity +from ._quaternion_magnitude import quaternion_magnitude +from ._quaternion_mean import quaternion_mean +from ._quaternion_slerp import quaternion_slerp +from ._quaternion_to_euler_angle import ( + quaternion_to_euler_angle, +) +from ._quaternion_to_rotation_matrix import ( + quaternion_to_rotation_matrix, +) +from ._quaternion_to_rotation_vector import ( + quaternion_to_rotation_vector, +) +from ._random_euler_angle import random_euler_angle +from ._random_quaternion import random_quaternion +from ._random_rotation_matrix import random_rotation_matrix +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_mean import rotation_matrix_mean +from ._rotation_matrix_to_euler_angle import rotation_matrix_to_euler_angle +from ._rotation_matrix_to_quaternion import ( + rotation_matrix_to_quaternion, +) +from ._rotation_matrix_to_rotation_vector import ( + rotation_matrix_to_rotation_vector, +) +from ._rotation_vector_identity import rotation_vector_identity +from ._rotation_vector_magnitude import rotation_vector_magnitude +from ._rotation_vector_mean import rotation_vector_mean +from ._rotation_vector_to_euler_angle import rotation_vector_to_euler_angle +from ._rotation_vector_to_quaternion import ( + rotation_vector_to_quaternion, +) +from ._rotation_vector_to_rotation_matrix import ( + rotation_vector_to_rotation_matrix, +) +from ._translation_identity import translation_identity diff --git a/src/beignet/ops/_geometry/_transformations/_apply_euler_angle.py b/src/beignet/ops/_geometry/_transformations/_apply_euler_angle.py new file mode 100644 index 0000000000..3cb9f09079 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_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 + ------- + output : 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/ops/_geometry/_transformations/_apply_quaternion.py b/src/beignet/ops/_geometry/_transformations/_apply_quaternion.py new file mode 100644 index 0000000000..1410453b88 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_apply_quaternion.py @@ -0,0 +1,52 @@ +from torch import Tensor + +from ._apply_rotation_matrix import apply_rotation_matrix +from ._quaternion_to_rotation_matrix import ( + quaternion_to_rotation_matrix, +) + + +def apply_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. + """ + return apply_rotation_matrix( + input, + quaternion_to_rotation_matrix( + rotation, + ), + inverse, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_apply_rotation_matrix.py b/src/beignet/ops/_geometry/_transformations/_apply_rotation_matrix.py new file mode 100644 index 0000000000..fcac443401 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_apply_rotation_matrix.py @@ -0,0 +1,47 @@ +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: + output = torch.einsum("ikj, ik -> ij", rotation, input) + else: + output = torch.einsum("ijk, ik -> ij", rotation, input) + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_apply_rotation_vector.py b/src/beignet/ops/_geometry/_transformations/_apply_rotation_vector.py new file mode 100644 index 0000000000..0682bd4ac6 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_apply_rotation_vector.py @@ -0,0 +1,63 @@ +from torch import Tensor + +from ._apply_rotation_matrix import apply_rotation_matrix +from ._quaternion_to_rotation_matrix import ( + quaternion_to_rotation_matrix, +) +from ._rotation_vector_to_quaternion import ( + rotation_vector_to_quaternion, +) + + +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 + ------- + output : Tensor, shape=(..., 3) + Rotated vectors. + """ + return apply_rotation_matrix( + input, + quaternion_to_rotation_matrix( + rotation_vector_to_quaternion( + rotation, + degrees, + ), + ), + inverse, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_compose_euler_angle.py b/src/beignet/ops/_geometry/_transformations/_compose_euler_angle.py new file mode 100644 index 0000000000..3427d724dc --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_compose_euler_angle.py @@ -0,0 +1,53 @@ +from torch import Tensor + +from ._compose_quaternion import compose_quaternion +from ._euler_angle_to_quaternion import euler_angle_to_quaternion +from ._quaternion_to_euler_angle import quaternion_to_euler_angle + + +def compose_euler_angle( + input: Tensor, + other: Tensor, + axes: str, + degrees: bool | None = False, +) -> Tensor: + r""" + Compose rotation quaternions. + + Parameters + ---------- + input : Tensor, shape=(..., 3) + Euler angles. + + other : Tensor, shape=(..., 3) + Euler angles. + + axes : str + Axes. One to three characters belonging to the set :math:`\{X, Y, Z\}` + for intrinsic rotations, or :math:`\{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) + Composed Euler angles. + """ + return quaternion_to_euler_angle( + compose_quaternion( + euler_angle_to_quaternion( + input, + axes, + degrees, + ), + euler_angle_to_quaternion( + other, + axes, + degrees, + ), + ), + axes, + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_compose_quaternion.py b/src/beignet/ops/_geometry/_transformations/_compose_quaternion.py new file mode 100644 index 0000000000..35d0b30ae9 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_compose_quaternion.py @@ -0,0 +1,73 @@ +import torch +from torch import Tensor + + +def compose_quaternion( + input: Tensor, + other: Tensor, + canonical: bool = False, +) -> Tensor: + r""" + Compose 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) + Composed rotation quaternions. + """ + output = torch.empty( + [max(input.shape[0], other.shape[0]), 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j in range(max(input.shape[0], other.shape[0])): + a = input[j, 0] + b = input[j, 1] + c = input[j, 2] + d = input[j, 3] + + p = other[j, 0] + q = other[j, 1] + r = other[j, 2] + s = other[j, 3] + + t = output[j, 0] + u = output[j, 1] + v = output[j, 2] + w = output[j, 3] + + output[j, 0] = d * p + s * a + b * r - c * q + output[j, 1] = d * q + s * b + c * p - a * r + output[j, 2] = d * r + s * c + a * q - b * p + output[j, 3] = d * s - a * p - b * q - c * r + + x = torch.sqrt(t**2.0 + u**2.0 + v**2.0 + w**2.0) + + if x == 0.0: + output[j] = torch.nan + + output[j] = output[j] / x + + if canonical: + if w == 0 and (t == 0 and (u == 0 and v < 0 or u < 0) or t < 0) or w < 0: + output[j] = -output[j] + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_compose_rotation_matrix.py b/src/beignet/ops/_geometry/_transformations/_compose_rotation_matrix.py new file mode 100644 index 0000000000..10b6ce1582 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_compose_rotation_matrix.py @@ -0,0 +1,33 @@ +from torch import Tensor + +from ._compose_quaternion import compose_quaternion +from ._quaternion_to_rotation_matrix import quaternion_to_rotation_matrix +from ._rotation_matrix_to_quaternion import rotation_matrix_to_quaternion + + +def compose_rotation_matrix( + input: Tensor, + other: Tensor, +) -> Tensor: + r""" + Compose rotation matrices. + + Parameters + ---------- + input : Tensor, shape=(..., 3, 3) + Rotation matrices. + + other : Tensor, shape=(..., 3, 3) + Rotation matrices. + + Returns + ------- + output : Tensor, shape=(..., 3, 3) + Composed rotation matrices. + """ + return quaternion_to_rotation_matrix( + compose_quaternion( + rotation_matrix_to_quaternion(input), + rotation_matrix_to_quaternion(other), + ), + ) diff --git a/src/beignet/ops/_geometry/_transformations/_compose_rotation_vector.py b/src/beignet/ops/_geometry/_transformations/_compose_rotation_vector.py new file mode 100644 index 0000000000..d15223e489 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_compose_rotation_vector.py @@ -0,0 +1,45 @@ +from torch import Tensor + +from ._compose_quaternion import compose_quaternion +from ._quaternion_to_rotation_vector import quaternion_to_rotation_vector +from ._rotation_vector_to_quaternion import rotation_vector_to_quaternion + + +def compose_rotation_vector( + input: Tensor, + other: Tensor, + degrees: bool | None = False, +) -> Tensor: + r""" + Compose rotation vectors. + + Parameters + ---------- + input : Tensor, shape=(..., 4) + Rotation vectors. + + other : Tensor, shape=(..., 4) + Rotation vectors. + + degrees : bool, optional + If `True`, rotation vector magnitudes are assumed to be in degrees. + Default, `False`. + + Returns + ------- + output : Tensor, shape=(..., 4) + Composed rotation vectors. + """ + return quaternion_to_rotation_vector( + compose_quaternion( + rotation_vector_to_quaternion( + input, + degrees, + ), + rotation_vector_to_quaternion( + other, + degrees, + ), + ), + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_euler_angle_identity.py b/src/beignet/ops/_geometry/_transformations/_euler_angle_identity.py new file mode 100644 index 0000000000..40c716cf3a --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_euler_angle_identity.py @@ -0,0 +1,70 @@ +import torch +from torch import Tensor + +from ._quaternion_identity import quaternion_identity +from ._quaternion_to_euler_angle import ( + 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: + r""" + 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 quaternion_to_euler_angle( + quaternion_identity( + size, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ), + axes, + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_euler_angle_magnitude.py b/src/beignet/ops/_geometry/_transformations/_euler_angle_magnitude.py new file mode 100644 index 0000000000..3429062d09 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_euler_angle_magnitude.py @@ -0,0 +1,41 @@ +from torch import Tensor + +from ._euler_angle_to_quaternion import ( + euler_angle_to_quaternion, +) +from ._quaternion_magnitude import 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 quaternion_magnitude( + euler_angle_to_quaternion( + input, + axes, + degrees, + ), + ) diff --git a/src/beignet/ops/_geometry/_transformations/_euler_angle_mean.py b/src/beignet/ops/_geometry/_transformations/_euler_angle_mean.py new file mode 100644 index 0000000000..5e711958bc --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_euler_angle_mean.py @@ -0,0 +1,49 @@ +from torch import Tensor + +from ._euler_angle_to_quaternion import euler_angle_to_quaternion +from ._quaternion_mean import quaternion_mean +from ._quaternion_to_euler_angle import quaternion_to_euler_angle + + +def euler_angle_mean( + input: Tensor, + weight: Tensor | None = None, + axes: str | None = None, + degrees: bool | None = False, +) -> Tensor: + r""" + Euler angle mean. + + Parameters + ---------- + input : Tensor, shape=(..., 3) + Euler angles. + + weight : Tensor, shape=(..., 4), optional + Relative importance of rotation quaternions. + + axes : str + Axes. One to three characters belonging to the set :math:`\{X, Y, Z\}` + for intrinsic rotations, or :math:`\{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 angle mean. + """ + return quaternion_to_euler_angle( + quaternion_mean( + euler_angle_to_quaternion( + input, + axes, + degrees, + ), + weight, + ), + axes, + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_euler_angle_to_quaternion.py b/src/beignet/ops/_geometry/_transformations/_euler_angle_to_quaternion.py new file mode 100644 index 0000000000..7736d8171b --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_euler_angle_to_quaternion.py @@ -0,0 +1,170 @@ +import re + +import torch +from torch import Tensor + + +def euler_angle_to_quaternion( + input: Tensor, + axes: str, + degrees: bool = False, + canonical: bool | None = False, +) -> Tensor: + r""" + Convert Euler angles to rotation quaternions. + + Parameters + ---------- + input : Tensor, shape=(..., 3) + Euler angles. + + axes : str + Axes. One to three characters belonging to the set :math:`\{X, Y, Z\}` + for intrinsic rotations, or :math:`\{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 + ------- + output : 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, 1] * y[k, 2] - z[k, 2] * y[k, 1] + c[1] = z[k, 2] * y[k, 0] - z[k, 0] * y[k, 2] + c[2] = z[k, 0] * y[k, 1] - z[k, 1] * y[k, 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, 1] * z[k, 2] - y[k, 2] * z[k, 1] + c[1] = y[k, 2] * z[k, 0] - y[k, 0] * z[k, 2] + c[2] = y[k, 0] * z[k, 1] - y[k, 1] * z[k, 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/ops/_geometry/_transformations/_euler_angle_to_rotation_matrix.py b/src/beignet/ops/_geometry/_transformations/_euler_angle_to_rotation_matrix.py new file mode 100644 index 0000000000..2d1baf1c28 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_euler_angle_to_rotation_matrix.py @@ -0,0 +1,80 @@ +import torch +from torch import Tensor + + +def euler_angle_to_rotation_matrix( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + r""" + Convert Euler angles to rotation matrices. + + Parameters + ---------- + input : Tensor, shape=(..., 3) + Euler angles. + + axes : str + Axes. One to three characters belonging to the set :math:`\{X, Y, Z\}` + for intrinsic rotations, or :math:`\{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, 3) + Rotation matrices. + """ + if degrees: + input = torch.deg2rad(input) + + output = torch.empty( + [input.shape[0], 3, 3], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j, axis in enumerate(axes): + a = torch.cos(input[..., j]) + b = torch.sin(input[..., j]) + + p = torch.full_like(a, 1.0) + q = torch.full_like(a, 0.0) + + match axis.lower(): + case "x": + x = [ + torch.stack([+p, +q, +q], dim=-1), + torch.stack([+q, +a, -b], dim=-1), + torch.stack([+q, +b, +a], dim=-1), + ] + case "y": + x = [ + torch.stack([+a, +q, +b], dim=-1), + torch.stack([+q, +p, +q], dim=-1), + torch.stack([-b, +q, +a], dim=-1), + ] + case "z": + x = [ + torch.stack([+a, -b, +q], dim=-1), + torch.stack([+b, +a, +q], dim=-1), + torch.stack([+q, +q, +p], dim=-1), + ] + case _: + raise ValueError + + x = torch.stack(x, dim=-2) + + if j == 0: + output = x + else: + if axes.islower(): + output = x @ output + else: + output = output @ x + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_euler_angle_to_rotation_vector.py b/src/beignet/ops/_geometry/_transformations/_euler_angle_to_rotation_vector.py new file mode 100644 index 0000000000..19dd27cc48 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_euler_angle_to_rotation_vector.py @@ -0,0 +1,253 @@ +import re + +import torch +from torch import Tensor + + +def euler_angle_to_rotation_vector( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + r""" + Convert Euler angles to rotation vectors. + + Parameters + ---------- + input : Tensor, shape=(..., 3) + Euler angles. + + axes : str + Axes. One to three characters belonging to the set :math:`\{X, Y, Z\}` + for intrinsic rotations, or :math:`\{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 + ------- + output : 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 + + if degrees: + input = torch.deg2rad(input) + + if len(axes.lower()) == 1: + match input.ndim: + case 0: + input = torch.reshape(input, [1, 1]) + case 1: + input = input[:, None] + case 2 if input.shape[-1] != 1: + raise ValueError + case _: + raise ValueError + + else: + if input.ndim not in [1, 2] or input.shape[-1] != len(axes.lower()): + raise ValueError + + if input.ndim == 1: + input = input[None, :] + + if input.ndim != 2 or input.shape[-1] != len(axes.lower()): + raise ValueError + + x = torch.zeros( + [input.shape[0], 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + match axes.lower()[0]: + case "x": + m = 0 + case "y": + m = 1 + case "z": + m = 2 + case _: + raise ValueError + + for j in range(input[:, 0].shape[0]): + x[j, 3] = torch.cos(input[:, 0][j] / 2) + x[j, m] = torch.sin(input[:, 0][j] / 2) + + for j in range(1, len(axes)): + y = torch.zeros( + [input.shape[0], 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + z = torch.empty( + [max(y.shape[0], x.shape[0]), 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + match axes.lower()[j]: + case "x": + m = 0 + case "y": + m = 1 + case "z": + m = 2 + case _: + raise ValueError + + for k in range(input[:, j].shape[0]): + y[k, 3] = torch.cos(input[:, j][k] / 2) + y[k, m] = torch.sin(input[:, j][k] / 2) + + if intrinsic: + if x.shape[0] == 1: + for k in range(max(x.shape[0], y.shape[0])): + q = y[k, 1] + r = y[k, 2] + s = y[k, 3] + p = y[k, 0] + + t = x[0, 0] + u = x[0, 1] + v = x[0, 2] + w = x[0, 3] + + z[k, 0] = w * p + s * t + u * r - v * q + z[k, 1] = w * q + s * u + v * p - t * r + z[k, 2] = w * r + s * v + t * q - u * p + z[k, 3] = w * s - t * p - u * q - v * r + elif y.shape[0] == 1: + for k in range(max(x.shape[0], y.shape[0])): + p = y[0, 0] + q = y[0, 1] + r = y[0, 2] + s = y[0, 3] + + t = x[k, 0] + u = x[k, 1] + v = x[k, 2] + w = x[k, 3] + + z[k, 0] = w * p + s * t + u * r - v * q + z[k, 1] = w * q + s * u + v * p - t * r + z[k, 2] = w * r + s * v + t * q - u * p + z[k, 3] = w * s - t * p - u * q - v * r + else: + for k in range(max(x.shape[0], y.shape[0])): + p = y[k, 0] + q = y[k, 1] + r = y[k, 2] + s = y[k, 3] + + t = x[k, 0] + u = x[k, 1] + v = x[k, 2] + w = x[k, 3] + + z[k, 0] = w * p + s * t + u * r - v * q + z[k, 1] = w * q + s * u + v * p - t * r + z[k, 2] = w * r + s * v + t * q - u * p + z[k, 3] = w * s - t * p - u * q - v * r + + x = z + else: + if y.shape[0] == 1: + for k in range(max(y.shape[0], x.shape[0])): + p = y[0, 0] + q = y[0, 1] + r = y[0, 2] + s = y[0, 3] + + t = x[k, 0] + u = x[k, 1] + v = x[k, 2] + w = x[k, 3] + + z[k, 0] = s * t + w * p + q * v - r * u + z[k, 1] = s * u + w * q + r * t - p * v + z[k, 2] = s * v + w * r + p * u - q * t + z[k, 3] = s * w - p * t - q * u - r * v + elif x.shape[0] == 1: + for k in range(max(y.shape[0], x.shape[0])): + t = x[0, 0] + u = x[0, 1] + v = x[0, 2] + w = x[0, 3] + + p = y[k, 0] + q = y[k, 1] + r = y[k, 2] + s = y[k, 3] + + z[k, 0] = s * t + w * p + q * v - r * u + z[k, 1] = s * u + w * q + r * t - p * v + z[k, 2] = s * v + w * r + p * u - q * t + z[k, 3] = s * w - p * t - q * u - r * v + else: + for k in range(max(y.shape[0], x.shape[0])): + p = y[k, 0] + q = y[k, 1] + r = y[k, 2] + s = y[k, 3] + + t = x[k, 0] + u = x[k, 1] + v = x[k, 2] + w = x[k, 3] + + z[k, 0] = s * t + w * p + q * v - r * u + z[k, 1] = s * u + w * q + r * t - p * v + z[k, 2] = s * v + w * r + p * u - q * t + z[k, 3] = s * w - p * t - q * u - r * v + + x = z + + output = torch.empty( + [x.shape[0], 3], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for j in range(x.shape[0]): + a = x[j, 0] + b = x[j, 1] + c = x[j, 2] + d = x[j, 3] + + if d == 0 and (a == 0 and (b == 0 and c < 0 or b < 0) or a < 0) or d < 0: + x[j] = -x[j] + + y = 2.0 * torch.atan2(torch.sqrt(a**2.0 + b**2.0 + c**2.0), d**1.0) + + if y < 0.001: + y = 2.0 + y**2.0 / 12.0 + 7.0 * y**2.0 * y**2.0 / 2880.0 + else: + y = y / torch.sin(y / 2.0) + + output[j] = x[j, :-1] * y + + if degrees: + output = torch.rad2deg(output) + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_invert_euler_angle.py b/src/beignet/ops/_geometry/_transformations/_invert_euler_angle.py new file mode 100644 index 0000000000..9a7c3c6a6d --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_invert_euler_angle.py @@ -0,0 +1,48 @@ +from torch import Tensor + +from ._euler_angle_to_quaternion import ( + euler_angle_to_quaternion, +) +from ._invert_quaternion import invert_quaternion +from ._quaternion_to_euler_angle import ( + quaternion_to_euler_angle, +) + + +def invert_euler_angle( + input: Tensor, + axes: str, + degrees: bool | None = False, +) -> Tensor: + r""" + 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 quaternion_to_euler_angle( + invert_quaternion( + euler_angle_to_quaternion( + input, + axes, + degrees, + ), + ), + axes, + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_invert_quaternion.py b/src/beignet/ops/_geometry/_transformations/_invert_quaternion.py new file mode 100644 index 0000000000..9762327f75 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_invert_quaternion.py @@ -0,0 +1,31 @@ +from torch import Tensor + + +def invert_quaternion( + input: Tensor, + canonical: bool = False, +) -> Tensor: + r""" + 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_quaternions : Tensor, shape (..., 4) + Inverted rotation quaternions. + """ + input[:, :3] = -input[:, :3] + + return input diff --git a/src/beignet/ops/_geometry/_transformations/_invert_rotation_matrix.py b/src/beignet/ops/_geometry/_transformations/_invert_rotation_matrix.py new file mode 100644 index 0000000000..4ceba7d74a --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_invert_rotation_matrix.py @@ -0,0 +1,19 @@ +import torch +from torch import Tensor + + +def invert_rotation_matrix(input: Tensor) -> Tensor: + r""" + 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/ops/_geometry/_transformations/_invert_rotation_vector.py b/src/beignet/ops/_geometry/_transformations/_invert_rotation_vector.py new file mode 100644 index 0000000000..c592cdbb91 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_invert_rotation_vector.py @@ -0,0 +1,20 @@ +from torch import Tensor + + +def invert_rotation_vector( + input: Tensor, +) -> Tensor: + r""" + 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/ops/_geometry/_transformations/_quaternion_identity.py b/src/beignet/ops/_geometry/_transformations/_quaternion_identity.py new file mode 100644 index 0000000000..9474cb9669 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_quaternion_identity.py @@ -0,0 +1,64 @@ +import torch +from torch import Tensor + + +def 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: + r""" + 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_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/ops/_geometry/_transformations/_quaternion_magnitude.py b/src/beignet/ops/_geometry/_transformations/_quaternion_magnitude.py new file mode 100644 index 0000000000..7d2179c094 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_quaternion_magnitude.py @@ -0,0 +1,37 @@ +import torch +from torch import Tensor + + +def quaternion_magnitude(input: Tensor, canonical=False) -> Tensor: + r""" + Rotation quaternion magnitudes. + + Parameters + ---------- + input : Tensor, shape=(..., 4) + Rotation quaternions. + + Returns + ------- + output : Tensor, shape=(...) + Angles in radians. Magnitudes will be in the range :math:`[0, \pi]`. + """ + output = torch.empty( + input.shape[0], + dtype=input.dtype, + layout=input.layout, + device=input.device, + requires_grad=input.requires_grad, + ) + + for j in range(input.shape[0]): + a = input[j, 0] + b = input[j, 1] + c = input[j, 2] + d = input[j, 3] + + x = torch.atan2(torch.sqrt(a**2 + b**2 + c**2), torch.abs(d)) + + output[j] = x * 2.0 + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_quaternion_mean.py b/src/beignet/ops/_geometry/_transformations/_quaternion_mean.py new file mode 100644 index 0000000000..c6c0d369be --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_quaternion_mean.py @@ -0,0 +1,34 @@ +import torch +from torch import Tensor + + +def quaternion_mean( + input: Tensor, + weight: Tensor | None = None, +) -> Tensor: + r""" + 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. + + Returns + ------- + output : Tensor, shape=(..., 4) + Rotation quaternions mean. + """ + if weight is None: + weight = torch.ones(input.shape[0]) + + _, output = torch.linalg.eigh((input.T * weight) @ input) + + output = output[:, -1] + + output = torch.unsqueeze(output, dim=0) + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_quaternion_slerp.py b/src/beignet/ops/_geometry/_transformations/_quaternion_slerp.py new file mode 100644 index 0000000000..694052a637 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_quaternion_slerp.py @@ -0,0 +1,93 @@ +import torch +import torch.testing +from torch import Tensor + + +def quaternion_slerp( + input: Tensor, + time: Tensor, + rotation: Tensor, +) -> 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. + """ # noqa: E501 + if time.shape[-1] != rotation.shape[-2]: + raise ValueError + + output = torch.empty( + [*input.shape, 4], + dtype=input.dtype, + layout=input.layout, + device=input.device, + ) + + for index, t in enumerate(input): + b = torch.min(torch.nonzero(torch.greater_equal(time, t))) + + if b > 0: + a = b - 1 + else: + a = b + + if time[b] == t or b == a: + output[index] = rotation[b] + + continue + + p, q = time[a], time[b] + + r = (t - p) / (q - p) + + t = rotation[a] + u = rotation[b] + + v = torch.dot(t, u) + + if v < 0.0: + u = -u + v = -v + + if v > 0.9995: + z = (1.0 - r) * t + r * u + else: + x = torch.sqrt(1.0 - v**2.0) + + y = torch.atan2(x, v) + + z = t * torch.sin((1.0 - r) * y) / x + u * torch.sin(r * y) / x + + output[index] = z / torch.linalg.norm(z) + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_quaternion_to_euler_angle.py b/src/beignet/ops/_geometry/_transformations/_quaternion_to_euler_angle.py new file mode 100644 index 0000000000..e66d7a6459 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_quaternion_to_euler_angle.py @@ -0,0 +1,143 @@ +import math +import re + +import torch +from torch import Tensor + + +def quaternion_to_euler_angle( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + r""" + 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/ops/_geometry/_transformations/_quaternion_to_rotation_matrix.py b/src/beignet/ops/_geometry/_transformations/_quaternion_to_rotation_matrix.py new file mode 100644 index 0000000000..20d753af6b --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_quaternion_to_rotation_matrix.py @@ -0,0 +1,43 @@ +import torch +from torch import Tensor + + +def quaternion_to_rotation_matrix(input: Tensor) -> Tensor: + r""" + 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/ops/_geometry/_transformations/_quaternion_to_rotation_vector.py b/src/beignet/ops/_geometry/_transformations/_quaternion_to_rotation_vector.py new file mode 100644 index 0000000000..6827c672e8 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_quaternion_to_rotation_vector.py @@ -0,0 +1,57 @@ +import torch +from torch import Tensor + + +def quaternion_to_rotation_vector( + input: Tensor, + degrees: bool | None = False, +) -> Tensor: + r""" + 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 = 2.0 + y**2.0 / 12 + 7 * y**2.0 * y**2.0 / 2880 + 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/ops/_geometry/_transformations/_random_euler_angle.py b/src/beignet/ops/_geometry/_transformations/_random_euler_angle.py new file mode 100644 index 0000000000..f41bea8995 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_random_euler_angle.py @@ -0,0 +1,91 @@ +import torch +from torch import Generator, Tensor + +from ._quaternion_to_euler_angle import ( + quaternion_to_euler_angle, +) +from ._random_quaternion import random_quaternion + + +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: + r""" + 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 quaternion_to_euler_angle( + random_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/ops/_geometry/_transformations/_random_quaternion.py b/src/beignet/ops/_geometry/_transformations/_random_quaternion.py new file mode 100644 index 0000000000..468e4c3e8e --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_random_quaternion.py @@ -0,0 +1,95 @@ +import torch +from torch import Generator, Tensor + + +def random_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: + r""" + 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_quaternions : Tensor, shape (..., 4) + Random rotation quaternions. + """ + 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(quaternions.size(0)): + if ( + (quaternions[index][3] < 0) + or (quaternions[index][3] == 0 and quaternions[index][0] < 0) + or ( + quaternions[index][3] == 0 + and quaternions[index][0] == 0 + and quaternions[index][1] < 0 + ) + or ( + quaternions[index][3] == 0 + and quaternions[index][0] == 0 + and quaternions[index][1] == 0 + and quaternions[index][2] < 0 + ) + ): + quaternions[index][0] *= -1.0 + quaternions[index][1] *= -1.0 + quaternions[index][2] *= -1.0 + quaternions[index][3] *= -1.0 + + return quaternions diff --git a/src/beignet/ops/_geometry/_transformations/_random_rotation_matrix.py b/src/beignet/ops/_geometry/_transformations/_random_rotation_matrix.py new file mode 100644 index 0000000000..8d9368a1c4 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_random_rotation_matrix.py @@ -0,0 +1,69 @@ +import torch +from torch import Generator, Tensor + +from ._quaternion_to_rotation_matrix import ( + quaternion_to_rotation_matrix, +) +from ._random_quaternion import random_quaternion + + +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: + r""" + 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 quaternion_to_rotation_matrix( + random_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/ops/_geometry/_transformations/_random_rotation_vector.py b/src/beignet/ops/_geometry/_transformations/_random_rotation_vector.py new file mode 100644 index 0000000000..1cfb578a42 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_random_rotation_vector.py @@ -0,0 +1,73 @@ +import torch +from torch import Generator, Tensor + +from ._quaternion_to_rotation_vector import ( + quaternion_to_rotation_vector, +) +from ._random_quaternion import random_quaternion + + +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: + r""" + 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 quaternion_to_rotation_vector( + random_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/ops/_geometry/_transformations/_rotation_matrix_identity.py b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_identity.py new file mode 100644 index 0000000000..4e646a5d85 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_identity.py @@ -0,0 +1,58 @@ +import torch +from torch import Tensor + +from ._quaternion_identity import quaternion_identity +from ._quaternion_to_rotation_matrix import ( + 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: + r""" + 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 quaternion_to_rotation_matrix( + quaternion_identity( + size, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ), + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_matrix_magnitude.py b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_magnitude.py new file mode 100644 index 0000000000..74c39721c7 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_magnitude.py @@ -0,0 +1,27 @@ +from torch import Tensor + +from ._quaternion_magnitude import quaternion_magnitude +from ._rotation_matrix_to_quaternion import ( + rotation_matrix_to_quaternion, +) + + +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 quaternion_magnitude( + rotation_matrix_to_quaternion( + input, + ), + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_matrix_mean.py b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_mean.py new file mode 100644 index 0000000000..0fb8bce5a5 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_mean.py @@ -0,0 +1,32 @@ +from torch import Tensor + +from ._quaternion_mean import quaternion_mean +from ._quaternion_to_rotation_matrix import quaternion_to_rotation_matrix +from ._rotation_matrix_to_quaternion import rotation_matrix_to_quaternion + + +def rotation_matrix_mean( + input: Tensor, + weight: Tensor | None = None, +) -> Tensor: + r""" + Parameters + ---------- + input : Tensor, shape=(..., 3, 3) + Rotation matrices. + + weight : Tensor, shape=(..., 4), optional + Relative importance of rotation matrices. + + Returns + ------- + output : Tensor, shape=(..., 3, 3) + """ + return quaternion_to_rotation_matrix( + quaternion_mean( + rotation_matrix_to_quaternion( + input, + ), + weight, + ), + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_euler_angle.py b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_euler_angle.py new file mode 100644 index 0000000000..cf6515a939 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_euler_angle.py @@ -0,0 +1,51 @@ +from torch import Tensor + +from ._quaternion_to_euler_angle import ( + quaternion_to_euler_angle, +) +from ._rotation_matrix_to_quaternion import ( + rotation_matrix_to_quaternion, +) + + +def rotation_matrix_to_euler_angle( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + r""" + 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 quaternion_to_euler_angle( + rotation_matrix_to_quaternion( + input, + ), + axes, + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_quaternion.py b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_quaternion.py new file mode 100644 index 0000000000..c415a6f888 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_quaternion.py @@ -0,0 +1,88 @@ +import torch +from torch import Tensor + + +def rotation_matrix_to_quaternion( + input: Tensor, + canonical: bool | None = False, +) -> Tensor: + r""" + 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 & (b == 0 & c < 0 | b < 0) | a < 0) | d < 0: + output[j] = -output[j] + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_rotation_vector.py b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_rotation_vector.py new file mode 100644 index 0000000000..75409a01b0 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_matrix_to_rotation_vector.py @@ -0,0 +1,37 @@ +from torch import Tensor + +from ._quaternion_to_rotation_vector import ( + quaternion_to_rotation_vector, +) +from ._rotation_matrix_to_quaternion import ( + rotation_matrix_to_quaternion, +) + + +def rotation_matrix_to_rotation_vector( + input: Tensor, + degrees: bool = False, +) -> Tensor: + r""" + 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 quaternion_to_rotation_vector( + rotation_matrix_to_quaternion( + input, + ), + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_vector_identity.py b/src/beignet/ops/_geometry/_transformations/_rotation_vector_identity.py new file mode 100644 index 0000000000..29a0d0ad38 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_vector_identity.py @@ -0,0 +1,64 @@ +import torch +from torch import Tensor + +from ._quaternion_identity import quaternion_identity +from ._quaternion_to_rotation_vector import ( + 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: + r""" + 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 quaternion_to_rotation_vector( + quaternion_identity( + size, + out=out, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ), + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_vector_magnitude.py b/src/beignet/ops/_geometry/_transformations/_rotation_vector_magnitude.py new file mode 100644 index 0000000000..ceb1725235 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_vector_magnitude.py @@ -0,0 +1,34 @@ +from torch import Tensor + +from ._quaternion_magnitude import quaternion_magnitude +from ._rotation_vector_to_quaternion import ( + rotation_vector_to_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 quaternion_magnitude( + rotation_vector_to_quaternion( + input, + degrees, + ), + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_vector_mean.py b/src/beignet/ops/_geometry/_transformations/_rotation_vector_mean.py new file mode 100644 index 0000000000..8fb8c2a968 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_vector_mean.py @@ -0,0 +1,42 @@ +from torch import Tensor + +from ._quaternion_mean import quaternion_mean +from ._quaternion_to_rotation_vector import quaternion_to_rotation_vector +from ._rotation_vector_to_quaternion import rotation_vector_to_quaternion + + +def rotation_vector_mean( + input: Tensor, + weight: Tensor | None = None, + degrees: bool | None = False, +) -> Tensor: + r""" + Compose rotation vectors. + + Parameters + ---------- + input : Tensor, shape=(..., 4) + Rotation vectors. + + weight : Tensor, shape=(..., 4), optional + Relative importance of rotation matrices. + + degrees : bool, optional + If `True`, rotation vector magnitudes are assumed to be in degrees. + Default, `False`. + + Returns + ------- + output : Tensor, shape=(..., 4) + Rotation vectors mean. + """ + return quaternion_to_rotation_vector( + quaternion_mean( + rotation_vector_to_quaternion( + input, + degrees, + ), + weight, + ), + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_euler_angle.py b/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_euler_angle.py new file mode 100644 index 0000000000..8ae5503d5c --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_euler_angle.py @@ -0,0 +1,48 @@ +from torch import Tensor + +from ._quaternion_to_euler_angle import ( + quaternion_to_euler_angle, +) +from ._rotation_vector_to_quaternion import ( + rotation_vector_to_quaternion, +) + + +def rotation_vector_to_euler_angle( + input: Tensor, + axes: str, + degrees: bool = False, +) -> Tensor: + r""" + 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 quaternion_to_euler_angle( + rotation_vector_to_quaternion( + input, + degrees, + ), + axes, + degrees, + ) diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_quaternion.py b/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_quaternion.py new file mode 100644 index 0000000000..92bd81b65a --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_quaternion.py @@ -0,0 +1,71 @@ +import torch +from torch import Tensor + + +def rotation_vector_to_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 & (b == 0 & c < 0 | b < 0) | a < 0) | d < 0: + output[j] = -output[j] + + return output diff --git a/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_rotation_matrix.py b/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_rotation_matrix.py new file mode 100644 index 0000000000..9259f8b233 --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_rotation_vector_to_rotation_matrix.py @@ -0,0 +1,37 @@ +from torch import Tensor + +from ._quaternion_to_rotation_matrix import ( + quaternion_to_rotation_matrix, +) +from ._rotation_vector_to_quaternion import ( + rotation_vector_to_quaternion, +) + + +def rotation_vector_to_rotation_matrix( + input: Tensor, + degrees: bool | None = False, +) -> Tensor: + r""" + 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 quaternion_to_rotation_matrix( + rotation_vector_to_quaternion( + input, + degrees, + ), + ) diff --git a/src/beignet/ops/_geometry/_transformations/_translation_identity.py b/src/beignet/ops/_geometry/_transformations/_translation_identity.py new file mode 100644 index 0000000000..1fa2b6966a --- /dev/null +++ b/src/beignet/ops/_geometry/_transformations/_translation_identity.py @@ -0,0 +1,51 @@ +import torch +from torch import Tensor + + +def translation_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/ops/_geometry/_transformations/test__apply_euler_angle.py b/tests/beignet/ops/_geometry/_transformations/test__apply_euler_angle.py new file mode 100644 index 0000000000..86aa0a1296 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__apply_euler_angle.py @@ -0,0 +1,94 @@ +import beignet.ops +import hypothesis.extra.numpy +import hypothesis.strategies +import numpy +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = function( + hypothesis.extra.numpy.arrays( + numpy.float64, + (size, 3), + elements={ + "allow_infinity": False, + "min_value": numpy.finfo(numpy.float32).min, + "max_value": numpy.finfo(numpy.float32).max, + }, + ), + ) + + rotation = Rotation.random( + 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_euler_angle(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.apply_euler_angle( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__apply_quaternion.py b/tests/beignet/ops/_geometry/_transformations/test__apply_quaternion.py new file mode 100644 index 0000000000..83a874f815 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__apply_quaternion.py @@ -0,0 +1,53 @@ +import beignet.ops +import hypothesis.extra.numpy +import hypothesis.strategies +import numpy +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = function( + hypothesis.extra.numpy.arrays( + numpy.float64, + (size, 3), + elements={ + "allow_infinity": False, + "min_value": numpy.finfo(numpy.float32).min, + "max_value": numpy.finfo(numpy.float32).max, + }, + ), + ) + + rotation = Rotation.random(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_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.apply_quaternion(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__apply_rotation_matrix.py b/tests/beignet/ops/_geometry/_transformations/test__apply_rotation_matrix.py new file mode 100644 index 0000000000..4ce23b05ac --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__apply_rotation_matrix.py @@ -0,0 +1,65 @@ +import beignet.ops +import hypothesis.extra.numpy +import hypothesis.strategies +import numpy +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = function( + hypothesis.extra.numpy.arrays( + numpy.float64, + (size, 3), + elements={ + "allow_infinity": False, + "min_value": numpy.finfo(numpy.float32).min, + "max_value": numpy.finfo(numpy.float32).max, + }, + ), + ) + + rotation = Rotation.random(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.ops.apply_rotation_matrix( + **parameters, + ), + expected, + equal_nan=True, + atol=1e-06, + rtol=1e-06, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__apply_rotation_vector.py b/tests/beignet/ops/_geometry/_transformations/test__apply_rotation_vector.py new file mode 100644 index 0000000000..33a5aa31ec --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__apply_rotation_vector.py @@ -0,0 +1,54 @@ +import beignet.ops +import hypothesis.extra.numpy +import hypothesis.strategies +import numpy +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=1, + max_value=8, + ), + ) + + input = function( + hypothesis.extra.numpy.arrays( + numpy.float64, + (size, 3), + elements={ + "allow_infinity": False, + "min_value": numpy.finfo(numpy.float32).min, + "max_value": numpy.finfo(numpy.float32).max, + }, + ), + ) + + rotation = Rotation.random(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.ops.apply_rotation_vector(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__compose_euler_angle.py b/tests/beignet/ops/_geometry/_transformations/test__compose_euler_angle.py new file mode 100644 index 0000000000..814e397bdc --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__compose_euler_angle.py @@ -0,0 +1,60 @@ +import operator + +import beignet.ops +import hypothesis.strategies +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ) + + input = Rotation.random(size) + other = Rotation.random(size) + + 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(input.as_euler(axes, degrees)), + "other": torch.from_numpy(other.as_euler(axes, degrees)), + "axes": axes, + "degrees": degrees, + }, + torch.from_numpy(operator.mul(input, other).as_euler(axes, degrees)), + ) + + +@hypothesis.given(_strategy()) +def test_compose_euler_angle(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.compose_euler_angle(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__compose_quaternion.py b/tests/beignet/ops/_geometry/_transformations/test__compose_quaternion.py new file mode 100644 index 0000000000..a0f51f686b --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__compose_quaternion.py @@ -0,0 +1,40 @@ +import operator + +import beignet.ops +import hypothesis.strategies +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ) + + 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(operator.mul(input, other).as_quat(canonical))), + ) + + +@hypothesis.given(_strategy()) +def test_compose_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs(beignet.ops.compose_quaternion(**parameters)), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__compose_rotation_matrix.py b/tests/beignet/ops/_geometry/_transformations/test__compose_rotation_matrix.py new file mode 100644 index 0000000000..a6e042d524 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__compose_rotation_matrix.py @@ -0,0 +1,37 @@ +import operator + +import beignet.ops +import hypothesis.strategies +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ) + + input = Rotation.random(size) + other = Rotation.random(size) + + return ( + { + "input": torch.from_numpy(input.as_matrix()), + "other": torch.from_numpy(other.as_matrix()), + }, + torch.from_numpy(operator.mul(input, other).as_matrix()), + ) + + +@hypothesis.given(_strategy()) +def test_compose_rotation_matrix(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.compose_rotation_matrix(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__compose_rotation_vector.py b/tests/beignet/ops/_geometry/_transformations/test__compose_rotation_vector.py new file mode 100644 index 0000000000..2a6f3e09ae --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__compose_rotation_vector.py @@ -0,0 +1,40 @@ +import operator + +import beignet.ops +import hypothesis.strategies +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ) + + input = Rotation.random(size) + other = Rotation.random(size) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy(input.as_rotvec(degrees)), + "other": torch.from_numpy(other.as_rotvec(degrees)), + "degrees": degrees, + }, + torch.from_numpy(operator.mul(input, other).as_rotvec(degrees)), + ) + + +@hypothesis.given(_strategy()) +def test_compose_rotation_vector(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.compose_rotation_vector(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__euler_angle_identity.py b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_identity.py new file mode 100644 index 0000000000..f9706d7bd8 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/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/ops/_geometry/_transformations/test__euler_angle_magnitude.py b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_magnitude.py new file mode 100644 index 0000000000..169376b4c5 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_magnitude.py @@ -0,0 +1,56 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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.ops.euler_angle_magnitude(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__euler_angle_mean.py b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_mean.py new file mode 100644 index 0000000000..31a568d598 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_mean.py @@ -0,0 +1,61 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ) + + input = Rotation.random(size) + + 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(input.as_euler(axes, degrees)), + "axes": axes, + "degrees": degrees, + }, + torch.unsqueeze( + torch.from_numpy( + input.mean().as_euler(axes, degrees), + ), + dim=0, + ), + ) + + +@hypothesis.given(_strategy()) +def test_euler_angle_mean(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.euler_angle_mean(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_quaternion.py b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_quaternion.py new file mode 100644 index 0000000000..1043a994c2 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_quaternion.py @@ -0,0 +1,75 @@ +import beignet.ops +import hypothesis.extra.numpy +import hypothesis.strategies +import torch +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_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.ops.euler_angle_to_quaternion( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_rotation_matrix.py b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_rotation_matrix.py new file mode 100644 index 0000000000..4e63942c86 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_rotation_matrix.py @@ -0,0 +1,65 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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.ops.euler_angle_to_rotation_matrix( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_rotation_vector.py b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_rotation_vector.py new file mode 100644 index 0000000000..20f2adea53 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__euler_angle_to_rotation_vector.py @@ -0,0 +1,63 @@ +import beignet.ops +import hypothesis.strategies +import torch +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.ops.euler_angle_to_rotation_vector( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__invert_euler_angle.py b/tests/beignet/ops/_geometry/_transformations/test__invert_euler_angle.py new file mode 100644 index 0000000000..d6505e2b31 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__invert_euler_angle.py @@ -0,0 +1,56 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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.ops.invert_euler_angle(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__invert_quaternion.py b/tests/beignet/ops/_geometry/_transformations/test__invert_quaternion.py new file mode 100644 index 0000000000..78cc7f5195 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__invert_quaternion.py @@ -0,0 +1,36 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.invert_quaternion(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__invert_rotation_matrix.py b/tests/beignet/ops/_geometry/_transformations/test__invert_rotation_matrix.py new file mode 100644 index 0000000000..9267f65c22 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__invert_rotation_matrix.py @@ -0,0 +1,33 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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.ops.invert_rotation_matrix(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__invert_rotation_vector.py b/tests/beignet/ops/_geometry/_transformations/test__invert_rotation_vector.py new file mode 100644 index 0000000000..a215243901 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__invert_rotation_vector.py @@ -0,0 +1,35 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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.ops.invert_rotation_vector(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__quaternion_identity.py b/tests/beignet/ops/_geometry/_transformations/test__quaternion_identity.py new file mode 100644 index 0000000000..f7f69c6f77 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__quaternion_identity.py @@ -0,0 +1,37 @@ +import beignet.ops +import hypothesis.strategies +import torch +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_quaternion_identity(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.quaternion_identity(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__quaternion_magnitude.py b/tests/beignet/ops/_geometry/_transformations/test__quaternion_magnitude.py new file mode 100644 index 0000000000..1bfe4e8a24 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__quaternion_magnitude.py @@ -0,0 +1,36 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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_quaternion_magnitude(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.quaternion_magnitude(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__quaternion_mean.py b/tests/beignet/ops/_geometry/_transformations/test__quaternion_mean.py new file mode 100644 index 0000000000..01d3b1baac --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__quaternion_mean.py @@ -0,0 +1,50 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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, + ), + ), + ) + + return ( + { + "input": torch.from_numpy( + rotation.as_quat( + canonical=False, + ), + ), + }, + torch.unsqueeze( + torch.abs( + torch.from_numpy( + rotation.mean().as_quat( + canonical=False, + ), + ), + ), + dim=0, + ), + ) + + +@hypothesis.given(_strategy()) +def test_quaternion_mean(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.ops.quaternion_mean( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_euler_angle.py b/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_euler_angle.py new file mode 100644 index 0000000000..9f2b6129e4 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_euler_angle.py @@ -0,0 +1,67 @@ +import beignet.ops +import hypothesis.strategies +import torch +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_quaternion_to_euler_angle(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.quaternion_to_euler_angle( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_rotation_matrix.py b/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_rotation_matrix.py new file mode 100644 index 0000000000..ebea9a70ed --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_rotation_matrix.py @@ -0,0 +1,41 @@ +import beignet.ops +import hypothesis.strategies +import torch +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_quaternion_to_rotation_matrix(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.quaternion_to_rotation_matrix( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_rotation_vector.py b/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_rotation_vector.py new file mode 100644 index 0000000000..5c69a9e1b3 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__quaternion_to_rotation_vector.py @@ -0,0 +1,46 @@ +import beignet.ops +import hypothesis.strategies +import torch +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_quaternion_to_rotation_vector(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.quaternion_to_rotation_vector( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__random_euler_angle.py b/tests/beignet/ops/_geometry/_transformations/test__random_euler_angle.py new file mode 100644 index 0000000000..ca0ef5fb2c --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__random_euler_angle.py @@ -0,0 +1,51 @@ +import beignet.ops +import hypothesis.strategies + + +@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.ops.random_euler_angle( + **parameters, + ).shape == (parameters["size"], 3) diff --git a/tests/beignet/ops/_geometry/_transformations/test__random_quaternion.py b/tests/beignet/ops/_geometry/_transformations/test__random_quaternion.py new file mode 100644 index 0000000000..938e79339c --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__random_quaternion.py @@ -0,0 +1,31 @@ +import beignet.ops +import hypothesis.strategies + + +@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_quaternion(data): + parameters, _ = data + + assert beignet.ops.random_quaternion( + **parameters, + ).shape == (parameters["size"], 4) diff --git a/tests/beignet/ops/_geometry/_transformations/test__random_rotation_matrix.py b/tests/beignet/ops/_geometry/_transformations/test__random_rotation_matrix.py new file mode 100644 index 0000000000..7e007b2b2f --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__random_rotation_matrix.py @@ -0,0 +1,28 @@ +import beignet.ops +import hypothesis.strategies + + +@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.ops.random_rotation_matrix( + **parameters, + ).shape == (parameters["size"], 3, 3) diff --git a/tests/beignet/ops/_geometry/_transformations/test__random_rotation_vector.py b/tests/beignet/ops/_geometry/_transformations/test__random_rotation_vector.py new file mode 100644 index 0000000000..c214441eb9 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__random_rotation_vector.py @@ -0,0 +1,31 @@ +import beignet.ops +import hypothesis.strategies + + +@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.ops.random_rotation_vector( + **parameters, + ).shape == (parameters["size"], 3) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_identity.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_identity.py new file mode 100644 index 0000000000..13d163f692 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_identity.py @@ -0,0 +1,34 @@ +import beignet.ops +import hypothesis.strategies +import torch +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.ops.rotation_matrix_identity(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_magnitude.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_magnitude.py new file mode 100644 index 0000000000..85fb6d94d9 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_magnitude.py @@ -0,0 +1,33 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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.ops.rotation_matrix_magnitude(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_mean.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_mean.py new file mode 100644 index 0000000000..22af4a1bbf --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_mean.py @@ -0,0 +1,33 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ) + + input = Rotation.random(size) + + return ( + { + "input": torch.from_numpy(input.as_matrix()), + }, + torch.unsqueeze(torch.from_numpy(input.mean().as_matrix()), dim=0), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_matrix_mean(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.rotation_matrix_mean(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_euler_angle.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_euler_angle.py new file mode 100644 index 0000000000..b143e9aa4f --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_euler_angle.py @@ -0,0 +1,65 @@ +import beignet.ops +import hypothesis.strategies +import torch +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.ops.rotation_matrix_to_euler_angle( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_quaternion.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_quaternion.py new file mode 100644 index 0000000000..fde6a951ff --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_quaternion.py @@ -0,0 +1,49 @@ +import beignet.ops +import hypothesis.extra.numpy +import hypothesis.strategies +import torch +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_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.ops.rotation_matrix_to_quaternion( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_rotation_vector.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_rotation_vector.py new file mode 100644 index 0000000000..491690a13a --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_matrix_to_rotation_vector.py @@ -0,0 +1,45 @@ +import beignet.ops +import hypothesis.extra.numpy +import hypothesis.strategies +import torch +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.ops.rotation_matrix_to_rotation_vector( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_identity.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_identity.py new file mode 100644 index 0000000000..d84b97d9fa --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_identity.py @@ -0,0 +1,36 @@ +import beignet.ops +import hypothesis.strategies +import torch +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.ops.rotation_vector_identity(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_magnitude.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_magnitude.py new file mode 100644 index 0000000000..a065bbed5f --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_magnitude.py @@ -0,0 +1,36 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +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.ops.rotation_vector_magnitude(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_mean.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_mean.py new file mode 100644 index 0000000000..6ae68d159f --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_mean.py @@ -0,0 +1,39 @@ +import beignet.ops +import hypothesis.strategies +import torch.testing +from scipy.spatial.transform import Rotation + + +@hypothesis.strategies.composite +def _strategy(function): + size = function( + hypothesis.strategies.integers( + min_value=16, + max_value=32, + ), + ) + + input = Rotation.random(size) + + degrees = function(hypothesis.strategies.booleans()) + + return ( + { + "input": torch.from_numpy(input.as_rotvec(degrees)), + "degrees": degrees, + }, + torch.unsqueeze( + torch.from_numpy(input.mean().as_rotvec(degrees)), + dim=0, + ), + ) + + +@hypothesis.given(_strategy()) +def test_rotation_vector_mean(data): + parameters, expected = data + + torch.testing.assert_close( + beignet.ops.rotation_vector_mean(**parameters), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_euler_angle.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_euler_angle.py new file mode 100644 index 0000000000..f49ab3ed1b --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_euler_angle.py @@ -0,0 +1,67 @@ +import beignet.ops +import hypothesis.strategies +import torch +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.ops.rotation_vector_to_euler_angle( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_quaternion.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_quaternion.py new file mode 100644 index 0000000000..9017271dbc --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_quaternion.py @@ -0,0 +1,54 @@ +import beignet.ops +import hypothesis.extra.numpy +import hypothesis.strategies +import torch +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_quaternion(data): + parameters, expected = data + + torch.testing.assert_close( + torch.abs( + beignet.ops.rotation_vector_to_quaternion( + **parameters, + ), + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_rotation_matrix.py b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_rotation_matrix.py new file mode 100644 index 0000000000..0c8e699e2a --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__rotation_vector_to_rotation_matrix.py @@ -0,0 +1,42 @@ +import beignet.ops +import hypothesis.strategies +import torch +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.ops.rotation_vector_to_rotation_matrix( + **parameters, + ), + expected, + ) diff --git a/tests/beignet/ops/_geometry/_transformations/test__slerp.py b/tests/beignet/ops/_geometry/_transformations/test__slerp.py new file mode 100644 index 0000000000..5b09726995 --- /dev/null +++ b/tests/beignet/ops/_geometry/_transformations/test__slerp.py @@ -0,0 +1,151 @@ +import beignet.ops +import hypothesis.strategies +import numpy +import torch +from scipy.spatial.transform import Rotation, Slerp + + +def test_slerp(): + # t = 0 + torch.testing.assert_close( + beignet.ops.quaternion_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.ops.quaternion_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.ops.quaternion_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.ops.quaternion_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.ops.quaternion_slerp(*parameters), expected_rotations + )