robosim.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 get_closest_square(n: int) -> tuple[int, int]:
 53    """
 54    Given a stack of N images, find the closest square X>=N*N and return that.
 55    Note: There are only 2 rows possible between x^2 and (x+1)^2 because (x+1)^2 = x^2 + 2*x + 1, thus we can add two
 56    columns at most. If a 3rd column is needed, then closest lower bound is (x+1)^2 and we must use that.
 57    Example: 9: 3*3; 12 -> 3*3 -> 3*4 (3 rows). 65 -> 8*8 -> 8*9. 73 -> 8*8 -> 8*9 -> 9*9
 58    """
 59    x = int(math.sqrt(n))
 60    r, c = x, x
 61    c = c + 1 if c * r < n else c
 62    r = r + 1 if c * r < n else r
 63    assert (c + 1) * r > n and c * (r + 1) > n
 64    return r, c
 65
 66def make_radius(size: Point3D) -> float:
 67    """creates a radius from a size and a scale"""
 68    return math.sqrt(size[0]**2 + size[1]**2 + size[2]**2) / 2
 69
 70
 71def is_rot(R: Rot3x3, tol: float = 20) -> bool:
 72    r"""
 73    Test if matrix belongs to SO(n).
 74    Checks orthogonality, ie. :math:`{\bf R} {\bf R}^T = {\bf I}` and :math:`\det({\bf R}) > 0`.
 75    For the first test we check that the norm of the residual is less than ``tol * eps``.
 76        >>> is_rot(np.eye(3))
 77        >>> is_rot(np.zeros((3,3)))
 78    """
 79    return bool(
 80        np.linalg.norm(R @ R.T - np.eye(R.shape[0])) < tol * _eps
 81        and np.linalg.det(R) > 0
 82    )
 83
 84"""
 85Below: trimmed, standalone port of the `spatialmath.base` functions robosim actually uses.
 86
 87Only 5 public functions, in the exact shapes the codebase calls them:
 88- is_rot(R[, tol])   -> mathutils.rot_from_forward_up, renormalize_rotation_matrix
 89- tr2rpy(pose4x4) -> mathutils.pose_to_trans_euler
 90- rpy2tr(rpy3)    -> mathutils.pose_from_trans_euler
 91- bskewa(twist6)   -> physics.motion._integrate_velocity_into_pose
 92- btrexp(se3 4x4)  -> physics.motion._integrate_velocity_into_pose
 93
 94Everything is self-contained (helpers are inner `_`-prefixed functions), radians only,
 95zyx order only, and float64 throughout. Callers downcast to float32 at their boundary.
 96"""
 97
 98def tr2rpy(T: Pose4x4) -> Point3D:
 99    """
100    Convert SO(3)/SE(3) to roll-pitch-yaw angles (radians, zyx order).
101        >>> tr2rpy(rpy2tr([0.2, 0.3, 0.5]))
102    The angles [roll, pitch, yaw] are sequential rotations about the X, Y, Z axes.
103    If the input is SE(3), the translation component is ignored. There is a singularity
104    at pitch = ±90° where roll is arbitrarily set to 0.
105    """
106    R = T[0:3, 0:3]
107
108    rpy = np.zeros(3)
109    if abs(abs(R[2, 0]) - 1) < 20 * _eps:  # |R31| == 1 -> gimbal-lock singularity
110        rpy[0] = 0.0  # roll is zero
111        if R[2, 0] < 0:
112            rpy[2] = -math.atan2(R[0, 1], R[0, 2])
113        else:
114            rpy[2] = math.atan2(-R[0, 1], -R[0, 2])
115        rpy[1] = -math.asin(np.clip(R[2, 0], -1.0, 1.0))
116    else:
117        rpy[0] = math.atan2(R[2, 1], R[2, 2])  # roll
118        rpy[2] = math.atan2(R[1, 0], R[0, 0])  # yaw
119        # pick the numerically largest term to solve pitch, avoiding division by ~0
120        k = int(np.argmax(np.abs([R[0, 0], R[1, 0], R[2, 1], R[2, 2]])))
121        if k == 0:
122            rpy[1] = -math.atan(R[2, 0] * math.cos(rpy[2]) / R[0, 0])
123        elif k == 1:
124            rpy[1] = -math.atan(R[2, 0] * math.sin(rpy[2]) / R[1, 0])
125        elif k == 2:
126            rpy[1] = -math.atan(R[2, 0] * math.sin(rpy[0]) / R[2, 1])
127        else:
128            rpy[1] = -math.atan(R[2, 0] * math.cos(rpy[0]) / R[2, 2])
129
130    return rpy
131
132def rpy2tr(rpy: Point3D) -> Pose4x4:
133    r"""
134    Create an SE(3) matrix from roll-pitch-yaw angles (radians, zyx order).
135        >>> rpy2tr([0.1, 0.2, 0.3])
136    zyx: rotate by yaw about z, then pitch about the new y, then roll about the new x.
137    The translation component is zero.
138    """
139    def _rotx(t: float) -> np.ndarray:
140        ct, st = math.cos(t), math.sin(t)
141        return np.array([[1.0, 0.0, 0.0], [0.0, ct, -st], [0.0, st, ct]], dtype=np.float64)
142
143    def _roty(t: float) -> np.ndarray:
144        ct, st = math.cos(t), math.sin(t)
145        return np.array([[ct, 0.0, st], [0.0, 1.0, 0.0], [-st, 0.0, ct]], dtype=np.float64)
146
147    def _rotz(t: float) -> np.ndarray:
148        ct, st = math.cos(t), math.sin(t)
149        return np.array([[ct, -st, 0.0], [st, ct, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64)
150
151    roll, pitch, yaw = float(rpy[0]), float(rpy[1]), float(rpy[2])
152    T = np.eye(4)
153    T[0:3, 0:3] = _rotz(yaw) @ _roty(pitch) @ _rotx(roll)
154    return T
155
156def skew(w: Point3D) -> Mat3x3:
157    """create a skew symmetric matrix from a 3-vector"""
158    return np.array([
159        [0.0,   -w[2],  w[1]],
160        [w[2],   0.0,  -w[0]],
161        [-w[1],  w[0],  0.0]], dtype=np.float64) # (4, 4)
162
163def bskew(v: Point3D) -> Mat4x4:
164    """batched skew"""
165    assert len(v.shape) == 2 and v.shape[1] == 3, v.shape
166    zero = np.zeros((len(v), ), "float32")
167    res = np.array([
168        [ zero,     -v[:, 2],    v[:, 1]],
169        [ v[:, 2],   zero,      -v[:, 0]],
170        [-v[:, 1],   v[:, 0],    zero,  ]], dtype=np.float64) # (3, 3, N)
171    return np.permute_dims(res, (2, 0, 1)) # (N, 3, 3)
172
173def bvexa(S: Mat4x4) -> Point6D:
174    """recover the twist [v, w] from the se(3) matrix (inverse of bskewa)"""
175    return np.array([
176        S[:, 0, 3], S[:, 1, 3], S[:, 2, 3],
177        (S[:, 2, 1] - S[:, 1, 2]) / 2,
178        (S[:, 0, 2] - S[:, 2, 0]) / 2,
179        (S[:, 1, 0] - S[:, 0, 1]) / 2], dtype=np.float64).T # (N, 6)
180
181def bskewa(v: Point6D) -> Mat4x4:
182    """batched skewa"""
183    assert len(v.shape) == 2 and v.shape[1] == 6, v.shape
184    zero = np.zeros((len(v), ), "float32")
185    res = np.array([
186        [ zero,     -v[:, 5],    v[:, 4],  v[:, 0]],
187        [ v[:, 5],   zero,      -v[:, 3],  v[:, 1]],
188        [-v[:, 4],   v[:, 3],    zero,     v[:, 2]],
189        [ zero,      zero,       zero,     zero]], dtype=np.float64) # (4, 4, N)
190    return np.permute_dims(res, (2, 0, 1)) # (N, 4, 4)
191
192def btrexp(S: Mat4x4) -> Pose4x4:
193    """
194    Batched matrix exponential of an se(3) element -> SE(3) (a screw motion).
195    ``S`` is the batched (Nx)4x4 augmented skew-symmetric matrix produced by ``bskewa``.
196        >>> btrexp(bskewa(np.array([[1, 0, 0, 0, 0, 0]])))
197    """
198    n = len(S)
199    tw = bvexa(S) # (N, 6)
200    v = tw[:, 0:3] # (3, )
201    w = tw[:, 3:6] # (3, )
202
203    w_norm = np.linalg.norm(w, axis=1) # (N, )
204    v_norm = np.linalg.norm(v, axis=1) # (N, )
205    # theta is the screw magnitude: rotation dominates, else it's a pure translation
206    theta = np.where(w_norm >= 20 * _eps, w_norm, v_norm)[:, None] # (N, 1)
207
208    with np.errstate(divide="ignore", invalid="ignore"):   # 0/0 on zero-twist is expected; nan_to_num handles it
209        v = np.nan_to_num(v / theta, nan=0) # (N, 3)
210        w = np.nan_to_num(w / theta, nan=0) # (N, 3)
211
212    skw = bskew(w)                                                      # (N, 3, 3)
213    theta_n = theta[..., None]                                          # (N, 1, 1)
214    st, ct = np.sin(theta_n), np.cos(theta_n)                           # (N, 1, 1), (N, 1, 1)
215    eye = np.eye(3)[None].repeat(n, axis=0)                             # (N, 3, 3)
216    R = eye + st * skw + (1.0 - ct) * skw @ skw                         # (N, 3, 3) Rodrigues
217    V = eye * theta_n + (1.0 - ct) * skw + (theta_n - st) * skw @ skw   # (N, 3, 3) translation map
218
219    T = np.eye(4)[None].repeat(n, axis=0)                               # (N, 4, 4)
220    T[:, 0:3, 0:3] = R
221    T[:, 0:3, 3] = np.einsum("bij,bj->bi", V, v)
222    return T