"""Array API backend for the `Rotation` class. This module provides generic, functional implementations of the `Rotation` class methods that work with any Array API-compatible backend. """ # Parts of the implementation are adapted from the cython backend and # https://github.com/jax-ml/jax/blob/d695aa4c63ffcebefce52794427c46bad576680c/jax/_src/scipy/spatial/transform.py. import re import warnings from types import EllipsisType import numpy as np from scipy._lib._array_api import ( array_namespace, Array, ArrayLike, is_lazy_array, xp_vector_norm, xp_result_type, xp_promote, is_jax, ) from scipy._lib._util import broadcastable from scipy._lib.array_api_compat import device as xp_device from scipy._lib.array_api_compat import is_array_api_obj import scipy._lib.array_api_extra as xpx def from_quat( quat: Array, normalize: bool = True, copy: bool = True, *, scalar_first: bool = False, ) -> Array: xp = array_namespace(quat) if scalar_first: quat = xp.roll(quat, -1, axis=-1) # Normalize will always copy, so we avoid the extra copy if we normalize if copy and not normalize: quat = xp.asarray(quat, copy=True) if normalize: quat = _normalize_quaternion(quat) return quat def from_matrix(matrix: Array, assume_valid: bool = False) -> Array: xp = array_namespace(matrix) device = xp_device(matrix) if not assume_valid: mask = xp.linalg.det(matrix) <= 0 lazy = is_lazy_array(mask) # Only non-lazy backends raise an error for non-positive determinants. if not lazy and xp.any(mask): ind = int(xp.nonzero(xpx.atleast_nd(mask, ndim=1, xp=xp))[0][0]) raise ValueError( "Non-positive determinant (left-handed or null coordinate frame) in " f"rotation matrix {ind}: {matrix[ind, ...]}." ) elif lazy: matrix = xp.where(mask[..., None, None], xp.nan, matrix) gramians = matrix @ xp.matrix_transpose(matrix) eye = xp.eye(3, dtype=matrix.dtype, device=device) is_orthogonal = xp.all( xpx.isclose(gramians, eye, atol=1e-12, xp=xp), axis=(-2, -1) ) if lazy: # Lazy backends do not support non-concrete boolean indexing or any form of # computation without statically known shapes, so we always compute SVD and # use xp.where to select the result. U, _, Vt = xp.linalg.svd(matrix, full_matrices=False) matrix = xp.where(is_orthogonal[..., None, None], matrix, U @ Vt) elif not xp.all(is_orthogonal): # For eager frameworks, only compute SVD if needed. is_not_orthogonal = ~is_orthogonal U, _, Vt = xp.linalg.svd(matrix[is_not_orthogonal], full_matrices=False) matrix = xpx.at(matrix)[is_not_orthogonal].set(U @ Vt) return _from_matrix_orthogonal(matrix) def _from_matrix_orthogonal(matrix: Array) -> Array: """Convert known orthogonal rotation matrix to quaternion""" xp = array_namespace(matrix) device = xp_device(matrix) matrix_trace = matrix[..., 0, 0] + matrix[..., 1, 1] + matrix[..., 2, 2] decision = xp.stack( [matrix[..., 0, 0], matrix[..., 1, 1], matrix[..., 2, 2], matrix_trace], axis=-1, ) choice = xp.argmax(decision, axis=-1, keepdims=True) quat = xp.empty((*matrix.shape[:-2], 4), dtype=matrix.dtype, device=device) # The Array API does not support mixing integer indexing with ellipsis, so we cannot # follow the same pattern as the cython backend. Instead, we compute each case # explicitly and assemble the final result with `xp.where`. # TODO: Revisit this implementation if the array API supports mixed integer and # ellipsis indexing. # Case 0 quat_0 = xp.stack( [ 1 - matrix_trace[...] + 2 * matrix[..., 0, 0], matrix[..., 1, 0] + matrix[..., 0, 1], matrix[..., 2, 0] + matrix[..., 0, 2], matrix[..., 2, 1] - matrix[..., 1, 2], ], axis=-1, ) quat = xp.where(choice == 0, quat_0, quat) # Case 1 quat_1 = xp.stack( [ matrix[..., 1, 0] + matrix[..., 0, 1], 1 - matrix_trace[...] + 2 * matrix[..., 1, 1], matrix[..., 2, 1] + matrix[..., 1, 2], matrix[..., 0, 2] - matrix[..., 2, 0], ], axis=-1, ) quat = xp.where(choice == 1, quat_1, quat) # Case 2 quat_2 = xp.stack( [ matrix[..., 2, 0] + matrix[..., 0, 2], matrix[..., 2, 1] + matrix[..., 1, 2], 1 - matrix_trace[...] + 2 * matrix[..., 2, 2], matrix[..., 1, 0] - matrix[..., 0, 1], ], axis=-1, ) quat = xp.where(choice == 2, quat_2, quat) # Case 3 quat_3 = xp.stack( [ matrix[..., 2, 1] - matrix[..., 1, 2], matrix[..., 0, 2] - matrix[..., 2, 0], matrix[..., 1, 0] - matrix[..., 0, 1], 1 + matrix_trace[...], ], axis=-1, ) quat = xp.where(choice == 3, quat_3, quat) return _normalize_quaternion(quat) def from_rotvec(rotvec: Array, degrees: bool = False) -> Array: xp = array_namespace(rotvec) if rotvec.shape[-1] != 3: raise ValueError( f"Expected `rot_vec` to have shape (..., 3), got {rotvec.shape}" ) rotvec = _deg2rad(rotvec) if degrees else rotvec angle = xp_vector_norm(rotvec, axis=-1, keepdims=True, xp=xp) small_angle = angle <= 1e-3 angle2 = angle**2 small_scale = 0.5 - angle2 / 48 + angle2**2 / 3840 # We need to handle the case where angle is 0 to avoid division by zero. We use the # value of the Taylor series approximation, but non-branching operations require # that we still divide by the angle. Since we do not use the result where the angle # is close to 0, this is safe. div_angle = angle + xp.asarray(small_angle, dtype=angle.dtype) large_scale = xp.sin(angle / 2) / div_angle scale = xp.where(small_angle, small_scale, large_scale) quat = xp.concat([rotvec * scale, xp.cos(angle / 2)], axis=-1) return quat def from_mrp(mrp: Array) -> Array: xp = array_namespace(mrp) if mrp.shape[-1] != 3: raise ValueError(f"Expected `mrp` to have shape (..., 3), got {mrp.shape}") mrp2_plus_1 = xp.linalg.vecdot(mrp, mrp, axis=-1)[..., None] + 1 q_no_norm = xp.concat([2 * mrp[..., :3], (2 - mrp2_plus_1)], axis=-1) quat = q_no_norm / mrp2_plus_1 return quat def from_euler(seq: str, angles: Array, degrees: bool = False) -> Array: xp = array_namespace(angles) num_axes = len(seq) if num_axes < 1 or num_axes > 3: raise ValueError( "Expected axis specification to be a non-empty " f"string of up to 3 characters, got {seq}" ) intrinsic = re.match(r"^[XYZ]{1,3}$", seq) is not None extrinsic = re.match(r"^[xyz]{1,3}$", seq) is not None if not (intrinsic or extrinsic): raise ValueError( "Expected axes from `seq` to be from ['x', 'y', " f"'z'] or ['X', 'Y', 'Z'], got {seq}" ) if any(seq[i] == seq[i + 1] for i in range(num_axes - 1)): raise ValueError(f"Expected consecutive axes to be different, got {seq}") if degrees: angles = _deg2rad(angles) angles = xpx.atleast_nd(angles, ndim=1, xp=xp) if angles.shape[-1] != num_axes: raise ValueError( "Expected last dimension of `angles` to match number of sequence axes " f"specified, got {angles.shape[-1]}." ) axes = [_elementary_basis_index(x) for x in seq.lower()] q = _elementary_quat_compose(axes, angles, intrinsic) return q def from_davenport( axes: Array, order: str, angles: Array | float, degrees: bool = False ) -> Array: xp = array_namespace(axes) device = xp_device(axes) if order in ["e", "extrinsic"]: # Must be static, cannot be jitted extrinsic = True elif order in ["i", "intrinsic"]: extrinsic = False else: raise ValueError( "order should be 'e'/'extrinsic' for extrinsic sequences or 'i'/" f"'intrinsic' for intrinsic sequences, got {order}" ) if axes.shape[-1] != 3: raise ValueError("Axes must be vectors of length 3.") axes = xpx.atleast_nd(axes, ndim=2, xp=xp) angles = xpx.atleast_nd(angles, ndim=1, xp=xp) num_axes = axes.shape[-2] if num_axes < 1 or num_axes > 3: raise ValueError(f"Expected up to 3 axes, got {num_axes}") axes = axes / xp_vector_norm(axes, axis=-1, keepdims=True, xp=xp) # Check if axes are orthogonal. Shape checks also work for lazy backends. axes_not_orthogonal = xp.zeros(axes.shape[:-2], dtype=xp.bool, device=device) if num_axes > 1: # Cannot be True yet, so we do not need to use xp.logical_or axes_not_orthogonal = axes_not_orthogonal | ( xp.abs(xp.vecdot(axes[..., 0, :], axes[..., 1, :])) > 1e-7 ) if num_axes > 2: axes_not_orthogonal = axes_not_orthogonal | ( xp.abs(xp.vecdot(axes[..., 1, :], axes[..., 2, :])) > 1e-7 ) if not is_lazy_array(axes_not_orthogonal) and xp.any(axes_not_orthogonal): raise ValueError("Consecutive axes must be orthogonal.") else: axes = xp.where(axes_not_orthogonal[..., None, None], xp.nan, axes) if degrees: angles = _deg2rad(angles) if ( not broadcastable(axes.shape[:-1], angles.shape) or axes.shape[-2] != angles.shape[-1] ): raise ValueError( f"Expected `angles` to match number of axes, got {angles.shape} angles " f"and {axes.shape} axes." ) q_shape = angles.shape[:-1] + (4,) q = xp.zeros(q_shape, dtype=angles.dtype, device=xp_device(angles)) q = xpx.at(q)[..., 3].set(1) for i in range(num_axes): qi = from_rotvec(angles[..., i, None] * axes[..., i, :]) q = compose_quat(qi, q) if extrinsic else compose_quat(q, qi) return q def as_quat( quat: Array, canonical: bool = False, *, scalar_first: bool = False ) -> Array: xp = array_namespace(quat) if canonical: quat = _quat_canonical(quat) if scalar_first: quat = xp.roll(quat, 1, axis=-1) return quat def as_matrix(quat: Array) -> Array: xp = array_namespace(quat) x = quat[..., 0] y = quat[..., 1] z = quat[..., 2] w = quat[..., 3] x2 = x * x y2 = y * y z2 = z * z w2 = w * w xy = x * y zw = z * w xz = x * z yw = y * w yz = y * z xw = x * w matrix_elements = [ x2 - y2 - z2 + w2, 2 * (xy - zw), 2 * (xz + yw), 2 * (xy + zw), -x2 + y2 - z2 + w2, 2 * (yz - xw), 2 * (xz - yw), 2 * (yz + xw), -x2 - y2 + z2 + w2, ] matrix = xp.reshape(xp.stack(matrix_elements, axis=-1), (*quat.shape[:-1], 3, 3)) return matrix def as_rotvec(quat: Array, degrees: bool = False) -> Array: xp = array_namespace(quat) quat = _quat_canonical(quat) ax_norm = xp_vector_norm(quat[..., :3], axis=-1, keepdims=True, xp=xp) angle = 2 * xp.atan2(ax_norm, quat[..., 3][..., None]) small_angle = angle <= 1e-3 angle2 = angle**2 small_scale = 2 + angle2 / 12 + 7 * angle2**2 / 2880 # We need to handle the case where sin(angle/2) is 0 to avoid division by zero. We # use the value of the Taylor series approximation, but non-branching operations # require that we still divide by the sin. Since we do not use the result where the # angle is close to 0, adding one to the sin where we discard the result is safe. div_sin = xp.sin(angle / 2.0) + xp.asarray(small_angle, dtype=angle.dtype) large_scale = angle / div_sin scale = xp.where(small_angle, small_scale, large_scale) if degrees: scale = _rad2deg(scale) rotvec = scale * quat[..., :3] return rotvec def as_mrp(quat: Array) -> Array: xp = array_namespace(quat) one = xp.asarray(1.0, device=xp_device(quat), dtype=quat.dtype) sign = xp.where(quat[..., 3, None] < 0, -1, one) denominator = 1.0 + sign * quat[..., 3, None] return sign * quat[..., :3] / denominator def as_euler( quat: Array, seq: str, degrees: bool = False, *, suppress_warnings: bool = False ) -> Array: xp = array_namespace(quat) # Sanitize the sequence if len(seq) != 3: raise ValueError(f"Expected 3 axes, got {seq}.") intrinsic = re.match(r"^[XYZ]{1,3}$", seq) is not None extrinsic = re.match(r"^[xyz]{1,3}$", seq) is not None if not (intrinsic or extrinsic): raise ValueError( "Expected axes from `seq` to be from ['x', 'y', 'z'] or ['X', 'Y', 'Z'], " f"got {seq}" ) if any(seq[i] == seq[i + 1] for i in range(2)): raise ValueError(f"Expected consecutive axes to be different, got {seq}") device = xp_device(quat) axes = [_elementary_basis_index(x) for x in seq.lower()] axes = axes if extrinsic else axes[::-1] i, j, k = axes symmetric = i == k k = 3 - i - j if symmetric else k mask = xp.asarray(symmetric, device=device) sign = xp.asarray((i - j) * (j - k) * (k - i) // 2, dtype=quat.dtype, device=device) # Permute quaternion elements a = xp.where(mask, quat[..., 3], quat[..., 3] - quat[..., j]) b = xp.where(mask, quat[..., i], quat[..., i] + quat[..., k] * sign) c = xp.where(mask, quat[..., j], quat[..., j] + quat[..., 3]) d = xp.where(mask, quat[..., k] * sign, quat[..., k] * sign - quat[..., i]) angles = _get_angles( extrinsic, symmetric, sign, xp.pi / 2, a, b, c, d, suppress_warnings ) return _rad2deg(angles) if degrees else angles def as_davenport( quat: Array, axes: ArrayLike, order: str, degrees: bool = False, *, suppress_warnings: bool = False, ) -> Array: xp = array_namespace(quat) # Check argument validity if order in ["e", "extrinsic"]: extrinsic = True elif order in ["i", "intrinsic"]: extrinsic = False else: raise ValueError( "order should be 'e'/'extrinsic' for extrinsic sequences or 'i'/'intrinsic'" f" for intrinsic sequences, got {order}" ) if axes.shape[-2] != 3: raise ValueError(f"Expected 3 axes, got {axes.shape}.") if axes.shape[-1] != 3: raise ValueError("Axes must be vectors of length 3.") if not broadcastable(axes.shape[:-2], quat.shape[:-1]): raise ValueError( f"Expected `axes` to match number of rotations, got {axes.shape} axes " f"and {quat.shape} rotations." ) # normalize axes axes = axes / xp_vector_norm(axes, axis=-1, keepdims=True, xp=xp) vdot_ax0_ax1 = xp.vecdot(axes[..., 0, :], axes[..., 1, :]) vdot_ax1_ax2 = xp.vecdot(axes[..., 1, :], axes[..., 2, :]) is_invalid = (vdot_ax0_ax1 >= 1e-7) | (vdot_ax1_ax2 >= 1e-7) if is_lazy_array(is_invalid): axes = xp.where(is_invalid[..., None, None], xp.nan, axes) elif xp.any(is_invalid): raise ValueError("Consecutive axes must be orthogonal.") angles = _compute_davenport_from_quat( quat, axes[..., 0, :], axes[..., 1, :], axes[..., 2, :], extrinsic, suppress_warnings, ) if degrees: angles = _rad2deg(angles) return angles def inv(quat: Array) -> Array: return xpx.at(quat)[..., :3].multiply(-1, copy=True) def magnitude(quat: Array) -> Array: xp = array_namespace(quat) sin_q = xp_vector_norm(quat[..., :3], axis=-1, xp=xp) cos_q = xp.abs(quat[..., 3]) angles = 2 * xp.atan2(sin_q, cos_q) return angles def approx_equal( quat: Array, other: Array, atol: float | None = None, degrees: bool = False ) -> Array: if atol is None: if degrees: warnings.warn( "atol must be set to use the degrees flag, defaulting to 1e-8 radians.", stacklevel=2, ) atol = 1e-8 elif degrees: atol = _deg2rad(atol) if not broadcastable(quat.shape, other.shape): raise ValueError( f"Expected broadcastable shapes in both rotations, got {quat.shape[:-1]} " f"rotations in first and {other.shape[:-1]} rotations in second object." ) quat_result = compose_quat(other, inv(quat)) angles = magnitude(quat_result) return angles < atol def mean( quat: Array, weights: ArrayLike | None = None, axis: None | int | tuple[int, ...] = None, ) -> Array: xp = array_namespace(quat) device = xp_device(quat) dtype = xp_result_type(quat, force_floating=True, xp=xp) if quat.shape[0] == 0: raise ValueError("Mean of an empty rotation set is undefined.") # Axis logic: For None, we reduce over all axes. For int, we only reduce over that # axis. For tuple, we reduce over all specified axes. all_axes = tuple(range(quat.ndim - 1)) if axis is None: axis = all_axes elif isinstance(axis, int): axis = (axis,) if not isinstance(axis, tuple): raise ValueError("`axis` must be None, int, or tuple of ints.") # Ensure all axes are within bounds if axis != () and (min(axis) < -(quat.ndim - 1) or max(axis) > (quat.ndim - 2)): raise ValueError( f"axis {axis} is out of bounds for rotation with shape {quat.shape[:-1]}." ) # Ensure all axes are positive and unique axis = tuple(sorted(set(x % (quat.ndim - 1) for x in axis))) lazy = is_lazy_array(quat) # Branching code is okay for checks that include meta info such as shapes and types quat_expand = quat[..., None, :] if weights is None: K = xp.matrix_transpose(quat_expand) @ quat_expand else: weights = xp.asarray(weights, dtype=dtype, device=device) neg_weights = weights < 0 if not lazy and xp.any(neg_weights): raise ValueError("`weights` must be non-negative.") elif lazy: # We cannot check for negative weights because jit code needs to be # non-branching. We return NaN instead weights = xp.where(neg_weights, xp.nan, weights) if not broadcastable(quat.shape[:-1], weights.shape): raise ValueError( "Expected `weights` to be broadcastable to rotation shape, got shape " f"{weights.shape} for {quat.shape[:-1]} rotations." ) # Make sure we can transpose quat weighted_quat = weights[..., None, None] * quat_expand K = xp.matrix_transpose(weighted_quat) @ quat_expand # Move reduction axes to the end keep_axes = tuple(i for i in all_axes if i not in axis) axes_order = keep_axes + axis K_reordered = xp.moveaxis(K, axes_order, all_axes) # Reshape to flatten reduction axes new_shape = K_reordered.shape[: len(keep_axes)] + (-1, 4, 4) K = xp.mean(xp.reshape(K_reordered, new_shape), axis=-3) _, v = xp.linalg.eigh(K) return v[..., -1] def reduce( quat: Array, left: Array | None = None, right: Array | None = None, ) -> tuple[Array, Array | None, Array | None]: if left is None and right is None: return quat, None, None # DECISION: We cannot have variable number of return arguments for jit compiled # functions. We therefore always return the indices, and filter out later. # TOOD: Properly support broadcasting. xp = array_namespace(quat) quat = xpx.atleast_nd(quat, ndim=2, xp=xp) if left is None: left = xp.ones_like(quat) if right is None: right = xp.ones_like(quat) # We want to calculate the real components of q = l * p * r. It can # be shown that: # qs = ls * ps * rs - ls * dot(pv, rv) - ps * dot(lv, rv) # - rs * dot(lv, pv) - dot(cross(lv, pv), rv) # where ls and lv denote the scalar and vector components of l. p = quat ps, pv = _split_rotation(p, xp) ls, lv = _split_rotation(left, xp) rs, rv = _split_rotation(right, xp) # Compute each term without einsum (not accessible in the Array API) # First term: np.einsum("i,j,k", ls, ps, rs) term1 = ls[..., :, None, None] * ps[..., None, :, None] * rs[..., None, None, :] # Second term: np.einsum('i,jx,kx', ls, pv, rv) prv = xp.sum(pv[..., :, None, :] * rv[..., None, :, :], axis=-1) term2 = ls[..., :, None, None] * prv[..., None, :, :] # Third term: np.einsum('ix,j,kx', lv, ps, rv) lrv = xp.sum(lv[..., :, None, :] * rv[..., None, :, :], axis=-1) term3 = ps[..., None, :, None] * lrv[..., :, None, :] # Fourth term: np.einsum('ix,jx,k', lv, pv, rs) lpv = xp.sum(lv[..., :, None, :] * pv[..., None, :, :], axis=-1) term4 = rs[..., None, None, :] * lpv[..., :, :, None] # Fifth term: np.einsum('xyz,ix,jy,kz', e, lv, pv, rv). We want to avoid expanding # the einsum into a 6D tensor to avoid excessive memory usage. Instead, we compute # the cross product between lv and pv and then compute the dot product with rv. # First compute cross products between lv and pv lv_expanded = lv[..., :, None, :] pv_expanded = pv[..., None, :, :] cross_lp = xp.linalg.cross(lv_expanded, pv_expanded) # Then compute dot product with rv term5 = xp.sum(cross_lp[..., :, :, None, :] * rv[..., None, None, :, :], axis=-1) # Combine all terms with proper shape alignment qs = xp.abs(term1 - term2 - term3 - term4 - term5) qs = xp.reshape(xp.moveaxis(qs, 1, 0), (qs.shape[1], -1)) # Find best indices from scalar components max_ind = xp.argmax(xp.reshape(qs, (qs.shape[0], -1)), axis=1) left_best = max_ind // rv.shape[0] right_best = max_ind % rv.shape[0] # Array API limitation: Integer index arrays are only allowed with integer indices # TODO: Can we somehow avoid this? all_idx = xp.reshape(xp.arange(left.shape[-1]), (1, -1)) left_idx = xp.reshape(left_best, (-1, 1)) left = left[left_idx, all_idx] right_idx = xp.reshape(right_best, (-1, 1)) right = right[right_idx, all_idx] # Reduce the rotation using the best indices reduced = compose_quat(left, compose_quat(p, right)) if left is None: left_best = None if right is None: right_best = None return reduced, left_best, right_best def apply(quat: Array, points: Array, inverse: bool = False) -> Array: xp = array_namespace(quat) mat = as_matrix(quat) # We do not have access to einsum. To avoid broadcasting issues, we add a singleton # dimension to the points array and remove it after the operation. points = points[..., None] if not broadcastable(mat.shape, points.shape): raise ValueError( f"Cannot broadcast {quat.shape[:-1]} rotations to {points.shape[:-1]} " "vectors." ) if inverse: # TODO: Replace with .mT once numpy 2.0 is the minimum supported version return (xp.matrix_transpose(mat) @ points)[..., 0] return (mat @ points)[..., 0] def setitem( quat: Array, value: Array, indexer: int | slice | EllipsisType | None ) -> Array: return xpx.at(quat)[indexer, ...].set(value) def align_vectors( a: Array, b: Array, weights: Array | None = None, return_sensitivity: bool = False ) -> tuple[Array, Array, Array]: xp = array_namespace(a) # Check input vectors dtype = xp_result_type(a, b, force_floating=True, xp=xp) a_original = xp.asarray(a, dtype=dtype) b_original = xp.asarray(b, dtype=dtype) a = xpx.atleast_nd(a_original, ndim=2, xp=xp) b = xpx.atleast_nd(b_original, ndim=2, xp=xp) if a.shape[-1] != 3: raise ValueError( f"Expected input `a` to have shape (3,) or (N, 3), got {a_original.shape}" ) if b.shape[-1] != 3: raise ValueError( f"Expected input `b` to have shape (3,) or (N, 3), got {b_original.shape}" ) if a.shape != b.shape: raise ValueError( "Expected inputs `a` and `b` to have same shapes" f", got {a_original.shape} and {b_original.shape} respectively." ) if a.ndim > 2 or b.ndim > 2: # This function does not support broadcasting raise ValueError( "Expected inputs `a` and `b` to have shape (3,) or (N, 3), got " f"{a_original.shape} and {b_original.shape} respectively." ) N = a.shape[0] # Check weights if weights is None: weights = xp.ones(N, device=xp_device(a), dtype=a.dtype) else: weights = xp.asarray(weights, device=xp_device(a), dtype=a.dtype) if weights.ndim != 1: raise ValueError( f"Expected `weights` to be 1 dimensional, got shape {weights.shape}." ) if N > 1 and (weights.shape[0] != N): raise ValueError( "Expected `weights` to have number of values equal to number of input " f"vectors, got {weights.shape[0]} values and {N} vectors." ) # We can only check for negative weights in eager execution models. Lazy # backends will return NaNs instead. negative_weights = weights < 0 if not is_lazy_array(negative_weights) and xp.any(negative_weights): raise ValueError("`weights` may not contain negative values") weights = xp.where(negative_weights, xp.nan, weights) # For the special case of a single vector pair, we use the infinite # weight code path weight_is_inf = xp.asarray([True]) if N == 1 else weights == xp.inf n_inf = xp.sum(xp.astype(weight_is_inf, a.dtype)) # We can only error out on multiple infinite weights or sensitivity return with # infinite weights in eager execution models. Lazy backends will return NaNs. if not is_lazy_array(n_inf): if n_inf > 1: raise ValueError("Only one infinite weight is allowed") if n_inf == 1 and return_sensitivity: raise ValueError( "Cannot return sensitivity matrix with an " "infinite weight or one vector pair" ) weights = xp.where(n_inf > 1, xp.nan, weights) inf_branch = xp.any(weight_is_inf, axis=-1) # DECISION: We cannot compute both branches for all frameworks. There are two main # reasons: # 1. Computing both for eager execution models is expensive. # 2. Some operations will fail when running the unused branch because of numerical # and algorithmical issues. Numpy e.g. will raise an exception when trying to # compute the svd of a matrix with infinite weights. To prevent this, we only # compute the branch that is needed. Lazy backends however require us to take the # full compute graph. Therefore, we use xp.where for lazy backends and a branching # version for eager frameworks. # # Note that we could also solve this by exploiting the externals of xpx.apply_where. # However, we'd have to rely on the implementation details of apply_where, which is # something we should avoid. # See https://github.com/scipy/scipy/pull/22777#discussion_r2028868364 if is_lazy_array(inf_branch): q_opt, rssd, sensitivity = _align_vectors(a, b, weights) q_opt_inf, rssd_inf, sensitivity_inf = _align_vectors_fixed(a, b, weights) q_opt = xp.where(inf_branch, q_opt_inf, q_opt) rssd = xp.where(inf_branch, rssd_inf, rssd) sensitivity = xp.where(inf_branch, sensitivity_inf, sensitivity) else: if xp.any(inf_branch): q_opt, rssd, sensitivity = _align_vectors_fixed(a, b, weights) else: q_opt, rssd, sensitivity = _align_vectors(a, b, weights) return q_opt, rssd, sensitivity def _align_vectors(a: Array, b: Array, weights: Array) -> tuple[Array, Array, Array]: xp = array_namespace(a) device = xp_device(a) B = (weights[:, None] * a).mT @ b u, s, vh = xp.linalg.svd(B) # Correct improper rotation if necessary (as in Kabsch algorithm) neg_det = xp.linalg.det(u @ vh) < 0 s = xpx.at(s)[..., -1].set(xp.where(neg_det, -s[..., -1], s[..., -1])) u = xpx.at(u)[..., :, -1].set(xp.where(neg_det, -u[..., :, -1], u[..., :, -1])) C = u @ vh # DECISION: We cannot branch on the condition because jit code needs to be # non-branching. Hence, we omit the check for uniqueness # (s[1] + s[2] < 1e-16 * s[0]) ssd = xp.sum(weights * xp.sum(b**2 + a**2, axis=-1), axis=-1) - 2 * xp.sum( s, axis=-1 ) rssd = xp.sqrt(xp.maximum(ssd, xp.zeros(1, device=device)))[..., 0] # TODO: We currently need to always compute the sensitivity matrix because lazy code # needs to be non-branching. We should check if compilers can optimize the where # statement (e.g. in jax) and check if we can have an eager version that only # evaluates the branch that is needed. # See xpx.apply_where, issue: https://github.com/data-apis/array-api-extra/pull/141 zeta = (s[..., 0] + s[..., 1]) * (s[..., 1] + s[..., 2]) * (s[..., 2] + s[..., 0]) kappa = s[..., 0] * s[..., 1] + s[..., 1] * s[..., 2] + s[..., 2] * s[..., 0] eye = xp.eye(3, dtype=a.dtype, device=device) sensitivity = xp.mean(weights) / zeta * (kappa * eye + B @ B.mT) q_opt = _from_matrix_orthogonal(C) return q_opt, rssd, sensitivity def _align_vectors_fixed( a: Array, b: Array, weights: Array ) -> tuple[Array, Array, Array]: xp = array_namespace(a) device = xp_device(a) N = a.shape[0] weight_is_inf = xp.asarray([True], device=device) if N == 1 else weights == xp.inf # We cannot use boolean masks for indexing because of jax. For the same reason, we # also cannot use dynamic slices. As a workaround, we roll the array so that the # infinitely weighted vector is at index 0. We then use static slices to get the # primary and secondary vectors. # # Note that argmax fulfils a secondary purpose here: # When we trace this function with jax, this function might get executed even if # weight_is_inf does not have a single valid entry. Argmax will still give us a # valid index which allows us to proceed with the function (the results are # discarded anyways), whereas boolean indexing would result in invalid, zero-shaped # arrays. inf_idx = xp.argmax(xp.astype(weight_is_inf, xp.uint8)) # xp.argmax returns an Array, but the specification of xp.roll does not support # Arrays as shifts. For lazy execution models we cannot convert to int because this # raises a concretization error. However, jax does accept Arrays as shifts which # allows jit compiling the function. Hence we do not convert to int for jax. # See https://github.com/data-apis/array-api/issues/914#issuecomment-2918959918 if not is_jax(xp): inf_idx = int(inf_idx) a_sorted = xp.roll(a, shift=-inf_idx, axis=0) b_sorted = xp.roll(b, shift=-inf_idx, axis=0) weights_sorted = xp.roll(weights, shift=-inf_idx, axis=0) a_pri = a_sorted[0, ...][None, ...] # Effectively [[0], ...] b_pri = b_sorted[0, ...][None, ...] a_pri_norm = xp_vector_norm(a_pri, axis=-1, keepdims=True, xp=xp) b_pri_norm = xp_vector_norm(b_pri, axis=-1, keepdims=True, xp=xp) if not is_lazy_array(a_pri_norm): if xp.any(a_pri_norm == 0) or xp.any(b_pri_norm == 0): raise ValueError("Cannot align zero length primary vectors") # We cannot error out on zero length vectors. We set the norm to NaN to avoid # division by zero and mark the result as invalid. a_pri_norm = xpx.at(a_pri_norm)[a_pri_norm == 0].set(xp.nan) b_pri_norm = xpx.at(b_pri_norm)[b_pri_norm == 0].set(xp.nan) a_pri = a_pri / a_pri_norm b_pri = b_pri / b_pri_norm # We first find the minimum angle rotation between the primary # vectors. cross = xp.linalg.cross(b_pri[..., 0, :], a_pri[..., 0, :]) cross_norm = xp_vector_norm(cross, axis=-1, xp=xp) theta = xp.atan2(cross_norm, xp.sum(a_pri[..., 0, :] * b_pri[..., 0, :], axis=-1)) tolerance = 1e-3 # tolerance for small angle approximation (rad) q_flip = xp.asarray([0.0, 0.0, 0.0, 1.0], dtype=a_pri.dtype, device=device) # Near pi radians, the Taylor series approximation of x/sin(x) diverges, so for # numerical stability we flip pi and then rotate back by the small angle pi - theta flip = xp.pi - theta < tolerance # For antiparallel vectors, cross = [0, 0, 0] so we need to manually set an # arbitrary orthogonal axis of rotation i = xp.argmin(xp.abs(a_pri[..., 0, :])) # TODO: Array API does not support fancy indexing __setitem__. The code is # equivalent to doing the following: # r_components = xp.zeros(3) # r_components = xpx.at(r_components)[i - 1].set(a_pri[0, i - 2]) # r_components = xpx.at(r_components)[i - 2].set(-a_pri[0, i - 1]) r_components = xp.stack( [ xp.where(i == 1, a_pri[0, 2], xp.where(i == 2, -a_pri[0, 1], 0)), xp.where(i == 2, a_pri[0, 0], xp.where(i == 0, -a_pri[0, 2], 0)), xp.where(i == 0, a_pri[0, 1], xp.where(i == 1, -a_pri[0, 0], 0)), ] ) r = xp.where(cross_norm == 0, r_components, cross) q_flip = xp.where(flip, from_rotvec(r / xp_vector_norm(r, xp=xp) * xp.pi), q_flip) theta = xp.where(flip, xp.pi - theta, theta) cross = xp.where(flip, -cross, cross) # Small angle Taylor series approximation for numerical stability theta2 = theta * theta small_scale = xp.abs(theta) < tolerance r_small_scale = cross * (1 + theta2 / 6 + theta2 * theta2 * 7 / 360) # We need to handle the case where theta is 0 to avoid division by zero. We use the # value of the Taylor series approximation, but non-branching operations require # that we still divide by the angle. Since we do not use the result where the angle # is close to 0, this is safe. theta = theta + xp.asarray(small_scale, dtype=theta.dtype) r_large_scale = cross * theta / xp.sin(theta) r = xp.where(small_scale, r_small_scale, r_large_scale) q_pri = compose_quat(from_rotvec(r), q_flip) # Case 1): No secondary vectors, q_opt is q_pri -> Immediately done # Case 2): Secondary vectors exist # We cannot use boolean masks here because of jax a_sec = a_sorted[1:, ...] b_sec = b_sorted[1:, ...] weights_sec = weights_sorted[1:] # We apply the first rotation to the b vectors to align the # primary vectors, resulting in vectors c. The secondary # vectors must now be rotated about that primary vector to best # align c to a. c_sec = apply(q_pri, b_sec) # Calculate vector components for the angle calculation. The # angle phi to rotate a single 2D vector C to align to 2D # vector A in the xy plane can be found with the equation # phi = atan2(cross(C, A), dot(C, A)) # = atan2(|C|*|A|*sin(phi), |C|*|A|*cos(phi)) # The below equations perform the same operation, but with the # 3D rotation restricted to the 2D plane normal to a_pri, where # the secondary vectors are projected into that plane. We then # find the composite angle from the weighted sum of the # axial components in that plane, minimizing the 2D alignment # problem. sin_term = xp.linalg.vecdot(xp.linalg.cross(c_sec, a_sec), a_pri) cos_term = xp.linalg.vecdot(c_sec, a_sec) - ( xp.linalg.vecdot(c_sec, a_pri) * xp.linalg.vecdot(a_sec, a_pri) ) phi = xp.atan2(xp.sum(weights_sec * sin_term), xp.sum(weights_sec * cos_term)) q_sec = from_rotvec(phi * a_pri[0, ...]) # Compose these to get the optimal rotation q_opt = q_pri if N == 1 else compose_quat(q_sec, q_pri) # Calculate the root sum squared distance. We force the error to # be zero for the infinite weight vectors since they will align # exactly. mask = ((N > 1) | (weights[0] == xp.inf)) & weight_is_inf # Skip non-infinite weight single vectors pairs, we used the # infinite weight code path but don't want to zero that weight weights_inf_zero = xpx.at(weights, mask).set(0, copy=True, xp=xp) a_est = apply(q_opt, b) rssd = xp.sqrt(xp.sum(weights_inf_zero @ (a - a_est) ** 2)) mask = xp.any(xp.isnan(weights), axis=-1) q_opt = xp.where(mask, xp.nan, q_opt) return q_opt, rssd, xp.asarray(xp.nan, device=device) def pow(quat: Array, n: float | Array) -> Array: xp = array_namespace(quat) device = xp_device(quat) # If n is an array, we sanitize it to a scalar and promote quat and n to # the same dtype. if is_array_api_obj(n): if n.shape == (1,): n = n[0] elif n.ndim != 0: raise ValueError("Array exponent must be a scalar") quat, n = xp_promote(quat, n, force_floating=True, xp=xp) # If n is a lazy array, we cannot take fast paths for special cases. if is_lazy_array(n): result = from_rotvec(n * as_rotvec(quat)) # general scaling of rotation angle # Special cases 0 -> identity, -1 -> inv, 1 -> copy identity = xp.zeros((*quat.shape[:-1], 4), dtype=quat.dtype, device=device) identity = xpx.at(identity)[..., 3].set(1) result = xp.where(n == 0, identity, result) result = xp.where(n == -1, inv(quat), result) result = xp.where(n == 1, quat, result) return result if n == 0: identity = xp.zeros((*quat.shape[:-1], 4), dtype=quat.dtype, device=device) return xpx.at(identity)[..., 3].set(1) if n == -1: return inv(quat) if n == 1: return quat return from_rotvec(n * as_rotvec(quat)) def _normalize_quaternion(quat: Array) -> Array: xp = array_namespace(quat) quat_norm = xp_vector_norm(quat, axis=-1, keepdims=True, xp=xp) zero_norm = quat_norm == 0 if is_lazy_array(quat_norm): quat = xp.where(zero_norm, xp.nan, quat) elif xp.any(zero_norm): raise ValueError("Found zero norm quaternions in `quat`.") return quat / quat_norm def _quat_canonical(quat: Array) -> Array: xp = array_namespace(quat) mask = quat[..., 3] < 0 zero_w = quat[..., 3] == 0 mask = mask | (zero_w & (quat[..., 0] < 0)) zero_wx = zero_w & (quat[..., 0] == 0) mask = mask | (zero_wx & (quat[..., 1] < 0)) zero_wxy = zero_wx & (quat[..., 1] == 0) mask = mask | (zero_wxy & (quat[..., 2] < 0)) return xp.where(mask[..., None], -quat, quat) def _elementary_basis_index(axis: str) -> int: if axis == "x": return 0 elif axis == "y": return 1 elif axis == "z": return 2 raise ValueError(f"Expected axis to be from ['x', 'y', 'z'], got {axis}") def _compute_davenport_from_quat( quat: Array, n1: Array, n2: Array, n3: Array, extrinsic: bool, suppress_warnings: bool, ) -> Array: # The algorithm assumes extrinsic frame transformations. The algorithm # in the paper is formulated for rotation quaternions, which are stored # directly by Rotation. # Adapt the algorithm for our case by reversing both axis sequence and # angles for intrinsic rotations when needed xp = array_namespace(quat) n1, n3 = (n1, n3) if extrinsic else (n3, n1) n_cross = xp.linalg.cross(n1, n2) lamb = xp.atan2(xp.vecdot(n3, n_cross), xp.vecdot(n3, n1)) # alternative set of angles compatible with as_euler implementation mask = lamb < 0 n2 = xp.where(mask[..., None], -n2, n2) lamb = xp.where(mask, -lamb, lamb) n_cross = xp.where(mask[..., None], -n_cross, n_cross) correct_set = mask quat_lamb = xp.concat( (xp.sin(lamb / 2)[..., None] * n2, xp.cos(lamb / 2)[..., None]), axis=-1 ) q_trans = compose_quat(quat_lamb, quat) a = q_trans[..., 3] b = xp.linalg.vecdot(q_trans[..., :3], n1) c = xp.linalg.vecdot(q_trans[..., :3], n2) d = xp.linalg.vecdot(q_trans[..., :3], n_cross) angles = _get_angles(extrinsic, False, 1, lamb, a, b, c, d, suppress_warnings) angles = xpx.at(angles)[..., 1].set( xp.where(correct_set, -angles[..., 1], angles[..., 1]) ) return angles def _elementary_quat_compose(axes: list[int], angles: Array, intrinsic: bool) -> Array: xp = array_namespace(angles) device = xp_device(angles) quat = _make_elementary_quat(axes[0], angles[..., 0], device=device, xp=xp) for i in range(1, len(axes)): ax_quat = _make_elementary_quat(axes[i], angles[..., i], device=device, xp=xp) quat = compose_quat(quat, ax_quat) if intrinsic else compose_quat(ax_quat, quat) return quat def _make_elementary_quat(axis: int, angle: Array, device, xp) -> Array: quat = xp.zeros((*angle.shape, 4), dtype=angle.dtype, device=device) quat = xpx.at(quat)[..., 3].set(xp.cos(angle / 2.0)) quat = xpx.at(quat)[..., axis].set(xp.sin(angle / 2.0)) return quat def _get_angles( extrinsic: bool, symmetric: bool, sign: int, lamb: float, a: Array, b: Array, c: Array, d: Array, suppress_warnings: bool, ) -> Array: xp = array_namespace(a) device = xp_device(a) eps = 1e-7 half_sum = xp.atan2(b, a) half_diff = xp.atan2(d, c) # We zero-initialize to automatically cover singular cases where the second angle is # not defined uniquely. angles = xp.zeros((*a.shape, 3), dtype=a.dtype, device=device) angles = xpx.at(angles)[..., 1].set(2 * xp.atan2(xp.hypot(c, d), xp.hypot(a, b))) angle_first = 0 if extrinsic else 2 angle_third = 2 if extrinsic else 0 # Check if the second angle is close to 0 or pi, causing a singularity. # - Case 0: Second angle is neither close to 0 nor pi. # - Case 1: Second angle is close to 0. # - Case 2: Second angle is close to pi. case1 = xp.abs(angles[..., 1]) <= eps case2 = xp.abs(angles[..., 1] - xp.pi) <= eps case0 = ~(case1 | case2) if not suppress_warnings and not is_lazy_array(case0) and xp.any(~case0): warnings.warn( "Gimbal lock detected. Setting third angle to zero " "since it is not possible to uniquely determine " "all angles.", stacklevel=3, ) # This writes case1 into a0 where True and case2 everywhere else. This is sound # because we later overwrite any values without singularity with the regular value # of case0. The second angle is covered by default since we zero-initialized the # second dimension. a0 = xp.where(case1, 2 * half_sum, 2 * half_diff * (-1 if extrinsic else 1)) angles = xpx.at(angles)[..., 0].set(a0) # We overwrite the values of angles without singularities (case0) a1 = xp.where(case0, half_sum - half_diff, angles[..., angle_first]) angles = xpx.at(angles)[..., angle_first].set(a1) # Same as above but for the third angle. We overwrite the non-singular case0 values a3 = xp.where(case0, half_sum + half_diff, angles[..., angle_third]) if not symmetric: a3 = a3 * sign angles = xpx.at(angles)[..., 1].set(angles[..., 1] - lamb) angles = xpx.at(angles)[..., angle_third].set(a3) angles = (angles + xp.pi) % (2 * xp.pi) - xp.pi return angles def compose_quat(p: Array, q: Array) -> Array: xp = array_namespace(p) cross = xp.linalg.cross(p[..., :3], q[..., :3]) qx = p[..., 3] * q[..., 0] + q[..., 3] * p[..., 0] + cross[..., 0] qy = p[..., 3] * q[..., 1] + q[..., 3] * p[..., 1] + cross[..., 1] qz = p[..., 3] * q[..., 2] + q[..., 3] * p[..., 2] + cross[..., 2] qw = ( p[..., 3] * q[..., 3] - p[..., 0] * q[..., 0] - p[..., 1] * q[..., 1] - p[..., 2] * q[..., 2] ) quat = xp.stack([qx, qy, qz, qw], axis=-1) return quat def _split_rotation(q: Array, xp) -> tuple[Array, Array]: q = xpx.atleast_nd(q, ndim=2, xp=xp) return q[..., -1], q[..., :-1] def _deg2rad(angles: Array) -> Array: return angles * (np.pi / 180.0) def _rad2deg(angles: Array) -> Array: return angles * (180.0 / np.pi)