import math import pytest import numpy as np from numpy.testing import assert_equal from scipy.spatial.transform import Rotation, Slerp import scipy.spatial.transform._rotation_cy as cython_backend import scipy.spatial.transform._rotation_xp as xp_backend from scipy.stats import special_ortho_group from itertools import permutations, product from contextlib import contextmanager import warnings from scipy._lib._array_api import ( xp_assert_equal, is_numpy, is_lazy_array, xp_vector_norm, xp_assert_close, eager_warns, xp_default_dtype, make_xp_test_case, make_xp_pytest_marks, xp_device_type, ) import scipy._lib.array_api_extra as xpx import pickle import copy lazy_xp_modules = [Rotation, Slerp] # from_quat and as_quat are used in almost all tests, so we mark them module-wide pytestmark = make_xp_pytest_marks(Rotation.as_quat, Rotation.from_quat) def basis_vec(axis): if axis == 'x': return [1, 0, 0] elif axis == 'y': return [0, 1, 0] elif axis == 'z': return [0, 0, 1] def rotation_to_xp(r: Rotation, xp): dtype = xpx.default_dtype(xp) return Rotation.from_quat(xp.asarray(r.as_quat(), dtype=dtype)) def test_init_non_array(): Rotation((0, 0, 0, 1)) Rotation([0, 0, 0, 1]) Rotation([[[0, 0, 0, 1]]]) def test_cython_backend_selection(): r = Rotation.from_quat(np.array([0, 0, 0, 1])) assert r._backend is cython_backend r = Rotation.from_quat(np.array([[0, 0, 0, 1]])) assert r._backend is cython_backend r = Rotation.from_quat(np.array([[[0, 0, 0, 1]]])) assert r._backend is xp_backend def test_numpy_float32_inputs(): Rotation.from_quat(np.array([1, 0, 0, 0], dtype=np.float32)) def test_generic_quat_matrix(xp): x = xp.asarray([[3.0, 4, 0, 0], [5, 12, 0, 0]]) r = Rotation.from_quat(x) expected_quat = x / xp.asarray([[5.0], [13.0]]) xp_assert_close(r.as_quat(), expected_quat) @pytest.mark.parametrize("ndim", range(1, 6)) def test_from_single_nd_quaternion(xp, ndim: int): x = xp.asarray([3.0, 4, 0, 0]) x = xp.reshape(x, (1,) * (ndim - 1) + (4,)) r = Rotation.from_quat(x) expected_quat = x / 5 xp_assert_close(r.as_quat(), expected_quat) @make_xp_test_case(Rotation.as_matrix) def test_from_quat_scalar_first(xp): rng = np.random.RandomState(0) r = Rotation.from_quat(xp.asarray([1, 0, 0, 0]), scalar_first=True) xp_assert_close(r.as_matrix(), xp.eye(3), rtol=1e-15, atol=1e-16) q = xp.tile(xp.asarray([1, 0, 0, 0]), (10, 1)) r = Rotation.from_quat(q, scalar_first=True) xp_assert_close( r.as_matrix(), xp.tile(xp.eye(3), (10, 1, 1)), rtol=1e-15, atol=1e-16 ) q = xp.asarray(rng.randn(100, 4)) q /= xp_vector_norm(q, axis=1)[:, None] for i in range(q.shape[0]): # Array API conforming loop qi = q[i, ...] r = Rotation.from_quat(qi, scalar_first=True) xp_assert_close(xp.roll(r.as_quat(), 1), qi, rtol=1e-15) r = Rotation.from_quat(q, scalar_first=True) xp_assert_close(xp.roll(r.as_quat(), 1, axis=1), q, rtol=1e-15) def test_from_quat_array_like(): rng = np.random.default_rng(123) # Single rotation r_expected = Rotation.random(rng=rng) r = Rotation.from_quat(r_expected.as_quat().tolist()) assert r_expected.approx_equal(r, atol=1e-12) # Multiple rotations r_expected = Rotation.random(3, rng=rng) r = Rotation.from_quat(r_expected.as_quat().tolist()) assert np.all(r_expected.approx_equal(r, atol=1e-12)) # Tensor of rotations q = rng.normal(size=(3, 4, 5, 4)) q /= xp_vector_norm(q, axis=-1)[..., None] r_expected = Rotation.from_quat(q) r = Rotation.from_quat(q) assert np.all(r_expected.approx_equal(r, atol=1e-12)) def test_from_quat_int_dtype(xp): r = Rotation.from_quat(xp.asarray([1, 0, 0, 0])) assert r.as_quat().dtype == xp_default_dtype(xp) def test_quat_canonical(xp): # Case 0: w < 0 q = xp.asarray([0.0, 0, 0, -1]) xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), -q) # Case 1: w == 0, x < 0 q = xp.asarray([-1.0, 0, 0, 0]) xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), -q) # Case 2: w == 0, x == 0, y < 0 q = xp.asarray([0.0, -1, 0, 0]) xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), -q) # Case 3: w == 0, x == 0, y == 0, z < 0 q = xp.asarray([0.0, 0, -1, 0]) xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), -q) # Other cases: w > 0, y < 0 q = xp.asarray([0.0, -0.1, 0, 0.9]) q = q / xp_vector_norm(q) xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), q) # Other cases: w > 0, z < 0 q = xp.asarray([0.0, 0.0, -0.1, 0.9]) q = q / xp_vector_norm(q) xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), q) @make_xp_test_case(Rotation.from_euler) def test_as_quat_scalar_first(xp): rng = np.random.RandomState(0) r = Rotation.from_euler('xyz', xp.zeros(3)) xp_assert_close(r.as_quat(scalar_first=True), xp.asarray([1.0, 0, 0, 0]), rtol=1e-15, atol=1e-16) r = Rotation.from_euler('xyz', xp.zeros((10, 3))) xp_assert_close(r.as_quat(scalar_first=True), xp.tile(xp.asarray([1.0, 0, 0, 0]), (10, 1)), rtol=1e-15, atol=1e-16) q = xp.asarray(rng.randn(100, 4)) q /= xp_vector_norm(q, axis=1)[:, None] for i in range(q.shape[0]): # Array API conforming loop qi = q[i, ...] r = Rotation.from_quat(qi) xp_assert_close(r.as_quat(scalar_first=True), xp.roll(qi, 1), rtol=1e-15) xp_assert_close(r.as_quat(canonical=True, scalar_first=True), xp.roll(r.as_quat(canonical=True), 1), rtol=1e-15) r = Rotation.from_quat(q) xp_assert_close(r.as_quat(scalar_first=True), xp.roll(q, 1, axis=1), rtol=1e-15) xp_assert_close(r.as_quat(canonical=True, scalar_first=True), xp.roll(r.as_quat(canonical=True), 1, axis=1), rtol=1e-15) def test_from_square_quat_matrix(xp): # Ensure proper norm array broadcasting x = xp.asarray([ [3.0, 0, 0, 4], [5, 0, 12, 0], [0, 0, 0, 1], [-1, -1, -1, 1], [0, 0, 0, -1], # Check double cover [-1, -1, -1, -1] # Check double cover ]) r = Rotation.from_quat(x) expected_quat = x / xp.asarray([[5.0], [13], [1], [2], [1], [2]]) xp_assert_close(r.as_quat(), expected_quat) def test_quat_double_to_canonical_single_cover(xp): x = xp.asarray([ [-1.0, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, -1], [-1, -1, -1, -1] ]) r = Rotation.from_quat(x) expected_quat = xp.abs(x) / xp_vector_norm(x, axis=1)[:, None] xp_assert_close(r.as_quat(canonical=True), expected_quat) @make_xp_test_case(Rotation.inv, Rotation.__mul__) def test_quat_double_cover(xp): # See the Rotation.from_quat() docstring for scope of the quaternion # double cover property. # Check from_quat and as_quat(canonical=False) q = xp.asarray([0.0, 0, 0, -1]) r = Rotation.from_quat(q) xp_assert_equal(q, r.as_quat(canonical=False)) # Check composition and inverse q = xp.asarray([1.0, 0, 0, 1])/math.sqrt(2) # 90 deg rotation about x r = Rotation.from_quat(q) r3 = r*r*r xp_assert_close(r.as_quat(canonical=False)*math.sqrt(2), xp.asarray([1.0, 0, 0, 1])) xp_assert_close(r.inv().as_quat(canonical=False)*math.sqrt(2), xp.asarray([-1.0, 0, 0, 1])) xp_assert_close(r3.as_quat(canonical=False)*math.sqrt(2), xp.asarray([1.0, 0, 0, -1])) xp_assert_close(r3.inv().as_quat(canonical=False)*math.sqrt(2), xp.asarray([-1.0, 0, 0, -1])) # More sanity checks xp_assert_close((r*r.inv()).as_quat(canonical=False), xp.asarray([0.0, 0, 0, 1]), atol=2e-16) xp_assert_close((r3*r3.inv()).as_quat(canonical=False), xp.asarray([0.0, 0, 0, 1]), atol=2e-16) xp_assert_close((r*r3).as_quat(canonical=False), xp.asarray([0.0, 0, 0, -1]), atol=2e-16) xp_assert_close((r.inv() * r3.inv()).as_quat(canonical=False), xp.asarray([0.0, 0, 0, -1]), atol=2e-16) @pytest.mark.parametrize("ndim", range(1, 6)) def test_from_quat_wrong_shape(xp, ndim: int): quat = xp.zeros((*((1,) * ndim), 5)) with pytest.raises(ValueError, match="Expected `quat` to have shape"): Rotation.from_quat(quat) def test_zero_norms_from_quat(xp): x = xp.asarray([ [3, 4, 0, 0], [0, 0, 0, 0], [5, 0, 12, 0] ]) if is_lazy_array(x): assert xp.all(xp.isnan(Rotation.from_quat(x).as_quat()[1, ...])) else: with pytest.raises(ValueError): Rotation.from_quat(x) @make_xp_test_case(Rotation.as_matrix) @pytest.mark.parametrize("ndim", range(1, 6)) def test_as_matrix_single_nd_quaternion(xp, ndim: int): quat = xp.asarray([0, 0, 1, 1]) quat = xp.reshape(quat, (1,) * (ndim - 1) + (4,)) mat = Rotation.from_quat(quat).as_matrix() expected_mat = xp.asarray([ [0.0, -1, 0], [1, 0, 0], [0, 0, 1] ]) expected_mat = xp.reshape(expected_mat, (1,) * (ndim - 1) + (3, 3)) xp_assert_close(mat, expected_mat, atol=1e-16) @make_xp_test_case(Rotation.as_matrix) def test_as_matrix_from_square_input(xp): quats = xp.asarray([ [0, 0, 1, 1], [0, 1, 0, 1], [0, 0, 0, 1], [0, 0, 0, -1] ]) mat = Rotation.from_quat(quats).as_matrix() assert_equal(mat.shape, (4, 3, 3)) expected0 = xp.asarray([ [0.0, -1, 0], [1, 0, 0], [0, 0, 1] ]) xp_assert_close(mat[0, ...], expected0, atol=1e-16) expected1 = xp.asarray([ [0.0, 0, 1], [0, 1, 0], [-1, 0, 0] ]) xp_assert_close(mat[1, ...], expected1, atol=1e-16) xp_assert_close(mat[2, ...], xp.eye(3)) xp_assert_close(mat[3, ...], xp.eye(3)) @make_xp_test_case(Rotation.as_matrix) def test_as_matrix_from_generic_input(xp): quats = xp.asarray([ [0, 0, 1, 1], [0, 1, 0, 1], [1, 2, 3, 4] ]) mat = Rotation.from_quat(quats).as_matrix() assert_equal(mat.shape, (3, 3, 3)) expected0 = xp.asarray([ [0.0, -1, 0], [1, 0, 0], [0, 0, 1] ]) xp_assert_close(mat[0, ...], expected0, atol=1e-16) expected1 = xp.asarray([ [0.0, 0, 1], [0, 1, 0], [-1, 0, 0] ]) xp_assert_close(mat[1, ...], expected1, atol=1e-16) expected2 = xp.asarray([ [0.4, -2, 2.2], [2.8, 1, 0.4], [-1, 2, 2] ]) / 3 xp_assert_close(mat[2, ...], expected2) @make_xp_test_case(Rotation.from_matrix) @pytest.mark.parametrize("ndim", range(1, 6)) def test_from_single_nd_matrix(xp, ndim: int): mat = xp.asarray([ [0, 0, 1], [1, 0, 0], [0, 1, 0] ]) mat = xp.reshape(mat, (1,) * (ndim - 1) + (3, 3)) expected_quat = xp.asarray([0.5, 0.5, 0.5, 0.5]) expected_quat = xp.reshape(expected_quat, (1,) * (ndim - 1) + (4,)) xp_assert_close(Rotation.from_matrix(mat).as_quat(), expected_quat) @make_xp_test_case(Rotation.from_matrix) def test_from_matrix_calculation(xp): atol = 1e-8 expected_quat = xp.asarray([1.0, 1, 6, 1]) / math.sqrt(39) mat = xp.asarray([ [-0.8974359, -0.2564103, 0.3589744], [0.3589744, -0.8974359, 0.2564103], [0.2564103, 0.3589744, 0.8974359] ]) xp_assert_close(Rotation.from_matrix(mat).as_quat(), expected_quat, atol=atol) xp_assert_close(Rotation.from_matrix(xp.reshape(mat, (1, 3, 3))).as_quat(), xp.reshape(expected_quat, (1, 4)), atol=atol) @make_xp_test_case(Rotation.from_matrix, Rotation.as_matrix) def test_matrix_calculation_pipeline(xp): mat = xp.asarray(special_ortho_group.rvs(3, size=10, random_state=0)) xp_assert_close(Rotation.from_matrix(mat).as_matrix(), mat) @make_xp_test_case(Rotation.from_matrix, Rotation.as_matrix) def test_from_matrix_ortho_output(xp): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-6 rnd = np.random.RandomState(0) mat = xp.asarray(rnd.random_sample((100, 3, 3)), dtype=dtype) dets = xp.linalg.det(mat) for i in range(dets.shape[0]): # Make sure we have a right-handed rotation matrix if dets[i] < 0: mat = xpx.at(mat)[i, ...].set(-mat[i, ...]) ortho_mat = Rotation.from_matrix(mat).as_matrix() mult_result = xp.matmul(ortho_mat, xp.matrix_transpose(ortho_mat)) eye3d = xp.zeros((100, 3, 3)) + xp.eye(3) xp_assert_close(mult_result, eye3d, atol=atol) @make_xp_test_case(Rotation.from_matrix, Rotation.as_matrix) def test_from_matrix_normalize(xp): mat = xp.asarray([ [1, 1, 0], [0, 1, 0], [0, 0, 1]]) expected = xp.asarray([[ 0.894427, 0.447214, 0.0], [-0.447214, 0.894427, 0.0], [ 0.0, 0.0, 1.0]]) xp_assert_close(Rotation.from_matrix(mat).as_matrix(), expected, atol=1e-6) mat = xp.asarray([ [0, -0.5, 0 ], [0.5, 0 , 0 ], [0, 0 , 0.5]]) expected = xp.asarray([[0.0, -1, 0], [ 1, 0, 0], [ 0, 0, 1]]) xp_assert_close(Rotation.from_matrix(mat).as_matrix(), expected, atol=1e-6) # Test a mix of normalized and non-normalized matrices mat = xp.stack([mat, xp.eye(3)]) expected = xp.stack([expected, xp.eye(3)]) xp_assert_close(Rotation.from_matrix(mat).as_matrix(), expected, atol=1e-6) @make_xp_test_case(Rotation.from_matrix, Rotation.as_matrix) def test_from_matrix_assume_valid(xp): rng = np.random.default_rng(0) dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-6 # Test that normal matrices remain unchanged rot = rotation_to_xp(Rotation.from_quat(rng.normal(size=(10, 4))), xp) rot_no_norm = Rotation.from_matrix(rot.as_matrix(), assume_valid=True) assert xp.all(rot.approx_equal(rot_no_norm, atol=atol)) # We make no guarantees about matrices that are not orthogonal or do not # have unit norm @make_xp_test_case(Rotation.from_matrix, Rotation.as_matrix) def test_from_matrix_non_positive_determinant(xp): mat = xp.eye(3) mat = xpx.at(mat)[0, 0].set(0) if is_lazy_array(mat): assert xp.all(xp.isnan(Rotation.from_matrix(mat).as_matrix())) else: with pytest.raises(ValueError, match="Non-positive determinant"): Rotation.from_matrix(mat) mat = xpx.at(mat)[0, 0].set(-1) if is_lazy_array(mat): assert xp.all(xp.isnan(Rotation.from_matrix(mat).as_matrix())) else: with pytest.raises(ValueError, match="Non-positive determinant"): Rotation.from_matrix(mat) def test_from_matrix_array_like(): rng = np.random.default_rng(123) # Single rotation r_expected = Rotation.random(rng=rng) r = Rotation.from_matrix(r_expected.as_matrix().tolist()) assert r_expected.approx_equal(r, atol=1e-12) # Multiple rotations r_expected = Rotation.random(3, rng=rng) r = Rotation.from_matrix(r_expected.as_matrix().tolist()) assert np.all(r_expected.approx_equal(r, atol=1e-12)) @make_xp_test_case(Rotation.from_matrix) def test_from_matrix_int_dtype(xp): mat = xp.asarray([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) r = Rotation.from_matrix(mat) assert r.as_quat().dtype == xp_default_dtype(xp) @make_xp_test_case(Rotation.from_rotvec) @pytest.mark.parametrize("ndim", range(1, 6)) def test_from_nd_single_rotvec(xp, ndim: int): atol = 1e-7 rotvec = xp.asarray([1, 0, 0]) rotvec = xp.reshape(rotvec, (1,) * (ndim - 1) + (3,)) expected_quat = xp.asarray([0.4794255, 0, 0, 0.8775826]) expected_quat = xp.reshape(expected_quat, (1,) * (ndim - 1) + (4,)) result = Rotation.from_rotvec(rotvec) xp_assert_close(result.as_quat(), expected_quat, atol=atol) @make_xp_test_case(Rotation.from_rotvec) def test_from_generic_rotvec(xp): atol = 1e-7 rotvec = xp.asarray([ [1, 2, 2], [1, -1, 0.5], [0, 0, 0]]) expected_quat = xp.asarray([ [0.3324983, 0.6649967, 0.6649967, 0.0707372], [0.4544258, -0.4544258, 0.2272129, 0.7316889], [0, 0, 0, 1] ]) xp_assert_close(Rotation.from_rotvec(rotvec).as_quat(), expected_quat, atol=atol) @make_xp_test_case(Rotation.from_rotvec) def test_from_rotvec_small_angle(xp): rotvec = xp.asarray([ [5e-4 / math.sqrt(3), -5e-4 / math.sqrt(3), 5e-4 / math.sqrt(3)], [0.2, 0.3, 0.4], [0, 0, 0] ]) quat = Rotation.from_rotvec(rotvec).as_quat() # cos(theta/2) ~~ 1 for small theta xp_assert_close(quat[0, 3], xp.asarray(1.0)[()]) # sin(theta/2) / theta ~~ 0.5 for small theta xp_assert_close(quat[0, :3], rotvec[0, ...] * 0.5) xp_assert_close(quat[1, 3], xp.asarray(0.9639685)[()]) xp_assert_close(quat[1, :3], xp.asarray([ 0.09879603932153465, 0.14819405898230198, 0.19759207864306931])) xp_assert_equal(quat[2, ...], xp.asarray([0.0, 0, 0, 1])) @make_xp_test_case(Rotation.from_rotvec, Rotation.as_rotvec) def test_from_rotvec_array_like(): rng = np.random.default_rng(123) # Single rotation r_expected = Rotation.random(rng=rng) r = Rotation.from_rotvec(r_expected.as_rotvec().tolist()) assert r_expected.approx_equal(r, atol=1e-12) # Multiple rotations r_expected = Rotation.random(3, rng=rng) r = Rotation.from_rotvec(r_expected.as_rotvec().tolist()) assert np.all(r_expected.approx_equal(r, atol=1e-12)) @make_xp_test_case(Rotation.from_rotvec) def test_from_rotvec_int_dtype(xp): rotvec = xp.asarray([1, 0, 0]) r = Rotation.from_rotvec(rotvec) assert r.as_quat().dtype == xp_default_dtype(xp) @make_xp_test_case(Rotation.from_rotvec) def test_degrees_from_rotvec(xp): rotvec1 = xp.asarray([1 / 3 ** (1/3)] * 3) rot1 = Rotation.from_rotvec(rotvec1, degrees=True) quat1 = rot1.as_quat() # deg2rad is not implemented in Array API -> / 180 * xp.pi rotvec2 = xp.asarray(rotvec1 / 180 * xp.pi) rot2 = Rotation.from_rotvec(rotvec2) quat2 = rot2.as_quat() xp_assert_close(quat1, quat2) @make_xp_test_case(Rotation.from_rotvec) def test_malformed_1d_from_rotvec(xp): with pytest.raises(ValueError, match='Expected `rot_vec` to have shape'): Rotation.from_rotvec(xp.asarray([1, 2])) @make_xp_test_case(Rotation.from_rotvec) @pytest.mark.parametrize("ndim", range(1, 6)) def test_malformed_nd_from_rotvec(xp, ndim: int): shape = (1,) * (ndim - 1) + (2,) with pytest.raises(ValueError, match='Expected `rot_vec` to have shape'): Rotation.from_rotvec(xp.ones(shape)) @make_xp_test_case(Rotation.as_rotvec) @pytest.mark.skip_xp_backends("dask.array", reason="missing required linalg.cross function") def test_as_generic_rotvec(xp): dtype = xpx.default_dtype(xp) atol = 1e-15 if dtype == xp.float64 else 1e-7 quat = xp.asarray([ [1, 2, -1, 0.5], [1, -1, 1, 0.0003], [0, 0, 0, 1] ]) quat /= xp_vector_norm(quat, axis=-1, keepdims=True) rotvec = Rotation.from_quat(quat).as_rotvec() angle = xp_vector_norm(rotvec, axis=-1) xp_assert_close(quat[:, 3], xp.cos(angle / 2)) xp_assert_close(xp.linalg.cross(rotvec, quat[:, :3]), xp.zeros((3, 3)), atol=atol) @make_xp_test_case(Rotation.as_rotvec) @pytest.mark.parametrize("ndim", range(1, 4)) def test_as_rotvec_single_nd_input(xp, ndim: int): quat = xp.asarray([1, 2, -3, 2]) quat = xp.reshape(quat, (1,) * (ndim - 1) + (4,)) expected_rotvec = xp.asarray([0.5772381, 1.1544763, -1.7317144]) expected_rotvec = xp.reshape(expected_rotvec, (1,) * (ndim - 1) + (3,)) actual_rotvec = Rotation.from_quat(quat).as_rotvec() assert_equal(actual_rotvec.shape, expected_rotvec.shape) xp_assert_close(actual_rotvec, expected_rotvec) @make_xp_test_case(Rotation.from_matrix, Rotation.as_rotvec) def test_as_rotvec_degrees(xp): # x->y, y->z, z->x mat = xp.asarray([[0, 0, 1], [1, 0, 0], [0, 1, 0]]) rot = Rotation.from_matrix(mat) rotvec = rot.as_rotvec(degrees=True) angle = xp_vector_norm(rotvec, axis=-1) xp_assert_close(angle, xp.asarray(120.0)[()]) xp_assert_close(rotvec[0], rotvec[1]) xp_assert_close(rotvec[1], rotvec[2]) @make_xp_test_case(Rotation.from_rotvec, Rotation.as_rotvec) def test_rotvec_calc_pipeline(xp): # Include small angles rotvec = xp.asarray([ [0, 0, 0], [1, -1, 2], [-3e-4, 3.5e-4, 7.5e-5] ]) xp_assert_close(Rotation.from_rotvec(rotvec).as_rotvec(), rotvec) xp_assert_close(Rotation.from_rotvec(rotvec, degrees=True).as_rotvec(degrees=True), rotvec) @make_xp_test_case(Rotation.from_mrp) @pytest.mark.parametrize("ndim", range(1, 4)) def test_from_mrp_single_nd_input(xp, ndim: int): mrp = xp.asarray([0, 0, 1.0]) mrp = xp.reshape(mrp, (1,) * (ndim - 1) + (3,)) expected_quat = xp.asarray([0.0, 0, 1, 0]) expected_quat = xp.reshape(expected_quat, (1,) * (ndim - 1) + (4,)) result = Rotation.from_mrp(mrp) xp_assert_close(result.as_quat(), expected_quat, atol=1e-12) def test_from_mrp_array_like(): rng = np.random.default_rng(123) # Single rotation r_expected = Rotation.random(rng=rng) r = Rotation.from_mrp(r_expected.as_mrp().tolist()) assert r_expected.approx_equal(r, atol=1e-12) # Multiple rotations r_expected = Rotation.random(3, rng=rng) r = Rotation.from_mrp(r_expected.as_mrp().tolist()) assert np.all(r_expected.approx_equal(r, atol=1e-12)) @make_xp_test_case(Rotation.from_mrp) def test_from_mrp_int_dtype(xp): mrp = xp.asarray([0, 0, 1]) r = Rotation.from_mrp(mrp) assert r.as_quat().dtype == xp_default_dtype(xp) @make_xp_test_case(Rotation.from_mrp) def test_from_generic_mrp(xp): mrp = xp.asarray([ [1, 2, 2], [1, -1, 0.5], [0, 0, 0]]) expected_quat = xp.asarray([ [0.2, 0.4, 0.4, -0.8], [0.61538462, -0.61538462, 0.30769231, -0.38461538], [0, 0, 0, 1]]) xp_assert_close(Rotation.from_mrp(mrp).as_quat(), expected_quat) @make_xp_test_case(Rotation.from_mrp) @pytest.mark.parametrize("ndim", range(1, 4)) def test_malformed_nd_from_mrp(xp, ndim: int): shape = (1,) * (ndim - 1) + (2,) with pytest.raises(ValueError, match='Expected `mrp` to have shape'): Rotation.from_mrp(xp.ones(shape)) @make_xp_test_case(Rotation.as_mrp) def test_as_generic_mrp(xp): quat = xp.asarray([ [1, 2, -1, 0.5], [1, -1, 1, 0.0003], [0, 0, 0, 1]]) quat /= xp_vector_norm(quat, axis=1)[:, None] expected_mrp = xp.asarray([ [0.33333333, 0.66666667, -0.33333333], [0.57725028, -0.57725028, 0.57725028], [0, 0, 0]]) xp_assert_close(Rotation.from_quat(quat).as_mrp(), expected_mrp) @make_xp_test_case(Rotation.from_euler, Rotation.as_mrp) def test_past_180_degree_rotation(xp): # ensure that a > 180 degree rotation is returned as a <180 rotation in MRPs # in this case 270 should be returned as -90 expected_mrp = xp.asarray([-math.tan(xp.pi / 2 / 4), 0.0, 0]) xp_assert_close( Rotation.from_euler('xyz', xp.asarray([270, 0, 0]), degrees=True).as_mrp(), expected_mrp, ) @make_xp_test_case(Rotation.as_mrp) @pytest.mark.parametrize("ndim", range(1, 4)) def test_as_mrp_single_nd_input(xp, ndim: int): quat = xp.asarray([1, 2, -3, 2]) quat = xp.reshape(quat, (1,) * (ndim - 1) + (4,)) expected_mrp = xp.asarray([0.16018862, 0.32037724, -0.48056586]) expected_mrp = xp.reshape(expected_mrp, (1,) * (ndim - 1) + (3,)) actual_mrp = Rotation.from_quat(quat).as_mrp() assert_equal(actual_mrp.shape, expected_mrp.shape) xp_assert_close(actual_mrp, expected_mrp) @make_xp_test_case(Rotation.from_mrp, Rotation.as_mrp) def test_mrp_calc_pipeline(xp): actual_mrp = xp.asarray([ [0, 0, 0], [1, -1, 2], [0.41421356, 0, 0], [0.1, 0.2, 0.1]]) expected_mrp = xp.asarray([ [0, 0, 0], [-0.16666667, 0.16666667, -0.33333333], [0.41421356, 0, 0], [0.1, 0.2, 0.1]]) xp_assert_close(Rotation.from_mrp(actual_mrp).as_mrp(), expected_mrp) @make_xp_test_case(Rotation.from_euler) def test_from_euler_single_rotation(xp): quat = Rotation.from_euler("z", xp.asarray(90), degrees=True).as_quat() expected_quat = xp.asarray([0.0, 0, 1, 1]) / math.sqrt(2) xp_assert_close(quat, expected_quat) @make_xp_test_case(Rotation.from_euler) def test_from_euler_input_validation(xp): # Single sequence with multiple angles with pytest.raises(ValueError, match="Expected last dimension of `angles` to"): Rotation.from_euler("X", xp.asarray([0, 90])) # Multiple sequences with single angle with pytest.raises(ValueError, match="Expected last dimension of `angles` to"): Rotation.from_euler("XYZ", xp.asarray([90])) @make_xp_test_case(Rotation.from_euler) @pytest.mark.parametrize("ndim", range(1, 4)) def test_from_euler_nd_rotation(xp, ndim: int): angles = xp.reshape(xp.asarray([0, 0, 90]), (1,) * (ndim - 1) + (3,)) quat = Rotation.from_euler("xyz", angles, degrees=True).as_quat() expected_quat = xp.asarray([0.0, 0, 1, 1]) / math.sqrt(2) expected_quat = xp.reshape(expected_quat, (1,) * (ndim - 1) + (4,)) xp_assert_close(quat, expected_quat) @make_xp_test_case(Rotation.from_euler, Rotation.as_matrix) def test_single_intrinsic_extrinsic_rotation(xp): extrinsic = Rotation.from_euler('z', xp.asarray(90), degrees=True).as_matrix() intrinsic = Rotation.from_euler('Z', xp.asarray(90), degrees=True).as_matrix() xp_assert_close(extrinsic, intrinsic) @make_xp_test_case(Rotation.from_euler) def test_from_euler_rotation_order(xp): # Intrinsic rotation is same as extrinsic with order reversed rnd = np.random.RandomState(0) a = xp.asarray(rnd.randint(low=0, high=180, size=(6, 3))) b = xp.flip(a, axis=-1) x = Rotation.from_euler('xyz', a, degrees=True).as_quat() y = Rotation.from_euler('ZYX', b, degrees=True).as_quat() xp_assert_close(x, y) @make_xp_test_case(Rotation.from_euler, Rotation.as_matrix) def test_from_euler_elementary_extrinsic_rotation(xp): atol = 1e-12 # Simple test to check if extrinsic rotations are implemented correctly mat = Rotation.from_euler('zx', xp.asarray([90, 90]), degrees=True).as_matrix() expected_mat = xp.asarray([ [0.0, -1, 0], [0, 0, -1], [1, 0, 0] ]) xp_assert_close(mat, expected_mat, atol=atol) @make_xp_test_case(Rotation.from_euler, Rotation.as_matrix) def test_from_euler_intrinsic_rotation_312(xp): atol = 1e-7 angles = xp.asarray([ [30, 60, 45], [30, 60, 30], [45, 30, 60] ]) mat = Rotation.from_euler('ZXY', angles, degrees=True).as_matrix() xp_assert_close(mat[0, ...], xp.asarray([ [0.3061862, -0.2500000, 0.9185587], [0.8838835, 0.4330127, -0.1767767], [-0.3535534, 0.8660254, 0.3535534] ]), atol=atol) xp_assert_close(mat[1, ...], xp.asarray([ [0.5334936, -0.2500000, 0.8080127], [0.8080127, 0.4330127, -0.3995191], [-0.2500000, 0.8660254, 0.4330127] ]), atol=atol) xp_assert_close(mat[2, ...], xp.asarray([ [0.0473672, -0.6123725, 0.7891491], [0.6597396, 0.6123725, 0.4355958], [-0.7500000, 0.5000000, 0.4330127] ]), atol=atol) @make_xp_test_case(Rotation.from_euler, Rotation.as_matrix) def test_from_euler_intrinsic_rotation_313(xp): angles = xp.asarray([ [30, 60, 45], [30, 60, 30], [45, 30, 60] ]) mat = Rotation.from_euler('ZXZ', angles, degrees=True).as_matrix() xp_assert_close(mat[0, ...], xp.asarray([ [0.43559574, -0.78914913, 0.4330127], [0.65973961, -0.04736717, -0.750000], [0.61237244, 0.61237244, 0.500000] ])) xp_assert_close(mat[1, ...], xp.asarray([ [0.6250000, -0.64951905, 0.4330127], [0.64951905, 0.1250000, -0.750000], [0.4330127, 0.750000, 0.500000] ])) xp_assert_close(mat[2, ...], xp.asarray([ [-0.1767767, -0.91855865, 0.35355339], [0.88388348, -0.30618622, -0.35355339], [0.4330127, 0.25000000, 0.8660254] ])) @make_xp_test_case(Rotation.from_euler, Rotation.as_matrix) def test_from_euler_extrinsic_rotation_312(xp): angles = xp.asarray([ [30, 60, 45], [30, 60, 30], [45, 30, 60] ]) mat = Rotation.from_euler('zxy', angles, degrees=True).as_matrix() xp_assert_close(mat[0, ...], xp.asarray([ [0.91855865, 0.1767767, 0.35355339], [0.25000000, 0.4330127, -0.8660254], [-0.30618622, 0.88388348, 0.35355339] ])) xp_assert_close(mat[1, ...], xp.asarray([ [0.96650635, -0.0580127, 0.2500000], [0.25000000, 0.4330127, -0.8660254], [-0.0580127, 0.89951905, 0.4330127] ])) xp_assert_close(mat[2, ...], xp.asarray([ [0.65973961, -0.04736717, 0.7500000], [0.61237244, 0.61237244, -0.5000000], [-0.43559574, 0.78914913, 0.4330127] ])) @make_xp_test_case(Rotation.from_euler, Rotation.as_matrix) def test_from_euler_extrinsic_rotation_313(xp): angles = xp.asarray([ [30, 60, 45], [30, 60, 30], [45, 30, 60] ]) mat = Rotation.from_euler('zxz', angles, degrees=True).as_matrix() xp_assert_close(mat[0, ...], xp.asarray([ [0.43559574, -0.65973961, 0.61237244], [0.78914913, -0.04736717, -0.61237244], [0.4330127, 0.75000000, 0.500000] ])) xp_assert_close(mat[1, ...], xp.asarray([ [0.62500000, -0.64951905, 0.4330127], [0.64951905, 0.12500000, -0.750000], [0.4330127, 0.75000000, 0.500000] ])) xp_assert_close(mat[2, ...], xp.asarray([ [-0.1767767, -0.88388348, 0.4330127], [0.91855865, -0.30618622, -0.250000], [0.35355339, 0.35355339, 0.8660254] ])) def test_from_euler_array_like(): rng = np.random.default_rng(123) order = "xyz" # Single rotation r_expected = Rotation.random(rng=rng) r = Rotation.from_euler(order, r_expected.as_euler(order).tolist()) assert r_expected.approx_equal(r, atol=1e-12) # Multiple rotations r_expected = Rotation.random(3, rng=rng) r = Rotation.from_euler(order, r_expected.as_euler(order).tolist()) assert np.all(r_expected.approx_equal(r, atol=1e-12)) def test_from_euler_scalar(): rng = np.random.default_rng(123) deg = rng.uniform(low=-180, high=180) r_expected = Rotation.from_euler("x", deg, degrees=True) r = Rotation.from_euler("x", float(deg), degrees=True) assert r_expected.approx_equal(r, atol=1e-12) @make_xp_test_case(Rotation.from_euler, Rotation.as_euler) @pytest.mark.parametrize("seq_tuple", permutations("xyz")) @pytest.mark.parametrize("intrinsic", (False, True)) def test_as_euler_asymmetric_axes(xp, seq_tuple, intrinsic): # helper function for mean error tests def test_stats(error, mean_max, rms_max): mean = xp.mean(error, axis=0) std = xp.std(error, axis=0) rms = xp.hypot(mean, std) assert xp.all(xp.abs(mean) < mean_max) assert xp.all(rms < rms_max) rnd = np.random.RandomState(0) n = 1000 angles = np.empty((n, 3)) angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles[:, 1] = rnd.uniform(low=-np.pi / 2, high=np.pi / 2, size=(n,)) angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles = xp.asarray(angles) seq = "".join(seq_tuple) if intrinsic: # Extrinsic rotation (wrt to global world) at lower case # intrinsic (WRT the object itself) lower case. seq = seq.upper() rotation = Rotation.from_euler(seq, angles) angles_quat = rotation.as_euler(seq) xp_assert_close(angles, angles_quat, atol=0, rtol=1e-12) test_stats(angles_quat - angles, 1e-15, 1e-14) @make_xp_test_case(Rotation.from_euler, Rotation.as_euler) @pytest.mark.parametrize("seq_tuple", permutations("xyz")) @pytest.mark.parametrize("intrinsic", (False, True)) def test_as_euler_symmetric_axes(xp, seq_tuple, intrinsic): # helper function for mean error tests def test_stats(error, mean_max, rms_max): mean = xp.mean(error, axis=0) std = xp.std(error, axis=0) rms = xp.hypot(mean, std) assert xp.all(xp.abs(mean) < mean_max) assert xp.all(rms < rms_max) rnd = np.random.RandomState(0) n = 1000 angles = np.empty((n, 3)) angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles[:, 1] = rnd.uniform(low=0, high=np.pi, size=(n,)) angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles = xp.asarray(angles) # Rotation of the form A/B/A are rotation around symmetric axes seq = "".join([seq_tuple[0], seq_tuple[1], seq_tuple[0]]) if intrinsic: seq = seq.upper() rotation = Rotation.from_euler(seq, angles) angles_quat = rotation.as_euler(seq) xp_assert_close(angles, angles_quat, atol=0, rtol=1.1e-13) test_stats(angles_quat - angles, 1e-16, 1e-14) @contextmanager def maybe_warn_gimbal_lock(should_warn, xp): if should_warn: # We can only warn on non-lazy backends because we'd need to condition on # traced booleans with eager_warns(UserWarning, match="Gimbal lock", xp=xp): yield else: with warnings.catch_warnings(): warnings.simplefilter("error") yield @make_xp_test_case(Rotation.from_euler, Rotation.as_matrix, Rotation.as_euler) @pytest.mark.parametrize("seq_tuple", permutations("xyz")) @pytest.mark.parametrize("intrinsic", (False, True)) @pytest.mark.parametrize("suppress_warnings", (False, True)) def test_as_euler_degenerate_asymmetric_axes( xp, seq_tuple, intrinsic, suppress_warnings ): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-6 # Since we cannot check for angle equality, we check for rotation matrix # equality angles = xp.asarray([ [45, 90, 35], [35, -90, 20], [35, 90, 25], [25, -90, 15]]) seq = "".join(seq_tuple) if intrinsic: # Extrinsic rotation (wrt to global world) at lower case # Intrinsic (WRT the object itself) upper case. seq = seq.upper() rotation = Rotation.from_euler(seq, angles, degrees=True) mat_expected = rotation.as_matrix() with maybe_warn_gimbal_lock(not suppress_warnings, xp): angle_estimates = rotation.as_euler( seq, degrees=True, suppress_warnings=suppress_warnings ) mat_estimated = Rotation.from_euler(seq, angle_estimates, degrees=True).as_matrix() xp_assert_close(mat_expected, mat_estimated, atol=atol) @make_xp_test_case(Rotation.from_euler, Rotation.as_matrix, Rotation.as_euler) @pytest.mark.parametrize("seq_tuple", permutations("xyz")) @pytest.mark.parametrize("intrinsic", (False, True)) @pytest.mark.parametrize("suppress_warnings", (False, True)) def test_as_euler_degenerate_symmetric_axes( xp, seq_tuple, intrinsic, suppress_warnings ): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-6 # Since we cannot check for angle equality, we check for rotation matrix # equality angles = xp.asarray([ [15, 0, 60], [35, 0, 75], [60, 180, 35], [15, -180, 25]]) # Rotation of the form A/B/A are rotation around symmetric axes seq = "".join([seq_tuple[0], seq_tuple[1], seq_tuple[0]]) if intrinsic: # Extrinsic rotation (wrt to global world) at lower case # Intrinsic (WRT the object itself) upper case. seq = seq.upper() rotation = Rotation.from_euler(seq, angles, degrees=True) mat_expected = rotation.as_matrix() with maybe_warn_gimbal_lock(not suppress_warnings, xp): angle_estimates = rotation.as_euler( seq, degrees=True, suppress_warnings=suppress_warnings ) mat_estimated = Rotation.from_euler(seq, angle_estimates, degrees=True).as_matrix() xp_assert_close(mat_expected, mat_estimated, atol=atol) @make_xp_test_case(Rotation.from_euler, Rotation.as_euler) @pytest.mark.parametrize("ndim", range(1, 4)) def test_as_euler_nd_rotation(xp, ndim: int): mat = xp.asarray([ [0.0, -1, 0], [1, 0, 0], [0, 0, 1] ]) mat = xp.reshape(mat, (1,) * (ndim - 1) + (3, 3)) angles = Rotation.from_matrix(mat).as_euler("xyz", degrees=True) expected_angles = xp.asarray([0, 0, 90.0]) expected_angles = xp.reshape(expected_angles, (1,) * (ndim - 1) + (3,)) xp_assert_close(angles, expected_angles, atol=1e-12) @make_xp_test_case(Rotation.as_matrix, Rotation.inv) def test_inv(xp): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-7 rnd = np.random.RandomState(0) n = 10 # preserve use of old random_state during SPEC 7 transition p = Rotation.random(num=n, random_state=rnd) p = rotation_to_xp(p, xp) p_mat = p.as_matrix() q_mat = p.inv().as_matrix() result1 = p_mat @ q_mat result2 = q_mat @ p_mat eye3d = xp.empty((n, 3, 3)) eye3d = xpx.at(eye3d)[..., :3, :3].set(xp.eye(3)) xp_assert_close(result1, eye3d, atol=atol) xp_assert_close(result2, eye3d, atol=atol) # Batched version batch_shape = (10, 3, 7) atol = 1e-12 if dtype == xp.float64 else 1e-6 quat = xp.asarray(rnd.normal(size=batch_shape + (4,)), dtype=dtype) r = Rotation.from_quat(quat) p_mat = r.as_matrix() q_mat = r.inv().as_matrix() result1 = p_mat @ q_mat result2 = q_mat @ p_mat eye_nd = xp.empty(batch_shape + (3, 3)) eye_nd = xpx.at(eye_nd)[..., :3, :3].set(xp.eye(3)) xp_assert_close(result1, eye_nd, atol=atol) xp_assert_close(result2, eye_nd, atol=atol) @make_xp_test_case(Rotation.inv, Rotation.as_matrix) def test_inv_single_rotation(xp): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-7 rng = np.random.default_rng(146972845698875399755764481408308808739) p = rotation_to_xp(Rotation.random(rng=rng), xp) q = p.inv() p_mat = p.as_matrix() q_mat = q.as_matrix() res1 = xp.matmul(p_mat, q_mat) res2 = xp.matmul(q_mat, p_mat) eye = xp.eye(3) xp_assert_close(res1, eye, atol=atol) xp_assert_close(res2, eye, atol=atol) x = rotation_to_xp(Rotation.random(num=1, rng=rng), xp) y = x.inv() x_matrix = x.as_matrix() y_matrix = y.as_matrix() result1 = xp.linalg.matmul(x_matrix, y_matrix) result2 = xp.linalg.matmul(y_matrix, x_matrix) eye3d = xp.empty((1, 3, 3)) eye3d = xpx.at(eye3d)[..., :3, :3].set(xp.eye(3)) xp_assert_close(result1, eye3d, atol=atol) xp_assert_close(result2, eye3d, atol=atol) @make_xp_test_case(Rotation.magnitude, Rotation.inv) def test_identity_magnitude(xp): n = 10 r = rotation_to_xp(Rotation.identity(n), xp) expected = xp.zeros(n) xp_assert_close(r.magnitude(), expected) xp_assert_close(r.inv().magnitude(), expected) @make_xp_test_case(Rotation.magnitude, Rotation.inv) def test_single_identity_magnitude(xp): r = rotation_to_xp(Rotation.identity(), xp) assert r.magnitude() == 0 assert r.inv().magnitude() == 0 @make_xp_test_case(Rotation.inv, Rotation.magnitude) def test_identity_invariance(xp): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-7 n = 10 p = rotation_to_xp(Rotation.random(n, rng=0), xp) q = rotation_to_xp(Rotation.identity(n), xp) result = p * q xp_assert_close(p.as_quat(), result.as_quat()) result = result * p.inv() xp_assert_close(result.magnitude(), xp.zeros(n), atol=atol) @make_xp_test_case(Rotation.inv, Rotation.magnitude) def test_single_identity_invariance(xp): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-7 n = 10 p = rotation_to_xp(Rotation.random(n, rng=0), xp) q = rotation_to_xp(Rotation.identity(), xp) result = p * q xp_assert_close(p.as_quat(), result.as_quat()) result = result * p.inv() xp_assert_close(result.magnitude(), xp.zeros(n), atol=atol) def test_identity_shape(): # Not an xp test, identity is using numpy only for now r = Rotation.identity(shape=()) assert r.as_quat().shape == (4,) r = Rotation.identity(shape=5) # Shape can be int assert r.as_quat().shape == (5, 4) r = Rotation.identity(shape=(2, 3)) assert r.as_quat().shape == (2, 3, 4) # Test values r = Rotation.identity(shape=(2, 2, 3)) xp_assert_equal(r.as_quat().reshape(-1, 4), np.tile(np.eye(4)[-1], (2 * 2 * 3, 1))) # Errors with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"): Rotation.identity(shape=2.5) with pytest.raises(ValueError, match="Only one of `num` or `shape` can be"): Rotation.identity(num=3, shape=(2, 2)) with pytest.raises(TypeError, match="takes from 0 to 1 positional arguments"): Rotation.identity(3, 3) @make_xp_test_case(Rotation.magnitude) @pytest.mark.parametrize("ndim", range(1, 4)) def test_magnitude(xp, ndim: int): quat_shape = (1,) * (ndim - 1) + (4,) quat = xp.reshape(xp.eye(4), quat_shape + (4,)) r = Rotation.from_quat(quat) result = r.magnitude() expected_result = xp.asarray([xp.pi, xp.pi, xp.pi, 0]) expected_result = xp.reshape(expected_result, quat_shape) xp_assert_close(result, expected_result) r = Rotation.from_quat(-quat) result = r.magnitude() xp_assert_close(result, expected_result) @make_xp_test_case(Rotation.magnitude) def test_magnitude_single_rotation(xp): r = Rotation.from_quat(xp.eye(4)) result1 = r[0].magnitude() xp_assert_close(result1, xp.asarray(xp.pi)[()]) result2 = r[3].magnitude() xp_assert_close(result2, xp.asarray(0.0)[()]) @make_xp_test_case(Rotation.inv, Rotation.magnitude, Rotation.approx_equal) def test_approx_equal(xp): rng = np.random.default_rng(146972845698875399755764481408308808739) p = Rotation.random(10, rng=rng) q = Rotation.random(10, rng=rng) r_mag = (p * q.inv()).magnitude() p = rotation_to_xp(p, xp) q = rotation_to_xp(q, xp) # ensure we get mix of Trues and Falses atol = xp.asarray(np.median(r_mag)) xp_assert_equal(p.approx_equal(q, atol), (xp.asarray(r_mag) < atol)) @make_xp_test_case(Rotation.from_rotvec, Rotation.approx_equal) def test_approx_equal_single_rotation(xp): # also tests passing single argument to approx_equal p = Rotation.from_rotvec(xp.asarray([0, 0, 1e-9])) # less than default atol of 1e-8 q = Rotation.from_quat(xp.eye(4)) assert p.approx_equal(q[3]) assert not p.approx_equal(q[0]) # test passing atol and using degrees assert not p.approx_equal(q[3], atol=1e-10) assert not p.approx_equal(q[3], atol=1e-8, degrees=True) with pytest.warns(UserWarning, match="atol must be set"): assert p.approx_equal(q[3], degrees=True) @make_xp_test_case(Rotation.inv, Rotation.magnitude, Rotation.approx_equal) def test_approx_equal_batched(xp): # Same shapes batch_shape = (2, 10, 3) rng = np.random.default_rng(0) p = Rotation.from_quat(rng.normal(size=batch_shape + (4,))) q = Rotation.from_quat(rng.normal(size=batch_shape + (4,))) r_mag = (p * q.inv()).magnitude() # Must be computed as numpy array for np.median p = rotation_to_xp(p, xp) q = rotation_to_xp(q, xp) assert r_mag.shape == batch_shape # ensure we get mix of Trues and Falses atol = xp.asarray(np.median(r_mag)) xp_assert_equal(p.approx_equal(q, atol), (xp.asarray(r_mag) < atol)) # Broadcastable shapes of same length p = Rotation.from_quat(rng.normal(size=batch_shape + (4,))) q = Rotation.from_quat(rng.normal(size=(1, 10, 1, 4))) r_mag = (p * q.inv()).magnitude() p = rotation_to_xp(p, xp) q = rotation_to_xp(q, xp) assert r_mag.shape == batch_shape atol = xp.asarray(np.median(r_mag)) xp_assert_equal(p.approx_equal(q, atol), (xp.asarray(r_mag) < atol)) # Broadcastable shapes of different length p = Rotation.from_quat(rng.normal(size=batch_shape + (4,))) q = Rotation.from_quat(rng.normal(size=(1, 3, 4))) r_mag = (p * q.inv()).magnitude() p = rotation_to_xp(p, xp) q = rotation_to_xp(q, xp) assert r_mag.shape == batch_shape atol = xp.asarray(np.median(r_mag)) xp_assert_equal(p.approx_equal(q, atol), (xp.asarray(r_mag) < atol)) @make_xp_test_case(Rotation.approx_equal) def test_approx_equal_batched_input_validation(xp): p = Rotation.from_quat(xp.ones((2, 3, 4))) q = Rotation.from_quat(xp.ones((3, 2, 4))) with pytest.raises(ValueError, match="broadcastable shapes"): p.approx_equal(q) p = Rotation.from_quat(xp.ones((2, 4))) q = Rotation.from_quat(xp.ones((3, 4))) with pytest.raises(ValueError, match="broadcastable shapes"): p.approx_equal(q) @make_xp_test_case(Rotation.from_rotvec, Rotation.mean, Rotation.magnitude) @pytest.mark.parametrize("ndim", range(1, 4)) def test_mean(xp, ndim: int): axes = xp.concat((-xp.eye(3), xp.eye(3))) axes = xp.reshape(axes, (1,) * (ndim - 1) + (6, 3)) thetas = xp.linspace(0, xp.pi / 2, 100) desired = xp.asarray(0.0)[()] atol = 1e-6 if xp_default_dtype(xp) is xp.float32 else 1e-10 for t in thetas: r_mean = Rotation.from_rotvec(t * axes).mean() assert r_mean.shape == () xp_assert_close(r_mean.magnitude(), desired, atol=atol) @make_xp_test_case(Rotation.from_rotvec, Rotation.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) # Test mean over last axis desired = xp.full(axes.shape[:-2], 0.0) if ndim == 1: desired = desired[()] atol = 1e-6 if xp_default_dtype(xp) is xp.float32 else 1e-10 xp_assert_close(r.mean(axis=-1).magnitude(), desired, atol=atol) # Test tuple axes desired = xp.full(axes.shape[1:-2], 0.0) if ndim < 3: desired = desired[()] xp_assert_close(r.mean(axis=(0, -1)).magnitude(), desired, atol=atol) # Empty axis tuple should return Rotation unchanged r_mean = r.mean(axis=()) xp_assert_close(r_mean.as_quat(canonical=True), r.as_quat(canonical=True), atol=atol) @make_xp_test_case(Rotation.mean, Rotation.magnitude) def test_mean_compare_axis(xp): # Create a random set of rotations and compare the mean over an axis with the # mean without axis of the sliced quaternion 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) mean_0 = r.mean(axis=0) for i in range(q.shape[1]): for j in range(q.shape[2]): mean_slice = Rotation.from_quat(q[:, i, j, ...]).mean() xp_assert_close((mean_0[i][j] * mean_slice.inv()).magnitude(), xp.asarray(0.0)[()], atol=atol) mean_1_2 = r.mean(axis=(1, 2)) for i in range(q.shape[0]): mean_slice = Rotation.from_quat(q[i, ...]).mean() xp_assert_close((mean_1_2[i] * mean_slice.inv()).magnitude(), xp.asarray(0.0)[()], atol=atol) @make_xp_test_case(Rotation.from_rotvec, Rotation.mean, Rotation.inv, Rotation.magnitude) @pytest.mark.parametrize("ndim", range(1, 4)) def test_weighted_mean(xp, ndim: int): # test that doubling a weight is equivalent to including a rotation twice. thetas = xp.linspace(0, xp.pi / 2, 100) # Create batched copies of the same setup batch_shape = (ndim,) * (ndim - 1) axes = xp.asarray([[0.0, 0, 0], [1, 0, 0], [1, 0, 0]]) weights = xp.asarray([1, 2]) axes = xp.tile(axes, batch_shape + (1, 1)) weights = xp.tile(weights, batch_shape + (1,)) expected = xp.asarray(0.0)[()] for t in thetas: rw = Rotation.from_rotvec(t * axes[..., :2, :]) mw = rw.mean(weights=weights) r = Rotation.from_rotvec(t * axes) m = r.mean() assert m.shape == () xp_assert_close((m * mw.inv()).magnitude(), expected, atol=1e-6) @make_xp_test_case(Rotation.mean) def test_mean_input_validation(xp): r = Rotation.from_quat(xp.eye(4)) if is_lazy_array(r.as_quat()): m = r.mean(weights=-xp.ones(4)) assert xp.all(xp.isnan(m._quat)) else: with pytest.raises(ValueError, match="non-negative"): r.mean(weights=-xp.ones(4)) # Test weight shape mismatch r = Rotation.from_quat(xp.ones((3, 4))) with pytest.raises(ValueError, match="Expected `weights` to"): r.mean(weights=xp.ones((2,))) r = Rotation.from_quat(xp.ones((2, 3, 4))) with pytest.raises(ValueError, match="Expected `weights` to"): r.mean(weights=xp.ones((2, 2))) # Test wrong axis with pytest.raises(ValueError, match=r"axis .* is out of bounds"): r.mean(axis=3) with pytest.raises(ValueError, match=r"axis .* is out of bounds"): r.mean(axis=(-1, 2)) with pytest.raises(ValueError, match="`axis` must be None, int, or tuple of ints."): r.mean(axis="0") with pytest.raises(ValueError, match=r"axis .* is out of bounds"): r.mean(axis=-12) @make_xp_test_case(Rotation.reduce) def test_reduction_no_indices(xp): r = Rotation.from_quat(xp.asarray([0.0, 0.0, 0.0, 1.0])) result = r.reduce(return_indices=False) assert isinstance(result, Rotation) @make_xp_test_case(Rotation.reduce) def test_reduction_none_indices(xp): r = Rotation.from_quat(xp.asarray([0.0, 0.0, 0.0, 1.0])) result = r.reduce(return_indices=True) assert type(result) is tuple assert len(result) == 3 reduced, left_best, right_best = result assert left_best is None assert right_best is None @make_xp_test_case(Rotation.reduce, Rotation.inv, Rotation.magnitude) def test_reduction_scalar_calculation(xp): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-6 rng = np.random.default_rng(146972845698875399755764481408308808739) l_np = Rotation.random(5, rng=rng) r_np = Rotation.random(10, rng=rng) p_np = Rotation.random(7, rng=rng) l = rotation_to_xp(l_np, xp) r = rotation_to_xp(r_np, xp) p = rotation_to_xp(p_np, xp) reduced, left_best, right_best = p.reduce(l, r, return_indices=True) # Loop implementation of the vectorized calculation in Rotation.reduce scalars = np.zeros((len(l_np), len(p_np), len(r_np))) for i, li in enumerate(l_np): for j, pj in enumerate(p_np): for k, rk in enumerate(r_np): scalars[i, j, k] = np.abs((li * pj * rk).as_quat()[3]) scalars = np.reshape(np.moveaxis(scalars, 1, 0), (scalars.shape[1], -1)) max_ind = np.argmax(np.reshape(scalars, (len(p), -1)), axis=1) left_best_check = xp.asarray(max_ind // len(r)) right_best_check = xp.asarray(max_ind % len(r)) assert xp.all(left_best == left_best_check) assert xp.all(right_best == right_best_check) reduced_check = l[left_best_check] * p * r[right_best_check] mag = (reduced.inv() * reduced_check).magnitude() xp_assert_close(mag, xp.zeros(len(p)), atol=atol) @make_xp_test_case(Rotation.from_matrix, Rotation.apply) def test_apply_single_rotation_single_point(xp): dtype = xpx.default_dtype(xp) mat = xp.asarray([ [0, -1, 0], [1, 0, 0], [0, 0, 1] ]) r_1d = Rotation.from_matrix(mat) r_2d = Rotation.from_matrix(xp.expand_dims(mat, axis=0)) v_1d = xp.asarray([1.0, 2, 3], dtype=dtype) v_2d = xp.expand_dims(v_1d, axis=0) v1d_rotated = xp.asarray([-2.0, 1, 3], dtype=dtype) v2d_rotated = xp.expand_dims(v1d_rotated, axis=0) xp_assert_close(r_1d.apply(v_1d), v1d_rotated) xp_assert_close(r_1d.apply(v_2d), v2d_rotated) xp_assert_close(r_2d.apply(v_1d), v2d_rotated) xp_assert_close(r_2d.apply(v_2d), v2d_rotated) v1d_inverse = xp.asarray([2.0, -1, 3], dtype=dtype) v2d_inverse = xp.expand_dims(v1d_inverse, axis=0) xp_assert_close(r_1d.apply(v_1d, inverse=True), v1d_inverse) xp_assert_close(r_1d.apply(v_2d, inverse=True), v2d_inverse) xp_assert_close(r_2d.apply(v_1d, inverse=True), v2d_inverse) xp_assert_close(r_2d.apply(v_2d, inverse=True), v2d_inverse) @make_xp_test_case(Rotation.from_matrix, Rotation.apply) @pytest.mark.parametrize("ndim", range(1, 4)) def test_apply_single_rotation_multiple_points(xp, ndim: int): dtype = xpx.default_dtype(xp) mat = xp.asarray([ [0, -1, 0], [1, 0, 0], [0, 0, 1] ]) r1 = Rotation.from_matrix(mat) r2 = Rotation.from_matrix(xp.expand_dims(mat, axis=0)) rng = np.random.default_rng(0) batch_shape = (ndim,) * (ndim - 1) v = xp.asarray(rng.normal(size=batch_shape + (2, 3)), dtype=dtype) v_rotated = xp.stack([-v[..., 1], v[..., 0], v[..., 2]], axis=-1) xp_assert_close(r1.apply(v), v_rotated) xp_assert_close(r2.apply(v), v_rotated) v_inverse = xp.stack([v[..., 1], -v[..., 0], v[..., 2]], axis=-1) xp_assert_close(r1.apply(v, inverse=True), v_inverse) xp_assert_close(r2.apply(v, inverse=True), v_inverse) @make_xp_test_case(Rotation.from_matrix, Rotation.apply) @pytest.mark.parametrize("ndim", range(1, 5)) def test_apply_multiple_rotations_single_point(xp, ndim: int): dtype = xpx.default_dtype(xp) mat = np.empty((2, 3, 3)) mat[0] = np.array([ [0, -1, 0], [1, 0, 0], [0, 0, 1] ]) mat[1] = np.array([ [1, 0, 0], [0, 0, -1], [0, 1, 0] ]) mat = xp.asarray(mat, dtype=dtype) batch_shape = (ndim,) * (ndim - 1) mat = xp.tile(mat, batch_shape + (1, 1, 1)) r = Rotation.from_matrix(mat) v1 = xp.asarray([1, 2, 3]) v2 = xp.expand_dims(v1, axis=0) v_rotated = xp.asarray([[-2.0, 1, 3], [1, -3, 2]]) v_rotated = xp.tile(v_rotated, batch_shape + (1, 1)) xp_assert_close(r.apply(v1), v_rotated) xp_assert_close(r.apply(v2), v_rotated) v_inverse = xp.asarray([[2.0, -1, 3], [1, 3, -2]]) v_inverse = xp.tile(v_inverse, batch_shape + (1, 1)) xp_assert_close(r.apply(v1, inverse=True), v_inverse) xp_assert_close(r.apply(v2, inverse=True), v_inverse) @make_xp_test_case(Rotation.from_matrix, Rotation.apply) @pytest.mark.parametrize("ndim", range(1, 5)) def test_apply_multiple_rotations_multiple_points(xp, ndim: int): dtype = xpx.default_dtype(xp) mat = np.empty((2, 3, 3)) mat[0] = np.array([ [0, -1, 0], [1, 0, 0], [0, 0, 1] ]) mat[1] = np.array([ [1, 0, 0], [0, 0, -1], [0, 1, 0] ]) mat = xp.asarray(mat, dtype=dtype) batch_shape = (ndim,) * (ndim - 1) mat = xp.tile(mat, batch_shape + (1, 1, 1)) r = Rotation.from_matrix(mat) v = xp.asarray([[1, 2, 3], [4, 5, 6]], dtype=dtype) v_rotated = xp.asarray([[-2.0, 1, 3], [4, -6, 5]], dtype=dtype) v_rotated = xp.tile(v_rotated, batch_shape + (1, 1)) xp_assert_close(r.apply(v), v_rotated) v_inverse = xp.asarray([[2.0, -1, 3], [4, 6, -5]], dtype=dtype) v_inverse = xp.tile(v_inverse, batch_shape + (1, 1)) xp_assert_close(r.apply(v, inverse=True), v_inverse) @make_xp_test_case(Rotation.apply) def test_apply_shapes(xp): rng = np.random.default_rng(0) # Broadcast shape: (6, 5, 4, 2) ( + (3,) for vectors, + (4,) for rotations) vector_shapes = [(), (1,), (2,), (1, 2), (5, 1, 2)] rot_shapes = [(), (1,), (2,), (1, 2), (4, 2), (1, 4, 2), (5, 4, 2), (6, 5, 4, 2)] for q_shape, v_shape in product(rot_shapes, vector_shapes): v = xp.asarray(rng.normal(size=v_shape + (3,))) q = xp.asarray(rng.normal(size=q_shape + (4,))) r = Rotation.from_quat(q) shape = np.broadcast_shapes(q_shape, v_shape) + (3,) x = r.apply(v) assert x.shape == shape x = r.apply(v, inverse=True) assert x.shape == shape def test_apply_array_like(): rng = np.random.default_rng(123) # Single vector r = Rotation.random(rng=rng) t = rng.uniform(-100, 100, size=(3,)) v = r.apply(t.tolist()) v_expected = r.apply(t) xp_assert_close(v, v_expected, atol=1e-12) # Multiple vectors t = rng.uniform(-100, 100, size=(2, 3)) v = r.apply(t.tolist()) v_expected = r.apply(t) xp_assert_close(v, v_expected, atol=1e-12) @make_xp_test_case(Rotation.apply) def test_apply_input_validation(xp): r = Rotation.from_quat(xp.ones(4)) with pytest.raises(ValueError, match="Expected input of shape"): r.apply(xp.ones(2)) with pytest.raises(ValueError, match="Expected input of shape"): r.apply(xp.ones((2, 2))) r = Rotation.from_quat(xp.ones((2, 4))) with pytest.raises(ValueError, match="Cannot broadcast"): r.apply(xp.ones((3, 3))) r = Rotation.from_quat(xp.ones((1, 7, 2, 4))) with pytest.raises(ValueError, match="Cannot broadcast"): r.apply(xp.ones((2, 2, 3))) @make_xp_test_case(Rotation.from_matrix, Rotation.as_matrix, Rotation.__getitem__) @pytest.mark.parametrize("ndim", range(1, 4)) def test_getitem(xp, ndim: int): rng = np.random.default_rng(0) quat = rng.normal(size=(2, ) + (ndim,) * (ndim - 1) + (4,)) mat = xp.asarray(Rotation.from_quat(quat).as_matrix()) r = Rotation.from_matrix(mat) xp_assert_close(r[0].as_matrix(), mat[0, ...], atol=1e-15) xp_assert_close(r[1].as_matrix(), mat[1, ...], atol=1e-15) xp_assert_close(r[:-1].as_matrix(), xp.expand_dims(mat[0, ...], axis=0), atol=1e-15) @make_xp_test_case(Rotation.__getitem__) def test_getitem_single(xp): with pytest.raises(TypeError, match='not subscriptable'): Rotation.from_quat(xp.asarray([0, 0, 0, 1]))[0] @make_xp_test_case(Rotation.from_matrix, Rotation.__getitem__, Rotation.as_matrix) def test_getitem_array_like(): mat = np.array([[[0.0, -1, 0], [1, 0, 0], [0, 0, 1]], [[1, 0, 0], [0, 0, -1], [0, 1, 0]]]) r = Rotation.from_matrix(mat) xp_assert_close(r[[0]].as_matrix(), mat[[0]], atol=1e-15) xp_assert_close(r[[0, 1]].as_matrix(), mat[[0, 1]], atol=1e-15) @make_xp_test_case(Rotation.__setitem__) def test_setitem_single(xp): r = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) with pytest.raises(TypeError, match='not subscriptable'): r[0] = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) @make_xp_test_case(Rotation.__setitem__) def test_setitem_slice(xp): rng = np.random.default_rng(146972845698875399755764481408308808739) r1 = rotation_to_xp(Rotation.random(10, rng=rng), xp) r2 = rotation_to_xp(Rotation.random(5, rng=rng), xp) r1[1:6] = r2 xp_assert_equal(r1[1:6].as_quat(), r2.as_quat()) # Multiple dimensions r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(3, 5, 4)))) r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 5, 4)))) r1[1:3] = r2 xp_assert_equal(r1[1:3].as_quat(), r2.as_quat()) @make_xp_test_case(Rotation.__setitem__) def test_setitem_integer(xp): rng = np.random.default_rng(146972845698875399755764481408308808739) r1 = rotation_to_xp(Rotation.random(10, rng=rng), xp) r2 = rotation_to_xp(Rotation.random(rng=rng), xp) r1[1] = r2 xp_assert_equal(r1[1].as_quat(), r2.as_quat()) # Multiple dimensions r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(3, 5, 4)))) r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(5, 4)))) r1[1] = r2 xp_assert_equal(r1[1].as_quat(), r2.as_quat()) @make_xp_test_case(Rotation.__setitem__) def test_setitem_wrong_type(xp): r = rotation_to_xp(Rotation.random(10, rng=0), xp) with pytest.raises(TypeError, match='Rotation object'): r[0] = 1 @make_xp_test_case(Rotation.from_matrix) def test_n_rotations(xp): mat = np.empty((2, 3, 3)) mat[0] = np.array([ [0, -1, 0], [1, 0, 0], [0, 0, 1] ]) mat[1] = np.array([ [1, 0, 0], [0, 0, -1], [0, 1, 0] ]) mat = xp.asarray(mat) r = Rotation.from_matrix(mat) assert_equal(len(r), 2) assert_equal(len(r[:-1]), 1) def test_random_rotation(): # No xp testing since random rotations are always using NumPy rng = np.random.default_rng(0) assert_equal(Rotation.random(rng=rng).as_quat().shape, (4,)) assert_equal(Rotation.random(None, rng=rng).as_quat().shape, (4,)) assert_equal(Rotation.random(1, rng=rng).as_quat().shape, (1, 4)) assert_equal(Rotation.random(5, rng=rng).as_quat().shape, (5, 4)) # Shape argument assert_equal(Rotation.random(rng=rng, shape=()).as_quat().shape, (4,)) assert_equal(Rotation.random(rng=rng, shape=(3,)).as_quat().shape, (3, 4)) assert_equal(Rotation.random(rng=rng, shape=(2, 3)).as_quat().shape, (2, 3, 4)) # Values should be the same for num=prod(shape) rng1, rng2 = np.random.default_rng(42), np.random.default_rng(42) r_num = Rotation.random(6, rng=rng1) r_shape = Rotation.random(rng=rng2, shape=(2, 3)) xp_assert_close(r_num.as_quat(), r_shape.as_quat().reshape(6, 4), atol=1e-12) # Errors with pytest.raises(ValueError, match="Only one of `num` or `shape` can be"): Rotation.random(num=3,rng=rng, shape=(2, 2)) with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"): Rotation.random(rng=rng, shape=2.5) with pytest.raises(TypeError, match="takes from 0 to 2 positional arguments"): Rotation.random(1, rng, None) # Shape should be kwarg only @make_xp_test_case(Rotation.align_vectors, Rotation.as_matrix) def test_align_vectors_no_rotation(xp): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-5 x = xp.asarray([[1, 2, 3], [4, 5, 6]], dtype=dtype) y = xp.asarray(x, copy=True) r, rssd = Rotation.align_vectors(x, y) xp_assert_close(r.as_matrix(), xp.eye(3), atol=atol) xp_assert_close(rssd, xp.asarray(0.0)[()], check_shape=False, atol=1e-6) @make_xp_test_case(Rotation.apply, Rotation.align_vectors) def test_align_vectors_no_noise(xp): dtype = xpx.default_dtype(xp) atol = 1e-7 if dtype == xp.float64 else 2e-3 rng = np.random.default_rng(14697284569885399755764481408308808739) c = rotation_to_xp(Rotation.random(rng=rng), xp) b = xp.asarray(rng.normal(size=(5, 3)), dtype=dtype) a = c.apply(b) est, rssd = Rotation.align_vectors(a, b) xp_assert_close(c.as_quat(), est.as_quat()) xp_assert_close(rssd, xp.asarray(0.0)[()], check_shape=False, atol=atol) @make_xp_test_case(Rotation.align_vectors, Rotation.apply) def test_align_vectors_improper_rotation(xp): dtype = xpx.default_dtype(xp) atol = 1e-7 if dtype == xp.float64 else 1e-3 # Tests correct logic for issue #10444 x = xp.asarray([[0.89299824, -0.44372674, 0.0752378], [0.60221789, -0.47564102, -0.6411702]]) y = xp.asarray([[0.02386536, -0.82176463, 0.5693271], [-0.27654929, -0.95191427, -0.1318321]]) est, rssd = Rotation.align_vectors(x, y) xp_assert_close(x, est.apply(y), atol=1e-6) xp_assert_close(rssd, xp.asarray(0.0)[()], check_shape=False, atol=atol) @make_xp_test_case(Rotation.align_vectors) def test_align_vectors_rssd_sensitivity(xp): rssd_expected = xp.asarray(0.141421356237308)[()] sens_expected = xp.asarray([[0.2, 0. , 0.], [0. , 1.5, 1.], [0. , 1. , 1.]]) atol = 1e-6 a = xp.asarray([[0, 1, 0], [0, 1, 1], [0, 1, 1]]) b = xp.asarray([[1, 0, 0], [1, 1.1, 0], [1, 0.9, 0]]) rot, rssd, sens = Rotation.align_vectors(a, b, return_sensitivity=True) xp_assert_close(rssd, rssd_expected, atol=atol) xp_assert_close(sens, sens_expected, atol=atol) @make_xp_test_case(Rotation.align_vectors, Rotation.as_matrix) def test_align_vectors_scaled_weights(xp): n = 10 a = xp.asarray(Rotation.random(n, rng=0).apply([1, 0, 0])) b = xp.asarray(Rotation.random(n, rng=1).apply([1, 0, 0])) scale = 2 est1, rssd1, cov1 = Rotation.align_vectors(a, b, xp.ones(n), True) est2, rssd2, cov2 = Rotation.align_vectors(a, b, scale * xp.ones(n), True) xp_assert_close(est1.as_matrix(), est2.as_matrix()) xp_assert_close(math.sqrt(scale) * rssd1, rssd2, atol=1e-6) xp_assert_close(cov1, cov2) @make_xp_test_case(Rotation.apply, Rotation.from_rotvec, Rotation.align_vectors) def test_align_vectors_noise(xp): dtype = xpx.default_dtype(xp) rng = np.random.default_rng(146972845698875399755764481408308808739) n_vectors = 100 rot = rotation_to_xp(Rotation.random(rng=rng), xp) vectors = xp.asarray(rng.normal(size=(n_vectors, 3)), dtype=dtype) result = rot.apply(vectors) # The paper adds noise as independently distributed angular errors sigma = np.deg2rad(1) tolerance = 1.5 * sigma noise = Rotation.from_rotvec( xp.asarray(rng.normal(size=(n_vectors, 3), scale=sigma), dtype=dtype) ) # Attitude errors must preserve norm. Hence apply individual random # rotations to each vector. noisy_result = noise.apply(result) est, rssd, cov = Rotation.align_vectors(noisy_result, vectors, return_sensitivity=True) # Use rotation compositions to find out closeness error_vector = (rot * est.inv()).as_rotvec() xp_assert_close(error_vector[0], xp.asarray(0.0)[()], atol=tolerance) xp_assert_close(error_vector[1], xp.asarray(0.0)[()], atol=tolerance) xp_assert_close(error_vector[2], xp.asarray(0.0)[()], atol=tolerance) # Check error bounds using covariance matrix cov *= xp.asarray(sigma) xp_assert_close(cov[0, 0], xp.asarray(0.0)[()], atol=tolerance) xp_assert_close(cov[1, 1], xp.asarray(0.0)[()], atol=tolerance) xp_assert_close(cov[2, 2], xp.asarray(0.0)[()], atol=tolerance) rssd_check = xp.sum((noisy_result - est.apply(vectors)) ** 2) ** 0.5 xp_assert_close(rssd, rssd_check, check_shape=False) @make_xp_test_case(Rotation.align_vectors) def test_align_vectors_invalid_input(xp): with pytest.raises(ValueError, match="Expected input `a` to have shape"): a, b = xp.asarray([1, 2, 3, 4]), xp.asarray([1, 2, 3]) Rotation.align_vectors(a, b) with pytest.raises(ValueError, match="Expected input `b` to have shape"): a, b = xp.asarray([1, 2, 3]), xp.asarray([1, 2, 3, 4]) Rotation.align_vectors(a, b) with pytest.raises(ValueError, match="Expected inputs `a` and `b` " "to have same shapes"): a, b = xp.asarray([[1, 2, 3], [4, 5, 6]]), xp.asarray([[1, 2, 3]]) Rotation.align_vectors(a, b) with pytest.raises(ValueError, match="Expected `weights` to be 1 dimensional"): a, b = xp.asarray([[1, 2, 3]]), xp.asarray([[1, 2, 3]]) weights = xp.asarray([[1]]) Rotation.align_vectors(a, b, weights) with pytest.raises(ValueError, match="Expected `weights` to have number of values"): a, b = xp.asarray([[1, 2, 3], [4, 5, 6]]), xp.asarray([[1, 2, 3], [4, 5, 6]]) weights = xp.asarray([1, 2, 3]) Rotation.align_vectors(a, b, weights) a, b = xp.asarray([[1, 2, 3]]), xp.asarray([[1, 2, 3]]) weights = xp.asarray([-1]) if is_lazy_array(weights): r, rssd = Rotation.align_vectors(a, b, weights) assert xp.all(xp.isnan(r.as_quat())), "Quaternion should be nan" assert xp.isnan(rssd), "RSSD should be nan" else: with pytest.raises(ValueError, match="`weights` may not contain negative values"): Rotation.align_vectors(a, b, weights) a, b = xp.asarray([[1, 2, 3], [4, 5, 6]]), xp.asarray([[1, 2, 3], [4, 5, 6]]) weights = xp.asarray([xp.inf, xp.inf]) if is_lazy_array(weights): r, rssd = Rotation.align_vectors(a, b, weights) assert xp.all(xp.isnan(r.as_quat())), "Quaternion should be nan" assert xp.isnan(rssd), "RSSD should be nan" else: with pytest.raises(ValueError, match="Only one infinite weight is allowed"): Rotation.align_vectors(a, b, weights) a, b = xp.asarray([[0, 0, 0]]), xp.asarray([[1, 2, 3]]) if is_lazy_array(a): r, rssd = Rotation.align_vectors(a, b) assert xp.all(xp.isnan(r.as_quat())), "Quaternion should be nan" assert xp.isnan(rssd), "RSSD should be nan" else: with pytest.raises(ValueError, match="Cannot align zero length primary vectors"): Rotation.align_vectors(a, b) a, b = xp.asarray([[1, 2, 3], [4, 5, 6]]), xp.asarray([[1, 2, 3], [4, 5, 6]]) weights = xp.asarray([xp.inf, 1]) if is_lazy_array(a): r, rssd, sens = Rotation.align_vectors(a, b, weights, return_sensitivity=True) assert xp.all(xp.isnan(sens)), "Sensitivity matrix should be nan" else: with pytest.raises(ValueError, match="Cannot return sensitivity matrix"): Rotation.align_vectors(a, b, weights, return_sensitivity=True) a, b = xp.asarray([[1, 2, 3]]), xp.asarray([[1, 2, 3]]) if is_lazy_array(a): r, rssd, sens = Rotation.align_vectors(a, b, return_sensitivity=True) assert xp.all(xp.isnan(sens)), "Sensitivity matrix should be nan" else: with pytest.raises(ValueError, match="Cannot return sensitivity matrix"): Rotation.align_vectors(a, b, return_sensitivity=True) # No broadcast support for align_vectors a, b = xp.asarray([[[1, 2, 3]]]), xp.asarray([[[1, 2, 3]]]) with pytest.raises(ValueError, match="Expected inputs `a` and `b` to have shape"): Rotation.align_vectors(a, b) @make_xp_test_case(Rotation.align_vectors, Rotation.as_matrix, Rotation.apply) def test_align_vectors_align_constrain(xp): # Align the primary +X B axis with the primary +Y A axis, and rotate about # it such that the +Y B axis (residual of the [1, 1, 0] secondary b vector) # is aligned with the +Z A axis (residual of the [0, 1, 1] secondary a # vector) dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-6 b = xp.asarray([[1, 0, 0], [1, 1, 0]]) a = xp.asarray([[0.0, 1, 0], [0, 1, 1]]) m_expected = xp.asarray([[0.0, 0, 1], [1, 0, 0], [0, 1, 0]]) R, rssd = Rotation.align_vectors(a, b, weights=xp.asarray([xp.inf, 1])) xp_assert_close(R.as_matrix(), m_expected, atol=atol) xp_assert_close(R.apply(b), a, atol=atol) # Pri and sec align exactly xp_assert_close(rssd, xp.asarray(0.0)[()], atol=atol) # Do the same but with an inexact secondary rotation b = xp.asarray([[1, 0, 0], [1, 2, 0]]) rssd_expected = 1.0 R, rssd = Rotation.align_vectors(a, b, weights=xp.asarray([xp.inf, 1])) xp_assert_close(R.as_matrix(), m_expected, atol=atol) xp_assert_close(R.apply(b)[0, ...], a[0, ...], atol=atol) # Only pri aligns exactly assert xpx.isclose(rssd, rssd_expected, atol=atol, xp=xp) a_expected = xp.asarray([[0.0, 1, 0], [0, 1, 2]]) xp_assert_close(R.apply(b), a_expected, atol=atol) # Check random vectors b = xp.asarray([[1, 2, 3], [-2, 3, -1]]) a = xp.asarray([[-1.0, 3, 2], [1, -1, 2]]) rssd_expected = 1.3101595297515016 R, rssd = Rotation.align_vectors(a, b, weights=xp.asarray([xp.inf, 1])) xp_assert_close(R.apply(b)[0, ...], a[0, ...], atol=atol) # Only pri aligns exactly assert xpx.isclose(rssd, rssd_expected, atol=atol, xp=xp) @make_xp_test_case(Rotation.align_vectors, Rotation.as_matrix) def test_align_vectors_near_inf(xp): # align_vectors should return near the same result for high weights as for # infinite weights. rssd will be different with floating point error on the # exactly aligned vector being multiplied by a large non-infinite weight dtype = xpx.default_dtype(xp) if dtype == xp.float32: pytest.skip("Align vectors near inf is numerically unstable in float32") n = 100 mats = [] for i in range(6): mats.append(Rotation.random(n, rng=10 + i).as_matrix()) for i in range(n): # Get random pairs of 3-element vectors a = xp.asarray(np.array([1 * mats[0][i][0], 2 * mats[1][i][0]]), dtype=dtype) b = xp.asarray(np.array([3 * mats[2][i][0], 4 * mats[3][i][0]]), dtype=dtype) R, _ = Rotation.align_vectors(a, b, weights=[1e10, 1]) R2, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) xp_assert_close(R.as_matrix(), R2.as_matrix(), atol=1e-4) for i in range(n): # Get random triplets of 3-element vectors a = xp.asarray(np.array([1*mats[0][i][0], 2*mats[1][i][0], 3*mats[2][i][0]]), dtype=dtype) b = xp.asarray(np.array([4*mats[3][i][0], 5*mats[4][i][0], 6*mats[5][i][0]]), dtype=dtype) R, _ = Rotation.align_vectors(a, b, weights=[1e10, 2, 1]) R2, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 2, 1]) xp_assert_close(R.as_matrix(), R2.as_matrix(), atol=1e-4) @make_xp_test_case(Rotation.align_vectors, Rotation.as_matrix) def test_align_vectors_parallel(xp): atol = 1e-12 a = xp.asarray([[1.0, 0, 0], [0, 1, 0]]) b = xp.asarray([[0.0, 1, 0], [0, 1, 0]]) m_expected = xp.asarray([[0.0, 1, 0], [-1, 0, 0], [0, 0, 1]]) R, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) xp_assert_close(R.as_matrix(), m_expected, atol=atol) R, _ = Rotation.align_vectors(a[0, ...], b[0, ...]) xp_assert_close(R.as_matrix(), m_expected, atol=atol) xp_assert_close(R.apply(b[0, ...]), a[0, ...], atol=atol) b = xp.asarray([[1, 0, 0], [1, 0, 0]]) m_expected = xp.asarray([[1.0, 0, 0], [0, 1, 0], [0, 0, 1]]) R, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) xp_assert_close(R.as_matrix(), m_expected, atol=atol) R, _ = Rotation.align_vectors(a[0, ...], b[0, ...]) xp_assert_close(R.as_matrix(), m_expected, atol=atol) xp_assert_close(R.apply(b[0, ...]), a[0, ...], atol=atol) @make_xp_test_case(Rotation.align_vectors, Rotation.magnitude, Rotation.apply, Rotation.from_rotvec, Rotation.as_rotvec, Rotation.as_matrix) def test_align_vectors_antiparallel(xp): dtype = xpx.default_dtype(xp) # Test exact 180 deg rotation atol = 1e-12 if dtype == xp.float64 else 1e-7 as_to_test = np.array([[[1.0, 0, 0], [0, 1, 0]], [[0, 1, 0], [1, 0, 0]], [[0, 0, 1], [0, 1, 0]]]) bs_to_test = np.array([[-a[0], a[1]] for a in as_to_test]) for a, b in zip(as_to_test, bs_to_test): a, b = xp.asarray(a, dtype=dtype), xp.asarray(b, dtype=dtype) R, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) xp_assert_close(R.magnitude(), xp.asarray(xp.pi)[()], atol=atol) xp_assert_close(R.apply(b[0, ...]), a[0, ...], atol=atol) # Test exact rotations near 180 deg Rs = Rotation.random(100, rng=0) dRs = Rotation.from_rotvec(Rs.as_rotvec()*1e-4) # scale down to small angle a = [[ 1, 0, 0], [0, 1, 0]] b = [[-1, 0, 0], [0, 1, 0]] as_to_test = [] for dR in dRs: as_to_test.append(np.array([dR.apply(a[0]), a[1]])) # GPU computations may be less accurate. See e.g. # https://github.com/jax-ml/jax/issues/18934 and # https://docs.pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html # We currently have no unified way to check which device is being used. Cupy will # always run on the GPU regardless of SCIPY_DEVICE, hence the explicit check for # cupy. Note that the current implementation lets other frameworks, e.g. numpy, run # on the CPU regardless of SCIPY_DEVICE but with increased GPU tolerances. if xp_device_type(xp.asarray(0)) == "cuda": atol = 1e-7 for a in as_to_test: a, b = xp.asarray(a), xp.asarray(b) R, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) R2, _ = Rotation.align_vectors(a, b, weights=[1e10, 1]) xp_assert_close(R.as_matrix(), R2.as_matrix(), atol=atol) @make_xp_test_case(Rotation.align_vectors, Rotation.apply) def test_align_vectors_primary_only(xp): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-5 mats_a = Rotation.random(100, rng=0).as_matrix() mats_b = Rotation.random(100, rng=1).as_matrix() for mat_a, mat_b in zip(mats_a, mats_b): # Get random 3-element unit vectors a = xp.asarray(mat_a[0], dtype=dtype) b = xp.asarray(mat_b[0], dtype=dtype) # Compare to align_vectors with primary only R, rssd = Rotation.align_vectors(a, b) xp_assert_close(R.apply(b), a, atol=atol) xp_assert_close(rssd, xp.asarray(0.0)[()], atol=atol) def test_align_vectors_array_like(): rng = np.random.default_rng(123) c = Rotation.random(rng=rng) b = rng.normal(size=(5, 3)) a = c.apply(b) est_expected, rssd_expected = Rotation.align_vectors(a, b) est, rssd = Rotation.align_vectors(a.tolist(), b.tolist()) xp_assert_close(est_expected.as_quat(), est.as_quat()) xp_assert_close(rssd, rssd_expected) @make_xp_test_case(Rotation.apply, Rotation.align_vectors) def test_align_vectors_mixed_dtypes(xp): dtype = xpx.default_dtype(xp) rng = np.random.default_rng(123) c = rotation_to_xp(Rotation.random(rng=rng), xp) b = xp.asarray(rng.normal(size=(5, 3)), dtype=dtype) a = xp.asarray(c.apply(b), dtype=xp.float32) # Intentionally float32 # Check that the dtype of the output is the result type of a and b est, _ = Rotation.align_vectors(a, b) xp_assert_close(est.as_quat(), c.as_quat()) def test_repr_single_rotation(xp): q = xp.asarray([0, 0, 0, 1]) actual = repr(Rotation.from_quat(q)) if is_numpy(xp): expected = """\ Rotation.from_matrix(array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))""" assert actual == expected else: assert actual.startswith("Rotation.from_matrix(") def test_repr_rotation_sequence(xp): q = xp.asarray([[0.0, 1, 0, 1], [0, 0, 1, 1]]) / math.sqrt(2) actual = f"{Rotation.from_quat(q)!r}" if is_numpy(xp): expected = """\ Rotation.from_matrix(array([[[ 0., 0., 1.], [ 0., 1., 0.], [-1., 0., 0.]], [[ 0., -1., 0.], [ 1., 0., 0.], [ 0., 0., 1.]]]))""" def stripped(s: str) -> str: # don't fail due to leading whitespace differences return "\n".join(map(str.lstrip, s.splitlines())) assert stripped(actual) == stripped(expected) else: assert actual.startswith("Rotation.from_matrix(") @make_xp_test_case(Slerp.__init__, Slerp.__call__) def test_slerp(xp): rnd = np.random.RandomState(0) key_rots = Rotation.from_quat(xp.asarray(rnd.uniform(size=(5, 4)))) key_quats = key_rots.as_quat() key_times = [0, 1, 2, 3, 4] interpolator = Slerp(key_times, key_rots) assert isinstance(interpolator.times, type(xp.asarray(0))) times = [0, 0.5, 0.25, 1, 1.5, 2, 2.75, 3, 3.25, 3.60, 4] interp_rots = interpolator(times) interp_quats = interp_rots.as_quat() # Dot products are affected by sign of quaternions mask = (interp_quats[:, -1] < 0)[:, None] interp_quats = xp.where(mask, -interp_quats, interp_quats) # Checking for quaternion equality, perform same operation mask = (key_quats[:, -1] < 0)[:, None] key_quats = xp.where(mask, -key_quats, key_quats) # Equality at keyframes, including both endpoints xp_assert_close(interp_quats[0, ...], key_quats[0, ...]) xp_assert_close(interp_quats[3, ...], key_quats[1, ...]) xp_assert_close(interp_quats[5, ...], key_quats[2, ...]) xp_assert_close(interp_quats[7, ...], key_quats[3, ...]) xp_assert_close(interp_quats[10, ...], key_quats[4, ...]) # Constant angular velocity between keyframes. Check by equating # cos(theta) between quaternion pairs with equal time difference. cos_theta1 = xp.sum(interp_quats[0, ...] * interp_quats[2, ...]) cos_theta2 = xp.sum(interp_quats[2, ...] * interp_quats[1, ...]) xp_assert_close(cos_theta1, cos_theta2) cos_theta4 = xp.sum(interp_quats[3, ...] * interp_quats[4, ...]) cos_theta5 = xp.sum(interp_quats[4, ...] * interp_quats[5, ...]) xp_assert_close(cos_theta4, cos_theta5) # theta1: 0 -> 0.25, theta3 : 0.5 -> 1 # Use double angle formula for double the time difference cos_theta3 = xp.sum(interp_quats[1, ...] * interp_quats[3, ...]) xp_assert_close(cos_theta3, 2 * (cos_theta1**2) - 1) # Miscellaneous checks assert_equal(len(interp_rots), len(times)) @make_xp_test_case(Slerp.__init__) def test_slerp_rot_is_rotation(xp): with pytest.raises(TypeError, match="must be a `Rotation` instance"): r = xp.asarray([[1,2,3,4], [0,0,0,1]]) t = xp.asarray([0, 1]) Slerp(t, r) SLERP_EXCEPTION_MESSAGE = "must be a sequence of at least 2 rotations" @make_xp_test_case(Slerp.__init__) def test_slerp_single_rot(xp): r = Rotation.from_quat(xp.asarray([[1.0, 2, 3, 4]])) with pytest.raises(ValueError, match=SLERP_EXCEPTION_MESSAGE): Slerp([1], r) @make_xp_test_case(Slerp.__init__) def test_slerp_rot_len0(xp): r = Rotation.random() r = Rotation.from_quat(xp.asarray(r.as_quat())) with pytest.raises(ValueError, match=SLERP_EXCEPTION_MESSAGE): Slerp([], r) @make_xp_test_case(Slerp.__init__) def test_slerp_rot_len1(xp): r = Rotation.random(1) r = Rotation.from_quat(xp.asarray(r.as_quat())) with pytest.raises(ValueError, match=SLERP_EXCEPTION_MESSAGE): Slerp([1], r) @make_xp_test_case(Slerp.__init__) def test_slerp_tensor_rot(xp): r = Rotation.from_quat(xp.ones((2, 2, 4))) with pytest.raises(ValueError, match="Rotations with more than 1 leading"): Slerp([1, 2], r) @make_xp_test_case(Slerp.__init__) def test_slerp_time_dim_mismatch(xp): with pytest.raises(ValueError, match="times to be specified in a 1 dimensional array"): rnd = np.random.RandomState(0) r = Rotation.from_quat(xp.asarray(rnd.uniform(size=(2, 4)))) t = xp.asarray([[1], [2]]) Slerp(t, r) @make_xp_test_case(Slerp.__init__) def test_slerp_num_rotations_mismatch(xp): with pytest.raises(ValueError, match="number of rotations to be equal to " "number of timestamps"): rnd = np.random.RandomState(0) r = Rotation.from_quat(xp.asarray(rnd.uniform(size=(5, 4)))) t = xp.arange(7) Slerp(t, r) @make_xp_test_case(Slerp.__init__) def test_slerp_equal_times(xp): rnd = np.random.RandomState(0) q = xp.asarray(rnd.uniform(size=(5, 4))) r = Rotation.from_quat(q) t = [0, 1, 2, 2, 4] if is_lazy_array(q): s = Slerp(t, r) assert xp.all(xp.isnan(s.times)) else: with pytest.raises(ValueError, match="strictly increasing order"): Slerp(t, r) @make_xp_test_case(Slerp.__init__) def test_slerp_decreasing_times(xp): rnd = np.random.RandomState(0) q = xp.asarray(rnd.uniform(size=(5, 4))) r = Rotation.from_quat(q) t = [0, 1, 3, 2, 4] if is_lazy_array(q): s = Slerp(t, r) assert xp.all(xp.isnan(s.times)) else: with pytest.raises(ValueError, match="strictly increasing order"): Slerp(t, r) @make_xp_test_case(Slerp.__init__, Slerp.__call__) def test_slerp_call_time_dim_mismatch(xp): rnd = np.random.RandomState(0) r = Rotation.from_quat(xp.asarray(rnd.uniform(size=(5, 4)))) t = xp.arange(5) s = Slerp(t, r) with pytest.raises(ValueError, match="`times` must be at most 1-dimensional."): interp_times = xp.asarray([[3.5], [4.2]]) s(interp_times) @make_xp_test_case(Slerp.__init__, Slerp.__call__) def test_slerp_call_time_out_of_range(xp): rnd = np.random.RandomState(0) r = Rotation.from_quat(xp.asarray(rnd.uniform(size=(5, 4)))) t = xp.arange(5) + 1 s = Slerp(t, r) times_low = xp.asarray([0, 1, 2]) times_high = xp.asarray([1, 2, 6]) if is_lazy_array(times_low): q = s(times_low).as_quat() in_range = xp.logical_and(times_low >= xp.min(t), times_low <= xp.max(t)) assert xp.all(xp.isnan(q[~in_range, ...])) assert xp.all(~xp.isnan(q[in_range, ...])) q = s(times_high).as_quat() in_range = xp.logical_and(times_high >= xp.min(t), times_high <= xp.max(t)) assert xp.all(xp.isnan(q[~in_range, ...])) assert xp.all(~xp.isnan(q[in_range, ...])) else: with pytest.raises(ValueError, match="times must be within the range"): s(times_low) with pytest.raises(ValueError, match="times must be within the range"): s(times_high) @make_xp_test_case(Slerp.__init__, Slerp.__call__, Rotation.from_euler, Rotation.inv, Rotation.magnitude) def test_slerp_call_scalar_time(xp): dtype = xpx.default_dtype(xp) atol = 1e-16 if dtype == xp.float64 else 1e-7 r = Rotation.from_euler('X', xp.asarray([[0], [80]]), degrees=True) s = Slerp([0, 1], r) r_interpolated = s(0.25) r_interpolated_expected = Rotation.from_euler('X', xp.asarray(20), degrees=True) delta = r_interpolated * r_interpolated_expected.inv() xp_assert_close(delta.magnitude(), xp.asarray(0.0)[()], atol=atol) @make_xp_test_case(Rotation.__mul__) def test_multiplication(xp): r1 = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) r2 = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) r3 = r1 * r2 xp_assert_close(r3.as_quat(), xp.asarray([0.0, 0, 0, 1])) # Check that multiplication with other types fails with pytest.raises(TypeError, match="unsupported operand type"): r1 * 2 # Check that __mul__ returns NotImplemented so that other types can implement # __rmul__. See https://github.com/scipy/scipy/issues/21541 assert r1.__mul__(1) is NotImplemented @make_xp_test_case(Rotation.__mul__) def test_multiplication_nd(xp): # multiple dimensions rng = np.random.default_rng(0) r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 3, 4)))) r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 3, 4)))) r3 = r1 * r2 assert r3.as_quat().shape == (2, 3, 4) # same shape len, different dimensions r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(1, 3, 4)))) r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 1, 4)))) r3 = r1 * r2 assert r3.as_quat().shape == (2, 3, 4) # different shape len, different dimensions r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(3, 1, 4, 4)))) r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 4, 4)))) r3 = r1 * r2 assert r3.as_quat().shape == (3, 2, 4, 4) # transition between 2D and 3D with 2D rotation as first argument. This needs to # choose the xp_backend even though r1's backend is cython r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 4)))) r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 2, 4)))) r3 = r1 * r2 assert r3.as_quat().shape == (2, 2, 4) @make_xp_test_case(Rotation.__mul__) def test_multiplication_errors(xp): rng = np.random.default_rng(0) r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 4)))) r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(1, 4, 4)))) with pytest.raises(ValueError, match="Cannot broadcast"): r1 * r2 @make_xp_test_case(Rotation.__mul__) def test_multiplication_stability(xp): qs = rotation_to_xp(Rotation.random(50, rng=0), xp) rs = rotation_to_xp(Rotation.random(1000, rng=1), xp) expected = xp.ones(len(rs)) for r in qs: rs = rs * r * rs xp_assert_close(xp_vector_norm(rs.as_quat(), axis=1), expected) @make_xp_test_case(Rotation.inv, Rotation.__pow__, Rotation.inv, Rotation.magnitude, Rotation.from_rotvec, Rotation.as_rotvec) @pytest.mark.parametrize("ndim", range(1, 4)) def test_pow(xp, ndim: int): dtype = xpx.default_dtype(xp) atol = 1e-14 if dtype == xp.float64 else 1e-6 rng = np.random.default_rng(0) batch_shape = (ndim,) * (ndim - 1) quat = rng.normal(size=batch_shape + (4,)) p = Rotation.from_quat(xp.asarray(quat)) p_inv = p.inv() # Test the short-cuts and other integers for n in [-5, -2, -1, 0, 1, 2, 5]: # Test accuracy q = p ** n q_identity = xp.asarray([0., 0, 0, 1]) r = Rotation.from_quat(xp.tile(q_identity, batch_shape + (1,))) for _ in range(abs(n)): if n > 0: r = r * p else: r = r * p_inv ang = (q * r.inv()).magnitude() assert xp.all(ang < atol) # Test shape preservation r = Rotation.from_quat(xp.tile(q_identity, batch_shape + (1,))) assert (r**n).as_quat().shape == batch_shape + (4,) # Large angle fractional for n in [-1.5, -0.5, -0.0, 0.0, 0.5, 1.5]: q = p ** n r = Rotation.from_rotvec(n * p.as_rotvec()) xp_assert_close(q.as_quat(), r.as_quat(), atol=atol) # Array exponent n = [-5, -2, -1.5, -1, -0.5, -0.0, 0, 0.0, 0.5, 1.0, 1.5, 2] for exponent in n: r = p ** exponent r_array = p ** xp.asarray([exponent]) # Test with 1D array xp_assert_close(r.as_quat(), r_array.as_quat()) r_array = p ** xp.asarray(exponent) # Test with scalar xp_assert_close(r.as_quat(), r_array.as_quat()) # Small angle rotvec = xp.zeros(batch_shape + (3,)) rotvec = xpx.at(rotvec)[..., 0].set(1e-12) p = Rotation.from_rotvec(rotvec) n = 3 q = p ** n r = Rotation.from_rotvec(n * p.as_rotvec()) xp_assert_close(q.as_quat(), r.as_quat(), atol=atol) # Array exponent q = p ** xp.asarray([n]) # Test with 1D array r = Rotation.from_rotvec(n * p.as_rotvec()) xp_assert_close(q.as_quat(), r.as_quat(), atol=atol) q = p ** xp.asarray(n) # Test with scalar r = Rotation.from_rotvec(n * p.as_rotvec()) xp_assert_close(q.as_quat(), r.as_quat(), atol=atol) @make_xp_test_case(Rotation.__pow__) def test_pow_errors(xp): p = rotation_to_xp(Rotation.random(rng=0), xp) with pytest.raises(NotImplementedError, match='modulus not supported'): pow(p, 1, 1) with pytest.raises(ValueError, match="Array exponent must be a scalar"): p ** xp.asarray([1, 2]) with pytest.raises(ValueError, match="Array exponent must be a scalar"): p ** xp.asarray([[1], [2]]) def test_rotation_within_numpy_array(): # TODO: Do we want to support this for all Array API frameworks? single = Rotation.random(rng=0) multiple = Rotation.random(2, rng=1) array = np.array(single) assert_equal(array.shape, ()) array = np.array(multiple) assert_equal(array.shape, (2,)) xp_assert_close(array[0].as_matrix(), multiple[0].as_matrix()) xp_assert_close(array[1].as_matrix(), multiple[1].as_matrix()) array = np.array([single]) assert_equal(array.shape, (1,)) assert_equal(array[0], single) array = np.array([multiple]) assert_equal(array.shape, (1, 2)) xp_assert_close(array[0, 0].as_matrix(), multiple[0].as_matrix()) xp_assert_close(array[0, 1].as_matrix(), multiple[1].as_matrix()) array = np.array([single, multiple], dtype=object) assert_equal(array.shape, (2,)) assert_equal(array[0], single) assert_equal(array[1], multiple) array = np.array([multiple, multiple, multiple]) assert_equal(array.shape, (3, 2)) @make_xp_test_case(Rotation.as_matrix) @pytest.mark.skip_xp_backends("array_api_strict", reason="array API doesn't support pickling") def test_pickling(xp): r = Rotation.from_quat(xp.asarray([0, 0, math.sin(np.pi/4), math.cos(np.pi/4)])) pkl = pickle.dumps(r) unpickled = pickle.loads(pkl) xp_assert_close(r.as_matrix(), unpickled.as_matrix(), atol=1e-15) @make_xp_test_case(Rotation.as_matrix) @pytest.mark.skip_xp_backends("array_api_strict", reason="array API doesn't support deepcopy") def test_deepcopy(xp): r = Rotation.from_quat(xp.asarray([0, 0, math.sin(np.pi/4), math.cos(np.pi/4)])) r1 = copy.deepcopy(r) xp_assert_close(r.as_matrix(), r1.as_matrix(), atol=1e-15) def test_as_euler_contiguous(): # The Array API does not specify contiguous arrays, so we can only check for NumPy r = Rotation.from_quat([0, 0, 0, 1]) e1 = r.as_euler('xyz') # extrinsic euler rotation e2 = r.as_euler('XYZ') # intrinsic assert e1.flags['C_CONTIGUOUS'] is True assert e2.flags['C_CONTIGUOUS'] is True assert all(i >= 0 for i in e1.strides) assert all(i >= 0 for i in e2.strides) @make_xp_test_case(Rotation.concatenate) def test_concatenate(xp): rotation = rotation_to_xp(Rotation.random(10, rng=0), xp) sizes = [1, 2, 3, 1, 3] starts = [0] + list(np.cumsum(sizes)) split = [rotation[i:i + n] for i, n in zip(starts, sizes)] result = Rotation.concatenate(split) xp_assert_equal(rotation.as_quat(), result.as_quat()) # Test Rotation input for multiple rotations result = Rotation.concatenate(rotation) xp_assert_equal(rotation.as_quat(), result.as_quat()) # Test that a copy is returned assert rotation is not result # Test Rotation input for single rotations rng = np.random.default_rng(0) quat = xp.asarray(rng.normal(size=(5, 2, 4))) rotation = Rotation.from_quat(quat) r1 = Rotation.from_quat(quat[:3, ...]) r2 = Rotation.from_quat(quat[3:, ...]) result = Rotation.concatenate([r1, r2]) xp_assert_equal(rotation.as_quat(), result.as_quat()) @make_xp_test_case(Rotation.concatenate) def test_concatenate_wrong_type(xp): with pytest.raises(TypeError, match='Rotation objects only'): rot = Rotation(xp.asarray(Rotation.identity().as_quat())) Rotation.concatenate([rot, 1, None]) @make_xp_test_case(Rotation.concatenate) def test_concatenate_wrong_shape(xp): r1 = Rotation.from_quat(xp.ones((5, 2, 4))) r2 = Rotation.from_quat(xp.ones((1, 4))) # Frameworks throw inconsistent error types on concat failures with pytest.raises((ValueError, RuntimeError, TypeError)): Rotation.concatenate([r1, r2]) # Regression test for gh-16663 @make_xp_test_case() def test_len_and_bool(xp): rotation_multi_one = Rotation(xp.asarray([[0, 0, 0, 1]])) rotation_multi = Rotation(xp.asarray([[0, 0, 0, 1], [0, 0, 0, 1]])) rotation_single = Rotation(xp.asarray([0, 0, 0, 1])) assert len(rotation_multi_one) == 1 assert len(rotation_multi) == 2 with pytest.raises(TypeError, match="Single rotation has no len()."): len(rotation_single) rotation_batched = Rotation.from_quat(xp.ones((3, 2, 4))) assert len(rotation_batched) == 3 # Rotation should always be truthy. See gh-16663 assert rotation_multi_one assert rotation_multi assert rotation_single @make_xp_test_case(Rotation.from_davenport) def test_from_davenport_single_rotation(xp): axis = xp.asarray([0, 0, 1]) quat = Rotation.from_davenport(axis, 'extrinsic', 90, degrees=True).as_quat() expected_quat = xp.asarray([0.0, 0, 1, 1]) / math.sqrt(2) xp_assert_close(quat, expected_quat) @make_xp_test_case(Rotation.from_rotvec, Rotation.from_davenport) def test_from_davenport_one_or_two_axes(xp): ez = xp.asarray([0.0, 0, 1]) ey = xp.asarray([0.0, 1, 0]) # Single rotation, single axis, axes.shape == (3, ) rot = Rotation.from_rotvec(ez * xp.pi/4) rot_dav = Rotation.from_davenport(ez, 'e', xp.pi/4) xp_assert_close(rot.as_quat(canonical=True), rot_dav.as_quat(canonical=True)) # Single rotation, single axis, axes.shape == (1, 3), angles.shape == (1, ) # -> Still single rotation axes = xp.reshape(ez, (1, 3)) # Torch can't create tensors from xp.asarray([ez]) rot = Rotation.from_rotvec(ez * xp.pi/4) rot_dav = Rotation.from_davenport(axes, 'e', [xp.pi/4]) xp_assert_close(rot.as_quat(canonical=True), rot_dav.as_quat(canonical=True)) # Single rotation, two axes, axes.shape == (2, 3) axes = xp.stack([ez, ey], axis=0) rot = Rotation.from_rotvec(axes * xp.asarray([[xp.pi/4], [xp.pi/6]])) rot = rot[0] * rot[1] axes_dav = xp.stack([ey, ez], axis=0) rot_dav = Rotation.from_davenport(axes_dav, 'e', [xp.pi/6, xp.pi/4]) xp_assert_close(rot.as_quat(canonical=True), rot_dav.as_quat(canonical=True)) # Two rotations, single axis, axes.shape == (3, ) axes = xp.stack([ez, ez], axis=0) rot = Rotation.from_rotvec(axes * xp.asarray([[xp.pi/6], [xp.pi/4]])) axes_dav = xp.reshape(ez, (1, 3)) rot_dav = Rotation.from_davenport(axes_dav, 'e', [[xp.pi/6], [xp.pi/4]]) xp_assert_close(rot.as_quat(canonical=True), rot_dav.as_quat(canonical=True)) @make_xp_test_case(Rotation.from_davenport) @pytest.mark.parametrize("ndim", range(1, 4)) def test_from_davenport_shapes(xp, ndim: int): # The shape rules for ND rotations are as follows: # axes.shape[-2] must be angles.shape[-1] # Resulting shape is np.broadcast_shapes(axes.shape[:-2], angles.shape[:-1]) + (4,) rng = np.random.default_rng(0) batch_shape = (ndim,) * (ndim - 1) # Create random, orthogonal axes r = Rotation.from_quat(xp.asarray(rng.normal(size=(4,)))) axes = r.as_matrix() # axes = (3,) angles = xp.asarray(rng.normal(size=batch_shape + (1,))) rot = Rotation.from_davenport(axes[0, ...], 'e', angles) assert rot.as_quat().shape == batch_shape + (4,) # axes = (1, 3) angles = xp.asarray(rng.normal(size=batch_shape + (1,))) rot = Rotation.from_davenport(axes[0, None, ...], 'e', angles) assert rot.as_quat().shape == batch_shape + (4,) # axes = (2, 3) angles = xp.asarray(rng.normal(size=batch_shape + (2,))) rot = Rotation.from_davenport(axes[:2, ...], 'e', angles) assert rot.as_quat().shape == batch_shape + (4,) # axes = (...,3, 3) r = Rotation.from_quat(xp.asarray(rng.normal(size=batch_shape + (4,)))) axes = r.as_matrix() angles = xp.asarray(rng.normal(size=batch_shape + (3,))) rot = Rotation.from_davenport(axes, 'e', angles) assert rot.as_quat().shape == batch_shape + (4,) @make_xp_test_case(Rotation.from_davenport) def test_from_davenport_broadcast(xp): rng = np.random.default_rng(0) # Create random, orthogonal axes r = Rotation.from_quat(xp.asarray(rng.normal(size=(4, 9, 1, 4)))) axes = r.as_matrix() angles = xp.asarray(rng.normal(size=(1, 4, 3))) rot = Rotation.from_davenport(axes, 'e', angles) # (4, 9, 1, 3) + (3,) axes, (1, 4, 3) angles -> (4, 9, 4) + (4,) for quaternion assert rot.as_quat().shape == (4, 9, 4, 4) @make_xp_test_case(Rotation.from_davenport) def test_from_davenport_invalid_input(xp): ez = [0, 0, 1] ey = [0, 1, 0] ezy = [0, 1, 1] # We can only raise in non-lazy frameworks. axes = xp.asarray([ez, ezy]) if is_lazy_array(axes): q = Rotation.from_davenport(axes, 'e', [0, 0]).as_quat() assert xp.all(xp.isnan(q)) else: with pytest.raises(ValueError, match="must be orthogonal"): Rotation.from_davenport(axes, 'e', [0, 0]) axes = xp.asarray([ez, ey, ezy]) if is_lazy_array(axes): q = Rotation.from_davenport(axes, 'e', [0, 0, 0]).as_quat() assert xp.all(xp.isnan(q)) else: with pytest.raises(ValueError, match="must be orthogonal"): Rotation.from_davenport(axes, 'e', [0, 0, 0]) with pytest.raises(ValueError, match="order should be"): Rotation.from_davenport(xp.asarray([ez]), 'xyz', [0]) with pytest.raises(ValueError, match="Expected `angles`"): Rotation.from_davenport(xp.asarray([ez, ey, ez]), 'e', [0, 1, 2, 3]) with pytest.raises(ValueError, match="Expected `angles`"): # Too many angles Rotation.from_davenport(xp.asarray(ez), 'e', [0, 1]) with pytest.raises(ValueError, match="Expected `angles`"): # Too few angles Rotation.from_davenport(xp.asarray([ez, ey, ez]), 'e', [0, 1]) def test_from_davenport_array_like(): rng = np.random.default_rng(123) # Single rotation e1 = np.array([1, 0, 0]) e2 = np.array([0, 1, 0]) e3 = np.array([0, 0, 1]) r_expected = Rotation.random(rng=rng) angles = r_expected.as_davenport([e1, e2, e3], 'e') r = Rotation.from_davenport([e1, e2, e3], 'e', angles.tolist()) assert r_expected.approx_equal(r, atol=1e-12) # Multiple rotations r_expected = Rotation.random(2, rng=rng) angles = r_expected.as_davenport([e1, e2, e3], 'e') r = Rotation.from_davenport([e1, e2, e3], 'e', angles.tolist()) assert np.all(r_expected.approx_equal(r, atol=1e-12)) @make_xp_test_case(Rotation.from_davenport, Rotation.as_davenport) def test_as_davenport(xp): dtype = xpx.default_dtype(xp) rnd = np.random.RandomState(0) n = 100 angles = np.empty((n, 3)) angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles_middle = rnd.uniform(low=0, high=np.pi, size=(n,)) angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) lambdas = rnd.uniform(low=0, high=np.pi, size=(20,)) e1 = xp.asarray([1.0, 0, 0]) e2 = xp.asarray([0.0, 1, 0]) for lamb in lambdas: e3 = xp.asarray(Rotation.from_rotvec(lamb*e2).apply(e1)) ax_lamb = xp.stack([e1, e2, e3], axis=0) angles[:, 1] = angles_middle - lamb for order in ['extrinsic', 'intrinsic']: ax = ax_lamb if order == "intrinsic" else xp.flip(ax_lamb, axis=0) rot = Rotation.from_davenport(ax, order, xp.asarray(angles, dtype=dtype)) angles_dav = rot.as_davenport(ax, order) xp_assert_close(angles_dav, xp.asarray(angles, dtype=dtype)) @make_xp_test_case(Rotation.from_davenport, Rotation.as_davenport) def test_as_davenport_nd(xp): rng = np.random.default_rng(0) r = Rotation.from_quat(xp.asarray(rng.normal(size=(4, 9, 1, 4)))) axes = r.as_matrix() # Get orthogonal axes angles = xp.asarray(rng.uniform(low=-np.pi, high=np.pi, size=(4, 9, 1, 3))) angles = xpx.at(angles)[..., 1].set(angles[..., 1] / 2) for order in ['extrinsic', 'intrinsic']: if order == "intrinsic": axes = xp.flip(axes, axis=-2) rot = Rotation.from_davenport(axes, order, angles) angles_dav = rot.as_davenport(axes, order) xp_assert_close(angles_dav, angles) @make_xp_test_case(Rotation.from_davenport, Rotation.as_davenport, Rotation.as_matrix) @pytest.mark.parametrize("suppress_warnings", (False, True)) def test_as_davenport_degenerate(xp, suppress_warnings): dtype = xpx.default_dtype(xp) atol = 1e-12 if dtype == xp.float64 else 1e-6 # Since we cannot check for angle equality, we check for rotation matrix # equality rnd = np.random.RandomState(0) n = 5 angles = np.empty((n, 3)) # symmetric sequences angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles_middle = [rnd.choice([0, np.pi]) for i in range(n)] angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) lambdas = rnd.uniform(low=0, high=np.pi, size=(5,)) e1 = xp.asarray([1.0, 0, 0]) e2 = xp.asarray([0.0, 1, 0]) for lamb in lambdas: e3 = xp.asarray(Rotation.from_rotvec(lamb*e2).apply(e1)) ax_lamb = xp.stack([e1, e2, e3], axis=0) angles[:, 1] = angles_middle - lamb for order in ['extrinsic', 'intrinsic']: ax = ax_lamb if order == 'intrinsic' else xp.flip(ax_lamb, axis=0) rot = Rotation.from_davenport(ax, order, xp.asarray(angles, dtype=dtype)) with maybe_warn_gimbal_lock(not suppress_warnings, xp): angles_dav = rot.as_davenport( ax, order, suppress_warnings=suppress_warnings ) mat_expected = rot.as_matrix() rot_estimated = Rotation.from_davenport(ax, order, angles_dav) mat_estimated = rot_estimated.as_matrix() xp_assert_close(mat_expected, mat_estimated, atol=atol) @make_xp_test_case(Rotation.from_euler, Rotation.from_davenport) def test_compare_from_davenport_from_euler(xp): dtype = xpx.default_dtype(xp) rnd = np.random.RandomState(0) n = 100 angles = np.empty((n, 3)) # symmetric sequences rtol = 1e-12 if dtype == xp.float64 else 1e-5 angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles[:, 1] = rnd.uniform(low=0, high=np.pi, size=(n,)) angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles = xp.asarray(angles, dtype=dtype) for order in ['extrinsic', 'intrinsic']: for seq_tuple in permutations('xyz'): seq = ''.join([seq_tuple[0], seq_tuple[1], seq_tuple[0]]) ax = xp.asarray([basis_vec(i) for i in seq], dtype=dtype) if order == 'intrinsic': seq = seq.upper() eul = Rotation.from_euler(seq, angles) dav = Rotation.from_davenport(ax, order, angles) xp_assert_close(eul.as_quat(canonical=True), dav.as_quat(canonical=True), rtol=rtol) # asymmetric sequences angles = xpx.at(angles)[:, 1].subtract(np.pi / 2) for order in ['extrinsic', 'intrinsic']: for seq_tuple in permutations('xyz'): seq = ''.join(seq_tuple) ax = xp.asarray([basis_vec(i) for i in seq], dtype=dtype) if order == 'intrinsic': seq = seq.upper() eul = Rotation.from_euler(seq, angles) dav = Rotation.from_davenport(ax, order, angles) xp_assert_close(eul.as_quat(), dav.as_quat(), rtol=rtol) @make_xp_test_case(Rotation.from_euler, Rotation.as_euler, Rotation.as_davenport) def test_compare_as_davenport_as_euler(xp): rnd = np.random.RandomState(0) n = 100 angles = np.empty((n, 3)) # symmetric sequences angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) angles[:, 1] = rnd.uniform(low=0, high=np.pi, size=(n,)) angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) for order in ['extrinsic', 'intrinsic']: for seq_tuple in permutations('xyz'): seq = ''.join([seq_tuple[0], seq_tuple[1], seq_tuple[0]]) ax = [basis_vec(i) for i in seq] if order == 'intrinsic': seq = seq.upper() rot = Rotation.from_euler(seq, xp.asarray(angles)) eul = rot.as_euler(seq) dav = rot.as_davenport(xp.asarray(ax), order) xp_assert_close(eul, dav, rtol=1e-12) # asymmetric sequences angles[:, 1] -= np.pi / 2 for order in ['extrinsic', 'intrinsic']: for seq_tuple in permutations('xyz'): seq = ''.join(seq_tuple) ax = [basis_vec(i) for i in seq] if order == 'intrinsic': seq = seq.upper() rot = Rotation.from_euler(seq, xp.asarray(angles)) eul = rot.as_euler(seq) dav = rot.as_davenport(xp.asarray(ax), order) xp_assert_close(eul, dav, rtol=1e-12) @make_xp_test_case(Rotation.from_matrix, Rotation.from_euler, Rotation.from_rotvec, Rotation.from_davenport, Rotation.from_mrp) def test_zero_rotation_construction(xp): r = Rotation.random(num=0) assert len(r) == 0 r_ide = Rotation.identity(num=0) assert len(r_ide) == 0 r_get = Rotation.random(num=3)[[]] assert len(r_get) == 0 r_quat = Rotation.from_quat(xp.zeros((0, 4))) assert len(r_quat) == 0 r_matrix = Rotation.from_matrix(xp.zeros((0, 3, 3))) assert len(r_matrix) == 0 r_euler = Rotation.from_euler("xyz", xp.zeros((0, 3))) assert len(r_euler) == 0 r_vec = Rotation.from_rotvec(xp.zeros((0, 3))) assert len(r_vec) == 0 r_dav = Rotation.from_davenport(xp.eye(3), "extrinsic", xp.zeros((0, 3))) assert len(r_dav) == 0 r_mrp = Rotation.from_mrp(xp.zeros((0, 3))) assert len(r_mrp) == 0 @make_xp_test_case(Rotation.as_matrix, Rotation.as_euler, Rotation.as_rotvec, Rotation.as_mrp, Rotation.as_davenport) def test_zero_rotation_representation(xp): r = Rotation.from_quat(xp.zeros((0, 4))) assert r.as_quat().shape == (0, 4) assert r.as_matrix().shape == (0, 3, 3) assert r.as_euler("xyz").shape == (0, 3) assert r.as_rotvec().shape == (0, 3) assert r.as_mrp().shape == (0, 3) assert r.as_davenport(xp.eye(3), "extrinsic").shape == (0, 3) @make_xp_test_case(Rotation.apply) def test_zero_rotation_array_rotation(xp): r = Rotation.from_quat(xp.zeros((0, 4))) v = xp.asarray([1, 2, 3]) v_rotated = r.apply(v) assert v_rotated.shape == (0, 3) v0 = xp.zeros((0, 3)) v0_rot = r.apply(v0) assert v0_rot.shape == (0, 3) v2 = xp.ones((2, 3)) with pytest.raises( ValueError, match="Cannot broadcast"): r.apply(v2) @make_xp_test_case(Rotation.__mul__) def test_zero_rotation_multiplication(xp): r = Rotation.from_quat(xp.zeros((0, 4))) r_single = Rotation.from_quat(xp.asarray([0.0, 0, 0, 1])) r_mult_left = r * r_single assert len(r_mult_left) == 0 r_mult_right = r_single * r assert len(r_mult_right) == 0 r0 = Rotation.from_quat(xp.zeros((0, 4))) r_mult = r * r0 assert len(r_mult) == 0 r2 = rotation_to_xp(Rotation.random(2), xp) with pytest.raises(ValueError, match="Cannot broadcast"): r0 * r2 with pytest.raises(ValueError, match="Cannot broadcast"): r2 * r0 @make_xp_test_case(Rotation.concatenate) def test_zero_rotation_concatentation(xp): r = Rotation.from_quat(xp.zeros((0, 4))) r0 = Rotation.concatenate([r, r]) assert len(r0) == 0 r1 = Rotation.from_quat(xp.asarray([0.0, 0, 0, 1])) r1 = r.concatenate([r1, r]) assert len(r1) == 1 r3 = rotation_to_xp(Rotation.random(3), xp) r3 = r.concatenate([r3, r]) assert len(r3) == 3 r4 = rotation_to_xp(Rotation.random(4), xp) r4 = r.concatenate([r, r4]) r4 = r.concatenate([r, r4]) assert len(r4) == 4 @make_xp_test_case(Rotation.__pow__) def test_zero_rotation_power(xp): r = Rotation.from_quat(xp.zeros((0, 4))) for pp in [-1.5, -1, 0, 1, 1.5]: pow0 = r**pp assert len(pow0) == 0 @make_xp_test_case(Rotation.inv) def test_zero_rotation_inverse(xp): r = Rotation.from_quat(xp.zeros((0, 4))) r_inv = r.inv() assert len(r_inv) == 0 @make_xp_test_case(Rotation.magnitude) def test_zero_rotation_magnitude(xp): r = Rotation.from_quat(xp.zeros((0, 4))) magnitude = r.magnitude() assert magnitude.shape == (0,) @make_xp_test_case(Rotation.mean) def test_zero_rotation_mean(xp): r = Rotation.from_quat(xp.zeros((0, 4))) with pytest.raises(ValueError, match="Mean of an empty rotation set is undefined."): r.mean() @make_xp_test_case(Rotation.approx_equal) def test_zero_rotation_approx_equal(xp): r = Rotation.from_quat(xp.zeros((0, 4))) r0 = Rotation.from_quat(xp.zeros((0, 4))) assert r.approx_equal(r0).shape == (0,) r1 = Rotation.from_quat(xp.asarray([0.0, 0, 0, 1])) assert r.approx_equal(r1).shape == (0,) r2 = rotation_to_xp(Rotation.random(), xp) assert r2.approx_equal(r).shape == (0,) approx_msg = "Expected broadcastable shapes in both rotations" r3 = rotation_to_xp(Rotation.random(2), xp) with pytest.raises(ValueError, match=approx_msg): r.approx_equal(r3) with pytest.raises(ValueError, match=approx_msg): r3.approx_equal(r) @pytest.mark.skip_xp_backends("jax.numpy", reason="JAX out-of-bounds indexing deviates from numpy") @pytest.mark.skip_xp_backends("dask.array", reason="zero-length arrays have nan-shapes") def test_zero_rotation_get_set(xp): r = Rotation.from_quat(xp.zeros((0, 4))) r_get = r[xp.asarray([], dtype=xp.bool)] assert len(r_get) == 0 r_slice = r[:0] assert len(r_slice) == 0 with pytest.raises(IndexError): r[xp.asarray([0])] with pytest.raises(IndexError): r[xp.asarray([True])] with pytest.raises(IndexError): r[0] = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) @make_xp_test_case(Rotation.__getitem__) def test_boolean_indexes(xp): r = rotation_to_xp(Rotation.random(3), xp) r0 = r[xp.asarray([False, False, False])] assert len(r0) == 0 r1 = r[xp.asarray([False, True, False])] assert len(r1) == 1 r3 = r[xp.asarray([True, True, True])] assert len(r3) == 3 # Multiple dimensions r = Rotation.from_quat(xp.ones((3, 2, 4))) r4 = r[xp.asarray([True, False, False])] assert len(r4) == 1 assert r4.as_quat().shape == (1, 2, 4) with pytest.raises(IndexError): r[xp.asarray([True, True])] @make_xp_test_case(Rotation.__iter__) def test_rotation_iter(xp): r = rotation_to_xp(Rotation.random(3), xp) for i, r_i in enumerate(r): assert isinstance(r_i, Rotation) xp_assert_equal(r_i.as_quat(), r[i].as_quat()) if i > len(r): raise RuntimeError("Iteration exceeded length of rotations") @pytest.mark.parametrize("ndim", range(1, 5)) def test_rotation_shape(xp, ndim: int): shape = tuple(range(2, 2 + ndim)[:ndim - 1]) quat = xp.ones(shape + (4,)) r = Rotation.from_quat(quat) assert r.shape == shape, f"Got {r.shape}, expected {shape}"