Fix small-angle approximations and update adjoint sign convention

fiber/_coefficients.py

@@ -0,0 +1,63 @@
+# Copyright 2026, Evan Palmer
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+from collections.abc import Callable
+
+import jax.numpy as jnp
+from jaxtyping import Array
+
+
+def series_or_closed(
+    theta2: Array, series: Callable[[Array], Array], closed: Callable[[Array], Array]
+) -> Array:
+    small = theta2 < 1e-2
+    return jnp.where(small, series(theta2), closed(jnp.where(small, 1.0, theta2)))
+
+
+def sinc(theta2: Array) -> Array:
+    return series_or_closed(
+        theta2,
+        lambda x: 1 - x / 6 + x**2 / 120 - x**3 / 5040 + x**4 / 362880,
+        lambda x: jnp.sin(jnp.sqrt(x)) / jnp.sqrt(x),
+    )
+
+
+def cosc(theta2: Array) -> Array:
+    return series_or_closed(
+        theta2,
+        lambda x: 1 / 2 - x / 24 + x**2 / 720 - x**3 / 40320 + x**4 / 3628800,
+        lambda x: 2 * jnp.sin(jnp.sqrt(x) / 2) ** 2 / x,
+    )
+
+
+def sinc3(theta2: Array) -> Array:
+    return series_or_closed(
+        theta2,
+        lambda x: 1 / 6 - x / 120 + x**2 / 5040 - x**3 / 362880 + x**4 / 39916800,
+        lambda x: (1 - sinc(x)) / x,
+    )
+
+
+def dlogc(theta2: Array) -> Array:
+    return series_or_closed(
+        theta2,
+        lambda x: 1 / 12 + x / 720 + x**2 / 30240 + x**3 / 1209600 + x**4 / 47900160,
+        lambda x: (1 - jnp.sqrt(x) / 2 / jnp.tan(jnp.sqrt(x) / 2)) / x,
+    )

fiber/_groups/se2/_element.py

@@ -222,7 +222,8 @@ def __sub__(self, other: Twist2d) -> Twist2d:  # type: ignore[reportRedeclaratio
     def __sub__(self, other: ArrayLike) -> Twist2d:
         return Twist2d(self.point, self.value - other)
 
-    __rsub__ = __sub__
+    def __rsub__(self, other: ArrayLike) -> Twist2d:
+        return Twist2d(self.point, other - self.value)
 
     @dispatch
     def __mul__(self, other: Twist2d) -> Twist2d:  # type: ignore[reportRedeclaration]
@@ -242,7 +243,8 @@ def __truediv__(self, other: Twist2d) -> Twist2d:  # type: ignore[reportRedeclar
     def __truediv__(self, other: ArrayLike) -> Twist2d:
         return Twist2d(self.point, self.value / other)
 
-    __rtruediv__ = __truediv__
+    def __rtruediv__(self, other: ArrayLike) -> Twist2d:
+        return Twist2d(self.point, other / self.value)
 
     def __repr__(self) -> str:
         repr = np.array2string(cast(np.ndarray, self.value), prefix="Twist2d(")
@@ -329,7 +331,8 @@ def __sub__(self, other: Wrench2d) -> Wrench2d:  # type: ignore[reportRedeclarat
     def __sub__(self, other: ArrayLike) -> Wrench2d:
         return Wrench2d(self.point, self.value - other)
 
-    __rsub__ = __sub__
+    def __rsub__(self, other: ArrayLike) -> Wrench2d:
+        return Wrench2d(self.point, other - self.value)
 
     @dispatch
     def __mul__(self, other: Wrench2d) -> Wrench2d:  # type: ignore[reportRedeclaration]
@@ -349,7 +352,8 @@ def __truediv__(self, other: Wrench2d) -> Wrench2d:  # type: ignore[reportRedecl
     def __truediv__(self, other: ArrayLike) -> Wrench2d:
         return Wrench2d(self.point, self.value / other)
 
-    __rtruediv__ = __truediv__
+    def __rtruediv__(self, other: ArrayLike) -> Wrench2d:
+        return Wrench2d(self.point, other / self.value)
 
     def __repr__(self) -> str:
         repr = np.array2string(cast(np.ndarray, self.value), prefix="Wrench2d(")

