Skip to content

Commit abe79e4

Browse files
petercorkeclaude
andauthored
perf(base): speed up trnorm and tr2adjoint (#214)
## Summary - `trnorm` (Gram-Schmidt re-orthonormalization used by e.g. `SE3.__matmul__`/`@`) and `tr2adjoint` (adjoint matrix, used by `SE3.Ad()`/`Twist3.Ad()`) both built small, fixed-size results (3x3/4x4/6x6) via generic NumPy helpers - `np.cross`/`unitvec`/`np.stack` in `trnorm`, `np.block` in `tr2adjoint` - whose dispatch overhead (built for arbitrary shapes/broadcasting) dominates cost at this size. Same pattern as the `isR`/`ishom` speedup in #213. - `trnorm`: replaced with explicit scalar cross-product/normalization arithmetic and direct in-place array construction. - `tr2adjoint`: replaced `np.block([[...]])` with a pre-allocated `np.zeros((6,6), dtype=T.dtype)` and direct slice assignment. `dtype=T.dtype` is preserved specifically so the function's documented `:SymPy: supported` contract still holds for symbolic (`dtype=object`) input. - Net effect (isolated `timeit`, `trnorm`/`tr2adjoint` called directly): `trnorm` ~12x faster (27.4us -> 2.3us), `tr2adjoint` ~2.7x faster (8.8us -> 3.3us). End to end, `SE3 @ SE3` (which normalizes via `trnorm`) drops from ~30us to ~4us. `SE3.Ad()`/`base.tr2adjoint()` calls on an already-known SE3/rotation matrix land around ~3.3us. - Note: `Twist3.Ad()` itself is not meaningfully faster from this change - its cost is dominated by the separate `self.SE3()` -> `trexp()` conversion step (matrix exponential), not by `tr2adjoint`. That's a different, inherent cost (same class as `Twist3 * Twist3`'s "product of exponentials"), not addressed here. ## Test plan - [x] `pytest tests/` - 343 passed - [x] Verified bit-for-bit numeric equivalence against the prior implementation over 200 random SO(3)/SE(3) trials (both `trnorm` and `tr2adjoint`, both 3x3 and 4x4 forms) - [x] Verified symbolic (SymPy, `dtype=object`) equivalence for `tr2adjoint` specifically, since it's the one with a documented `:SymPy: supported` contract 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 2829ced commit abe79e4

1 file changed

Lines changed: 36 additions & 16 deletions

File tree

‎spatialmath/base/transforms3d.py‎

Lines changed: 36 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1606,17 +1606,36 @@ def trnorm(T: SE3Array) -> SE3Array:
16061606
if not ishom(T) and not isrot(T):
16071607
raise ValueError("expecting SO(3) or SE(3)")
16081608

1609-
o = T[:3, 1]
1610-
a = T[:3, 2]
1609+
# explicit scalar arithmetic avoids the generic-dispatch overhead of
1610+
# np.cross/unitvec/np.stack (built for arbitrary shapes/broadcasting),
1611+
# which dominates cost for a fixed 3x3 rotation submatrix
1612+
o0, o1, o2 = T[0, 1], T[1, 1], T[2, 1]
1613+
a0, a1, a2 = T[0, 2], T[1, 2], T[2, 2]
16111614

1612-
n = np.cross(o, a) # N = O x A
1613-
o = np.cross(a, n) # (a)];
1614-
R = np.stack((unitvec(n), unitvec(o), unitvec(a)), axis=1)
1615+
# n = o x a
1616+
n0 = o1 * a2 - o2 * a1
1617+
n1 = o2 * a0 - o0 * a2
1618+
n2 = o0 * a1 - o1 * a0
16151619

1616-
if ishom(T):
1617-
return rt2tr(cast(SO3Array, R), T[:3, 3])
1618-
else:
1619-
return R
1620+
# o = a x n, to re-orthogonalize
1621+
o0, o1, o2 = a1 * n2 - a2 * n1, a2 * n0 - a0 * n2, a0 * n1 - a1 * n0
1622+
1623+
n_norm = (n0 * n0 + n1 * n1 + n2 * n2) ** 0.5
1624+
o_norm = (o0 * o0 + o1 * o1 + o2 * o2) ** 0.5
1625+
a_norm = (a0 * a0 + a1 * a1 + a2 * a2) ** 0.5
1626+
1627+
is_hom = ishom(T)
1628+
R = np.empty((4, 4) if is_hom else (3, 3))
1629+
R[0, 0], R[1, 0], R[2, 0] = n0 / n_norm, n1 / n_norm, n2 / n_norm
1630+
R[0, 1], R[1, 1], R[2, 1] = o0 / o_norm, o1 / o_norm, o2 / o_norm
1631+
R[0, 2], R[1, 2], R[2, 2] = a0 / a_norm, a1 / a_norm, a2 / a_norm
1632+
1633+
if is_hom:
1634+
R[3, :3] = 0.0
1635+
R[3, 3] = 1.0
1636+
R[:3, 3] = T[:3, 3]
1637+
1638+
return cast(SE3Array, R)
16201639

16211640

16221641
@overload
@@ -2733,20 +2752,21 @@ def tr2adjoint(T):
27332752
:SymPy: supported
27342753
"""
27352754

2736-
Z = np.zeros((3, 3), dtype=T.dtype)
27372755
if T.shape == (3, 3):
27382756
# SO(3) adjoint
27392757
R = T
27402758
return R
27412759
elif T.shape == (4, 4):
27422760
# SE(3) adjoint
27432761
(R, t) = tr2rt(T)
2744-
# fmt: off
2745-
return np.block([
2746-
[R, skew(t) @ R],
2747-
[Z, R]
2748-
])
2749-
# fmt: on
2762+
# direct block assignment avoids np.block's generic nested-list
2763+
# assembly overhead, which dominates cost for this fixed 6x6-from-
2764+
# four-3x3 layout; dtype preserved for SymPy support
2765+
A = np.zeros((6, 6), dtype=T.dtype)
2766+
A[:3, :3] = R
2767+
A[:3, 3:] = skew(t) @ R
2768+
A[3:, 3:] = R
2769+
return A
27502770
else:
27512771
raise ValueError("bad argument")
27522772

0 commit comments

Comments
 (0)