robolib.utils.mathutils

  1"""mathutils.py - math utilities for robosim. Should be only pure python + numpy."""
  2import math
  3import numpy as np
  4from .utils import logger, Rot3x3, Pose4x4, Point3D, Point6D, Mat4x4, Mat3x3
  5
  6_eps = np.finfo(np.float64).eps
  7
  8# math
  9
 10def norm(x: np.ndarray) -> np.ndarray:
 11    """L2 normalize a vector"""
 12    return (x / np.linalg.norm(x, ord=2)).astype("float32")
 13
 14def rot_from_forward_up(forward: Point3D, up: Point3D) -> Rot3x3:
 15    """right hand rule: x x y = z, y x z = x, z x x = y; z = forward. (3, ) + (3, ) -> (3, 3)"""
 16    z = norm(forward) # first, normalize forward, as we care about it's orientation only. z = ||z||
 17    y_init = np.float32(up) # up = Y
 18    x = norm(np.cross(y_init, z)) # second build x=||y_init x z|| as the cross of existing forward and up
 19    y = np.cross(z, x) # third, get the new y as the cross of the already built x and z.
 20    R = np.column_stack([x, y, z]) # NOTE: somehow it's left-handed and we make it right-handed here
 21    assert is_rot(R), f"\n-{R=}\n-{R@R.T=}"
 22    return R
 23
 24def pose_to_trans_euler(pose: Pose4x4) -> Point6D:
 25    """converts a pose to a 6DoF translation + euler vector, mostly for printing reasons (4, 4) -> (6, )"""
 26    return np.float32([*pose[0:3, 3], *tr2rpy(pose)])
 27
 28def pose_from_trans_euler(position_rpy: Point6D) -> Pose4x4:
 29    """converts a position + a euler vector to a Pose. (6, ) -> (4, 4)"""
 30    res = rpy2tr(position_rpy[-3:]).astype("float32")
 31    res[0:3, 3] = position_rpy[0:3]
 32    return res
 33
 34def renormalize_rotation_matrix(rot: Rot3x3, tol: float=1e10) -> Rot3x3:
 35    """uses SVD to renormalize a rotation matrix if it drifted too much (i.e. during trajectory updates)"""
 36    if is_rot(rot, tol=tol):
 37        return rot
 38    logger.debug(f"Renormalizing R as err is {np.linalg.norm((rot @ rot.T) - np.eye(3)):.7f}")
 39    U, _, Vt = np.linalg.svd(rot)
 40    D = np.diag([1, 1, np.linalg.det(U @ Vt)]) # force det=1 (not -1)
 41    new_rot = U @ D @ Vt
 42    assert is_rot(new_rot, tol=1e10), new_rot
 43    return new_rot
 44
 45def pose_from_position_target_up(position: Point3D, target: Point3D, up: Point3D) -> Pose4x4:
 46    """returns a 4x4 pose matrix from initial position+target+up (3, ) + (3, ) + (3, ) -> (4, 4)"""
 47    pose = np.eye(4, dtype="float32")
 48    pose[0:3, 0:3] = rot_from_forward_up(forward=np.float32(target) - position, up=up)
 49    pose[0:3, 3] = position
 50    return pose
 51
 52def make_radius(size: Point3D) -> float:
 53    """creates a radius from a size and a scale"""
 54    return math.sqrt(size[0]**2 + size[1]**2 + size[2]**2) / 2
 55
 56def is_rot(R: Rot3x3, tol: float = 20) -> bool:
 57    r"""
 58    Test if matrix belongs to SO(n).
 59    Checks orthogonality, ie. :math:`{\bf R} {\bf R}^T = {\bf I}` and :math:`\det({\bf R}) > 0`.
 60    For the first test we check that the norm of the residual is less than ``tol * eps``.
 61        >>> is_rot(np.eye(3))
 62        >>> is_rot(np.zeros((3,3)))
 63    """
 64    return bool(
 65        np.linalg.norm(R @ R.T - np.eye(R.shape[0])) < tol * _eps
 66        and np.linalg.det(R) > 0
 67    )
 68
 69"""
 70Below: trimmed, standalone port of the `spatialmath.base` functions robosim actually uses.
 71
 72Only 5 public functions, in the exact shapes the codebase calls them:
 73- is_rot(R[, tol])   -> mathutils.rot_from_forward_up, renormalize_rotation_matrix
 74- tr2rpy(pose4x4) -> mathutils.pose_to_trans_euler
 75- rpy2tr(rpy3)    -> mathutils.pose_from_trans_euler
 76- bskewa(twist6)   -> physics.motion._integrate_velocity_into_pose
 77- btrexp(se3 4x4)  -> physics.motion._integrate_velocity_into_pose
 78
 79Everything is self-contained (helpers are inner `_`-prefixed functions), radians only,
 80zyx order only, and float64 throughout. Callers downcast to float32 at their boundary.
 81"""
 82
 83def tr2rpy(T: Pose4x4) -> Point3D:
 84    """
 85    Convert SO(3)/SE(3) to roll-pitch-yaw angles (radians, zyx order).
 86        >>> tr2rpy(rpy2tr([0.2, 0.3, 0.5]))
 87    The angles [roll, pitch, yaw] are sequential rotations about the X, Y, Z axes.
 88    If the input is SE(3), the translation component is ignored. There is a singularity
 89    at pitch = ±90° where roll is arbitrarily set to 0.
 90    """
 91    R = T[0:3, 0:3]
 92
 93    rpy = np.zeros(3)
 94    if abs(abs(R[2, 0]) - 1) < 20 * _eps:  # |R31| == 1 -> gimbal-lock singularity
 95        rpy[0] = 0.0  # roll is zero
 96        if R[2, 0] < 0:
 97            rpy[2] = -math.atan2(R[0, 1], R[0, 2])
 98        else:
 99            rpy[2] = math.atan2(-R[0, 1], -R[0, 2])
