Fix Euler-Heun solver

fiber/solvers/__init__.py

@@ -20,7 +20,7 @@
 
 from ._euler import LieEuler
 from ._euler_heun import EulerHeun
-from ._term import ImplicitVariationalTerm, SharpTerm, VariationalDiffusionTerm
+from ._term import ImplicitVariationalTerm, SharpTerm, VariationalControlTerm
 from ._variational_integrator import LieSVI
 
 __all__ = [
@@ -29,5 +29,5 @@
     "LieEuler",
     "LieSVI",
     "SharpTerm",
-    "VariationalDiffusionTerm",
+    "VariationalControlTerm",
 ]

fiber/solvers/_euler_heun.py

@@ -25,7 +25,7 @@
 from diffrax import RESULTS, AbstractStratonovichSolver, AbstractTerm, MultiTerm
 from diffrax._term import WrapTerm
 
-from .._custom_types import VF, Args, BoolScalarLike, Control, DenseInfo, RealScalarLike
+from .._custom_types import VF, Args, BoolScalarLike, DenseInfo, RealScalarLike
 from .._groups import AbstractCotangentVector, AbstractTangentVector
 from .._local_interpolation import LocalLeftBundleInterpolation as LocalInterpolation
 from .._operations import rplus
@@ -38,7 +38,7 @@
 _Terms = MultiTerm[
     tuple[
         SharpTerm[AbstractCotangentVector],
-        AbstractTerm[AbstractCotangentVector, Control],
+        AbstractTerm,
     ]
 ]
 
@@ -85,10 +85,14 @@ def step(
         dw = diffusion.contr(t0, t1)
 
         f0 = drift.vf_prod(t0, y0, args, dt)
-        h0 = diffusion.vf_prod(t0, y0, args, dw)
+        cv = type(f0)
+
+        h0 = diffusion.prod(diffusion.vf(t0, y0, args), dw)
+        h0 = cv.from_vector(h0, point=y0.point)
 
         y_prime = y0 + dual_metric(y0, h0)  # type: ignore[reportOperatorIssue]
-        h_prime = diffusion.vf_prod(t0, y_prime, args, dw)
+        h_prime = diffusion.prod(diffusion.vf(t0, y_prime, args), dw)
+        h_prime = cv.from_vector(h_prime, point=y_prime.point)
 
         k = f0 + 0.5 * (h0 + h_prime)
         y1 = y0 + dual_metric(y0, k)  # type: ignore[reportOperatorIssue]

fiber/solvers/_term.py

@@ -70,7 +70,7 @@ def vf_prod(self, t: RealScalarLike, y: _V, args: Args, control: Control) -> _CV
         return self.vf(t, y, args, control)
 
 
-class VariationalDiffusionTerm(AbstractTerm[_CV, Control]):
+class VariationalControlTerm(AbstractTerm[_CV, Control]):
     vector_field: Callable[[RealScalarLike, _V, Args, Control], _CV]
     control: Control
 

fiber/solvers/_variational_integrator.py

@@ -36,7 +36,7 @@
 from .._custom_types import Args, BoolScalarLike, DenseInfo, RealScalarLike
 from .._groups._element import AbstractCotangentVector, AbstractTangentVector
 from .._local_interpolation import LocalLeftBundleInterpolation as LocalInterpolation
-from ._term import ImplicitVariationalTerm, VariationalDiffusionTerm
+from ._term import ImplicitVariationalTerm, VariationalControlTerm
 
 type _ErrorEstimate = None
 type _SolverState = None
@@ -46,7 +46,7 @@
 _Terms = MultiTerm[
     tuple[
         ImplicitVariationalTerm[AbstractCotangentVector],
-        VariationalDiffusionTerm[AbstractCotangentVector],
+        VariationalControlTerm[AbstractCotangentVector],
     ]
 ]
 

pyproject.toml

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