fiber/_groups/se2/_operations.py

@@ -20,10 +20,10 @@
 
 import functools
 
-import jax
 import jax.numpy as jnp
 from jaxtyping import Array
 
+from ..._coefficients import cosc, dlogc, sinc, sinc3
 from ..._vecfuncs import vex2
 from .. import so2
 
@@ -54,7 +54,7 @@ def dadj(w: Array, p: Array):
 
 @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)")
 def dadj_op(w: Array) -> Array:
-    return -adj_op(w).T
+    return adj_op(w).T
 
 
 @functools.partial(jnp.vectorize, signature="(n,n),(n)->(n)")
@@ -96,7 +96,7 @@ def dAdj(g: Array, p: Array) -> Array:
 
 @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)")
 def dAdj_op(g: Array) -> Array:
-    return Adj_op(inv(g)).T
+    return Adj_op(g).T
 
 
 @functools.partial(jnp.vectorize, signature="(n,n),(n)->(n)")
@@ -113,77 +113,32 @@ def dAdj_inv_op(g: Array) -> Array:
 def expm(w: Array) -> Array:
     lin, ang_hat = w[:2, 2], w[:2, :2]
     theta = vex2(ang_hat)
-    A = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: 1 - (theta**2 / 6) + (theta**4 / 120),
-        lambda: jnp.sin(theta) / theta,
-    )
-    B = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: theta * (0.5 - (theta**2 / 24) + (theta**4 / 720)),
-        lambda: (1 - jnp.cos(theta)) / theta,
-    )
-    V = A * jnp.eye(2) + B * _J
+    theta2 = theta**2
+    V = sinc(theta2) * jnp.eye(2) + theta * cosc(theta2) * _J
     return jnp.block([[so2.expm(ang_hat), (V @ lin).reshape(2, 1)], [jnp.zeros(2), 1]])
 
 
 @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)")
 def dexpm(w: Array) -> Array:
     A = adj_op(w)
-    theta = vex2(w[:2, :2])
-    a = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: 1 - (theta**2 / 6) + (theta**4 / 120),
-        lambda: jnp.sin(theta) / theta,
-    )
-    b = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: 0.5 - (theta**2 / 24) + (theta**4 / 720),
-        lambda: (1 - jnp.cos(theta)) / (theta**2),
-    )
-    c = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: (1 / 6) + (theta**2 / 120),
-        lambda: (1 - a) / (theta**2),
-    )
-    return jnp.eye(3) + b * A + c * (A @ A)
+    theta2 = vex2(w[:2, :2]) ** 2
+    return jnp.eye(3) + cosc(theta2) * A + sinc3(theta2) * (A @ A)
 
 
 @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)")
 def logm(g: Array) -> Array:
     pos, rot = g[:2, 2], g[:2, :2]
     ang_hat = so2.logm(rot)
     theta = vex2(ang_hat)
-    half = 0.5 * theta
-    c = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: 1 - (theta**2 / 12) - (theta**4 / 720),
-        lambda: half / jnp.tan(half),
-    )
-    W = c * jnp.eye(2) - half * _J
+    W = (1 - theta**2 * dlogc(theta**2)) * jnp.eye(2) - 0.5 * theta * _J
     return jnp.block([[ang_hat, (W @ pos).reshape(2, 1)], [jnp.zeros(3)]])
 
 
 @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)")
 def dlogm(w: Array) -> Array:
     A = adj_op(w)
-    theta = vex2(w[:2, :2])
-    a = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: 1 - (theta**2 / 6) + (theta**4 / 120),
-        lambda: jnp.sin(theta) / theta,
-    )
-    b = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: 0.5 - (theta**2 / 24) + (theta**4 / 720),
-        lambda: (1 - jnp.cos(theta)) / (theta**2),
-    )
-    e = jax.lax.cond(
-        jnp.isclose(theta, 0.0),  # type: ignore
-        lambda: (1 / 12) + (theta**2 / 720),
-        lambda: (b - 0.5 * a) / (1 - jnp.cos(theta)),
-    )
-    return jnp.eye(3) - 0.5 * A + e * (A @ A)
+    theta2 = vex2(w[:2, :2]) ** 2
+    return jnp.eye(3) - 0.5 * A + dlogc(theta2) * (A @ A)
 
 
 @functools.partial(jnp.vectorize, signature="(n,n),(n,n)->(n,n)")

