Fix Frechet mean objective function

fiber/_groups/se2/random/_gaussian.py

@@ -25,7 +25,6 @@
 import optimistix as optx
 from jaxtyping import Array, PRNGKeyArray
 
-from ...._vecfuncs import softnorm
 from .._element import _as_vector, _from_vector
 from .._operations import expm, logm, lplus, rminus, rplus
 
@@ -73,22 +72,22 @@ def mean(
     max_steps: int = 100,
     throw: bool = True,
 ) -> Array:
-    def residuals(mean, samples):
+    def update(mean, samples):
         errors = _as_vector(rminus(samples, mean))
-        return softnorm(jnp.sum(errors, axis=0))
+        return rplus(mean, _from_vector(jnp.mean(errors, axis=0)))
 
     # construct the initial guess to warm-start the solver
     exp_coords = _as_vector(logm(samples))
     init_mean = jnp.mean(exp_coords, axis=0)
     y0 = expm(_from_vector(init_mean))
 
-    # find the mean using a root-finder
-    sol = optx.root_find(
-        residuals,
-        optx.Newton(rtol, atol),
+    sol = optx.fixed_point(
+        update,
+        optx.FixedPointIteration(rtol, atol),
         y0,
         args=samples,
         max_steps=max_steps,
+        adjoint=optx.RecursiveCheckpointAdjoint(),
         throw=throw,
     )
 

fiber/_groups/se3/random/_gaussian.py

@@ -25,7 +25,6 @@
 import optimistix as optx
 from jaxtyping import Array, PRNGKeyArray
 
-from ...._vecfuncs import softnorm
 from .._element import _as_vector, _from_vector
 from .._operations import expm, logm, lplus, rminus, rplus
 
@@ -73,22 +72,22 @@ def mean(
     max_steps: int = 100,
     throw: bool = True,
 ) -> Array:
-    def residuals(mean, samples):
+    def update(mean, samples):
         errors = _as_vector(rminus(samples, mean))
-        return softnorm(jnp.sum(errors, axis=0))
+        return rplus(mean, _from_vector(jnp.mean(errors, axis=0)))
 
     # construct the initial guess to warm-start the solver
     exp_coords = _as_vector(logm(samples))
     init_mean = jnp.mean(exp_coords, axis=0)
     y0 = expm(_from_vector(init_mean))
 
-    # find the mean using a root-finder
-    sol = optx.root_find(
-        residuals,
-        optx.Newton(rtol, atol),
+    sol = optx.fixed_point(
+        update,
+        optx.FixedPointIteration(rtol, atol),
         y0,
         args=samples,
         max_steps=max_steps,
+        adjoint=optx.RecursiveCheckpointAdjoint(),
         throw=throw,
     )
 

fiber/_groups/so2/random/_gaussian.py

@@ -25,7 +25,7 @@
 import optimistix as optx
 from jaxtyping import Array, PRNGKeyArray
 
-from ...._vecfuncs import skew2, softnorm, vex2
+from ...._vecfuncs import skew2, vex2
 from .._operations import expm, logm, lplus, rminus, rplus
 
 
@@ -76,22 +76,22 @@ def mean(
     max_steps: int = 100,
     throw: bool = True,
 ) -> Array:
-    def residuals(mean, samples):
+    def update(mean, samples):
         errors = vex2(rminus(samples, mean))
-        return softnorm(jnp.sum(errors, axis=0))
+        return rplus(mean, skew2(jnp.mean(errors, axis=0)))
 
     # construct the initial guess to warm-start the solver
     exp_coords = vex2(logm(samples))
     init_mean = jnp.mean(exp_coords, axis=0)
     y0 = expm(skew2(init_mean))
 
-    # find the mean using a root-finder
-    sol = optx.root_find(
-        residuals,
-        optx.Newton(rtol, atol),
+    sol = optx.fixed_point(
+        update,
+        optx.FixedPointIteration(rtol, atol),
         y0,
         args=samples,
         max_steps=max_steps,
+        adjoint=optx.RecursiveCheckpointAdjoint(),
         throw=throw,
     )
 

fiber/_groups/so3/random/_gaussian.py

@@ -25,7 +25,7 @@
 import optimistix as optx
 from jaxtyping import Array, PRNGKeyArray
 
-from ...._vecfuncs import skew3, softnorm, vex3
+from ...._vecfuncs import skew3, vex3
 from .._operations import expm, logm, lplus, rminus, rplus
 
 
@@ -72,22 +72,22 @@ def mean(
     max_steps: int = 100,
     throw: bool = True,
 ) -> Array:
-    def residuals(mean, samples):
+    def update(mean, samples):
         errors = vex3(rminus(samples, mean))
-        return softnorm(jnp.sum(errors, axis=0))
+        return rplus(mean, skew3(jnp.mean(errors, axis=0)))
 
     # construct the initial guess to warm-start the solver
     exp_coords = vex3(logm(samples))
     init_mean = jnp.mean(exp_coords, axis=0)
     y0 = expm(skew3(init_mean))
 
-    # find the mean using a root-finder
-    sol = optx.root_find(
-        residuals,
-        optx.Newton(rtol, atol),
+    sol = optx.fixed_point(
+        update,
+        optx.FixedPointIteration(rtol, atol),
         y0,
         args=samples,
         max_steps=max_steps,
+        adjoint=optx.RecursiveCheckpointAdjoint(),
         throw=throw,
     )
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "fiber"
-version = "0.1.2"
+version = "0.1.3"
 description = "A JAX library for Lie group operations."
 readme = "README.md"
 requires-python = ">=3.14"