robomath 0.1.0__py3-none-any.whl

robomath/__init__.py

@@ -0,0 +1,221 @@
+"""Declare all public usable functions from robomath.
+
+Groups: SO(2), SO(3), SE(2), SE(3), Sim(3) (pose with a scale, the scale last in the
+tangent vector) and SE_2(3) (rotation, velocity and position). The Transform class
+wraps SE(3) only.
+
+Planned: kinematics (screw axes, product of exponentials, forward
+and inverse kinematics, Jacobians), trajectories on SO(3) and SE(3), and covariance
+propagation. Estimators, simulators and ROS code stay outside this package.
+"""
+
+from .convert import (
+    pq_to_trans,
+    quat_wxyz_to_xyzw,
+    quat_xyzw_to_wxyz,
+    swap_twist,
+    trans_to_pq,
+)
+from .euler import euler_zyx_to_matrix, matrix_to_euler_zyx
+from .geodesic import (
+    distance_so3,
+    interp_se2,
+    interp_se3,
+    interp_so3,
+    mean_se2,
+    mean_se3,
+    mean_so3,
+)
+from .quaternions import (
+    axis_angle_to_quat,
+    euler_zyx_to_quat,
+    matrix_to_quat,
+    quat_conjugate,
+    quat_multiply,
+    quat_normalize,
+    quat_to_axis_angle,
+    quat_to_euler_zyx,
+    quat_to_matrix,
+    slerp,
+)
+from .registration import kabsch, umeyama
+from .rotations import inverse, rotx, roty, rotz
+from .se2 import (
+    adjoint_se2,
+    exp_vec_se2,
+    hat_se2,
+    lminus_se2,
+    log_vec_se2,
+    lplus_se2,
+    rminus_se2,
+    rplus_se2,
+    vee_se2,
+)
+from .se3 import (
+    adjoint,
+    angular,
+    exp_vec_se3,
+    hat_se3,
+    linear,
+    lminus_se3,
+    log_vec_se3,
+    lplus_se3,
+    rminus_se3,
+    rp_to_trans,
+    rplus_se3,
+    trans_inv,
+    trans_to_rp,
+    twist,
+    vee_se3,
+)
+from .se23 import (
+    adjoint_se23,
+    exp_vec_se23,
+    hat_se23,
+    lminus_se23,
+    log_vec_se23,
+    lplus_se23,
+    rminus_se23,
+    rplus_se23,
+    rvp_to_se23,
+    se23_inv,
+    se23_to_rvp,
+    vee_se23,
+)
+from .sim3 import (
+    adjoint_sim3,
+    exp_vec_sim3,
+    hat_sim3,
+    lminus_sim3,
+    log_vec_sim3,
+    lplus_sim3,
+    rminus_sim3,
+    rplus_sim3,
+    sim3_inv,
+    sim3_to_srt,
+    srt_to_sim3,
+    vee_sim3,
+)
+from .so2 import exp_so2, log_so2, wrap_angle
+from .so3 import (
+    axis_angle_to_matrix,
+    exp_so3,
+    exp_vec,
+    hat,
+    jl,
+    jl_inv,
+    jr,
+    jr_inv,
+    lminus,
+    log_so3,
+    log_vec,
+    lplus,
+    matrix_to_axis_angle,
+    nearest_rotation,
+    rminus,
+    rplus,
+    vee,
+)
+from .transform import Transform
+
+__all__ = [
+    "axis_angle_to_matrix",
+    "axis_angle_to_quat",
+    "matrix_to_axis_angle",
+    "matrix_to_quat",
+    "quat_conjugate",
+    "quat_multiply",
+    "quat_normalize",
+    "quat_to_axis_angle",
+    "quat_to_matrix",
+    "slerp",
+    "euler_zyx_to_matrix",
+    "inverse",
+    "matrix_to_euler_zyx",
+    "rotx",
+    "roty",
+    "rotz",
+    "hat",
+    "vee",
+    "exp_vec",
+    "log_vec",
+    "exp_so3",
+    "log_so3",
+    "rplus",
+    "rminus",
+    "lplus",
+    "lminus",
+    "jr",
+    "jr_inv",
+    "jl",
+    "jl_inv",
+    "euler_zyx_to_quat",
+    "quat_to_euler_zyx",
+    "nearest_rotation",
+    "Transform",
+    "adjoint_sim3",
+    "exp_vec_sim3",
+    "hat_sim3",
+    "lminus_sim3",
+    "log_vec_sim3",
+    "lplus_sim3",
+    "rminus_sim3",
+    "rplus_sim3",
+    "sim3_inv",
+    "sim3_to_srt",
+    "srt_to_sim3",
+    "vee_sim3",
+    "adjoint_se23",
+    "exp_vec_se23",
+    "hat_se23",
+    "lminus_se23",
+    "log_vec_se23",
+    "lplus_se23",
+    "rminus_se23",
+    "rplus_se23",
+    "rvp_to_se23",
+    "se23_inv",
+    "se23_to_rvp",
+    "vee_se23",
+    "pq_to_trans",
+    "quat_wxyz_to_xyzw",
+    "quat_xyzw_to_wxyz",
+    "swap_twist",
+    "trans_to_pq",
+    "kabsch",
+    "umeyama",
+    "distance_so3",
+    "interp_se2",
+    "interp_se3",
+    "interp_so3",
+    "mean_se2",
+    "mean_se3",
+    "mean_so3",
+    "exp_so2",
+    "log_so2",
+    "wrap_angle",
+    "adjoint_se2",
+    "exp_vec_se2",
+    "hat_se2",
+    "lminus_se2",
+    "log_vec_se2",
+    "lplus_se2",
+    "rminus_se2",
+    "rplus_se2",
+    "vee_se2",
+    "adjoint",
+    "angular",
+    "exp_vec_se3",
+    "hat_se3",
+    "linear",
+    "lminus_se3",
+    "log_vec_se3",
+    "lplus_se3",
+    "rminus_se3",
+    "rp_to_trans",
+    "rplus_se3",
+    "trans_inv",
+    "trans_to_rp",
+    "twist",
+    "vee_se3",
+]

