Source code for skbot.transform.affine

```from numpy.typing import ArrayLike
from typing import List
import numpy as np

from ._utils import vector_project, angle_between
from .functions import translate, rotate, as_affine_matrix

"""A link representing a affine transformation.

This is an abstract class useful when writing links for affine
transformations. An affine transformation is a transformation that can be
described using the equation ``y = Ax+b``.

The main utility of this class is that it computes the corresponding
transformation matrix once ``self.transform`` is known.

Parameters
----------
parent : int
The frame in which vectors are specified.
child : int
The frame into which this link transforms vectors.
axis : int
The axis along which computation takes place. All other axes are considered
batch dimensions.

"""

def __init__(self, parent_dim: int, child_dim: int, *, axis: int = -1) -> None:

super().__init__(parent_dim, child_dim)

self._tf_matrix = None
self._axis = axis

@property
def affine_matrix(self) -> np.ndarray:
"""The transformation matrix mapping the parent to the child frame."""

cartesian_parent = Frame(self.parent_dim)
cartesian_child = Frame(self.child_dim)
self(cartesian_parent, cartesian_child)

affine_parent = AffineSpace(self.parent_dim, axis=self._axis)(cartesian_parent)
affine_child = AffineSpace(self.child_dim, axis=self._axis)(cartesian_child)

affine_matrix = as_affine_matrix(affine_parent, affine_child, axis=self._axis)

return affine_matrix

@property
def _inverse_tf_matrix(self):
cartesian_parent = Frame(self.parent_dim)
cartesian_child = Frame(self.child_dim)
self(cartesian_parent, cartesian_child)

affine_parent = AffineSpace(self.parent_dim, axis=self._axis)(cartesian_parent)
affine_child = AffineSpace(self.child_dim, axis=self._axis)(cartesian_child)

affine_matrix = as_affine_matrix(affine_child, affine_parent, axis=self._axis)

return affine_matrix

def invert(self) -> Frame:
return Inverse(self)

@property
def affine_matrix(self) -> np.ndarray:
"""The transformation matrix mapping the parent to the child frame."""

def transform(self, x: ArrayLike) -> np.ndarray:

return x

def __inverse_transform__(self, x: ArrayLike) -> np.ndarray:

return x

@property
def affine_matrix(self) -> np.ndarray:
"""The transformation matrix mapping the parent to the child frame."""

return matrix

def invert(self) -> Frame:

"""Planar rotation in N-D.

The plane of rotation is described by two vectors (u, v). The initial
angle of rotation is twice the angle between u and v (measured from u to v)
and can be modified by setting the angle explicitly, e.g. ``link.angle =
np.pi``. The angle is measured in radians.

Parameters
----------
u : ArrayLike
The first vector defining the plane of rotation. The angle is mesured
from here.
v : ArrayLike
The second vector defining the plane of rotation. The angle is measured
to here, i.e., from u to v.

Notes
-----
Implements __inverse_transform__.

If you wish to initialize a rotation that rotates u onto v, and both are of
same length, you can use ``tf.Rotation(u+v, v)``.

"""

def __init__(self, u: ArrayLike, v: ArrayLike, *, axis: int = -1) -> None:
u = np.asarray(u)
v = np.asarray(v)

u = np.moveaxis(u, axis, -1)
v = np.moveaxis(v, axis, -1)

frame_dim = u.shape[axis]

super().__init__(frame_dim, frame_dim)

self._u = u / np.linalg.norm(u)
self._v = v / np.linalg.norm(v)
self._angle = 2 * angle_between(u, v)

u_orthogonal = v - vector_project(v, u)
self._u_ortho = u_orthogonal / np.linalg.norm(u_orthogonal)

[docs]    def transform(self, x: ArrayLike) -> np.ndarray:
return rotate(x, self._u, self._v)

[docs]    def __inverse_transform__(self, x: ArrayLike) -> np.ndarray:
return rotate(x, self._v, self._u)

@property
def angle(self) -> float:
"""The magnitude of the rotation (in radians)."""
return self._angle

@angle.setter
def angle(self, angle: ArrayLike) -> None:
self._angle = angle

self._v = np.cos(angle / 2) * self._u - np.sin(angle / 2) * self._u_ortho

"""Translation in N-D.

.. versionchanged:: 0.6.0
``direction`` and ``scalar`` may now be broadcastable arrays

Parameters
----------
direction : ArrayLike
A vector, or batch of vectors, describing the translation.
amount : ArrayLike
A scalar indicating by how much to scale ``direction``. Default is 1.
axis : int
The axis along which computation takes place. All other axes are considered
batch dimensions.

"""

def __init__(
self, direction: ArrayLike, *, amount: ArrayLike = 1, axis: int = -1
) -> None:
direction = np.asarray(direction)

frame_dim = direction.shape[axis]

super().__init__(frame_dim, frame_dim, axis=axis)

self._axis = axis
self._amount = np.asarray(amount)
self._direction = np.moveaxis(direction, axis, -1)

@property
def direction(self) -> np.ndarray:
"""The direction in which vectors are translated."""
return self._direction

@direction.setter
def direction(self, direction: ArrayLike) -> None:
self._direction = np.asarray(direction)

@property
def amount(self) -> float:
"""The amount by which to scale the direction vector."""
return self._amount

@amount.setter
def amount(self, amount: ArrayLike) -> None:
self._amount = np.asarray(amount)

[docs]    def transform(self, x: ArrayLike) -> np.ndarray:
x = np.asarray(x)
x = np.moveaxis(x, self._axis, -1)
result = translate(x, self._amount[..., None] * self._direction)
return np.moveaxis(result, -1, self._axis)

[docs]    def __inverse_transform__(self, x: ArrayLike) -> np.ndarray:
x = np.asarray(x)
x = np.moveaxis(x, self._axis, -1)
result = translate(x, -self._amount[..., None] * self._direction)
return np.moveaxis(result, -1, self._axis)

"""Transform to affine space

Parameters
----------
ndim : int
The number of dimensions of the cartesian space.
axis : int
The axis along which computation takes place. All other axes are considered
batch dimensions.

"""

def __init__(self, ndim: int, *, axis=-1) -> None:
super().__init__(ndim, ndim + 1)
self._axis = axis

[docs]    def transform(self, x: ArrayLike) -> np.ndarray:
x = np.moveaxis(x, self._axis, -1)

shape = list(x.shape)
shape[-1] += 1

affine_vector = np.ones(shape, dtype=x.dtype)
affine_vector[..., :-1] = x

return np.moveaxis(affine_vector, -1, self._axis)

[docs]    def __inverse_transform__(self, x: ArrayLike) -> np.ndarray:
x = np.moveaxis(x, self._axis, -1)

values = x[..., :-1]
scaling = x[..., -1][..., None]

cartesian_vector = values / scaling

return np.moveaxis(cartesian_vector, -1, self._axis)
```