Add a Leapfrog solver for diffrax

src/galax/dynamics/_src/experimental/__init__.py

@@ -1,4 +1,5 @@
 """Experimental dynamics."""
 
 from .integrate import *
+from .leapfrog import *
 from .stream import *

src/galax/dynamics/_src/experimental/integrate.py

@@ -21,7 +21,6 @@
 import galax.dynamics.loop_strategies as lstrat
 import galax.potential as gp
 from galax.dynamics._src.orbit.field_base import AbstractOrbitField
-from galax.dynamics._src.orbit.field_hamiltonian import HamiltonianField
 
 BQParr: TypeAlias = tuple[Real[gdt.Qarr, "B"], Real[gdt.Parr, "B"]]
 
@@ -377,6 +376,11 @@ def integrate_orbit(
         evaluation of the solution.
 
     """
+    # Note: this is needed to prevent a circular import
+    from galax.dynamics._src.orbit.field_hamiltonian import (
+        HamiltonianField,
+    )
+
     field = pot if isinstance(pot, AbstractOrbitField) else HamiltonianField(pot)
     terms = field.terms(solver)
 

src/galax/dynamics/_src/experimental/leapfrog.py

@@ -0,0 +1,142 @@
+# ruff: noqa: ARG002
+"""
+Note: This module implements a diffrax solver for Leapfrog integration. There is a
+stalled PR to add a similar integrator to diffrax, so in the meantime we implement it
+here.
+"""  # noqa: D205
+
+__all__ = ["Leapfrog", "SymplecticSolverT"]
+
+from collections.abc import Callable
+from jaxtyping import ArrayLike, Float, PyTree
+from typing import Any, ClassVar, TypeAlias
+
+from diffrax import (
+    RESULTS,
+    AbstractSolver,
+    AbstractTerm,
+    LocalLinearInterpolation,
+    ODETerm,
+    SemiImplicitEuler,
+)
+from equinox.internal import ω  # noqa: PLC2403
+
+from galax.dynamics._src.orbit.field_base import AbstractOrbitField
+from galax.dynamics._src.orbit.field_hamiltonian import HamiltonianField
+
+# diffrax doesn't publicly export these; they're type hints only (no runtime
+# behavior depends on them), so redeclare as `Any` rather than importing
+# diffrax's private `_custom_types` module.
+VF: TypeAlias = Any
+Args: TypeAlias = Any
+BoolScalarLike: TypeAlias = Any
+DenseInfo: TypeAlias = Any
+RealScalarLike: TypeAlias = Any
+
+_ErrorEstimate: TypeAlias = None
+_SolverState: TypeAlias = None
+
+Ya: TypeAlias = PyTree[Float[ArrayLike, "?*y"], " Y"]
+Yb: TypeAlias = PyTree[Float[ArrayLike, "?*y"], " Y"]
+
+
+class Leapfrog(AbstractSolver):  # type: ignore[misc]
+    """Leapfrog (velocity Verlet) symplectic integrator.
+
+    This is a 2nd order symplectic integration method. This integrator does not support
+    adaptive step sizing: it provides no error estimate, so pair it with
+    `diffrax.ConstantStepSize` (e.g. via ``gd.OrbitSolver(Leapfrog(),
+    stepsize_controller=dfx.ConstantStepSize())``) rather than
+    `galax.dynamics.OrbitSolver`'s default adaptive `diffrax.PIDController`, which will
+    raise a `RuntimeError` for a solver without error estimates. This is either known as
+    kick-drift-kick leapfrog or velocity Verlet.
+
+    Assuming that:
+
+        x0, v0 = y0
+
+    and:
+
+        f, g = terms
+
+    This method computes the next step as:
+
+        v_half = v0 + h/2 * g(t0, x0)
+        x1 = x0 + h * f(t0, v_half)
+        v1 = v_half + h/2 * g(t1, x1)
+    """
+
+    term_structure: ClassVar = (AbstractTerm, AbstractTerm)
+    interpolation_cls: ClassVar[Callable[..., LocalLinearInterpolation]] = (
+        LocalLinearInterpolation
+    )
+
+    def order(self, _: Any) -> int:
+        return 2
+
+    def init(
+        self,
+        terms: tuple[AbstractTerm, AbstractTerm],
+        t0: RealScalarLike,
+        t1: RealScalarLike,
+        y0: tuple[Ya, Yb],
+        args: Args,
+    ) -> _SolverState:
+        return None
+
+    def step(
+        self,
+        terms: tuple[AbstractTerm, AbstractTerm],
+        t0: RealScalarLike,
+        t1: RealScalarLike,
+        y0: tuple[Ya, Yb],
+        args: Args,
+        solver_state: _SolverState,
+        made_jump: BoolScalarLike,
+    ) -> tuple[tuple[Ya, Yb], _ErrorEstimate, DenseInfo, _SolverState, RESULTS]:
+        del solver_state, made_jump
+
+        f, g = terms
+        q0, p0 = y0
+        h = t1 - t0
+
+        v_half = (p0**ω + 0.5 * h * g.vf(t0, q0, args) ** ω).ω
+        q1 = (q0**ω + h * f.vf(t0, v_half, args) ** ω).ω
+        p1 = (v_half**ω + 0.5 * h * g.vf(t1, q1, args) ** ω).ω
+
+        y1 = (q1, p1)
+        dense_info = {"y0": y0, "y1": y1}
+        return y1, None, dense_info, None, RESULTS.successful
+
+    def func(
+        self,
+        terms: tuple[AbstractTerm, AbstractTerm],
+        t0: RealScalarLike,
+        y0: tuple[Ya, Yb],
+        args: Args,
+    ) -> VF:
+        f, g = terms
+        q0, p0 = y0
+        qdot = f.vf(t0, p0, args)
+        pdot = g.vf(t0, q0, args)
+        return qdot, pdot
+
+
+Leapfrog.__init__.__doc__ = """**Arguments:** None"""
+
+
+SymplecticSolverT: TypeAlias = Leapfrog | SemiImplicitEuler
+
+
+# ===============================================
+# Terms dispatch
+#
+# Registered here (rather than in `field_hamiltonian.py`) so that core orbit
+# code doesn't need to import the experimental package just to support
+# `Leapfrog`. See `field_hamiltonian.py` for the `SemiImplicitEuler` dispatch.
+
+
+@AbstractOrbitField.terms.dispatch
+def terms(self: HamiltonianField, _: Leapfrog, /) -> tuple[ODETerm, ODETerm]:
+    """Return the AbstractTerm terms for the Leapfrog solver."""
+    return (ODETerm(self.dx_dt), ODETerm(self.dv_dt))

src/galax/dynamics/_src/orbit/field_hamiltonian.py

@@ -261,6 +261,11 @@ def terms(
 ) -> tuple[dfx.ODETerm, dfx.ODETerm]:
     r"""Return the AbstractTerm terms for the SemiImplicitEuler solver.
 
+    See also `galax.dynamics.experimental.Leapfrog`, another symplectic solver,
+    for which the analogous dispatch is registered in
+    `galax.dynamics._src.experimental.leapfrog` (to avoid this core module
+    importing the experimental package).
+
     Examples
     --------
     >>> import diffrax as dfx

tests/unit/dynamics/experimental/test_leapfrog.py

@@ -0,0 +1,111 @@
+# ruff: noqa: ARG005
+
+import diffrax as dfx
+import jax.numpy as jnp
+import pytest
+
+from galax.dynamics import experimental
+
+
+def shaped_allclose(x, y, **kwargs):
+    return jnp.shape(x) == jnp.shape(y) and jnp.allclose(x, y, **kwargs)
+
+
+@pytest.mark.parametrize("solver", [experimental.Leapfrog()])
+def test_symplectic_solvers(solver):
+    dq_dt = dfx.ODETerm(lambda t, p, args: p)
+    dp_dt = dfx.ODETerm(lambda t, q, args: -q)
+    y0 = (1.0, -0.5)
+    dt0 = 0.001
+    sol1 = dfx.diffeqsolve(
+        (dq_dt, dp_dt),
+        solver,
+        0,
+        1,
+        dt0,
+        y0,
+        max_steps=2000,
+    )
+    term_combined = dfx.ODETerm(lambda t, y, args: (y[1], -y[0]))
+    sol2 = dfx.diffeqsolve(term_combined, dfx.Tsit5(), 0, 1, 0.001, y0)
+    assert shaped_allclose(sol1.ys[0], sol2.ys[0])
+    assert shaped_allclose(sol1.ys[1], sol2.ys[1])
+
+
+def test_leapfrog_time_dependent_force():
+    """The second kick must evaluate the force at ``t1`` (post-drift), not ``t0``.
+
+    A time-independent force can't distinguish these -- galax's potentials are
+    usually static, so this guards against a future refactor silently swapping
+    the two.  With a time-dependent force, evaluating the second kick at the
+    wrong time turns the O(h^2) global error into O(h).
+    """
+    dq_dt = dfx.ODETerm(lambda t, p, args: p)
+    dp_dt = dfx.ODETerm(lambda t, q, args: -q + 0.3 * jnp.sin(2 * t))
+    y0 = (1.0, -0.5)
+    dt0 = 0.01
+
+    sol = dfx.diffeqsolve(
+        (dq_dt, dp_dt), experimental.Leapfrog(), 0, 1, dt0, y0, max_steps=1000
+    )
+
+    term_combined = dfx.ODETerm(lambda t, y, args: (y[1], -y[0] + 0.3 * jnp.sin(2 * t)))
+    ref = dfx.diffeqsolve(
+        term_combined,
+        dfx.Tsit5(),
+        0,
+        1,
+        0.0001,
+        y0,
+        stepsize_controller=dfx.PIDController(rtol=1e-10, atol=1e-10),
+        max_steps=200_000,
+    )
+
+    # Correct: ~1.5e-5. A t0/t1 swap in the 2nd kick: ~1e-3. 1e-4 cleanly
+    # separates the two with margin on both sides.
+    assert jnp.abs(sol.ys[0][-1] - ref.ys[0][-1]) < 1e-4
+    assert jnp.abs(sol.ys[1][-1] - ref.ys[1][-1]) < 1e-4
+
+
+def test_leapfrog_conserves_energy_over_many_periods():
+    """A symplectic integrator's energy error oscillates; it should not drift.
+
+    Uses a step size coarse enough (50 steps/period) that the per-step error
+    is easily visible, over enough periods that a bug breaking the symplectic
+    structure (e.g. the drift step using the wrong velocity) would show up as
+    unbounded growth rather than being masked by an overly accurate step.
+    """
+    dq_dt = dfx.ODETerm(lambda t, p, args: p)
+    dp_dt = dfx.ODETerm(lambda t, q, args: -q)
+    y0 = (1.0, 0.0)
+
+    period = 2 * jnp.pi
+    n_periods = 100
+    dt0 = period / 50
+    t1 = n_periods * period
+    saveat = dfx.SaveAt(ts=jnp.linspace(0, t1, n_periods * 10))
+
+    sol = dfx.diffeqsolve(
+        (dq_dt, dp_dt),
+        experimental.Leapfrog(),
+        0,
+        t1,
+        dt0,
+        y0,
+        saveat=saveat,
+        max_steps=1_000_000,
+    )
+    q, p = sol.ys
+    energy = 0.5 * p**2 + 0.5 * q**2
+    rel_err = jnp.abs(energy - energy[0]) / energy[0]
+
+    half = len(rel_err) // 2
+    first_half_max = rel_err[:half].max()
+    second_half_max = rel_err[half:].max()
+
+    # Bounded, not growing: 100 periods in, the error is still the same order
+    # as it was at the start (a non-symplectic or otherwise broken step blows
+    # up by many orders of magnitude over this many periods -- see e.g. a
+    # drift step using v0 instead of v_half).
+    assert second_half_max < 5 * first_half_max
+    assert second_half_max < 0.05