from types import EllipsisType
from scipy._lib._array_api import (
array_namespace,
Array,
ArrayLike,
is_lazy_array,
xp_vector_norm,
xp_result_type,
xp_promote,
is_array_api_obj,
)
import scipy._lib.array_api_extra as xpx
from scipy.spatial.transform._rotation_xp import (
as_matrix as quat_as_matrix,
from_matrix as quat_from_matrix,
_from_matrix_orthogonal as quat_from_matrix_orthogonal,
from_rotvec as quat_from_rotvec,
as_rotvec as quat_as_rotvec,
compose_quat,
from_quat,
inv as quat_inv,
mean as quat_mean,
)
from scipy._lib.array_api_compat import device as xp_device
from scipy._lib._util import broadcastable
def from_matrix(matrix: Array, normalize: bool = True, copy: bool = True) -> Array:
xp = array_namespace(matrix)
# Shape check should be done before calling this function
if normalize or copy:
matrix = xp.asarray(matrix, copy=True)
last_row_ok = xp.all(
matrix[..., 3, :] == xp.asarray([0, 0, 0, 1.0], device=xp_device(matrix)),
axis=-1,
)
lazy = is_lazy_array(matrix)
# We delay lazy branch checks until after normalization to avoid overwriting nans
# with the rotation matrix
if not lazy and xp.any(~last_row_ok):
if last_row_ok.shape == ():
idx = ()
else:
idx = tuple(int(i[0]) for i in xp.nonzero(~last_row_ok))
vals = matrix[idx + (3, ...)]
raise ValueError(
f"Expected last row of transformation matrix {idx} to be "
f"exactly [0, 0, 0, 1], got {vals}"
)
# The quat_from_matrix() method orthogonalizes the rotation
# component of the transformation matrix. While this does have some
# overhead in converting a rotation matrix to a quaternion and back, it
# allows for skipping singular value decomposition for near-orthogonal
# matrices, which is a computationally expensive operation.
if normalize:
rotmat = quat_as_matrix(quat_from_matrix(matrix[..., :3, :3]))
matrix = xpx.at(matrix)[..., :3, :3].set(rotmat)
# Lazy branch matrix invalidation
if lazy:
matrix = xp.where(last_row_ok[..., None, None], matrix, xp.nan)
return matrix
def from_rotation(quat: Array) -> Array:
xp = array_namespace(quat)
rotmat = quat_as_matrix(quat)
matrix = xp.zeros(
(*rotmat.shape[:-2], 4, 4), dtype=quat.dtype, device=xp_device(quat)
)
matrix = xpx.at(matrix)[..., :3, :3].set(rotmat)
matrix = xpx.at(matrix)[..., 3, 3].set(1)
return matrix
def from_translation(translation: Array) -> Array:
xp = array_namespace(translation)
if translation.shape[-1] != 3:
raise ValueError(
f"Expected `translation` to have shape (..., 3), got {translation.shape}."
)
device = xp_device(translation)
dtype = xp_result_type(translation, force_floating=True, xp=xp)
eye = xp.eye(4, dtype=dtype, device=device)
matrix = xpx.atleast_nd(eye, ndim=translation.ndim + 1, xp=xp)
matrix = xp.zeros(
(*translation.shape[:-1], 4, 4),
dtype=dtype,
device=device,
)
matrix = xpx.at(matrix)[...].set(xp.eye(4, dtype=dtype, device=device))
matrix = xpx.at(matrix)[..., :3, 3].set(
xp_promote(translation, force_floating=True, xp=xp)
)
return matrix
def from_components(translation: Array, quat: Array) -> Array:
return compose_transforms(from_translation(translation), from_rotation(quat))
def from_exp_coords(exp_coords: Array) -> Array:
rot_vec = exp_coords[..., :3]
rot_matrix = quat_as_matrix(quat_from_rotvec(rot_vec))
translation_transform = _compute_se3_exp_translation_transform(rot_vec)
translations = (translation_transform @ exp_coords[..., 3:, None])[..., 0]
return _create_transformation_matrix(translations, rot_matrix)
def from_dual_quat(dual_quat: Array, *, scalar_first: bool = False) -> Array:
xp = array_namespace(dual_quat)
if dual_quat.shape[-1] != 8:
raise ValueError(
f"Expected `dual_quat` to have shape (..., 8), got {dual_quat.shape}."
)
real_part = dual_quat[..., :4]
dual_part = dual_quat[..., 4:]
if scalar_first:
real_part = xp.roll(real_part, -1, axis=-1)
dual_part = xp.roll(dual_part, -1, axis=-1)
real_part, dual_part = _normalize_dual_quaternion(real_part, dual_part)
rot_quat = from_quat(real_part)
translation = 2.0 * compose_quat(dual_part, quat_inv(rot_quat))[..., :3]
matrix = _create_transformation_matrix(translation, quat_as_matrix(rot_quat))
return matrix
def as_exp_coords(matrix: Array) -> Array:
xp = array_namespace(matrix)
rot_vec = quat_as_rotvec(quat_from_matrix_orthogonal(matrix[..., :3, :3]))
translation_transform = _compute_se3_log_translation_transform(rot_vec)
translations = (translation_transform @ matrix[..., :3, 3][..., None])[..., 0]
exp_coords = xp.concat([rot_vec, translations], axis=-1)
return exp_coords
def as_dual_quat(matrix: Array, *, scalar_first: bool = False) -> Array:
xp = array_namespace(matrix)
real_parts = quat_from_matrix_orthogonal(matrix[..., :3, :3])
pure_translation_quats = xp.empty(
(*matrix.shape[:-2], 4), dtype=matrix.dtype, device=xp_device(matrix)
)
pure_translation_quats = xpx.at(pure_translation_quats)[..., :3].set(
matrix[..., :3, 3]
)
pure_translation_quats = xpx.at(pure_translation_quats)[..., 3].set(0.0)
dual_parts = 0.5 * compose_quat(pure_translation_quats, real_parts)
if scalar_first:
real_parts = xp.roll(real_parts, 1, axis=-1)
dual_parts = xp.roll(dual_parts, 1, axis=-1)
dual_quats = xp.concat([real_parts, dual_parts], axis=-1)
return dual_quats
def compose_transforms(tf_matrix: Array, other_tf_matrix: Array) -> Array:
if not broadcastable(tf_matrix.shape, other_tf_matrix.shape):
len_a = tf_matrix.shape[0] if tf_matrix.ndim == 3 else 1
len_b = other_tf_matrix.shape[0] if other_tf_matrix.ndim == 3 else 1
raise ValueError(
"Expected equal number of transforms in both or a "
"single transform in either object, got "
f"{len_a} transforms in first and {len_b}"
"transforms in second object."
)
return tf_matrix @ other_tf_matrix
def inv(matrix: Array) -> Array:
xp = array_namespace(matrix)
r_inv = xp.matrix_transpose(matrix[..., :3, :3])
# Matrix multiplication of r_inv and translation vector
t_inv = -(r_inv @ matrix[..., :3, 3][..., None])[..., 0]
matrix = xp.zeros(
(*matrix.shape[:-2], 4, 4), dtype=matrix.dtype, device=xp_device(matrix)
)
matrix = xpx.at(matrix)[..., :3, :3].set(r_inv)
matrix = xpx.at(matrix)[..., :3, 3].set(t_inv)
matrix = xpx.at(matrix)[..., 3, 3].set(1)
return matrix
def apply(matrix: Array, vector: Array, inverse: bool = False) -> Array:
xp = array_namespace(matrix)
if vector.shape[-1] != 3:
raise ValueError(f"Expected vector to have shape (..., 3), got {vector.shape}.")
vec = xp.empty(
(*vector.shape[:-1], 4), dtype=vector.dtype, device=xp_device(vector)
)
vec = xpx.at(vec)[..., :3].set(vector)
vec = xpx.at(vec)[..., 3].set(1)
vec = vec[..., None]
if inverse:
matrix = inv(matrix)
# We raise a ValueError manually here because letting the function run its course
# would raise heterogeneous error types and messages for different frameworks.
# However, the error only mimics numpy's error message and does not provide the
# same amount of context.
if not broadcastable(matrix.shape, vec.shape):
raise ValueError("operands could not be broadcast together")
return (matrix @ vec)[..., :3, 0]
def pow(matrix: Array, n: float | Array) -> Array:
xp = array_namespace(matrix)
device = xp_device(matrix)
# If n is an array, we sanitize it to a scalar and promote quat and n to
# the same dtype.
if is_array_api_obj(n):
if n.shape == (1,):
n = n[0]
elif n.ndim != 0:
raise ValueError("Array exponent must be a scalar")
matrix, n = xp_promote(matrix, n, force_floating=True, xp=xp)
# If n is a lazy array, we cannot take fast paths for special cases.
if is_lazy_array(n):
# Lazy execution. We compute all special cases and the general case
result = from_exp_coords(as_exp_coords(matrix) * n)
identity = xp.eye(4, dtype=matrix.dtype, device=device)
result = xp.where(n == 0, identity, result)
result = xp.where(n == -1, inv(matrix), result)
result = xp.where(n == 1, matrix, result)
return result
if n == 0:
identity = xp.eye(4, dtype=matrix.dtype, device=device)
identity = xpx.atleast_nd(identity, ndim=matrix.ndim, xp=xp)
return xp.tile(identity, (*matrix.shape[:-2], 1, 1))
elif n == -1:
return inv(matrix)
elif n == 1:
return matrix
return from_exp_coords(as_exp_coords(matrix) * n)
def mean(
matrix: Array,
weights: ArrayLike | None = None,
axis: None | int | tuple[int, ...] = None
) -> Array:
xp = array_namespace(matrix)
if matrix.shape[0] == 0:
raise ValueError("Mean of an empty rotation set is undefined.")
# Axis logic: For None, we reduce over all axes. For int, we only reduce over that
# axis. For tuple, we reduce over all specified axes.
all_axes = tuple(range(matrix.ndim - 2))
if axis is None:
axis = all_axes
elif isinstance(axis, int):
axis = (axis,)
if not isinstance(axis, tuple):
raise ValueError("`axis` must be None, int, or tuple of ints.")
# Ensure all axes are within bounds
if (axis != () and
(min(axis) < -(matrix.ndim - 2) or max(axis) > (matrix.ndim - 3))
):
raise ValueError(
f"axis {axis} is out of bounds for transform with shape "
f"{matrix.shape[:-2]}."
)
# Ensure all axes are positive and unique
axis = tuple(sorted(set(x % (matrix.ndim - 2) for x in axis)))
lazy = is_lazy_array(matrix)
quats = quat_from_matrix_orthogonal(matrix[..., :3, :3])
if weights is None:
quats_mean = quat_mean(quats, axis=axis)
else:
neg_weights = weights < 0
any_neg_weights = xp.any(neg_weights)
if not lazy and any_neg_weights:
raise ValueError("`weights` must be non-negative.")
if weights.shape != matrix.shape[:-2]:
raise ValueError(
f"Expected `weights` to match transform shape, got shape "
f"{weights.shape} for {matrix.shape[:-2]} transformations."
)
quats_mean = quat_mean(quats, weights=weights, axis=axis)
r_mean = quat_as_matrix(quats_mean)
t = matrix[..., :3, 3]
if weights is None:
t_mean = xp.mean(t, axis=axis)
else:
norm = xp.sum(weights[..., None], axis=axis)
wsum = xp.sum(t * weights[..., None], axis=axis)
t_mean = wsum / norm
tf = _create_transformation_matrix(t_mean, r_mean)
if weights is not None and lazy:
# We cannot raise on negative weights because jit code needs to be
# non-branching. We return NaN instead
mask = xp.where(any_neg_weights, xp.nan, 1.0)
tf = mask * tf
return tf
def setitem(
matrix: Array,
indexer: Array | int | tuple | slice | EllipsisType | None,
value: Array,
) -> Array:
xp = array_namespace(matrix)
if isinstance(indexer, EllipsisType):
return xpx.at(matrix)[indexer].set(value)
if is_array_api_obj(indexer) and indexer.dtype == xp.bool:
return xpx.at(matrix)[indexer].set(value)
return xpx.at(matrix)[indexer, ...].set(value)
def normalize_dual_quaternion(dual_quat: Array) -> Array:
"""Normalize dual quaternion."""
xp = array_namespace(dual_quat)
real, dual = _normalize_dual_quaternion(dual_quat[..., :4], dual_quat[..., 4:])
return xp.concat((real, dual), axis=-1)
def _create_transformation_matrix(
translations: Array, rotation_matrices: Array
) -> Array:
if not translations.shape[:-1] == rotation_matrices.shape[:-2]:
raise ValueError(
"The number of rotation matrices and translations must be the same."
)
xp = array_namespace(translations)
matrix = xp.empty(
(*translations.shape[:-1], 4, 4),
dtype=translations.dtype,
device=xp_device(translations),
)
matrix = xpx.at(matrix)[..., :3, :3].set(rotation_matrices)
matrix = xpx.at(matrix)[..., :3, 3].set(translations)
matrix = xpx.at(matrix)[..., 3, :3].set(0)
matrix = xpx.at(matrix)[..., 3, 3].set(1)
return matrix
def _compute_se3_exp_translation_transform(rot_vec: Array) -> Array:
"""Compute the transformation matrix from the se3 translation part to SE3
translation.
The transformation matrix depends on the rotation vector.
"""
xp = array_namespace(rot_vec)
device = xp_device(rot_vec)
dtype = rot_vec.dtype
angle = xp_vector_norm(rot_vec, axis=-1, keepdims=True, xp=xp)
small_scale = angle < 1e-3
k1_small = 0.5 - angle**2 / 24 + angle**4 / 720
# Avoid division by zero for non-branching computations. The value will get
# discarded in the xp.where selection.
safe_angle = angle + xp.asarray(small_scale, dtype=dtype, device=device)
k1 = (1.0 - xp.cos(angle)) / safe_angle**2
k1 = xp.where(small_scale, k1_small, k1)
k2_small = 1 / 6 - angle**2 / 120 + angle**4 / 5040
# Again, avoid division by zero by adding one to all near-zero angles.
safe_angle = angle + xp.asarray(small_scale, dtype=dtype, device=device)
k2 = (angle - xp.sin(angle)) / safe_angle**3
k2 = xp.where(small_scale, k2_small, k2)
s = _create_skew_matrix(rot_vec)
eye = xp.eye(3, dtype=dtype, device=device)
return eye + k1[..., None] * s + k2[..., None] * s @ s
def _compute_se3_log_translation_transform(rot_vec: Array) -> Array:
"""Compute the transformation matrix from SE3 translation to the se3
translation part.
It is the inverse of `_compute_se3_exp_translation_transform` in a closed
analytical form.
"""
xp = array_namespace(rot_vec)
dtype = rot_vec.dtype
device = xp_device(rot_vec)
angle = xp_vector_norm(rot_vec, axis=-1, keepdims=True, xp=xp)
mask = angle < 1e-3
k_small = 1 / 12 + angle**2 / 720 + angle**4 / 30240
safe_angle = angle + xp.asarray(mask, dtype=dtype, device=device)
k = (1 - 0.5 * angle / xp.tan(0.5 * safe_angle)) / safe_angle**2
k = xp.where(mask, k_small, k)
s = _create_skew_matrix(rot_vec)
return xp.eye(3, dtype=dtype, device=device) - 0.5 * s + k[..., None] * s @ s
def _create_skew_matrix(vec: Array) -> Array:
"""Create skew-symmetric (aka cross-product) matrix for stack of vectors."""
xp = array_namespace(vec)
result = xp.zeros((*vec.shape[:-1], 3, 3), dtype=vec.dtype, device=xp_device(vec))
result = xpx.at(result)[..., 0, 1].set(-vec[..., 2])
result = xpx.at(result)[..., 0, 2].set(vec[..., 1])
result = xpx.at(result)[..., 1, 0].set(vec[..., 2])
result = xpx.at(result)[..., 1, 2].set(-vec[..., 0])
result = xpx.at(result)[..., 2, 0].set(-vec[..., 1])
result = xpx.at(result)[..., 2, 1].set(vec[..., 0])
return result
def _normalize_dual_quaternion(
real_part: Array, dual_part: Array
) -> tuple[Array, Array]:
"""Ensure that the dual quaternion has unit norm.
The norm is a dual number and must be 1 + 0 * epsilon, which means that
the real quaternion must have unit norm and the dual quaternion must be
orthogonal to the real quaternion.
"""
xp = array_namespace(real_part)
real_norm = xp_vector_norm(real_part, axis=-1, keepdims=True, xp=xp)
# special case: real quaternion is 0, we set it to identity
zero_real_mask = real_norm == 0.0
unit_quat = xp.asarray(
[0.0, 0.0, 0.0, 1.0], dtype=real_part.dtype, device=xp_device(real_part)
)
real_part = xp.where(zero_real_mask, unit_quat, real_part)
real_norm = xp.where(zero_real_mask, 1.0, real_norm)
# 1. ensure unit real quaternion
real_part = real_part / real_norm
dual_part = dual_part / real_norm
# 2. ensure orthogonality of real and dual quaternion
dual_part -= xp.sum(real_part * dual_part, axis=-1, keepdims=True) * real_part
return real_part, dual_part