Skip to content

Commit af473db

Browse files
petercorkeclaude
andcommitted
perf(base): speed up trnorm and tr2adjoint
Both functions built small, fixed-size (3x3/4x4/6x6) results out of generic NumPy helpers (np.cross/unitvec/np.stack in trnorm, np.block in tr2adjoint) whose dispatch overhead - built for arbitrary shapes and broadcasting - dominates cost at this size. Same pattern as the isR/ishom speedup in #213. Replaced with explicit scalar arithmetic (trnorm) and direct pre-allocated slice-assignment (tr2adjoint), dtype preserved via np.zeros(..., dtype=T.dtype) so tr2adjoint's documented SymPy support is unaffected. ~12x faster trnorm, ~2.7x faster tr2adjoint standalone. End to end: SE3 @ SE3 (which normalizes via trnorm) drops from ~30us to ~4us. Verified bit-for-bit numeric equivalence against the prior implementation over 200 random SO(3)/SE(3) trials, plus symbolic (SymPy dtype=object) equivalence for tr2adjoint. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1 parent fff7009 commit af473db

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
@@ -1591,17 +1591,36 @@ def trnorm(T: SE3Array) -> SE3Array:
15911591
if not ishom(T) and not isrot(T):
15921592
raise ValueError("expecting SO(3) or SE(3)")
15931593

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

1597-
n = np.cross(o, a) # N = O x A
1598-
o = np.cross(a, n) # (a)];
1599-
R = np.stack((unitvec(n), unitvec(o), unitvec(a)), axis=1)
1600+
# n = o x a
1601+
n0 = o1 * a2 - o2 * a1
1602+
n1 = o2 * a0 - o0 * a2
1603+
n2 = o0 * a1 - o1 * a0
16001604

1601-
if ishom(T):
1602-
return rt2tr(cast(SO3Array, R), T[:3, 3])
1603-
else:
1604-
return R
1605+
# o = a x n, to re-orthogonalize
1606+
o0, o1, o2 = a1 * n2 - a2 * n1, a2 * n0 - a0 * n2, a0 * n1 - a1 * n0
1607+
1608+
n_norm = (n0 * n0 + n1 * n1 + n2 * n2) ** 0.5
1609+
o_norm = (o0 * o0 + o1 * o1 + o2 * o2) ** 0.5
1610+
a_norm = (a0 * a0 + a1 * a1 + a2 * a2) ** 0.5
1611+
1612+
is_hom = ishom(T)
1613+
R = np.empty((4, 4) if is_hom else (3, 3))
1614+
R[0, 0], R[1, 0], R[2, 0] = n0 / n_norm, n1 / n_norm, n2 / n_norm
1615+
R[0, 1], R[1, 1], R[2, 1] = o0 / o_norm, o1 / o_norm, o2 / o_norm
1616+
R[0, 2], R[1, 2], R[2, 2] = a0 / a_norm, a1 / a_norm, a2 / a_norm
1617+
1618+
if is_hom:
1619+
R[3, :3] = 0.0
1620+
R[3, 3] = 1.0
1621+
R[:3, 3] = T[:3, 3]
1622+
1623+
return cast(SE3Array, R)
16051624

16061625

16071626
@overload
@@ -2718,20 +2737,21 @@ def tr2adjoint(T):
27182737
:SymPy: supported
27192738
"""
27202739

2721-
Z = np.zeros((3, 3), dtype=T.dtype)
27222740
if T.shape == (3, 3):
27232741
# SO(3) adjoint
27242742
R = T
27252743
return R
27262744
elif T.shape == (4, 4):
27272745
# SE(3) adjoint
27282746
(R, t) = tr2rt(T)
2729-
# fmt: off
2730-
return np.block([
2731-
[R, skew(t) @ R],
2732-
[Z, R]
2733-
])
2734-
# fmt: on
2747+
# direct block assignment avoids np.block's generic nested-list
2748+
# assembly overhead, which dominates cost for this fixed 6x6-from-
2749+
# four-3x3 layout; dtype preserved for SymPy support
2750+
A = np.zeros((6, 6), dtype=T.dtype)
2751+
A[:3, :3] = R
2752+
A[:3, 3:] = skew(t) @ R
2753+
A[3:, 3:] = R
2754+
return A
27352755
else:
27362756
raise ValueError("bad argument")
27372757

0 commit comments

Comments
 (0)