From d05bf28f4480dab78175af286308ad9d11987293 Mon Sep 17 00:00:00 2001 From: Evan Palmer Date: Thu, 8 Oct 2026 17:04:32 -0700 Subject: [PATCH 1/3] Updated dadj convention to match equations used in papers --- fiber/_groups/se2/_operations.py | 4 ++-- fiber/_groups/se3/_operations.py | 6 +++--- fiber/_groups/so2/_operations.py | 4 ++-- fiber/_groups/so3/_operations.py | 4 ++-- pyproject.toml | 2 +- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/fiber/_groups/se2/_operations.py b/fiber/_groups/se2/_operations.py index a645b27..32abee2 100644 --- a/fiber/_groups/se2/_operations.py +++ b/fiber/_groups/se2/_operations.py @@ -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)") diff --git a/fiber/_groups/se3/_operations.py b/fiber/_groups/se3/_operations.py index de5b43e..2009b75 100644 --- a/fiber/_groups/se3/_operations.py +++ b/fiber/_groups/se3/_operations.py @@ -51,7 +51,7 @@ def dadj(w: Array, p: Array): @functools.partial(jnp.vectorize, signature="(n,n)->(m,m)") def dadj_op(w: Array) -> Array: - return -adj_op(w).T + return adj_op(w).T @functools.partial(jnp.vectorize, signature="(n,n),(m)->(m)") @@ -72,7 +72,7 @@ def Adj(g: Array, w: Array) -> Array: @functools.partial(jnp.vectorize, signature="(n,n)->(m,m)") def Adj_op(g: Array) -> Array: pos, rot = g[:3, 3], g[:3, :3] - return jnp.block([[rot, skew3(pos)], [jnp.zeros_like(rot), rot]]) + return jnp.block([[rot, skew3(pos) @ rot], [jnp.zeros_like(rot), rot]]) @functools.partial(jnp.vectorize, signature="(n,n),(n,n)->(n,n)") @@ -92,7 +92,7 @@ def dAdj(g: Array, p: Array) -> Array: @functools.partial(jnp.vectorize, signature="(n,n)->(m,m)") def dAdj_op(g: Array) -> Array: - return Adj_op(inv(g)).T + return Adj_op(g).T @functools.partial(jnp.vectorize, signature="(n,n),(m)->(m)") diff --git a/fiber/_groups/so2/_operations.py b/fiber/_groups/so2/_operations.py index 2b8c917..eda064e 100644 --- a/fiber/_groups/so2/_operations.py +++ b/fiber/_groups/so2/_operations.py @@ -50,7 +50,7 @@ def dadj(w: Array, p: RealScalarLike) -> RealScalarLike: @functools.partial(jnp.vectorize, signature="(n,n)->()") def dadj_op(w: Array) -> RealScalarLike: - return -adj_op(w) + return adj_op(w) @functools.partial(jnp.vectorize, signature="(n,n),()->()") @@ -91,7 +91,7 @@ def dAdj(g: Array, p: RealScalarLike) -> RealScalarLike: @functools.partial(jnp.vectorize, signature="(n,n)->()") def dAdj_op(g: Array) -> RealScalarLike: - return Adj_op(inv(g)) + return Adj_op(g) @functools.partial(jnp.vectorize, signature="(n,n),()->()") diff --git a/fiber/_groups/so3/_operations.py b/fiber/_groups/so3/_operations.py index 6c31e5a..bb32a76 100644 --- a/fiber/_groups/so3/_operations.py +++ b/fiber/_groups/so3/_operations.py @@ -49,7 +49,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)") @@ -89,7 +89,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)") diff --git a/pyproject.toml b/pyproject.toml index 35ddd0d..5c9a5ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi" [project] name = "fiber" -version = "0.1.3" +version = "0.1.4" description = "A JAX library for Lie group operations." readme = "README.md" requires-python = ">=3.14" From 7c4df186c2a1e034f29b597a4796ddb70ec1c57c Mon Sep 17 00:00:00 2001 From: Evan Palmer Date: Thu, 8 Oct 2026 17:45:40 -0700 Subject: [PATCH 2/3] Fixed small-angle approximations --- fiber/_coefficients.py | 63 ++++++++++++++++ fiber/_groups/se2/_element.py | 12 ++- fiber/_groups/se2/_operations.py | 61 ++------------- fiber/_groups/se2/random/_gaussian.py | 14 ++-- fiber/_groups/se3/_element.py | 12 ++- fiber/_groups/se3/_operations.py | 105 +++++--------------------- fiber/_groups/se3/random/_gaussian.py | 14 ++-- fiber/_groups/so2/_element.py | 12 ++- fiber/_groups/so2/random/_gaussian.py | 14 ++-- fiber/_groups/so3/_element.py | 12 ++- fiber/_groups/so3/_operations.py | 79 +++++++------------ fiber/_groups/so3/random/_gaussian.py | 14 ++-- fiber/_operations.py | 64 ++++++++-------- fiber/random/_gaussian.py | 36 ++++++--- 14 files changed, 241 insertions(+), 271 deletions(-) create mode 100644 fiber/_coefficients.py diff --git a/fiber/_coefficients.py b/fiber/_coefficients.py new file mode 100644 index 0000000..14e4b75 --- /dev/null +++ b/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, + ) diff --git a/fiber/_groups/se2/_element.py b/fiber/_groups/se2/_element.py index ac53e95..b73860d 100644 --- a/fiber/_groups/se2/_element.py +++ b/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(") diff --git a/fiber/_groups/se2/_operations.py b/fiber/_groups/se2/_operations.py index 32abee2..eef41f7 100644 --- a/fiber/_groups/se2/_operations.py +++ b/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 @@ -113,40 +113,16 @@ 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)") @@ -154,36 +130,15 @@ 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)") diff --git a/fiber/_groups/se2/random/_gaussian.py b/fiber/_groups/se2/random/_gaussian.py index c0eef70..3a229f7 100644 --- a/fiber/_groups/se2/random/_gaussian.py +++ b/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 diff --git a/fiber/_groups/se3/_element.py b/fiber/_groups/se3/_element.py index 478bb59..dd87a22 100644 --- a/fiber/_groups/se3/_element.py +++ b/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(") diff --git a/fiber/_groups/se3/_operations.py b/fiber/_groups/se3/_operations.py index 2009b75..a846b4b 100644 --- a/fiber/_groups/se3/_operations.py +++ b/fiber/_groups/se3/_operations.py @@ -20,11 +20,11 @@ import functools -import jax import jax.numpy as jnp from jaxtyping import Array -from ..._vecfuncs import skew3, softnorm, vex3 +from ..._coefficients import cosc, series_or_closed, sinc, sinc3 +from ..._vecfuncs import skew3, vex3 from .. import so3 @@ -108,39 +108,28 @@ def dAdj_inv_op(g: Array) -> Array: @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)") def expm(w: Array) -> Array: lin, ang = w[:3, 3], w[:3, :3] - theta = softnorm(vex3(ang)) - 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 = (1 - A) / theta**2 - V = jnp.eye(3) + B * ang + C * (ang @ ang) + V = so3.dexpm(ang) return jnp.block([[so3.expm(ang), (V @ lin).reshape(3, 1)], [jnp.zeros(3), 1]]) -def _se3_V(lin: Array, ang_hat: Array, theta, a, b, c) -> Array: +def _se3_V(lin: Array, ang_hat: Array) -> Array: lin_hat = skew3(lin) ang = vex3(ang_hat) - q1 = jax.lax.cond( - jnp.isclose(theta, 0.0), # type: ignore - lambda: -(1 / 12) + (theta**2 / 180), - lambda: (a - 2 * b) / theta**2, + theta2 = jnp.sum(ang**2) + q1 = series_or_closed( + theta2, + lambda x: -1 / 12 + x / 180 - x**2 / 6720 + x**3 / 453600 - x**4 / 47900160, + lambda x: (sinc(x) - 2 * cosc(x)) / x, ) - q2 = jax.lax.cond( - jnp.isclose(theta, 0.0), # type: ignore - lambda: -(1 / 60) + (theta**2 / 1260), - lambda: (b - 3 * c) / theta**2, + q2 = series_or_closed( + theta2, + lambda x: -1 / 60 + x / 1260 - x**2 / 60480 + x**3 / 4989600 - x**4 / 622702080, + lambda x: (cosc(x) - 3 * sinc3(x)) / x, ) Q = q1 * ang_hat + q2 * (ang_hat @ ang_hat) return ( - b * lin_hat - + c * (ang_hat @ lin_hat + lin_hat @ ang_hat) + cosc(theta2) * lin_hat + + sinc3(theta2) * (ang_hat @ lin_hat + lin_hat @ ang_hat) + jnp.dot(ang, lin) * Q ) @@ -148,25 +137,8 @@ def _se3_V(lin: Array, ang_hat: Array, theta, a, b, c) -> Array: @functools.partial(jnp.vectorize, signature="(n,n)->(m,m)") def dexpm(w: Array) -> Array: lin, ang_hat = w[:3, 3], w[:3, :3] - ang = vex3(ang_hat) - theta = softnorm(ang) - 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: (theta - jnp.sin(theta)) / theta**3, - ) - V = _se3_V(lin, ang_hat, theta, a, b, c) - D = jnp.eye(3) + b * ang_hat + c * (ang_hat @ ang_hat) + V = _se3_V(lin, ang_hat) + D = so3.dexpm(ang_hat) return jnp.block([[D, V], [jnp.zeros_like(V), D]]) @@ -174,52 +146,15 @@ def dexpm(w: Array) -> Array: def logm(g: Array) -> Array: pos, rot = g[:3, 3], g[:3, :3] ang_hat = so3.logm(rot) - ang = vex3(ang_hat) - jac = jnp.eye(3) - V = jax.lax.cond( - jnp.isclose(softnorm(ang), 0.0), - lambda: jac, - lambda: ( - jac - - 0.5 * ang_hat - + ( - (1 / (softnorm(ang) ** 2)) - - (1 + jnp.cos(softnorm(ang))) - / (2 * softnorm(ang) * jnp.sin(softnorm(ang))) - ) - * (ang_hat @ ang_hat) - ), - ) + V = so3.dlogm(ang_hat) return jnp.block([[ang_hat, (V @ pos).reshape(3, 1)], [jnp.zeros(4)]]) @functools.partial(jnp.vectorize, signature="(n,n)->(m,m)") def dlogm(w: Array) -> Array: lin, ang_hat = w[:3, 3], w[:3, :3] - ang = vex3(ang_hat) - theta = softnorm(ang) - 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: (theta - jnp.sin(theta)) / theta**3, - ) - 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)), - ) - B = _se3_V(lin, ang_hat, theta, a, b, c) - D = jnp.eye(3) - 0.5 * ang_hat + e * (ang_hat @ ang_hat) + B = _se3_V(lin, ang_hat) + D = so3.dlogm(ang_hat) return jnp.block([[D, -D @ B @ D], [jnp.zeros_like(D), D]]) diff --git a/fiber/_groups/se3/random/_gaussian.py b/fiber/_groups/se3/random/_gaussian.py index 2586b2c..541978a 100644 --- a/fiber/_groups/se3/random/_gaussian.py +++ b/fiber/_groups/se3/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 diff --git a/fiber/_groups/so2/_element.py b/fiber/_groups/so2/_element.py index d9e5a93..9eb5b1d 100644 --- a/fiber/_groups/so2/_element.py +++ b/fiber/_groups/so2/_element.py @@ -209,7 +209,8 @@ def __sub__(self, other: Spin2d) -> Spin2d: # type: ignore[reportRedeclaration] def __sub__(self, other: ArrayLike) -> Spin2d: return Spin2d(self.point, self.value - other) - __rsub__ = __sub__ + def __rsub__(self, other: ArrayLike) -> Spin2d: + return Spin2d(self.point, other - self.value) @dispatch def __mul__(self, other: Spin2d) -> Spin2d: # type: ignore[reportRedeclaration] @@ -229,7 +230,8 @@ def __truediv__(self, other: Spin2d) -> Spin2d: # type: ignore[reportRedeclarat def __truediv__(self, other: ArrayLike) -> Spin2d: return Spin2d(self.point, self.value / other) - __rtruediv__ = __truediv__ + def __rtruediv__(self, other: ArrayLike) -> Spin2d: + return Spin2d(self.point, other / self.value) def __repr__(self) -> str: repr = np.array2string(cast(np.ndarray, self.value), prefix="Spin2d(") @@ -320,7 +322,8 @@ def __sub__(self, other: Moment2d) -> Moment2d: # type: ignore[reportRedeclarat def __sub__(self, other: ArrayLike) -> Moment2d: return Moment2d(self.point, self.value - other) - __rsub__ = __sub__ + def __rsub__(self, other: ArrayLike) -> Moment2d: + return Moment2d(self.point, other - self.value) @dispatch def __mul__(self, other: Moment2d) -> Moment2d: # type: ignore[reportRedeclaration] @@ -340,7 +343,8 @@ def __truediv__(self, other: Moment2d) -> Moment2d: # type: ignore[reportRedecl def __truediv__(self, other: ArrayLike) -> Moment2d: return Moment2d(self.point, self.value / other) - __rtruediv__ = __truediv__ + def __rtruediv__(self, other: ArrayLike) -> Moment2d: + return Moment2d(self.point, other / self.value) def __repr__(self) -> str: repr = np.array2string(cast(np.ndarray, self.value), prefix="Moment2d(") diff --git a/fiber/_groups/so2/random/_gaussian.py b/fiber/_groups/so2/random/_gaussian.py index 04c473d..3487021 100644 --- a/fiber/_groups/so2/random/_gaussian.py +++ b/fiber/_groups/so2/random/_gaussian.py @@ -26,7 +26,7 @@ from jaxtyping import Array, PRNGKeyArray from ...._vecfuncs import skew2, vex2 -from .._operations import expm, logm, lplus, rminus, rplus +from .._operations import expm, lminus, logm, lplus, rminus, rplus def sample_lie_algebra( @@ -75,10 +75,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 = vex2(rminus(samples, mean)) - return rplus(mean, skew2(jnp.mean(errors, axis=0))) + errors = vex2(minus(samples, mean)) + return plus(mean, skew2(jnp.mean(errors, axis=0))) # construct the initial guess to warm-start the solver exp_coords = vex2(logm(samples)) @@ -98,7 +101,8 @@ def update(mean, samples): return sol.value -def cov(mean: Array, samples: Array) -> Array: - distances = vex2(rminus(samples, mean)) +def cov(mean: Array, samples: Array, left: bool = True) -> Array: + minus = rminus if left else lminus + distances = vex2(minus(samples, mean)) sigma = jnp.divide(jnp.einsum("n,n->", distances, distances), len(samples)) return sigma diff --git a/fiber/_groups/so3/_element.py b/fiber/_groups/so3/_element.py index f78e9cd..35a646e 100644 --- a/fiber/_groups/so3/_element.py +++ b/fiber/_groups/so3/_element.py @@ -221,7 +221,8 @@ def __sub__(self, other: Spin3d) -> Spin3d: # type: ignore[reportRedeclaration] def __sub__(self, other: ArrayLike) -> Spin3d: return Spin3d(self.point, self.value - other) - __rsub__ = __sub__ + def __rsub__(self, other: ArrayLike) -> Spin3d: + return Spin3d(self.point, other - self.value) @dispatch def __mul__(self, other: Spin3d) -> Spin3d: # type: ignore[reportRedeclaration] @@ -241,7 +242,8 @@ def __truediv__(self, other: Spin3d) -> Spin3d: # type: ignore[reportRedeclarat def __truediv__(self, other: ArrayLike) -> Spin3d: return Spin3d(self.point, self.value / other) - __rtruediv__ = __truediv__ + def __rtruediv__(self, other: ArrayLike) -> Spin3d: + return Spin3d(self.point, other / self.value) def __repr__(self) -> str: repr = np.array2string(cast(np.ndarray, self.value), prefix="Spin3d(") @@ -326,7 +328,8 @@ def __sub__(self, other: Moment3d) -> Moment3d: # type: ignore[reportRedeclarat def __sub__(self, other: ArrayLike) -> Moment3d: return Moment3d(self.point, self.value - other) - __rsub__ = __sub__ + def __rsub__(self, other: ArrayLike) -> Moment3d: + return Moment3d(self.point, other - self.value) @dispatch def __mul__(self, other: Moment3d) -> Moment3d: # type: ignore[reportRedeclaration] @@ -346,7 +349,8 @@ def __truediv__(self, other: Moment3d) -> Moment3d: # type: ignore[reportRedecl def __truediv__(self, other: ArrayLike) -> Moment3d: return Moment3d(self.point, self.value / other) - __rtruediv__ = __truediv__ + def __rtruediv__(self, other: ArrayLike) -> Moment3d: + return Moment3d(self.point, other / self.value) def __repr__(self) -> str: repr = np.array2string(cast(np.ndarray, self.value), prefix="Moment3d(") diff --git a/fiber/_groups/so3/_operations.py b/fiber/_groups/so3/_operations.py index bb32a76..04a5cbf 100644 --- a/fiber/_groups/so3/_operations.py +++ b/fiber/_groups/so3/_operations.py @@ -20,11 +20,11 @@ import functools -import jax import jax.numpy as jnp from jaxtyping import Array -from ..._vecfuncs import softclip, softnorm, vex3 +from ..._coefficients import cosc, dlogc, sinc, sinc3 +from ..._vecfuncs import skew3, vex3 @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)") @@ -104,70 +104,43 @@ def dAdj_inv_op(g: Array) -> Array: @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)") def expm(w: Array) -> Array: - theta = softnorm(vex3(w)) - sin = jax.lax.cond( - jnp.isclose(theta, 0.0), # type: ignore - lambda: 1 - (theta**2 / 6) + (theta**4 / 120), - lambda: jnp.sin(theta) / theta, - ) - cos = 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), - ) - return jnp.eye(3) + sin * w + cos * (w @ w) + theta2 = jnp.sum(vex3(w) ** 2) + return jnp.eye(3) + sinc(theta2) * w + cosc(theta2) * (w @ w) @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)") def dexpm(w: Array) -> Array: - theta = softnorm(vex3(w)) - 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 = (1 - a) / (theta**2) - return jnp.eye(3) + b * w + c * (w @ w) + theta2 = jnp.sum(vex3(w) ** 2) + return jnp.eye(3) + cosc(theta2) * w + sinc3(theta2) * (w @ w) @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)") def logm(g: Array) -> Array: - cos = (jnp.trace(g) - 1) / 2 - cos = softclip(cos, -1, 1) - theta = jnp.arccos(cos) - w_hat = jnp.zeros((3, 3)) - return jax.lax.cond( - jnp.sin(theta) < 1e-6, - lambda: w_hat, - lambda: 0.5 * theta / jnp.sin(theta) * (g - g.T), + cos = jnp.clip(0.5 * (jnp.trace(g) - 1), -1.0, 1.0) + sin_axis = 0.5 * vex3(g - g.T) + sin2 = jnp.sum(sin_axis**2) + + # guard the square root so its gradient stays finite at the identity + sin = jnp.where(sin2 > 0.0, jnp.sqrt(jnp.where(sin2 > 0.0, sin2, 1.0)), 0.0) + theta = jnp.arctan2(sin, cos) + + # past pi / 2 the skew part loses the axis, so recover it from the symmetric part + obtuse = cos < 0.0 + outer = (0.5 * (g + g.T) - cos * jnp.eye(3)) / jnp.where(obtuse, 1 - cos, 1.0) + i = jnp.argmax(jnp.diag(outer)) + axis = outer[:, i] / jnp.sqrt(jnp.where(obtuse, outer[i, i], 1.0)) + axis = jnp.where(axis @ sin_axis < 0.0, -axis, axis) + + w = jnp.where( + obtuse, theta * axis, sin_axis / sinc(jnp.where(obtuse, 0.0, theta**2)) ) + return skew3(w) @functools.partial(jnp.vectorize, signature="(n,n)->(n,n)") def dlogm(w: Array) -> Array: - theta = softnorm(vex3(w)) - cos = jnp.cos(theta) - 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 - cos) / (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 - cos), - ) - return jnp.eye(3) - 0.5 * w + e * (w @ w) + theta2 = jnp.sum(vex3(w) ** 2) + return jnp.eye(3) - 0.5 * w + dlogc(theta2) * (w @ w) @functools.partial(jnp.vectorize, signature="(n,n),(n,n)->(n,n)") diff --git a/fiber/_groups/so3/random/_gaussian.py b/fiber/_groups/so3/random/_gaussian.py index e0e3f69..884aebf 100644 --- a/fiber/_groups/so3/random/_gaussian.py +++ b/fiber/_groups/so3/random/_gaussian.py @@ -26,7 +26,7 @@ from jaxtyping import Array, PRNGKeyArray from ...._vecfuncs import skew3, vex3 -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 = vex3(rminus(samples, mean)) - return rplus(mean, skew3(jnp.mean(errors, axis=0))) + errors = vex3(minus(samples, mean)) + return plus(mean, skew3(jnp.mean(errors, axis=0))) # construct the initial guess to warm-start the solver exp_coords = vex3(logm(samples)) @@ -94,7 +97,8 @@ def update(mean, samples): return sol.value -def cov(mean: Array, samples: Array) -> Array: - distances = vex3(rminus(samples, mean)) +def cov(mean: Array, samples: Array, left: bool = True) -> Array: + minus = rminus if left else lminus + distances = vex3(minus(samples, mean)) sigma = jnp.divide(jnp.einsum("ni,nj->ij", distances, distances), len(samples)) return sigma diff --git a/fiber/_operations.py b/fiber/_operations.py index 59a5151..65d9ff1 100644 --- a/fiber/_operations.py +++ b/fiber/_operations.py @@ -58,22 +58,22 @@ def inv(w: Isometry3d) -> Isometry3d: @dispatch def adj(w: Spin2d, v: Spin2d) -> Spin2d: # type: ignore[reportRedeclaration] - return Spin2d.from_matrix(so2.adj(w.value, v.value)) + return Spin2d.from_matrix(so2.adj(w.value, v.value), w.point) @dispatch def adj(w: Spin3d, v: Spin3d) -> Spin3d: # type: ignore[reportRedeclaration] - return Spin3d.from_matrix(so3.adj(w.value, v.value)) + return Spin3d.from_matrix(so3.adj(w.value, v.value), w.point) @dispatch def adj(w: Twist2d, v: Twist2d) -> Twist2d: # type: ignore[reportRedeclaration] - return Twist2d.from_matrix(se2.adj(w.value, v.value)) + return Twist2d.from_matrix(se2.adj(w.value, v.value), w.point) @dispatch def adj(w: Twist3d, v: Twist3d) -> Twist3d: # type: ignore[reportRedeclaration] - return Twist3d.from_matrix(se3.adj(w.value, v.value)) + return Twist3d.from_matrix(se3.adj(w.value, v.value), w.point) @dispatch @@ -98,22 +98,22 @@ def adj_op(w: Twist3d) -> Array: # type: ignore[reportRedeclaration] @dispatch def dadj(w: Spin2d, p: Moment2d) -> Moment2d: # type: ignore[reportRedeclaration] - return Moment2d.from_vector(so2.dadj(w.value, p.value)) + return Moment2d.from_vector(so2.dadj(w.value, p.value), p.point) @dispatch def dadj(w: Spin3d, p: Moment3d) -> Moment3d: # type: ignore[reportRedeclaration] - return Moment3d.from_vector(so3.dadj(w.value, p.value)) + return Moment3d.from_vector(so3.dadj(w.value, p.value), p.point) @dispatch def dadj(w: Twist2d, p: Wrench2d) -> Wrench2d: # type: ignore[reportRedeclaration] - return Wrench2d.from_vector(se2.dadj(w.value, p.value)) + return Wrench2d.from_vector(se2.dadj(w.value, p.value), p.point) @dispatch def dadj(w: Twist3d, p: Wrench3d) -> Wrench3d: # type: ignore[reportRedeclaration] - return Wrench3d.from_vector(se3.dadj(w.value, p.value)) + return Wrench3d.from_vector(se3.dadj(w.value, p.value), p.point) @dispatch @@ -138,22 +138,22 @@ def dadj_op(w: Twist3d) -> Array: # type: ignore[reportRedeclaration] @dispatch def dadj_inv(w: Spin2d, p: Moment2d) -> Moment2d: # type: ignore[reportRedeclaration] - return Moment2d.from_vector(so2.dadj_inv(w.value, p.value)) + return Moment2d.from_vector(so2.dadj_inv(w.value, p.value), p.point) @dispatch def dadj_inv(w: Spin3d, p: Moment3d) -> Moment3d: # type: ignore[reportRedeclaration] - return Moment3d.from_vector(so3.dadj_inv(w.value, p.value)) + return Moment3d.from_vector(so3.dadj_inv(w.value, p.value), p.point) @dispatch def dadj_inv(w: Twist2d, p: Wrench2d) -> Wrench2d: # type: ignore[reportRedeclaration] - return Wrench2d.from_vector(se2.dadj_inv(w.value, p.value)) + return Wrench2d.from_vector(se2.dadj_inv(w.value, p.value), p.point) @dispatch def dadj_inv(w: Twist3d, p: Wrench3d) -> Wrench3d: # type: ignore[reportRedeclaration] - return Wrench3d.from_vector(se3.dadj_inv(w.value, p.value)) + return Wrench3d.from_vector(se3.dadj_inv(w.value, p.value), p.point) @dispatch @@ -178,22 +178,22 @@ def dadj_inv_op(w: Twist3d) -> Array: # type: ignore[reportRedeclaration] @dispatch def Adj(g: Rotation2d, w: Spin2d) -> Spin2d: # type: ignore[reportRedeclaration] - return Spin2d.from_matrix(so2.Adj(g.value, w.value)) + return Spin2d.from_matrix(so2.Adj(g.value, w.value), w.point) @dispatch def Adj(g: Rotation3d, w: Spin3d) -> Spin3d: # type: ignore[reportRedeclaration] - return Spin3d.from_matrix(so3.Adj(g.value, w.value)) + return Spin3d.from_matrix(so3.Adj(g.value, w.value), w.point) @dispatch def Adj(g: Isometry2d, w: Twist2d) -> Twist2d: # type: ignore[reportRedeclaration] - return Twist2d.from_matrix(se2.Adj(g.value, w.value)) + return Twist2d.from_matrix(se2.Adj(g.value, w.value), w.point) @dispatch def Adj(g: Isometry3d, w: Twist3d) -> Twist3d: # type: ignore[reportRedeclaration] - return Twist3d.from_matrix(se3.Adj(g.value, w.value)) + return Twist3d.from_matrix(se3.Adj(g.value, w.value), w.point) @dispatch @@ -218,22 +218,22 @@ def Adj_op(g: Isometry3d) -> Array: # type: ignore[reportRedeclaration] @dispatch def Adj_inv(g: Rotation2d, w: Spin2d) -> Spin2d: # type: ignore[reportRedeclaration] - return Spin2d.from_matrix(so2.Adj_inv(g.value, w.value)) + return Spin2d.from_matrix(so2.Adj_inv(g.value, w.value), w.point) @dispatch def Adj_inv(g: Rotation3d, w: Spin3d) -> Spin3d: # type: ignore[reportRedeclaration] - return Spin3d.from_matrix(so3.Adj_inv(g.value, w.value)) + return Spin3d.from_matrix(so3.Adj_inv(g.value, w.value), w.point) @dispatch def Adj_inv(g: Isometry2d, w: Twist2d) -> Twist2d: # type: ignore[reportRedeclaration] - return Twist2d.from_matrix(se2.Adj_inv(g.value, w.value)) + return Twist2d.from_matrix(se2.Adj_inv(g.value, w.value), w.point) @dispatch def Adj_inv(g: Isometry3d, w: Twist3d) -> Twist3d: # type: ignore[reportRedeclaration] - return Twist3d.from_matrix(se3.Adj_inv(g.value, w.value)) + return Twist3d.from_matrix(se3.Adj_inv(g.value, w.value), w.point) @dispatch @@ -258,22 +258,22 @@ def Adj_inv_op(g: Isometry3d) -> Array: # type: ignore[reportRedeclaration] @dispatch def dAdj(g: Rotation2d, p: Moment2d) -> Moment2d: # type: ignore[reportRedeclaration] - return Moment2d.from_vector(so2.dAdj(g.value, p.value)) + return Moment2d.from_vector(so2.dAdj(g.value, p.value), p.point) @dispatch def dAdj(g: Rotation3d, p: Moment3d) -> Moment3d: # type: ignore[reportRedeclaration] - return Moment3d.from_vector(so3.dAdj(g.value, p.value)) + return Moment3d.from_vector(so3.dAdj(g.value, p.value), p.point) @dispatch def dAdj(g: Isometry2d, p: Wrench2d) -> Wrench2d: # type: ignore[reportRedeclaration] - return Wrench2d.from_vector(se2.dAdj(g.value, p.value)) + return Wrench2d.from_vector(se2.dAdj(g.value, p.value), p.point) @dispatch def dAdj(g: Isometry3d, p: Wrench3d) -> Wrench3d: # type: ignore[reportRedeclaration] - return Wrench3d.from_vector(se3.dAdj(g.value, p.value)) + return Wrench3d.from_vector(se3.dAdj(g.value, p.value), p.point) @dispatch @@ -298,22 +298,22 @@ def dAdj_op(g: Isometry3d) -> Array: # type: ignore[reportRedeclaration] @dispatch def dAdj_inv(g: Rotation2d, p: Moment2d) -> Moment2d: # type: ignore[reportRedeclaration] - return Moment2d.from_vector(so2.dAdj_inv(g.value, p.value)) + return Moment2d.from_vector(so2.dAdj_inv(g.value, p.value), p.point) @dispatch def dAdj_inv(g: Rotation3d, p: Moment3d) -> Moment3d: # type: ignore[reportRedeclaration] - return Moment3d.from_vector(so3.dAdj_inv(g.value, p.value)) + return Moment3d.from_vector(so3.dAdj_inv(g.value, p.value), p.point) @dispatch def dAdj_inv(g: Isometry2d, p: Wrench2d) -> Wrench2d: # type: ignore[reportRedeclaration] - return Wrench2d.from_vector(se2.dAdj_inv(g.value, p.value)) + return Wrench2d.from_vector(se2.dAdj_inv(g.value, p.value), p.point) @dispatch def dAdj_inv(g: Isometry3d, p: Wrench3d) -> Wrench3d: # type: ignore[reportRedeclaration] - return Wrench3d.from_vector(se3.dAdj_inv(g.value, p.value)) + return Wrench3d.from_vector(se3.dAdj_inv(g.value, p.value), p.point) @dispatch @@ -478,22 +478,22 @@ def lminus(g: Isometry3d, h: Isometry3d) -> Twist3d: # type: ignore[reportRedec @dispatch def rminus(g: Rotation2d, h: Rotation2d) -> Spin2d: # type: ignore[reportRedeclaration] - return Spin2d.from_matrix(so2.rminus(g.value, h.value)) + return Spin2d.from_matrix(so2.rminus(g.value, h.value), h) @dispatch def rminus(g: Rotation3d, h: Rotation3d) -> Spin3d: # type: ignore[reportRedeclaration] - return Spin3d.from_matrix(so3.rminus(g.value, h.value)) + return Spin3d.from_matrix(so3.rminus(g.value, h.value), h) @dispatch def rminus(g: Isometry2d, h: Isometry2d) -> Twist2d: # type: ignore[reportRedeclaration] - return Twist2d.from_matrix(se2.rminus(g.value, h.value)) + return Twist2d.from_matrix(se2.rminus(g.value, h.value), h) @dispatch def rminus(g: Isometry3d, h: Isometry3d) -> Twist3d: # type: ignore[reportRedeclaration] - return Twist3d.from_matrix(se3.rminus(g.value, h.value)) + return Twist3d.from_matrix(se3.rminus(g.value, h.value), h) @dispatch diff --git a/fiber/random/_gaussian.py b/fiber/random/_gaussian.py index b341d94..3b3a966 100644 --- a/fiber/random/_gaussian.py +++ b/fiber/random/_gaussian.py @@ -94,8 +94,9 @@ def mean( # type: ignore[reportRedeclaration] atol=1e-6, max_steps: int = 100, throw: bool = True, + left: bool = True, ) -> Rotation2d: - est_mean = so2.random.mean(samples.value, rtol, atol, max_steps, throw) + est_mean = so2.random.mean(samples.value, rtol, atol, max_steps, throw, left) return Rotation2d.from_matrix(est_mean) @@ -106,8 +107,9 @@ def mean( # type: ignore[reportRedeclaration] atol=1e-6, max_steps: int = 100, throw: bool = True, + left: bool = True, ) -> Rotation3d: - est_mean = so3.random.mean(samples.value, rtol, atol, max_steps, throw) + est_mean = so3.random.mean(samples.value, rtol, atol, max_steps, throw, left) return Rotation3d.from_matrix(est_mean) @@ -118,8 +120,9 @@ def mean( # type: ignore[reportRedeclaration] atol=1e-6, max_steps: int = 100, throw: bool = True, + left: bool = True, ) -> Isometry2d: - est_mean = se2.random.mean(samples.value, rtol, atol, max_steps, throw) + est_mean = se2.random.mean(samples.value, rtol, atol, max_steps, throw, left) return Isometry2d.from_matrix(est_mean) @@ -130,26 +133,35 @@ def mean( # type: ignore[reportRedeclaration] atol=1e-6, max_steps: int = 100, throw: bool = True, + left: bool = True, ) -> Isometry3d: - est_mean = se3.random.mean(samples.value, rtol, atol, max_steps, throw) + est_mean = se3.random.mean(samples.value, rtol, atol, max_steps, throw, left) return Isometry3d.from_matrix(est_mean) @dispatch -def cov(mean: Rotation2d, samples: Rotation2d) -> Array: # type: ignore[reportRedeclaration] - return so2.random.cov(mean.value, samples.value) +def cov( # type: ignore[reportRedeclaration] + mean: Rotation2d, samples: Rotation2d, left: bool = True +) -> Array: + return so2.random.cov(mean.value, samples.value, left) @dispatch -def cov(mean: Rotation3d, samples: Rotation3d) -> Array: # type: ignore[reportRedeclaration] - return so3.random.cov(mean.value, samples.value) +def cov( # type: ignore[reportRedeclaration] + mean: Rotation3d, samples: Rotation3d, left: bool = True +) -> Array: + return so3.random.cov(mean.value, samples.value, left) @dispatch -def cov(mean: Isometry2d, samples: Isometry2d) -> Array: # type: ignore[reportRedeclaration] - return se2.random.cov(mean.value, samples.value) +def cov( # type: ignore[reportRedeclaration] + mean: Isometry2d, samples: Isometry2d, left: bool = True +) -> Array: + return se2.random.cov(mean.value, samples.value, left) @dispatch -def cov(mean: Isometry3d, samples: Isometry3d) -> Array: # type: ignore[reportRedeclaration] - return se3.random.cov(mean.value, samples.value) +def cov( # type: ignore[reportRedeclaration] + mean: Isometry3d, samples: Isometry3d, left: bool = True +) -> Array: + return se3.random.cov(mean.value, samples.value, left) From 1260f9bf4f749e2e06ee6260ee951ee894d622b7 Mon Sep 17 00:00:00 2001 From: Evan Palmer Date: Thu, 8 Oct 2026 17:48:03 -0700 Subject: [PATCH 3/3] update uv lock --- uv.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uv.lock b/uv.lock index c6031d6..5593bb6 100644 --- a/uv.lock +++ b/uv.lock @@ -46,7 +46,7 @@ wheels = [ [[package]] name = "fiber" -version = "0.1.3" +version = "0.1.4" source = { editable = "." } dependencies = [ { name = "diffrax" },