robomath/_util.py

@@ -0,0 +1,16 @@
+"""Robomath - small helpers shared by the other modules."""
+
+import numpy as np
+
+
+def norm(v: np.ndarray) -> np.ndarray:
+    """Length of the vectors along the last axis, shape (...).
+
+    Same result as np.linalg.norm(v, axis=-1)
+    """
+    return np.sqrt(np.einsum("...i,...i->...", v, v))
+
+
+def apply(a: np.ndarray, v: np.ndarray) -> np.ndarray:
+    """Matrix-vector product for stacks: a (..., n, n) times v (..., n)."""
+    return (a @ v[..., None])[..., 0]

robomath/convert.py

@@ -0,0 +1,68 @@
+"""Robomath - Convert between robomath's layouts and those of other libraries.
+
+Inside robomath a quaternion is (x, y, z, w) and a tangent vector is (translation,
+rotation). Other libraries differ, and a mix-up still runs but gives wrong results.
+Each function here changes only the layout, named by the layout and not the library.
+Every function accepts leading batch dimensions.
+
+Quaternion (x, y, z, w), scalar last: scipy, ROS messages, Eigen coeffs(), Pinocchio.
+Quaternion (w, x, y, z), scalar first: Drake, MuJoCo, Ceres.
+Tangent vector (translation, rotation): Sola's paper, Sophus, manif, Pinocchio.
+Tangent vector (rotation, translation): Modern Robotics, Drake, MuJoCo, GTSAM Pose3.
+
+Not here yet: poses with the camera convention
+(OpenCV and COLMAP store world-to-camera, so the inverse of ours), other Euler
+sequences, SE(2) layouts such as GTSAM's, and ROS message classes (this module takes
+plain arrays, so it needs no ROS install).
+"""
+
+import numpy as np
+
+from .quaternions import matrix_to_quat, quat_to_matrix
+from .se3 import angular, linear, rp_to_trans, trans_to_rp, twist
+
+
+def quat_xyzw_to_wxyz(q: np.ndarray) -> np.ndarray:
+    """Move the scalar part from last to first: (x, y, z, w) -> (w, x, y, z)."""
+    return np.roll(np.asarray(q, dtype=float), 1, axis=-1)
+
+
+def quat_wxyz_to_xyzw(q: np.ndarray) -> np.ndarray:
+    """Move the scalar part from first to last: (w, x, y, z) -> (x, y, z, w)."""
+    return np.roll(np.asarray(q, dtype=float), -1, axis=-1)
+
+
+def swap_twist(tau: np.ndarray) -> np.ndarray:
+    """Swap the halves of a tangent vector, (a, b) -> (b, a).
+
+    Turns (translation, rotation) into (rotation, translation) and back; the swap is
+    its own inverse. Works for a tangent vector or a spatial velocity.
+    """
+    tau = np.asarray(tau, dtype=float)
+    return twist(angular(tau), linear(tau))
+
+
+def pq_to_trans(p: np.ndarray, q: np.ndarray) -> np.ndarray:
+    """4x4 pose from a position and an (x, y, z, w) quaternion, as in a ROS Pose.
+
+    Args:
+        p (np.ndarray): Positions, shape (..., 3).
+        q (np.ndarray): Unit quaternions, shape (..., 4).
+
+    Returns:
+        np.ndarray: The poses, shape (..., 4, 4).
+    """
+    return rp_to_trans(quat_to_matrix(q), p)
+
+
+def trans_to_pq(m: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Position and (x, y, z, w) quaternion of a 4x4 pose, as in a ROS Pose.
+
+    Args:
+        m (np.ndarray): Poses, shape (..., 4, 4).
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: (p, q), shapes (..., 3) and (..., 4).
+    """
+    r, t = trans_to_rp(m)
+    return t, matrix_to_quat(r)

robomath/euler.py

@@ -0,0 +1,100 @@
+"""Robomath - ZYX Euler angles.
+
+Every function accepts leading batch dimensions.
+"""
+
+import math
+
+import numpy as np
+
+GIMBAL_TOL = 1e-15
+
+
+def euler_zyx_to_matrix(alpha: float, beta: float, gamma: float) -> np.ndarray:
+    """Rotation matrix for ZYX Euler angles.
+
+    Args:
+        alpha (float): Angle about z, in radians, scalar or shape (...).
+        beta (float): Angle about y, in radians; broadcasts with alpha.
+        gamma (float): Angle about x, in radians; broadcasts with alpha.
+
+    Returns:
+        np.ndarray: The rotation matrix, shape (..., 3, 3),
+            R = Rz(alpha) Ry(beta) Rx(gamma).
+    """
+    alpha, beta, gamma = (np.asarray(a, dtype=float) for a in (alpha, beta, gamma))
+    if alpha.ndim == beta.ndim == gamma.ndim == 0:
+        ca, sa = math.cos(alpha), math.sin(alpha)
+        cb, sb = math.cos(beta), math.sin(beta)
+        cg, sg = math.cos(gamma), math.sin(gamma)
+        return np.array(
+            [
+                [ca * cb, ca * sb * sg - sa * cg, ca * sb * cg + sa * sg],
+                [sa * cb, sa * sb * sg + ca * cg, sa * sb * cg - ca * sg],
+                [-sb, cb * sg, cb * cg],
+            ]
+        )
+    ca, sa = np.cos(alpha), np.sin(alpha)
+    cb, sb = np.cos(beta), np.sin(beta)
+    cg, sg = np.cos(gamma), np.sin(gamma)
+    r = np.empty(np.broadcast_shapes(alpha.shape, beta.shape, gamma.shape) + (3, 3))
+    r[..., 0, 0] = ca * cb
+    r[..., 0, 1] = ca * sb * sg - sa * cg
+    r[..., 0, 2] = ca * sb * cg + sa * sg
+    r[..., 1, 0] = sa * cb
+    r[..., 1, 1] = sa * sb * sg + ca * cg
+    r[..., 1, 2] = sa * sb * cg - ca * sg
+    r[..., 2, 0] = -sb
+    r[..., 2, 1] = cb * sg
+    r[..., 2, 2] = cb * cg
+    return r
+
+
+def matrix_to_euler_zyx(r: np.ndarray) -> tuple[float, float, float]:
+    """ZYX Euler angles of a rotation matrix.
+
+    At gimbal lock (beta = +-pi/2) only one combination of alpha and gamma is
+    fixed by the matrix, so alpha is set to 0 and gamma takes the rest.
+
+    Args:
+        r (np.ndarray): Rotation matrices, shape (..., 3, 3).
+
+    Returns:
+        tuple[float, float, float]: (alpha, beta, gamma) in radians, each of
+            shape (...), with alpha and gamma in (-pi, pi] and beta in
+            [-pi/2, pi/2].
+    """
+    r = np.asarray(r, dtype=float)
+    if r.ndim == 2:
+        (r11, r12, _), (r21, r22, _), (r31, r32, r33) = r.tolist()
+        if abs(r31 + 1) < GIMBAL_TOL:
+            return 0.0, math.pi / 2, math.atan2(r12, r22)
+        if abs(r31 - 1) < GIMBAL_TOL:
+            return 0.0, -math.pi / 2, -math.atan2(r12, r22)
+        return (
+            math.atan2(r21, r11),
+            math.atan2(-r31, math.hypot(r11, r21)),
+            math.atan2(r32, r33),
+        )
+    r11, r12 = r[..., 0, 0], r[..., 0, 1]
+    r21, r22 = r[..., 1, 0], r[..., 1, 1]
+    r31, r32, r33 = r[..., 2, 0], r[..., 2, 1], r[..., 2, 2]
+
+    # Both gimbal-lock cases and the general one are computed, then chosen by
+    # mask, so a batch can mix them.
+    lock_up = np.abs(r31 + 1) < GIMBAL_TOL
+    lock_down = np.abs(r31 - 1) < GIMBAL_TOL
+    locked_gamma = np.arctan2(r12, r22)
+
+    beta = np.where(
+        lock_up,
+        np.pi / 2,
+        np.where(lock_down, -np.pi / 2, np.arctan2(-r31, np.hypot(r11, r21))),
+    )
+    alpha = np.where(lock_up | lock_down, 0.0, np.arctan2(r21, r11))
+    gamma = np.where(
+        lock_up,
+        locked_gamma,
+        np.where(lock_down, -locked_gamma, np.arctan2(r32, r33)),
+    )
+    return alpha[()], beta[()], gamma[()]

robomath/geodesic.py

@@ -0,0 +1,148 @@
+"""Robomath - Distance, interpolation and mean on SO(3) and SE(3).
+
+All of them follow the shortest path (the geodesic) given by the right plus and
+minus of the group: the step from A to B is minus(B, A), and moving a fraction t
+of it is plus(A, t * minus(B, A)). Every function accepts leading batch dimensions.
+"""
+
+import warnings
+
+import numpy as np
+
+from ._util import norm
+from .quaternions import matrix_to_quat, quat_to_matrix
+from .se2 import rminus_se2, rplus_se2
+from .se3 import angular, rminus_se3, rp_to_trans, rplus_se3, trans_to_rp
+from .so2 import exp_so2, log_so2
+from .so3 import rminus, rplus
+
+MEAN_TOL = 1e-13
+MEAN_ITERATIONS = 100
+UNIQUE_ANGLE = np.pi / 2
+
+
+def distance_so3(r1: np.ndarray, r2: np.ndarray) -> np.ndarray:
+    """Angle in [0, pi] of the rotation that takes r1 to r2, shape (...)."""
+    return norm(rminus(r2, r1))
+
+
+def interp_so3(r1: np.ndarray, r2: np.ndarray, t: np.ndarray) -> np.ndarray:
+    """Rotation a fraction t of the way from r1 (t = 0) to r2 (t = 1)."""
+    t = np.asarray(t, dtype=float)
+    return rplus(r1, t[..., None] * rminus(r2, r1))
+
+
+def interp_se2(m1: np.ndarray, m2: np.ndarray, t: np.ndarray) -> np.ndarray:
+    """Planar pose a fraction t of the way from m1 (t = 0) to m2 (t = 1)."""
+    t = np.asarray(t, dtype=float)
+    return rplus_se2(m1, t[..., None] * rminus_se2(m2, m1))
+
+
+def interp_se3(m1: np.ndarray, m2: np.ndarray, t: np.ndarray) -> np.ndarray:
+    """Pose a fraction t of the way from m1 (t = 0) to m2 (t = 1)."""
+    t = np.asarray(t, dtype=float)
+    return rplus_se3(m1, t[..., None] * rminus_se3(m2, m1))
+
+
+def _weighted(w, x):
+    """Weighted sum over the N items: w (..., N) times x (..., N, k) -> (..., k)."""
+    return np.einsum("...n,...nk->...k", w, x)
+
+
+def _start_so3(rotations, w):
+    """Chordal mean: the top eigenvector of the weighted sum of q q^T (Markley).
+
+    No iteration and no local minima, and it does not depend on the item order.
+    Same mean as scipy's Rotation.mean, written here from the method.
+    """
+    q = matrix_to_quat(rotations)
+    k = np.einsum("...n,...ni,...nj->...ij", w, q, q)
+    return quat_to_matrix(np.linalg.eigh(k)[1][..., -1])
+
+
+def _start_se3(poses, w):
+    r, t = trans_to_rp(poses)
+    return rp_to_trans(_start_so3(r, w), _weighted(w, t))
+
+
+def _start_se2(poses, w):
+    r, t = trans_to_rp(poses)
+    theta = log_so2(r)
+    mean_angle = np.arctan2(
+        np.einsum("...n,...n->...", w, np.sin(theta)),
+        np.einsum("...n,...n->...", w, np.cos(theta)),
+    )
+    return rp_to_trans(exp_so2(mean_angle), _weighted(w, t))
+
+
+def _angle_se3(tau):
+    """Rotation angle of each SE(3) tangent vector."""
+    return norm(angular(tau))
+
+
+def _angle_se2(tau):
+    """Rotation angle of each SE(2) tangent vector."""
+    return np.abs(tau[..., 2])
+
+
+def _mean(items, weights, start, plus, minus, angle):
+    """Karcher mean: move the guess by the weighted mean of its offsets to the items.
+
+    items has shape (..., N, a, a); the N items are averaged. The first guess is
+    the chordal mean of the rotations, so the answer does not depend on the item
+    order. angle maps tangent offsets to the rotation angle of each. The mean is
+    unique when all items lie within 90 degrees of it. Beyond that the loop can
+    settle on a local answer, so it warns.
+    """
+    items = np.asarray(items, dtype=float)
+    n = items.shape[-3]
+    if n == 0:
+        raise ValueError("the mean of no items is undefined")
+    w = np.full(n, 1.0 / n) if weights is None else np.asarray(weights, dtype=float)
+    if np.any(w < 0) or np.any(np.sum(w, axis=-1) <= 0):
+        raise ValueError("weights must be non-negative and not all zero")
+    w = w / np.sum(w, axis=-1, keepdims=True)
+    mean = start(items, w)
+    for _ in range(MEAN_ITERATIONS):
+        step = np.sum(w[..., None] * minus(items, mean[..., None, :, :]), axis=-2)
+        mean = plus(mean, step)
+        if np.max(np.abs(step), initial=0.0) < MEAN_TOL:
+            offsets = minus(items, mean[..., None, :, :])
+            if np.max(angle(offsets), initial=0.0) >= UNIQUE_ANGLE:
+                warnings.warn(
+                    "an item is 90 degrees or more from the mean; "
+                    "the mean may not be unique",
+                    stacklevel=3,
+                )
+            return mean
+    raise ValueError("mean did not settle; the items are too spread out")
+
+
+def mean_so3(rotations: np.ndarray, weights: np.ndarray | None = None) -> np.ndarray:
+    """Mean of N rotations, shape (..., N, 3, 3) -> (..., 3, 3).
+
+    Args:
+        rotations (np.ndarray): The rotations to average.
+        weights (np.ndarray | None): Weights, shape (..., N). Equal if None.
+    """
+    return _mean(rotations, weights, _start_so3, rplus, rminus, norm)
+
+
+def mean_se3(poses: np.ndarray, weights: np.ndarray | None = None) -> np.ndarray:
+    """Mean of N poses, shape (..., N, 4, 4) -> (..., 4, 4).
+
+    Args:
+        poses (np.ndarray): The poses to average.
+        weights (np.ndarray | None): Weights, shape (..., N). Equal if None.
+    """
+    return _mean(poses, weights, _start_se3, rplus_se3, rminus_se3, _angle_se3)
+
+
+def mean_se2(poses: np.ndarray, weights: np.ndarray | None = None) -> np.ndarray:
+    """Mean of N planar poses, shape (..., N, 3, 3) -> (..., 3, 3).
+
+    Args:
+        poses (np.ndarray): The poses to average.
+        weights (np.ndarray | None): Weights, shape (..., N). Equal if None.
+    """
+    return _mean(poses, weights, _start_se2, rplus_se2, rminus_se2, _angle_se2)