100        rpy[1] = -math.asin(np.clip(R[2, 0], -1.0, 1.0))
101    else:
102        rpy[0] = math.atan2(R[2, 1], R[2, 2])  # roll
103        rpy[2] = math.atan2(R[1, 0], R[0, 0])  # yaw
104        # pick the numerically largest term to solve pitch, avoiding division by ~0
105        k = int(np.argmax(np.abs([R[0, 0], R[1, 0], R[2, 1], R[2, 2]])))
106        if k == 0:
107            rpy[1] = -math.atan(R[2, 0] * math.cos(rpy[2]) / R[0, 0])
108        elif k == 1:
109            rpy[1] = -math.atan(R[2, 0] * math.sin(rpy[2]) / R[1, 0])
110        elif k == 2:
111            rpy[1] = -math.atan(R[2, 0] * math.sin(rpy[0]) / R[2, 1])
112        else:
113            rpy[1] = -math.atan(R[2, 0] * math.cos(rpy[0]) / R[2, 2])
114
115    return rpy
116
117def rpy2tr(rpy: Point3D) -> Pose4x4:
118    r"""
119    Create an SE(3) matrix from roll-pitch-yaw angles (radians, zyx order).
120        >>> rpy2tr([0.1, 0.2, 0.3])
121    zyx: rotate by yaw about z, then pitch about the new y, then roll about the new x.
122    The translation component is zero.
123    """
124    def _rotx(t: float) -> np.ndarray:
125        ct, st = math.cos(t), math.sin(t)
126        return np.array([[1.0, 0.0, 0.0], [0.0, ct, -st], [0.0, st, ct]], dtype=np.float64)
127
128    def _roty(t: float) -> np.ndarray:
129        ct, st = math.cos(t), math.sin(t)
130        return np.array([[ct, 0.0, st], [0.0, 1.0, 0.0], [-st, 0.0, ct]], dtype=np.float64)
131
132    def _rotz(t: float) -> np.ndarray:
133        ct, st = math.cos(t), math.sin(t)
134        return np.array([[ct, -st, 0.0], [st, ct, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64)
135
136    roll, pitch, yaw = float(rpy[0]), float(rpy[1]), float(rpy[2])
137    T = np.eye(4)
138    T[0:3, 0:3] = _rotz(yaw) @ _roty(pitch) @ _rotx(roll)
139    return T
140
141def skew(w: Point3D) -> Mat3x3:
142    """create a skew symmetric matrix from a 3-vector"""
143    return np.array([
144        [0.0,   -w[2],  w[1]],
145        [w[2],   0.0,  -w[0]],
146        [-w[1],  w[0],  0.0]], dtype=np.float64) # (4, 4)
147
148def bskew(v: Point3D) -> Mat4x4:
149    """batched skew"""
150    assert len(v.shape) == 2 and v.shape[1] == 3, v.shape
151    zero = np.zeros((len(v), ), "float32")
152    res = np.array([
153        [ zero,     -v[:, 2],    v[:, 1]],
154        [ v[:, 2],   zero,      -v[:, 0]],
155        [-v[:, 1],   v[:, 0],    zero,  ]], dtype=np.float64) # (3, 3, N)
156    return np.permute_dims(res, (2, 0, 1)) # (N, 3, 3)
157
158def bvexa(S: Mat4x4) -> Point6D:
159    """recover the twist [v, w] from the se(3) matrix (inverse of bskewa)"""
160    return np.array([
161        S[:, 0, 3], S[:, 1, 3], S[:, 2, 3],
162        (S[:, 2, 1] - S[:, 1, 2]) / 2,
163        (S[:, 0, 2] - S[:, 2, 0]) / 2,
164        (S[:, 1, 0] - S[:, 0, 1]) / 2], dtype=np.float64).T # (N, 6)
165
166def bskewa(v: Point6D) -> Mat4x4:
167    """batched skewa"""
168    assert len(v.shape) == 2 and v.shape[1] == 6, v.shape
169    zero = np.zeros((len(v), ), "float32")
170    res = np.array([
171        [ zero,     -v[:, 5],    v[:, 4],  v[:, 0]],
172        [ v[:, 5],   zero,      -v[:, 3],  v[:, 1]],
173        [-v[:, 4],   v[:, 3],    zero,     v[:, 2]],
174        [ zero,      zero,       zero,     zero]], dtype=np.float64) # (4, 4, N)
175    return np.permute_dims(res, (2, 0, 1)) # (N, 4, 4)
176
177def btrexp(S: Mat4x4) -> Pose4x4:
178    """
179    Batched matrix exponential of an se(3) element -> SE(3) (a screw motion).
180    ``S`` is the batched (Nx)4x4 augmented skew-symmetric matrix produced by ``bskewa``.
181        >>> btrexp(bskewa(np.array([[1, 0, 0, 0, 0, 0]])))
182    """
183    n = len(S)
184    tw = bvexa(S) # (N, 6)
185    v = tw[:, 0:3] # (3, )
186    w = tw[:, 3:6] # (3, )
187
188    w_norm = np.linalg.norm(w, axis=1) # (N, )
189    v_norm = np.linalg.norm(v, axis=1) # (N, )
190    # theta is the screw magnitude: rotation dominates, else it's a pure translation
191    theta = np.where(w_norm >= 20 * _eps, w_norm, v_norm)[:, None] # (N, 1)
192
193    with np.errstate(divide="ignore", invalid="ignore"):   # 0/0 on zero-twist is expected; nan_to_num handles it
194        v = np.nan_to_num(v / theta, nan=0) # (N, 3)
195        w = np.nan_to_num(w / theta, nan=0) # (N, 3)
196
197    skw = bskew(w)                                                      # (N, 3, 3)
198    theta_n = theta[..., None]                                          # (N, 1, 1)
199    st, ct = np.sin(theta_n), np.cos(theta_n)                           # (N, 1, 1), (N, 1, 1)
200    eye = np.eye(3)[None].repeat(n, axis=0)                             # (N, 3, 3)
201    R = eye + st * skw + (1.0 - ct) * skw @ skw                         # (N, 3, 3) Rodrigues
202    V = eye * theta_n + (1.0 - ct) * skw + (theta_n - st) * skw @ skw   # (N, 3, 3) translation map
203
204    T = np.eye(4)[None].repeat(n, axis=0)                               # (N, 4, 4)
205    T[:, 0:3, 0:3] = R
206    T[:, 0:3, 3] = np.einsum("bij,bj->bi", V, v)
207    return T