fiber/_groups/se2/random/_gaussian.py

@@ -26,7 +26,7 @@
 from jaxtyping import Array, PRNGKeyArray
 
 from .._element import _as_vector, _from_vector
-from .._operations import expm, logm, lplus, rminus, rplus
+from .._operations import expm, lminus, logm, lplus, rminus, rplus
 
 
 def sample_lie_algebra(
@@ -71,10 +71,13 @@ def mean(
     atol=1e-6,
     max_steps: int = 100,
     throw: bool = True,
+    left: bool = True,
 ) -> Array:
+    minus, plus = (rminus, rplus) if left else (lminus, lplus)
+
     def update(mean, samples):
-        errors = _as_vector(rminus(samples, mean))
-        return rplus(mean, _from_vector(jnp.mean(errors, axis=0)))
+        errors = _as_vector(minus(samples, mean))
+        return plus(mean, _from_vector(jnp.mean(errors, axis=0)))
 
     # construct the initial guess to warm-start the solver
     exp_coords = _as_vector(logm(samples))
@@ -94,7 +97,8 @@ def update(mean, samples):
     return sol.value
 
 
-def cov(mean: Array, samples: Array) -> Array:
-    distances = _as_vector(rminus(samples, mean))
+def cov(mean: Array, samples: Array, left: bool = True) -> Array:
+    minus = rminus if left else lminus
+    distances = _as_vector(minus(samples, mean))
     sigma = jnp.divide(jnp.einsum("ni,nj->ij", distances, distances), len(samples))
     return sigma

fiber/_groups/se3/_element.py

@@ -232,7 +232,8 @@ def __sub__(self, other: Twist3d) -> Twist3d:  # type: ignore[reportRedeclaratio
     def __sub__(self, other: ArrayLike) -> Twist3d:
         return Twist3d(self.point, self.value - other)
 
-    __rsub__ = __sub__
+    def __rsub__(self, other: ArrayLike) -> Twist3d:
+        return Twist3d(self.point, other - self.value)
 
     @dispatch
     def __mul__(self, other: Twist3d) -> Twist3d:  # type: ignore[reportRedeclaration]
@@ -252,7 +253,8 @@ def __truediv__(self, other: Twist3d) -> Twist3d:  # type: ignore[reportRedeclar
     def __truediv__(self, other: ArrayLike) -> Twist3d:
         return Twist3d(self.point, self.value / other)
 
-    __rtruediv__ = __truediv__
+    def __rtruediv__(self, other: ArrayLike) -> Twist3d:
+        return Twist3d(self.point, other / self.value)
 
     def __repr__(self) -> str:
         repr = np.array2string(cast(np.ndarray, self.value), prefix="Twist3d(")
@@ -339,7 +341,8 @@ def __sub__(self, other: Wrench3d) -> Wrench3d:  # type: ignore[reportRedeclarat
     def __sub__(self, other: ArrayLike) -> Wrench3d:
         return Wrench3d(self.point, self.value - other)
 
-    __rsub__ = __sub__
+    def __rsub__(self, other: ArrayLike) -> Wrench3d:
+        return Wrench3d(self.point, other - self.value)
 
     @dispatch
     def __mul__(self, other: Wrench3d) -> Wrench3d:  # type: ignore[reportRedeclaration]
@@ -359,7 +362,8 @@ def __truediv__(self, other: Wrench3d) -> Wrench3d:  # type: ignore[reportRedecl
     def __truediv__(self, other: ArrayLike) -> Wrench3d:
         return Wrench3d(self.point, self.value / other)
 
-    __rtruediv__ = __truediv__
+    def __rtruediv__(self, other: ArrayLike) -> Wrench3d:
+        return Wrench3d(self.point, other / self.value)
 
     def __repr__(self) -> str:
         repr = np.array2string(cast(np.ndarray, self.value), prefix="Wrench3d(")