import pickle
from itertools import product
import pytest
import numpy as np
from scipy.spatial.transform import Rotation, RigidTransform
from scipy.spatial.transform._rigid_transform import normalize_dual_quaternion
from scipy._lib._array_api import (
is_lazy_array,
xp_vector_norm,
is_numpy,
xp_assert_close,
make_xp_test_case,
xp_assert_equal,
xp_promote
)
import scipy._lib.array_api_extra as xpx
lazy_xp_modules = [RigidTransform]
def rotation_to_xp(r: Rotation, xp):
dtype = xpx.default_dtype(xp)
return Rotation.from_quat(xp.asarray(r.as_quat(), dtype=dtype))
def rigid_transform_to_xp(r: RigidTransform, xp):
dtype = xpx.default_dtype(xp)
return RigidTransform.from_matrix(xp.asarray(r.as_matrix(), dtype=dtype))
@make_xp_test_case(RigidTransform.as_matrix)
def test_repr(xp):
actual = repr(RigidTransform.from_matrix(xp.eye(4)))
expected = """\
RigidTransform.from_matrix(array([[1., 0., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]]))"""
if is_numpy(xp):
assert actual == expected
else:
assert actual.startswith("RigidTransform.from_matrix(")
tf = RigidTransform.from_matrix(xp.asarray(RigidTransform.identity(2).as_matrix()))
actual = repr(tf)
expected = """\
RigidTransform.from_matrix(array([[[1., 0., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]],
[[1., 0., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]]]))"""
if is_numpy(xp):
assert actual == expected
else:
assert actual.startswith("RigidTransform.from_matrix(")
@make_xp_test_case(RigidTransform.from_rotation)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_rotation(xp, ndim: int):
atol = 1e-12
rng = np.random.default_rng(0)
shape = (ndim,) * (ndim - 1) + (4,)
r = rotation_to_xp(Rotation.from_quat(rng.normal(size=shape)), xp=xp)
tf = RigidTransform.from_rotation(r)
xp_assert_close(tf.as_matrix()[..., :3, :3], r.as_matrix(), atol=atol)
xp_assert_close(tf.as_matrix()[..., :3, 3], xp.zeros(shape[:-1] + (3,)), atol=atol)
xp_assert_close(tf.as_matrix()[..., 3, :3], xp.zeros(shape[:-1] + (3,)), atol=atol)
xp_assert_close(tf.as_matrix()[..., 3, 3], xp.ones(shape[:-1]), atol=atol)
assert tf.single == (ndim == 1)
@make_xp_test_case(RigidTransform.from_translation)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_translation(xp, ndim: int):
shape = (ndim,) * (ndim - 1)
t = xp.reshape(xp.arange(ndim ** (ndim-1) * 3), shape + (3,))
tf = RigidTransform.from_translation(t)
expected = xp.tile(xp.eye(4), shape + (1, 1))
t_float = xp_promote(t, force_floating=True, xp=xp)
expected = xpx.at(expected)[..., :3, 3].set(t_float)
xp_assert_close(tf.as_matrix(), expected)
assert tf.single == (ndim == 1)
def test_from_translation_array_like():
# Test single translation
t = [1, 2, 3]
tf = RigidTransform.from_translation(t)
tf_expected = RigidTransform.from_translation(np.array(t))
xp_assert_close(tf.as_matrix(), tf_expected.as_matrix())
assert tf.single
# Test multiple translations
t = [[1, 2, 3], [4, 5, 6]]
tf = RigidTransform.from_translation(t)
tf_expected = RigidTransform.from_translation(np.array(t))
xp_assert_close(tf.as_matrix(), tf_expected.as_matrix())
assert not tf.single
@make_xp_test_case(RigidTransform.from_matrix, RigidTransform.as_matrix)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_matrix(xp, ndim: int):
atol = 1e-12
shape = (ndim,) * (ndim - 1) + (4, 4)
dtype = xpx.default_dtype(xp)
matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
t = xp.reshape(xp.arange(ndim ** (ndim-1) * 3, dtype=dtype), shape[:-2] + (3,))
matrix = xpx.at(matrix)[..., :3, 3].set(t)
tf = RigidTransform.from_matrix(matrix)
xp_assert_close(tf.as_matrix(), matrix, atol=atol)
assert tf.single == (ndim == 1)
# Test non-1 determinant
matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
matrix = xpx.at(matrix)[..., :3, :3].set(xp.eye(3) * 2.0)
tf = RigidTransform.from_matrix(matrix)
expected = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
xp_assert_close(tf.as_matrix(), expected, atol=atol)
# Test non-orthogonal rotation matrix
matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
# matrix is equivalent to [[1, 1, 0, 0],
# [0, 1, 0, 0],
# [0, 0, 1, 0],
# [0, 0, 0, 1]]
matrix = xpx.at(matrix)[..., 0, 1].set(1.0)
tf = RigidTransform.from_matrix(matrix)
expected = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
expected = xpx.at(expected)[..., 0, 0].set(0.894427)
expected = xpx.at(expected)[..., 0, 1].set(0.447214)
expected = xpx.at(expected)[..., 1, 0].set(-0.447214)
expected = xpx.at(expected)[..., 1, 1].set(0.894427)
xp_assert_close(tf.as_matrix(), expected, atol=1e-6)
# Test invalid matrix
invalid = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
invalid = xpx.at(invalid)[..., 3, 3].set(2) # Invalid last row
if is_lazy_array(invalid):
tf = RigidTransform.from_matrix(invalid)
assert xp.all(xp.isnan(tf.as_matrix()))
else:
with pytest.raises(ValueError):
RigidTransform.from_matrix(invalid)
def test_from_matrix_array_like():
# Test single transform matrix
matrix = [[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]]
expected = np.eye(4)
tf = RigidTransform.from_matrix(matrix)
xp_assert_close(tf.as_matrix(), expected)
assert tf.single
# Test multiple transform matrices
matrices = [matrix, matrix]
tf = RigidTransform.from_matrix(matrices)
for i in range(len(matrices)):
xp_assert_close(tf.as_matrix()[i, ...], expected)
assert not tf.single
@make_xp_test_case(RigidTransform.from_components)
@pytest.mark.parametrize("r_ndim", range(1, 3))
@pytest.mark.parametrize("t_ndim", range(1, 3))
def test_from_components(xp, r_ndim: int, t_ndim: int):
atol = 1e-12
dims = (6, 5, 4, 3) # Common shape
q_shape = dims[:r_ndim - 1][::-1] + (4,)
t_shape = dims[:t_ndim - 1][::-1] + (3,)
tf_shape = np.broadcast_shapes(q_shape[:-1], t_shape[:-1]) + (4, 4)
rng = np.random.default_rng(0)
t = xp.reshape(xp.arange(np.prod(t_shape[:-1]) * 3), t_shape)
r = rotation_to_xp(Rotation.from_quat(rng.random(size=q_shape)), xp=xp)
tf = RigidTransform.from_components(t, r)
expected = xp.zeros(tf_shape)
expected = xpx.at(expected)[..., :3, :3].set(r.as_matrix())
t_float = xp_promote(t, force_floating=True, xp=xp)
expected = xpx.at(expected)[..., :3, 3].set(t_float)
expected = xpx.at(expected)[..., 3, 3].set(1)
xp_assert_close(tf.as_matrix(), expected, atol=atol)
assert tf.single == (r_ndim == 1 and t_ndim == 1)
def test_from_components_array_like():
rng = np.random.default_rng(123)
# Test single rotation and translation
t = [1, 2, 3]
r = Rotation.random(rng=rng)
tf = RigidTransform.from_components(t, r)
tf_expected = RigidTransform.from_components(np.array(t), r)
xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)
assert tf.single
# Test multiple rotations and translations
t = [[1, 2, 3], [4, 5, 6]]
r = Rotation.random(len(t), rng=rng)
tf = RigidTransform.from_components(t, r)
tf_expected = RigidTransform.from_components(np.array(t), r)
xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)
assert not tf.single
@make_xp_test_case(RigidTransform.as_components)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_as_components(xp, ndim):
dtype = xpx.default_dtype(xp)
atol = 1e-12 if dtype == xp.float64 else 1e-6
shape = (ndim,) * (ndim - 1)
rng = np.random.default_rng(123)
t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
r = rotation_to_xp(Rotation.from_quat(rng.random(shape + (4,))), xp=xp)
tf = RigidTransform.from_components(t, r)
new_t, new_r = tf.as_components()
assert xp.all(new_r.approx_equal(r, atol=atol))
xp_assert_close(new_t, t, atol=atol)
@make_xp_test_case(RigidTransform.from_exp_coords)
@pytest.mark.parametrize("dim", range(1, 4))
def test_from_exp_coords(xp, dim: int):
shape = (dim,) * (dim - 1)
# example from 3.3 of
# https://hades.mech.northwestern.edu/images/2/25/MR-v2.pdf
dtype = xpx.default_dtype(xp)
angle1 = np.deg2rad(30.0)
mat = xp.asarray([
[np.cos(angle1), -np.sin(angle1), 0.0, 1.0],
[np.sin(angle1), np.cos(angle1), 0.0, 2.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0]
], dtype=dtype)
mat = xp.tile(mat, shape + (1, 1))
tf1 = RigidTransform.from_matrix(mat)
angle2 = np.deg2rad(60.0)
mat = xp.asarray([
[np.cos(angle2), -np.sin(angle2), 0.0, 2.0],
[np.sin(angle2), np.cos(angle2), 0.0, 1.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0]
], dtype=dtype)
mat = xp.tile(mat, shape + (1, 1))
tf2 = RigidTransform.from_matrix(mat)
expected = tf2 * tf1.inv()
deg2rag = xp.asarray(np.deg2rad(30.0))
exp_coords = deg2rag * xp.asarray([0.0, 0.0, 1.0, 3.37, -3.37, 0.0])
exp_coords = xp.tile(exp_coords, shape + (1,))
actual = RigidTransform.from_exp_coords(exp_coords)
xp_assert_close(actual.as_matrix(), expected.as_matrix(), atol=1e-2)
# test cases generated by comparison to pytransform3d
exp_coords = xp.asarray([
[-2.01041204, -0.52983629, 0.65773501,
0.10386614, 0.05855009, 0.54959179],
[-0.22537438, -0.24132627, -2.4747121,
-0.09158594, 1.88075832, -0.03197204]
])
exp_coords = xp.tile(exp_coords, shape + (1,))
expected_matrix = xp.asarray([
[[0.76406621, 0.10504613, -0.63652819, -0.10209961],
[0.59956454, -0.47987325, 0.64050295, 0.40158789],
[-0.2381705, -0.87102639, -0.42963687, 0.19637636],
[0., 0., 0., 1.]],
[[-0.78446989, 0.61157488, 0.10287448, 1.33330055],
[-0.58017785, -0.78232107, 0.22664378, 0.52660831],
[0.21909052, 0.11810973, 0.96852952, -0.02968529],
[0., 0., 0., 1.]]
])
expected_matrix = xp.tile(expected_matrix, shape + (1, 1))
xp_assert_close(
RigidTransform.from_exp_coords(exp_coords).as_matrix(),
expected_matrix, atol=1e-8)
# identity
expected_matrix = xp.tile(xp.eye(4), shape + (1, 1))
exp_coords = xp.zeros(shape + (6,), dtype=dtype)
xp_assert_close(
RigidTransform.from_exp_coords(exp_coords).as_matrix(),
expected_matrix, atol=1e-12)
# only translation
expected_matrix = xp.asarray([
[[1.0, 0.0, 0.0, 3.0],
[0.0, 1.0, 0.0, -5.4],
[0.0, 0.0, 1.0, 100.2],
[0.0, 0.0, 0.0, 1.0]],
[[1.0, 0.0, 0.0, -3.0],
[0.0, 1.0, 0.0, 13.3],
[0.0, 0.0, 1.0, 1.3],
[0.0, 0.0, 0.0, 1.0]]
])
expected_matrix = xp.tile(expected_matrix, shape + (1, 1, 1))
exp_coords = xp.asarray([
[0.0, 0.0, 0.0, 3.0, -5.4, 100.2],
[0.0, 0.0, 0.0, -3.0, 13.3, 1.3],
])
exp_coords = xp.tile(exp_coords, shape + (1, 1))
actual = RigidTransform.from_exp_coords(exp_coords)
xp_assert_close(actual.as_matrix(), expected_matrix, atol=1e-12)
# only rotation
angles = xp.asarray([[34, -12, 0.5], [-102, -55, 30]])
angles = xp.tile(angles, shape + (1, 1))
rot = Rotation.from_euler('zyx', angles, degrees=True)
rotvec = rot.as_rotvec()
expected_matrix = xp.tile(xp.eye(4), shape + (2, 1, 1))
expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(rot.as_matrix())
exp_coords = xp.concat((rotvec, xp.zeros_like(rotvec)), axis=-1)
actual = RigidTransform.from_exp_coords(exp_coords)
xp_assert_close(actual.as_matrix(), expected_matrix, atol=1e-12)
def test_from_exp_coords_array_like():
rng = np.random.default_rng(123)
# Test single transform
t = np.array([1, 2, 3])
r = Rotation.random(rng=rng)
tf_expected = RigidTransform.from_components(t, r)
exp_coords = tf_expected.as_exp_coords().tolist()
assert isinstance(exp_coords, list)
tf = RigidTransform.from_exp_coords(exp_coords)
xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)
# Test multiple transforms
t = [[1, 2, 3], [4, 5, 6]]
r = Rotation.random(len(t), rng=rng)
tf_expected = RigidTransform.from_components(t, r)
exp_coords = tf_expected.as_exp_coords().tolist()
assert isinstance(exp_coords, list)
tf = RigidTransform.from_exp_coords(exp_coords)
xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)
@make_xp_test_case(RigidTransform.as_exp_coords)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_as_exp_coords(xp, ndim: int):
shape = (ndim,) * (ndim - 1)
# identity
expected = xp.zeros(shape + (6,))
actual = RigidTransform.from_exp_coords(expected).as_exp_coords()
xp_assert_close(actual, expected, atol=1e-12)
rng = np.random.default_rng(10)
# pure rotation
rot_vec = xp.asarray(rng.normal(scale=0.1, size=shape + (1000, 3)))
tf = RigidTransform.from_rotation(Rotation.from_rotvec(rot_vec))
exp_coords = tf.as_exp_coords()
xp_assert_close(exp_coords[..., :3], rot_vec, atol=1e-12)
expected = xp.zeros_like(rot_vec)
xp_assert_close(exp_coords[..., 3:], expected, atol=1e-16)
# pure translation
translation = xp.asarray(rng.normal(scale=100.0, size=shape + (1000, 3)))
tf = RigidTransform.from_translation(translation)
exp_coords = tf.as_exp_coords()
xp_assert_close(exp_coords[..., :3], expected, atol=1e-16)
xp_assert_close(exp_coords[..., 3:], translation, atol=1e-15)
@make_xp_test_case(RigidTransform.from_dual_quat)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_dual_quat(xp, ndim: int):
dtype = xpx.default_dtype(xp)
atol = 1e-12 if dtype == xp.float64 else 1e-7
shape = (ndim,) * (ndim - 1)
# identity
dq = xp.asarray([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype)
dq = xp.tile(dq, shape + (1,))
expected = xp.tile(xp.eye(4), shape + (1, 1))
xp_assert_close(RigidTransform.from_dual_quat(dq).as_matrix(), expected, atol=atol)
dq = xp.asarray([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype)
dq = xp.tile(dq, shape + (1,))
xp_assert_close(RigidTransform.from_dual_quat(dq, scalar_first=True).as_matrix(),
expected, atol=atol)
# only translation
dq = xp.asarray([0, 0, 0, 1, 0.25, 0.15, -0.7, 0], dtype=dtype)
dq = xp.tile(dq, shape + (1,))
actual = RigidTransform.from_dual_quat(dq)
expected_matrix = xp.asarray([
[1, 0, 0, 0.5],
[0, 1, 0, 0.3],
[0, 0, 1, -1.4],
[0, 0, 0, 1]
])
expected_matrix = xp.tile(expected_matrix, shape + (1, 1))
xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol)
dq = xp.asarray([1, 0, 0, 0, 0, 0.25, 0.15, -0.7], dtype=dtype)
dq = xp.tile(dq, shape + (1,))
actual = RigidTransform.from_dual_quat(dq, scalar_first=True)
xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol)
# only rotation
angles = xp.asarray([65, -13, 90], dtype=dtype)
angles = xp.tile(angles, shape + (1,))
actual_rot = Rotation.from_euler("xyz", angles, degrees=True)
qrot = actual_rot.as_quat()
dq = xp.concat((qrot, xp.zeros_like(qrot)), axis=-1)
actual = RigidTransform.from_dual_quat(dq)
expected_matrix = xp.tile(xp.eye(4), shape + (1, 1))
expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(actual_rot.as_matrix())
xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol)
qrot = actual_rot.as_quat(scalar_first=True)
dq = xp.concat((qrot, xp.zeros_like(qrot)), axis=-1)
actual = RigidTransform.from_dual_quat(dq, scalar_first=True)
expected_matrix = xp.tile(xp.eye(4), shape + (1, 1))
expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(actual_rot.as_matrix())
xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol)
# rotation and translation
# rtol is set to 1e-7 because xp_assert_close deviates from
# np.testing.assert_allclose in that it does not automatically default to 1e-7 for
# floating point inputs.
# See https://numpy.org/doc/2.2/reference/generated/numpy.testing.assert_allclose.html
dq = xp.asarray(
[[0.0617101, -0.06483886, 0.31432811, 0.94508498,
0.04985168, -0.26119618, 0.1691491, -0.07743254],
[0.19507259, 0.49404931, -0.06091285, 0.8450749,
0.65049656, -0.30782513, 0.16566752, 0.04174544]])
dq = xp.tile(dq, shape + (1, 1))
actual = RigidTransform.from_dual_quat(dq)
expected_matrix = xp.asarray(
[[[0.79398752, -0.60213598, -0.08376202, 0.24605262],
[0.58613113, 0.79477941, -0.15740392, -0.4932833],
[0.16135089, 0.07588122, 0.98397557, 0.34262676],
[0., 0., 0., 1.]],
[[0.50440981, 0.2957028, 0.81125249, 1.20934468],
[0.08979911, 0.91647262, -0.3898898, -0.70540077],
[-0.8587822, 0.26951399, 0.43572393, -0.47776265],
[0., 0., 0., 1.]]])
expected_matrix = xp.tile(expected_matrix, shape + (1, 1, 1))
xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol, rtol=1e-7)
dq = xp.asarray(
[[0.94508498, 0.0617101, -0.06483886, 0.31432811,
-0.07743254, 0.04985168, -0.26119618, 0.1691491],
[0.8450749, 0.19507259, 0.49404931, -0.06091285,
0.04174544, 0.65049656, -0.30782513, 0.16566752]])
dq = xp.tile(dq, shape + (1, 1))
actual = RigidTransform.from_dual_quat(dq, scalar_first=True)
xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol, rtol=1e-7)
# unnormalized dual quaternions
# invalid real quaternion with norm 0
dq = xp.zeros(shape + (8,))
actual = RigidTransform.from_dual_quat(dq)
expected = xp.tile(xp.eye(4), shape + (1, 1))
xp_assert_close(actual.as_matrix(), expected, atol=atol)
# real quaternion with norm != 1
unnormalized_dual_quat = xp.asarray(
[-0.2547655, 1.23506123, 0.20230088, 0.24247194, # norm 1.3
0.38559628, 0.08184063, 0.1755943, -0.1582222] # orthogonal
)
xp_assert_close(xp_vector_norm(unnormalized_dual_quat[:4]), xp.asarray(1.3)[()],
atol=atol)
xp_assert_close(xp.vecdot(unnormalized_dual_quat[:4],
unnormalized_dual_quat[4:])[()],
xp.asarray(0.0)[()], atol=1e-8)
dq = xp.tile(unnormalized_dual_quat, shape + (1,))
dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat()
expected_ones = xp.ones(shape) if shape != () else xp.asarray(1.0)[()]
expected_zeros = xp.zeros(shape) if shape != () else xp.asarray(0.0)[()]
xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones,
atol=1e-12)
vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
vecdot = vecdot[()] if vecdot.shape == () else vecdot
xp_assert_close(vecdot, expected_zeros, atol=atol)
# real and dual quaternion are not orthogonal
unnormalized_dual_quat = xp.asarray(
[0.20824458, 0.75098079, 0.54542913, -0.30849493, # unit norm
-0.16051025, 0.10742978, 0.21277201, 0.20596935] # not orthogonal
)
xp_assert_close(xp_vector_norm(unnormalized_dual_quat[:4]), xp.asarray(1.0)[()],
atol=atol)
assert xp.vecdot(unnormalized_dual_quat[:4], unnormalized_dual_quat[4:]) != 0.0
dq = xp.tile(unnormalized_dual_quat, shape + (1,))
dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat()
xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones,
atol=1e-12)
vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
vecdot = vecdot[()] if vecdot.shape == () else vecdot
xp_assert_close(vecdot, expected_zeros, atol=atol)
# invalid real quaternion with norm 0, non-orthogonal dual quaternion
unnormalized_dual_quat = xp.asarray(
[0.0, 0.0, 0.0, 0.0, -0.16051025, 0.10742978, 0.21277201, 0.20596935])
assert xp.vecdot(xp.asarray([0.0, 0, 0, 1]), unnormalized_dual_quat[4:]) != 0.0
dq = xp.tile(unnormalized_dual_quat, shape + (1,))
dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat()
xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones,
atol=1e-12)
vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
vecdot = vecdot[()] if vecdot.shape == () else vecdot
xp_assert_close(vecdot, expected_zeros, atol=atol)
# compensation for precision loss in real quaternion
rng = np.random.default_rng(1000)
t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)
r = Rotation.from_quat(q)
random_dual_quats = RigidTransform.from_components(t, r).as_dual_quat()
# ensure that random quaternions are not normalized
random_dual_quats = xpx.at(random_dual_quats)[..., :4].add(0.01)
assert not xp.any(xpx.isclose(xp_vector_norm(random_dual_quats[..., :4], axis=-1),
1.0, atol=0.0001))
dual_quat_norm = RigidTransform.from_dual_quat(
random_dual_quats).as_dual_quat()
xp_assert_close(xp_vector_norm(dual_quat_norm[..., :4], axis=-1), expected_ones,
atol=atol)
# compensation for precision loss in dual quaternion, results in violation
# of orthogonality constraint
t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)
r = Rotation.from_quat(q)
random_dual_quats = RigidTransform.from_components(t, r).as_dual_quat()
# ensure that random quaternions are not normalized
random_dual_quats = xpx.at(random_dual_quats)[..., 4:].add(0.1)
q_norm = xp.vecdot(random_dual_quats[..., :4], random_dual_quats[..., 4:])
assert not xp.any(xpx.isclose(q_norm, 0.0, atol=0.0001))
dual_quat_norm = RigidTransform.from_dual_quat(
random_dual_quats).as_dual_quat()
vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
vecdot = vecdot[()] if vecdot.shape == () else vecdot
xp_assert_close(vecdot, expected_zeros, atol=atol)
xp_assert_close(random_dual_quats[..., :4], dual_quat_norm[..., :4], atol=atol)
def test_from_dual_quat_array_like():
rng = np.random.default_rng(123)
# Test single transform
t = np.array([1, 2, 3])
r = Rotation.random(rng=rng)
tf_expected = RigidTransform.from_components(t, r)
dual_quat = tf_expected.as_dual_quat().tolist()
assert isinstance(dual_quat, list)
tf = RigidTransform.from_dual_quat(dual_quat)
xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)
# Test multiple transforms
t = [[1, 2, 3], [4, 5, 6]]
r = Rotation.random(len(t), rng=rng)
tf_expected = RigidTransform.from_components(t, r)
dual_quat = tf_expected.as_dual_quat().tolist()
assert isinstance(dual_quat, list)
tf = RigidTransform.from_dual_quat(dual_quat)
xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)
@make_xp_test_case(RigidTransform.as_dual_quat)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_as_dual_quat(xp, ndim: int):
dtype = xpx.default_dtype(xp)
shape = (ndim,) * (ndim - 1)
# identity
expected = xp.asarray([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype)
actual = rigid_transform_to_xp(RigidTransform.identity(), xp).as_dual_quat()
xp_assert_close(actual, expected, atol=1e-12)
expected = xp.asarray([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
tf = rigid_transform_to_xp(RigidTransform.identity(), xp)
actual = tf.as_dual_quat(scalar_first=True)
xp_assert_close(actual, expected, atol=1e-12)
rng = np.random.default_rng(10)
# only rotation
for _ in range(10):
q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)
real_part = Rotation.from_quat(q).as_quat()
dual_part = xp.zeros_like(real_part)
expected = xp.concat((real_part, dual_part), axis=-1)
actual = RigidTransform.from_dual_quat(expected).as_dual_quat()
# because of double cover:
actual = actual * xp.sign(actual[..., 0, None])
expected = expected * xp.sign(expected[..., 0, None])
xp_assert_close(actual, expected, atol=1e-12)
# only translation
for _ in range(10):
tf = 0.5 * xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
expected = xp.zeros(shape + (8,), dtype=dtype)
expected = xpx.at(expected)[..., 3].set(1.0)
expected = xpx.at(expected)[..., 4:7].set(tf)
actual = RigidTransform.from_dual_quat(expected).as_dual_quat()
# because of double cover:
actual = actual * xp.sign(actual[..., 0, None])
expected = expected * xp.sign(expected[..., 0, None])
xp_assert_close(actual, expected, atol=1e-12)
# rotation and translation
for _ in range(10):
t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
expected = RigidTransform.from_components(t, r).as_dual_quat()
actual = RigidTransform.from_dual_quat(expected).as_dual_quat()
# because of double cover:
actual = actual * xp.sign(actual[..., 0, None])
expected = expected * xp.sign(expected[..., 0, None])
xp_assert_close(actual, expected, atol=1e-12)
@make_xp_test_case(RigidTransform.from_components, RigidTransform.as_components,
RigidTransform.from_exp_coords, RigidTransform.as_exp_coords,
RigidTransform.from_matrix, RigidTransform.as_matrix,
RigidTransform.from_dual_quat, RigidTransform.as_dual_quat)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_as_internal_consistency(xp, ndim: int):
dtype = xpx.default_dtype(xp)
atol = 1e-12
n = 10
rng = np.random.default_rng(10)
shape = (n,) + (ndim,) * (ndim - 1)
t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
tf0 = RigidTransform.from_components(t, r)
tf1 = RigidTransform.from_components(*tf0.as_components())
xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)
tf1 = RigidTransform.from_components(tf0.translation, tf0.rotation)
xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)
tf1 = RigidTransform.from_exp_coords(tf0.as_exp_coords())
xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)
tf1 = RigidTransform.from_matrix(tf0.as_matrix())
xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)
tf1 = RigidTransform.from_dual_quat(tf0.as_dual_quat())
xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)
# exp_coords small rotation
t = xp.asarray(rng.normal(scale=1000.0, size=shape + (3,)), dtype=dtype)
rotvec = xp.asarray(rng.normal(scale=1e-10, size=shape + (3,)), dtype=dtype)
r = Rotation.from_rotvec(rotvec)
tf0 = RigidTransform.from_components(t, r)
tf1 = RigidTransform.from_exp_coords(tf0.as_exp_coords())
xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)
def test_identity():
# We do not use xp here because identity always returns numpy arrays
atol = 1e-12
# Test single identity
tf = RigidTransform.identity()
xp_assert_close(tf.as_matrix(), np.eye(4), atol=atol)
# Test multiple identities
tf = RigidTransform.identity(5)
xp_assert_close(tf.as_matrix(), np.array([np.eye(4)] * 5), atol=atol)
# Test shape
tf = RigidTransform.identity(shape=3)
expected = np.tile(np.eye(4), (3, 1, 1))
xp_assert_close(tf.as_matrix(), expected, atol=atol)
tf = RigidTransform.identity(shape=(2, 3))
expected = np.tile(np.eye(4), (2, 3, 1, 1))
xp_assert_close(tf.as_matrix(), expected, atol=atol)
# Test errors
with pytest.raises(ValueError, match="Only one of `num` and `shape` can be."):
RigidTransform.identity(10, shape=(2, 3))
with pytest.raises(TypeError, match="takes from 0 to 1 positional arguments"):
RigidTransform.identity(None, (-1, 3)) # Shape is kwarg only
with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"):
RigidTransform.identity(shape="invalid")
@make_xp_test_case(RigidTransform.apply)
def test_apply(xp):
atol = 1e-12
# Broadcast shape: (6, 5, 4, 2) ( + (3,) for vectors, + (4,) for rotations)
vector_shapes = [(), (1,), (2,), (1, 2), (5, 1, 2)]
tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (1, 4, 2), (5, 4, 2), (6, 5, 4, 2)]
rng = np.random.default_rng(123)
for tf_shape, v_shape in product(tf_shapes, vector_shapes):
# Random rotation and translation
t = xp.asarray(rng.normal(size=tf_shape + (3,)))
q = xp.asarray(rng.normal(size=tf_shape + (4,)))
r = Rotation.from_quat(q)
tf = RigidTransform.from_components(t, r)
vecs = xp.asarray(rng.normal(size=v_shape + (3,)))
expected = t + r.apply(vecs)
res = tf.apply(vecs)
assert res.shape == np.broadcast_shapes(tf_shape, v_shape) + (3,)
xp_assert_close(res, expected, atol=atol)
def test_apply_array_like():
rng = np.random.default_rng(123)
# Single vector
t = np.array([1, 2, 3])
r = Rotation.random(rng=rng)
tf = RigidTransform.from_components(t, r)
vec = [1, 0, 0]
expected = t + r.apply(vec)
xp_assert_close(tf.apply(vec), expected, atol=1e-12)
# Multiple vectors
t = np.array([[1, 2, 3], [4, 5, 6]])
r = Rotation.random(len(t), rng=rng)
tf = RigidTransform.from_components(t, r)
vec = [[1, 0, 0], [0, 1, 0]]
expected = t + r.apply(vec)
xp_assert_close(tf.apply(vec), expected, atol=1e-12)
@make_xp_test_case(RigidTransform.apply)
def test_inverse_apply(xp):
atol = 1e-12
# Broadcast shape: (6, 5, 4, 2) ( + (3,) for vectors, + (4,) for rotations)
vector_shapes = [(), (1,), (2,), (1, 2), (5, 1, 2)]
tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (1, 4, 2), (5, 4, 2), (6, 5, 4, 2)]
rng = np.random.default_rng(123)
for tf_shape, v_shape in product(tf_shapes, vector_shapes):
# Random rotation and translation
t = xp.asarray(rng.normal(size=tf_shape + (3,)))
q = xp.asarray(rng.normal(size=tf_shape + (4,)))
r = Rotation.from_quat(q)
tf = RigidTransform.from_components(t, r)
vecs = xp.asarray(rng.normal(size=v_shape + (3,)))
expected = tf.inv().apply(vecs)
res = tf.apply(vecs, inverse=True)
assert res.shape == np.broadcast_shapes(tf_shape, v_shape) + (3,)
xp_assert_close(res, expected, atol=atol)
@make_xp_test_case(RigidTransform.apply)
def test_rotation_alone(xp):
atol = 1e-12
r = Rotation.from_euler('z', xp.asarray(90), degrees=True)
tf = RigidTransform.from_rotation(r)
vec = xp.asarray([1, 0, 0])
expected = r.apply(vec)
xp_assert_close(tf.apply(vec), expected, atol=atol)
@make_xp_test_case(RigidTransform.apply)
def test_translation_alone(xp):
atol = 1e-12
t = xp.asarray([1.0, 2, 3])
tf = RigidTransform.from_translation(t)
vec = xp.asarray([5.0, 6, 7])
expected = t + vec
xp_assert_close(tf.apply(vec), expected, atol=atol)
@make_xp_test_case(RigidTransform.apply, RigidTransform.__mul__)
def test_composition(xp):
atol = 1e-12
tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (5, 4, 2)]
dtype = xpx.default_dtype(xp)
rng = np.random.default_rng(123)
for tf_shape1, tf_shape2 in product(tf_shapes, repeat=2):
# Random rotation and translation
t1 = xp.asarray(rng.normal(size=tf_shape1 + (3,)))
q1 = xp.asarray(rng.normal(size=tf_shape1 + (4,)))
r1 = Rotation.from_quat(q1)
tf1 = RigidTransform.from_components(t1, r1)
t2 = xp.asarray(rng.normal(size=tf_shape2 + (3,)))
q2 = xp.asarray(rng.normal(size=tf_shape2 + (4,)))
r2 = Rotation.from_quat(q2)
tf2 = RigidTransform.from_components(t2, r2)
composed = tf2 * tf1
vec = xp.asarray(rng.normal(size=(3,)), dtype=dtype)
expected = tf2.apply(tf1.apply(vec))
res = composed.apply(vec)
assert res.shape == np.broadcast_shapes(tf_shape1, tf_shape2) + (3,)
xp_assert_close(res, expected, atol=atol)
expected = t2 + r2.apply(t1 + r1.apply(vec))
xp_assert_close(composed.apply(vec), expected, atol=atol)
assert composed.single == (tf1.single and tf2.single)
@make_xp_test_case(RigidTransform.__pow__, RigidTransform.__mul__)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_pow(xp, ndim: int):
dtype = xpx.default_dtype(xp)
atol = 1e-12 if dtype == xp.float64 else 1e-6
num = 10
rng = np.random.default_rng(100)
shape = (num,) + (ndim,) * (ndim - 1)
t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)
r = Rotation.from_quat(q)
p = RigidTransform.from_components(t, r)
p_inv = p.inv()
# Test the short-cuts and other integers
for n in [-5, -2, -1, 0, 1, 2, 5]:
q = p**n
r = RigidTransform.from_matrix(xp.tile(xp.eye(4), shape + (1, 1)))
for _ in range(abs(n)):
if n > 0:
r = r * p
else:
r = r * p_inv
xp_assert_close(q.as_matrix(), r.as_matrix(), atol=atol)
# Test shape preservation of single
single_tf = RigidTransform.identity()
assert (single_tf**n).as_matrix().shape == (4, 4)
# Test fractional powers
q = p**0.5
xp_assert_close((q * q).as_matrix(), p.as_matrix(), atol=atol)
q = p**-0.5
xp_assert_close((q * q).as_matrix(), p.inv().as_matrix(), atol=atol)
q = p** 1.5
xp_assert_close((q * q).as_matrix(), (p**3).as_matrix(), atol=atol)
q = p** -1.5
xp_assert_close((q * q).as_matrix(), (p**-3).as_matrix(), atol=atol)
# pow function
tf = pow(RigidTransform.from_matrix(xp.eye(4)), 2)
xp_assert_close(tf.as_matrix(), xp.eye(4), atol=atol)
@make_xp_test_case(RigidTransform.__pow__)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_pow_equivalence_with_rotation(xp, ndim: int):
atol = 1e-12
num = 10
rng = np.random.default_rng(100)
dtype = xpx.default_dtype(xp)
shape = (num,) + (ndim,) * (ndim - 1)
r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
p = RigidTransform.from_rotation(r)
for n in [-5, -2, -1.5, -1, -0.5, 0.0, 0.5, 1, 1.5, 2, 5]:
xp_assert_close((p**n).rotation.as_matrix(), (r**n).as_matrix(), atol=atol)
@make_xp_test_case(RigidTransform.inv, RigidTransform.__mul__)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_inverse(xp, ndim: int):
dtype = xpx.default_dtype(xp)
atol = 1e-12 if dtype == xp.float64 else 1e-6
rng = np.random.default_rng(100)
shape = (ndim,) * (ndim - 1)
# Test inverse transform
r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
tf = RigidTransform.from_components(t, r)
# Test that tf * tf.inv() equals identity
tf_inv = tf.inv()
composed = tf * tf_inv
expected = xp.tile(xp.eye(4), shape + (1, 1))
xp_assert_close(composed.as_matrix(), expected, atol=atol)
@make_xp_test_case(RigidTransform.as_matrix)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_properties(xp, ndim: int):
atol = 1e-12 if xpx.default_dtype(xp) == xp.float64 else 1e-6
shape = (ndim,) * (ndim - 1)
dtype = xpx.default_dtype(xp)
rng = np.random.default_rng(100)
# Test rotation and translation properties
r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
tf = RigidTransform.from_components(t, r)
xp_assert_close(tf.rotation.as_matrix(), r.as_matrix(), atol=atol)
assert xp.all(tf.rotation.approx_equal(r, atol=atol))
xp_assert_close(tf.translation, t, atol=atol)
# Test that we don't return views that would modify the original array
xpx.at(tf.translation)[..., 0].set(0.0)
xp_assert_close(tf.translation, t, atol=atol)
assert tf.single == (shape == ())
@make_xp_test_case(RigidTransform.__getitem__)
def test_indexing(xp):
atol = 1e-12
# Test indexing for multiple transforms
r = Rotation.from_euler('zyx', xp.asarray([[90, 0, 0], [0, 90, 0]]), degrees=True)
t = xp.asarray([[1.0, 2, 3], [4, 5, 6]])
tf = RigidTransform.from_components(t, r)
# Test single index
xp_assert_close(tf[0].as_matrix()[:3, :3], r[0].as_matrix(), atol=atol)
xp_assert_close(tf[0].as_matrix()[:3, 3], t[0, ...], atol=atol)
# Test slice
tf_slice = tf[0:2]
xp_assert_close(tf_slice.as_matrix()[:, :3, :3], r[0:2].as_matrix(), atol=atol)
xp_assert_close(tf_slice.as_matrix()[:, :3, 3], t[0:2, ...], atol=atol)
# Test boolean indexing
tf_masked = tf[xp.asarray([True, True])]
xp_assert_close(tf_masked.as_matrix()[:, :3, :3], r.as_matrix(), atol=atol)
xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t, atol=atol)
tf_masked = tf[xp.asarray([False, True])]
xp_assert_close(tf_masked.as_matrix()[:, :3, :3],
r[xp.asarray([False, True])].as_matrix(), atol=atol)
xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t[xp.asarray([False, True])],
atol=atol)
tf_masked = tf[xp.asarray([False, False])]
assert len(tf_masked) == 0
# Test integer array indexing
idx = xp.asarray([0, 1])
xp_assert_close(tf[idx].as_matrix()[:, :3, :3], r[idx].as_matrix(), atol=atol)
xp_assert_close(tf[idx].as_matrix()[:, :3, 3], t, atol=atol)
def test_indexing_array_like():
atol = 1e-12
r = Rotation.from_euler('zyx', np.array([[90, 0, 0], [0, 90, 0]]), degrees=True)
t = np.array([[1.0, 2, 3], [4, 5, 6]])
tf = RigidTransform.from_components(t, r)
tf_masked = tf[[False, True]]
xp_assert_close(tf_masked.as_matrix()[:, :3, :3], r[[False, True]].as_matrix(),
atol=atol)
xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t[[False, True]], atol=atol)
tf_masked = tf[[False, False]]
assert len(tf_masked) == 0
@make_xp_test_case(RigidTransform.concatenate)
def test_concatenate(xp):
atol = 1e-12
# Test concatenation of transforms
t1 = xp.asarray([1, 0, 0])
r1 = Rotation.from_euler('z', xp.asarray(90), degrees=True)
tf1 = RigidTransform.from_components(t1, r1)
t2 = xp.asarray([0, 1, 0])
r2 = Rotation.from_euler('x', xp.asarray(90), degrees=True)
tf2 = RigidTransform.from_components(t2, r2)
# Concatenate single transforms
concatenated1 = RigidTransform.concatenate([tf1, tf2])
xp_assert_close(concatenated1[0].as_matrix(), tf1.as_matrix(), atol=atol)
xp_assert_close(concatenated1[1].as_matrix(), tf2.as_matrix(), atol=atol)
# Concatenate multiple transforms
concatenated2 = RigidTransform.concatenate([tf1, concatenated1])
xp_assert_close(concatenated2[0].as_matrix(), tf1.as_matrix(), atol=atol)
xp_assert_close(concatenated2[1].as_matrix(), tf1.as_matrix(), atol=atol)
xp_assert_close(concatenated2[2].as_matrix(), tf2.as_matrix(), atol=atol)
# Test ND concatenation
tf3 = RigidTransform.from_translation(xp.reshape(xp.arange(18), (3, 2, 3)))
tf4 = RigidTransform.from_translation(xp.reshape(xp.arange(18) + 18, (3, 2, 3)))
concatenated3 = RigidTransform.concatenate([tf3, tf4])
xp_assert_close(concatenated3.as_matrix()[:3, ...], tf3.as_matrix(), atol=atol)
xp_assert_close(concatenated3.as_matrix()[3:, ...], tf4.as_matrix(), atol=atol)
@make_xp_test_case(RigidTransform.from_matrix)
def test_input_validation(xp):
# Test invalid matrix shapes
inputs = [xp.eye(3), xp.zeros((4, 3)), []]
for input in inputs:
with pytest.raises(ValueError, match="Expected `matrix` to have shape"):
RigidTransform.from_matrix(input)
# Test invalid last row
for ndim in range(3):
shape = (ndim,) * (ndim - 1)
matrix = xp.zeros(shape + (4, 4))
matrix = xpx.at(matrix)[...].set(xp.eye(4))
matrix = xpx.at(matrix)[..., 3, :].set(xp.asarray([1.0, 0, 0, 1]))
if is_lazy_array(matrix):
matrix = RigidTransform.from_matrix(matrix).as_matrix()
assert xp.all(xp.isnan(matrix[..., 3, :]))
else:
with pytest.raises(ValueError, match="last row of transformation matrix"):
RigidTransform.from_matrix(matrix)
# Test left handed rotation matrix
matrix = xp.eye(4)
matrix = xpx.at(matrix)[0, 0].set(-1)
if is_lazy_array(matrix):
matrix = RigidTransform.from_matrix(matrix).as_matrix()
assert xp.all(xp.isnan(matrix[..., :3, :3]))
else:
with pytest.raises(ValueError, match="Non-positive determinant"):
RigidTransform(matrix, normalize=True)
# Test non-Rotation input
with pytest.raises(TypeError,
match="Expected `rotation` to be a `Rotation` instance"):
RigidTransform.from_rotation(xp.eye(3))
@make_xp_test_case(RigidTransform.mean)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_mean(xp, ndim: int):
atol = 1e-12
rng = np.random.default_rng(123)
dtype = xpx.default_dtype(xp)
t = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (3,)), dtype=dtype)
q = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (4,)), dtype=dtype)
r = Rotation.from_quat(q)
tf = RigidTransform.from_components(t, r)
# Unweighted mean
axis = tuple(range(t.ndim - 1))
t_mean = xp.mean(t, axis=axis)
r_mean = r.mean()
tf_mean = tf.mean()
assert tf_mean.shape == ()
xp_assert_close(tf_mean.as_matrix(),
RigidTransform.from_components(t_mean, r_mean).as_matrix(),
atol=atol)
# Weighted mean
if ndim == 1:
weights = None
t_mean = t
else:
weights = xp.asarray(rng.random(size=(ndim,) * (ndim - 1)), dtype=dtype)
norm = xp.sum(weights[..., None], axis=axis)
wsum = xp.sum(t * weights[..., None], axis=axis)
t_mean = wsum/norm
r_mean = r.mean(weights=weights)
tf_mean = tf.mean(weights=weights)
assert tf_mean.shape == ()
xp_assert_close(tf_mean.as_matrix(),
RigidTransform.from_components(t_mean, r_mean).as_matrix(),
atol=atol)
@make_xp_test_case(
RigidTransform.from_rotation, RigidTransform.mean, Rotation.magnitude
)
@pytest.mark.parametrize("ndim", range(1, 5))
def test_mean_axis(xp, ndim: int):
axes = xp.tile(xp.concat((-xp.eye(3), xp.eye(3))), (3,) * (ndim - 1) + (1, 1))
theta = xp.pi / 4
r = Rotation.from_rotvec(theta * axes)
tf = RigidTransform.from_rotation(r)
# Test mean over last axis
desired = xp.full(axes.shape[:-2], 0.0)
if ndim == 1:
desired = desired[()]
atol = 1e-6 if xpx.default_dtype(xp) is xp.float32 else 1e-10
xp_assert_close(tf.mean(axis=-1).rotation.magnitude(), desired, atol=atol)
# Test tuple axes
desired = xp.full(axes.shape[1:-2], 0.0)
if ndim < 3:
desired = desired[()]
xp_assert_close(tf.mean(axis=(0, -1)).rotation.magnitude(), desired, atol=atol)
# Empty axis tuple should return RigidTransform unchanged
tf_mean = tf.mean(axis=())
xp_assert_close(tf_mean.as_matrix(), tf.as_matrix(), atol=atol)
@make_xp_test_case(RigidTransform.mean, Rotation.magnitude)
def test_mean_compare_axis(xp):
# Create a random set of transforms and compare the mean over an axis with
# the mean without axis of the sliced transform
atol = 1e-10 if xpx.default_dtype(xp) == xp.float64 else 1e-6
rng = np.random.default_rng(0)
q = xp.asarray(rng.normal(size=(4, 5, 6, 4)), dtype=xpx.default_dtype(xp))
r = Rotation.from_quat(q)
t = xp.asarray(rng.normal(size=(4, 5, 6, 3)), dtype=xpx.default_dtype(xp))
tf = RigidTransform.from_components(t, r)
mean_0 = tf.mean(axis=0)
for i in range(q.shape[1]):
for j in range(q.shape[2]):
r_slice = Rotation.from_quat(q[:, i, j, ...])
t_slice = t[:, i, j, ...]
mean_slice_tf = RigidTransform.from_components(t_slice, r_slice).mean()
xp_assert_close(
(mean_0[i][j].rotation * mean_slice_tf.rotation.inv()).magnitude(),
xp.asarray(0.0)[()], atol=atol,
)
xp_assert_close(
mean_0[i][j].translation, mean_slice_tf.translation, atol=atol,
)
mean_1_2 = tf.mean(axis=(1, 2))
for i in range(q.shape[0]):
r_slice = Rotation.from_quat(q[i, ...])
t_slice = t[i, ...]
mean_slice_tf = RigidTransform.from_components(t_slice, r_slice).mean()
xp_assert_close(
(mean_1_2[i].rotation * mean_slice_tf.rotation.inv()).magnitude(),
xp.asarray(0.0)[()], atol=atol,
)
xp_assert_close(
mean_1_2[i].translation, mean_slice_tf.translation, atol=atol,
)
@make_xp_test_case(RigidTransform.mean)
def test_mean_invalid_weights(xp):
tf = RigidTransform.from_matrix(xp.tile(xp.eye(4), (4, 1, 1)))
if is_lazy_array(tf.as_matrix()):
m = tf.mean(weights=-xp.ones(4))
assert xp.all(xp.isnan(m.as_matrix()))
else:
with pytest.raises(ValueError, match="non-negative"):
tf.mean(weights=-xp.ones(4))
# Test weight shape mismatch
tf = RigidTransform.from_matrix(xp.eye(4))
with pytest.raises(ValueError, match="Expected `weights` to"):
tf.mean(weights=xp.ones((2,)))
tf = RigidTransform.from_matrix(xp.tile(xp.eye(4), (3, 2, 1, 1, 1)))
with pytest.raises(ValueError, match="Expected `weights` to"):
tf.mean(weights=xp.ones((2, 1)))
@make_xp_test_case(RigidTransform.from_translation)
def test_translation_validation(xp):
# Test invalid translation shapes
with pytest.raises(ValueError, match="Expected `translation` to have shape"):
RigidTransform.from_translation(xp.asarray([1, 2]))
with pytest.raises(ValueError, match="Expected `translation` to have shape"):
RigidTransform.from_translation(xp.zeros((2, 2)))
@make_xp_test_case(RigidTransform.apply)
def test_vector_validation(xp):
tf = rigid_transform_to_xp(RigidTransform.identity(2), xp=xp)
# Test invalid vector shapes
with pytest.raises(ValueError, match="Expected vector to have shape"):
tf.apply(xp.asarray([1, 2]))
with pytest.raises(ValueError, match="Expected vector to have shape"):
tf.apply(xp.zeros((2, 2)))
with pytest.raises(ValueError, match="operands could not be broadcast"):
tf.apply(xp.zeros((1, 4, 3)))
@make_xp_test_case(RigidTransform.__getitem__)
def test_indexing_validation(xp):
tf = RigidTransform.from_matrix(xp.eye(4))
# Test indexing on single transform
with pytest.raises(TypeError, match="Single transform is not subscriptable"):
tf[0]
with pytest.raises(TypeError, match="Single transform is not subscriptable"):
tf[0:1]
# Test length on single transform
with pytest.raises(TypeError, match="Single transform has no len"):
len(tf)
@make_xp_test_case(RigidTransform.__mul__)
def test_composition_validation(xp):
tf2 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6]]))
tf3 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
# Test incompatible shapes
with pytest.raises(ValueError, match="Cannot broadcast"):
tf2 * tf3
tf4 = RigidTransform.from_matrix(xp.tile(xp.eye(4), (1, 4, 1, 1)))
# Test invalid broadcasting shape
with pytest.raises(ValueError, match="Cannot broadcast"):
tf2 * tf4
@make_xp_test_case(RigidTransform.concatenate)
def test_concatenate_validation(xp):
tf = RigidTransform.from_matrix(xp.eye(4))
# Test invalid inputs
with pytest.raises(TypeError,
match="input must contain RigidTransform objects"):
RigidTransform.concatenate([tf, xp.eye(4)])
# Test incompatible shapes
tf2 = RigidTransform.from_translation(xp.ones((1, 1, 1, 3)))
# Frameworks have a highly heterogeneous way of reporting errors for this case
with pytest.raises((ValueError, TypeError, RuntimeError)):
RigidTransform.concatenate([tf, tf2])
@make_xp_test_case(RigidTransform.__setitem__)
def test_setitem(xp):
tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
single = RigidTransform.from_translation(xp.asarray([1, 1, 1]))
double = RigidTransform.from_translation(xp.asarray([[2, 2, 2], [3, 3, 3]]))
triple = RigidTransform.from_translation(xp.asarray([[3, 3, 3],
[4, 4, 4],
[5, 5, 5]]))
# Test indexing with integer index
tf[0] = single
xp_assert_close(tf.translation, xp.asarray([[1.0, 1, 1], [4, 5, 6], [7, 8, 9]]))
# Test indexing with slice
tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
tf[:2] = double
xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [3, 3, 3], [7, 8, 9]]))
# Test indexing with ellipsis
tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
tf[...] = triple
xp_assert_close(tf.translation, xp.asarray([[3.0, 3, 3], [4, 4, 4], [5, 5, 5]]))
# Test indexing with boolean array
tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
mask = xp.asarray([True, False, True])
tf[mask] = double
xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [4, 5, 6], [3, 3, 3]]))
@make_xp_test_case(RigidTransform.__setitem__)
@pytest.mark.skip_xp_backends("array_api_strict",
reason="doesn't support fancy indexing __setitem__")
def test_setitem_fancy_indexing(xp):
double = RigidTransform.from_translation(xp.asarray([[2, 2, 2], [3, 3, 3]]))
tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
idx = xp.asarray([0, 2])
tf[idx] = double
xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [4, 5, 6], [3, 3, 3]]))
@make_xp_test_case(RigidTransform.__setitem__)
def test_setitem_validation(xp):
tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6]]))
single = RigidTransform.from_matrix(xp.eye(4))
# Test setting item on single transform
with pytest.raises(TypeError, match="Single transform is not subscriptable"):
single[0] = tf
# Test invalid value type
with pytest.raises(TypeError, match="value must be a RigidTransform"):
tf[0] = xp.eye(4)
@pytest.mark.skip_xp_backends("jax.numpy",
reason="JAX does not support memory sharing")
@make_xp_test_case(RigidTransform.as_matrix)
def test_copy_flag(xp):
# Test that copy=True creates new memory
matrix = xp.eye(4)
tf = RigidTransform(matrix, normalize=False, copy=True)
matrix[0, 0] = 2
assert tf.as_matrix()[0, 0] == 1
# Test that copy=False shares memory
matrix = xp.eye(4)
tf = RigidTransform(matrix, normalize=False, copy=False)
matrix[0, 0] = 2
assert tf.as_matrix()[0, 0] == 2
@make_xp_test_case(normalize_dual_quaternion)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_normalize_dual_quaternion(xp, ndim: int):
dtype = xpx.default_dtype(xp)
atol = 1e-12 if dtype == xp.float64 else 1e-6
rng = np.random.default_rng(100)
shape = (ndim,) * (ndim - 1)
dual_quat = normalize_dual_quaternion(xp.zeros((1, 8)))
xp_assert_close(xp_vector_norm(dual_quat[0, :4], axis=-1), xp.asarray(1.0)[()],
atol=1e-12)
xp_assert_close(xp.vecdot(dual_quat[0, :4], dual_quat[0, 4:])[()],
xp.asarray(0.0)[()], atol=1e-12)
dual_quat = xp.asarray(rng.normal(size=shape + (8,)), dtype=dtype)
dual_quat = normalize_dual_quaternion(dual_quat)
expected = xp.ones(shape) if shape != () else xp.asarray(1.0)[()]
xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected, atol=atol)
expected = xp.zeros(shape) if shape != () else xp.asarray(0.0)[()]
vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
vecdot = vecdot[()] if vecdot.shape == () else vecdot
xp_assert_close(vecdot, expected, atol=atol)
@make_xp_test_case(RigidTransform.from_matrix, RigidTransform.from_rotation,
RigidTransform.from_translation, RigidTransform.from_components,
RigidTransform.from_exp_coords, RigidTransform.from_dual_quat)
def test_empty_transform_construction(xp):
tf = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
assert len(tf) == 0
assert not tf.single
tf = RigidTransform.from_rotation(Rotation.from_quat(xp.zeros((0, 4))))
assert len(tf) == 0
assert not tf.single
tf = RigidTransform.from_translation(xp.empty((0, 3)))
assert len(tf) == 0
assert not tf.single
empty_rot = Rotation.from_quat(xp.zeros((0, 4)))
tf = RigidTransform.from_components(xp.empty((0, 3)), empty_rot)
assert len(tf) == 0
assert not tf.single
tf = RigidTransform.from_exp_coords(xp.empty((0, 6)))
assert len(tf) == 0
assert not tf.single
tf = RigidTransform.from_dual_quat(xp.empty((0, 8)))
assert len(tf) == 0
assert not tf.single
tf = RigidTransform.identity(0)
assert len(tf) == 0
assert not tf.single
@make_xp_test_case(RigidTransform.from_matrix, RigidTransform.as_components,
RigidTransform.as_exp_coords, RigidTransform.as_dual_quat)
def test_empty_transform_representation(xp):
tf = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
assert len(tf.rotation) == 0
assert tf.translation.shape == (0, 3)
t, r = tf.as_components()
assert t.shape == (0, 3)
assert len(r) == 0
assert tf.as_matrix().shape == (0, 4, 4)
assert tf.as_exp_coords().shape == (0, 6)
assert tf.as_dual_quat().shape == (0, 8)
@make_xp_test_case(RigidTransform.from_matrix, RigidTransform.apply)
def test_empty_transform_application(xp):
tf = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
assert tf.apply(xp.zeros((3,))).shape == (0, 3)
assert tf.apply(xp.empty((0, 3))).shape == (0, 3)
with pytest.raises(ValueError, match="operands could not be broadcast together"):
tf.apply(xp.zeros((2, 3)))
@make_xp_test_case(RigidTransform.from_matrix, RigidTransform.__mul__)
def test_empty_transform_composition(xp):
tf_empty = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
tf_single = RigidTransform.from_matrix(xp.eye(4))
tf_many = rigid_transform_to_xp(RigidTransform.identity(3), xp=xp)
assert len(tf_empty * tf_empty) == 0
assert len(tf_empty * tf_single) == 0
assert len(tf_single * tf_empty) == 0
with pytest.raises(ValueError, match="Cannot broadcast"):
tf_many * tf_empty
with pytest.raises(ValueError, match="Cannot broadcast"):
tf_empty * tf_many
@make_xp_test_case(RigidTransform.from_matrix, RigidTransform.concatenate)
def test_empty_transform_concatenation(xp):
tf_empty = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
tf_single = RigidTransform.from_matrix(xp.eye(4))
tf_many = rigid_transform_to_xp(RigidTransform.identity(2), xp=xp)
assert len(RigidTransform.concatenate([tf_empty, tf_empty])) == 0
assert len(RigidTransform.concatenate([tf_empty, tf_single])) == 1
assert len(RigidTransform.concatenate([tf_single, tf_empty])) == 1
assert len(RigidTransform.concatenate([tf_empty, tf_many])) == 2
assert len(RigidTransform.concatenate([tf_many, tf_empty])) == 2
assert len(RigidTransform.concatenate([tf_many, tf_empty, tf_single])) == 3
@make_xp_test_case(RigidTransform.from_matrix, RigidTransform.inv,
RigidTransform.__pow__)
def test_empty_transform_inv_and_pow(xp):
tf = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
assert len(tf.inv()) == 0
assert len(tf ** 0) == 0
assert len(tf ** 1) == 0
assert len(tf ** -1) == 0
assert len(tf ** 0.5) == 0
@make_xp_test_case(RigidTransform.__getitem__)
def test_empty_transform_indexing(xp):
tf_many = rigid_transform_to_xp(RigidTransform.identity(3), xp=xp)
tf_zero = tf_many[xp.asarray([], dtype=xp.int32)]
assert len(tf_zero) == 0
# Array API does not specify out-of-bounds indexing. Only check for numpy.
if is_numpy(xp):
assert len(tf_zero[:5]) == 0 # Slices can go out of bounds.
with pytest.raises(IndexError):
tf_zero[0]
with pytest.raises(IndexError):
tf_zero[xp.asarray([0, 2])]
with pytest.raises(IndexError):
tf_zero[xp.asarray([False, True])]
@make_xp_test_case(RigidTransform.from_matrix)
@pytest.mark.skip_xp_backends("array_api_strict",
reason="array API doesn't support pickling")
def test_pickling(xp):
# Note: Array API makes no provision for arrays to be pickleable, so
# it's OK to skip this test for the backends that don't support it
mat = xp.eye(4)
mat = xpx.at(mat)[0, 3].set(2.0)
tf = RigidTransform.from_matrix(mat)
pkl = pickle.dumps(tf)
unpickled = pickle.loads(pkl)
xp_assert_close(tf.as_matrix(), unpickled.as_matrix(), atol=1e-15)
@make_xp_test_case(RigidTransform.as_matrix, RigidTransform.__iter__)
def test_rigid_transform_iter(xp):
r = rigid_transform_to_xp(RigidTransform.identity(3), xp)
for i, r_i in enumerate(r):
assert isinstance(r_i, RigidTransform)
xp_assert_equal(r_i.as_matrix(), r[i].as_matrix())
if i > len(r):
raise RuntimeError("Iteration exceeded length of transforms")
@make_xp_test_case()
@pytest.mark.parametrize("dim", range(1, 5))
def test_shape_property(xp, dim: int):
shape = (dim,) * (dim - 1)
tf = RigidTransform.from_translation(xp.zeros(shape + (3,)))
assert tf.shape == shape