beignet.func.space

docs/beignet.func.md

@@ -0,0 +1,3 @@
+# beignet.func
+
+::: beignet.func.space

pyproject.toml

@@ -10,7 +10,7 @@ requires = [
 authors = [{ email = "[email protected]", name = "Allen Goodman" }]
 dependencies = [
     "pooch",
-    "torch",
+    "torch==2.2.2",
     "tqdm",
 ]
 dynamic = ["version"]

src/beignet/__init__.py

@@ -12,6 +12,7 @@
 )
 from ._apply_rotation_matrix import apply_rotation_matrix
 from ._apply_rotation_vector import apply_rotation_vector
+from ._apply_transform import apply_transform
 from ._compose_euler_angle import compose_euler_angle
 from ._compose_quaternion import compose_quaternion
 from ._compose_rotation_matrix import compose_rotation_matrix
@@ -28,6 +29,7 @@
 from ._invert_quaternion import invert_quaternion
 from ._invert_rotation_matrix import invert_rotation_matrix
 from ._invert_rotation_vector import invert_rotation_vector
+from ._invert_transform import invert_transform
 from ._quaternion_identity import quaternion_identity
 from ._quaternion_magnitude import quaternion_magnitude
 from ._quaternion_mean import quaternion_mean
@@ -72,6 +74,7 @@
     "apply_quaternion",
     "apply_rotation_matrix",
     "apply_rotation_vector",
+    "apply_transform",
     "compose_euler_angle",
     "compose_quaternion",
     "compose_rotation_matrix",
@@ -86,9 +89,11 @@
     "invert_quaternion",
     "invert_rotation_matrix",
     "invert_rotation_vector",
+    "invert_transform",
     "quaternion_identity",
     "quaternion_magnitude",
     "quaternion_mean",
+    "quaternion_slerp",
     "quaternion_to_euler_angle",
     "quaternion_to_rotation_matrix",
     "quaternion_to_rotation_vector",
@@ -108,6 +113,5 @@
     "rotation_vector_to_euler_angle",
     "rotation_vector_to_quaternion",
     "rotation_vector_to_rotation_matrix",
-    "quaternion_slerp",
     "translation_identity",
 ]

src/beignet/_apply_transform.py

@@ -0,0 +1,101 @@
+import torch
+from torch import Tensor
+from torch.autograd import Function
+
+
+def _apply_transform(input: Tensor, transform: Tensor) -> Tensor:
+    """
+    Applies an affine transformation to the position vector.
+
+    Parameters
+    ----------
+    input : Tensor
+        Position, must have the shape `(..., dimension)`.
+
+    transform : Tensor
+        The affine transformation matrix, must be a scalar, a vector, or a
+        matrix with the shape `(dimension, dimension)`.
+
+    Returns
+    -------
+    Tensor
+        Affine transformed position vector, has the same shape as the
+        position vector.
+    """
+    match transform.ndim:
+        case 0:
+            return input * transform
+        case 1:
+            return torch.einsum("i,...i->...i", transform, input)
+        case 2:
+            return torch.einsum("ij,...j->...i", transform, input)
+        case _:
+            raise ValueError
+
+
+class _ApplyTransform(Function):
+    generate_vmap_rule = True
+
+    @staticmethod
+    def forward(transform: Tensor, input: Tensor) -> Tensor:
+        """
+        Return affine transformed position.
+
+        Parameters
+        ----------
+        transform : Tensor
+            Affine transformation matrix, must have shape
+            `(dimension, dimension)`.
+
+        input : Tensor
+            Position, must have shape `(..., dimension)`.
+
+        Returns
+        -------
+        Tensor
+            Affine transformed position of shape `(..., dimension)`.
+        """
+        return _apply_transform(input, transform)
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        transform, input = inputs
+
+        ctx.save_for_backward(transform, input, output)
+
+    @staticmethod
+    def jvp(ctx, grad_transform: Tensor, grad_input: Tensor) -> (Tensor, Tensor):
+        transform, input, _ = ctx.saved_tensors
+
+        output = _apply_transform(input, transform)
+
+        grad_output = grad_input + _apply_transform(input, grad_transform)
+
+        return output, grad_output
+
+    @staticmethod
+    def backward(ctx, grad_output: Tensor) -> (Tensor, Tensor):
+        _, _, output = ctx.saved_tensors
+
+        return output, grad_output
+
+
+def apply_transform(input: Tensor, transform: Tensor) -> Tensor:
+    """
+    Return affine transformed position.
+
+    Parameters
+    ----------
+    input : Tensor
+        Position, must have shape `(..., dimension)`.
+
+    transform : Tensor
+        Affine transformation matrix, must have shape
+        `(dimension, dimension)`.
+
+    Returns
+    -------
+    Tensor
+        Affine transformed position of shape `(..., dimension)`.
+    """
+    return _ApplyTransform.apply(transform, input)

src/beignet/_invert_transform.py

@@ -0,0 +1,25 @@
+import torch
+from torch import Tensor
+
+
+def invert_transform(transform: Tensor) -> Tensor:
+    """
+    Calculates the inverse of an affine transformation matrix.
+
+    Parameters
+    ----------
+    transform : Tensor
+        The affine transformation matrix to be inverted.
+
+    Returns
+    -------
+    Tensor
+        The inverse of the given affine transformation matrix.
+    """
+    if transform.ndim in {0, 1}:
+        return 1.0 / transform
+
+    if transform.ndim == 2:
+        return torch.linalg.inv(transform)
+
+    raise ValueError

src/beignet/func/__init__.py

@@ -0,0 +1 @@
+from ._space import space