Langevin PR

diffrax/__init__.py

@@ -71,6 +71,7 @@
     AbstractDIRK as AbstractDIRK,
     AbstractERK as AbstractERK,
     AbstractESDIRK as AbstractESDIRK,
+    AbstractFosterLangevinSRK as AbstractFosterLangevinSRK,
     AbstractImplicitSolver as AbstractImplicitSolver,
     AbstractItoSolver as AbstractItoSolver,
     AbstractRungeKutta as AbstractRungeKutta,
@@ -79,6 +80,7 @@
     AbstractSRK as AbstractSRK,
     AbstractStratonovichSolver as AbstractStratonovichSolver,
     AbstractWrappedSolver as AbstractWrappedSolver,
+    ALIGN as ALIGN,
     Bosh3 as Bosh3,
     ButcherTableau as ButcherTableau,
     CalculateJacobian as CalculateJacobian,
@@ -100,11 +102,13 @@
     LeapfrogMidpoint as LeapfrogMidpoint,
     Midpoint as Midpoint,
     MultiButcherTableau as MultiButcherTableau,
+    QUICSORT as QUICSORT,
     Ralston as Ralston,
     ReversibleHeun as ReversibleHeun,
     SEA as SEA,
     SemiImplicitEuler as SemiImplicitEuler,
     ShARK as ShARK,
+    ShOULD as ShOULD,
     Sil3 as Sil3,
     SlowRK as SlowRK,
     SPaRK as SPaRK,
@@ -125,6 +129,8 @@
     ControlTerm as ControlTerm,
     MultiTerm as MultiTerm,
     ODETerm as ODETerm,
+    UnderdampedLangevinDiffusionTerm as UnderdampedLangevinDiffusionTerm,
+    UnderdampedLangevinDriftTerm as UnderdampedLangevinDriftTerm,
     WeaklyDiagonalControlTerm as WeaklyDiagonalControlTerm,
 )
 

diffrax/_solver/__init__.py

@@ -1,3 +1,4 @@
+from .align import ALIGN as ALIGN
 from .base import (
     AbstractAdaptiveSolver as AbstractAdaptiveSolver,
     AbstractImplicitSolver as AbstractImplicitSolver,
@@ -12,6 +13,7 @@
 from .dopri8 import Dopri8 as Dopri8
 from .euler import Euler as Euler
 from .euler_heun import EulerHeun as EulerHeun
+from .foster_langevin_srk import AbstractFosterLangevinSRK as AbstractFosterLangevinSRK
 from .heun import Heun as Heun
 from .implicit_euler import ImplicitEuler as ImplicitEuler
 from .kencarp3 import KenCarp3 as KenCarp3
@@ -26,6 +28,7 @@
     ItoMilstein as ItoMilstein,
     StratonovichMilstein as StratonovichMilstein,
 )
+from .quicsort import QUICSORT as QUICSORT
 from .ralston import Ralston as Ralston
 from .reversible_heun import ReversibleHeun as ReversibleHeun
 from .runge_kutta import (
@@ -42,6 +45,7 @@
 from .semi_implicit_euler import SemiImplicitEuler as SemiImplicitEuler
 from .shark import ShARK as ShARK
 from .shark_general import GeneralShARK as GeneralShARK
+from .should import ShOULD as ShOULD
 from .sil3 import Sil3 as Sil3
 from .slowrk import SlowRK as SlowRK
 from .spark import SPaRK as SPaRK

diffrax/_solver/align.py

@@ -0,0 +1,188 @@
+import equinox as eqx
+import jax.numpy as jnp
+import jax.tree_util as jtu
+from equinox.internal import ω
+from jaxtyping import ArrayLike, PyTree
+
+from .._custom_types import (
+    AbstractSpaceTimeLevyArea,
+    RealScalarLike,
+)
+from .._local_interpolation import LocalLinearInterpolation
+from .._term import (
+    UnderdampedLangevinLeaf,
+    UnderdampedLangevinTuple,
+    UnderdampedLangevinX,
+)
+from .foster_langevin_srk import (
+    AbstractCoeffs,
+    AbstractFosterLangevinSRK,
+    UnderdampedLangevinArgs,
+)
+
+
+# For an explanation of the coefficients, see foster_langevin_srk.py
+class _ALIGNCoeffs(AbstractCoeffs):
+    beta: PyTree[ArrayLike]
+    a1: PyTree[ArrayLike]
+    b1: PyTree[ArrayLike]
+    aa: PyTree[ArrayLike]
+    chh: PyTree[ArrayLike]
+    dtype: jnp.dtype = eqx.field(static=True)
+
+    def __init__(self, beta, a1, b1, aa, chh):
+        self.beta = beta
+        self.a1 = a1
+        self.b1 = b1
+        self.aa = aa
+        self.chh = chh
+        all_leaves = jtu.tree_leaves([self.beta, self.a1, self.b1, self.aa, self.chh])
+        self.dtype = jnp.result_type(*all_leaves)
+
+
+_ErrorEstimate = UnderdampedLangevinTuple
+
+
+class ALIGN(AbstractFosterLangevinSRK[_ALIGNCoeffs, _ErrorEstimate]):
+    r"""The Adaptive Langevin via Interpolated Gradients and Noise method
+    designed by James Foster. This is a second order solver for the
+    Underdamped Langevin Diffusion, the terms for which can be created using
+    [`diffrax.make_underdamped_langevin_term`][]. Uses two evaluations of the vector
+    field per step, but is FSAL, so in practice it only requires one.
+
+    ??? cite "Reference"
+
+        This is a modification of the Strang-Splitting method from Definition 4.2 of
+
+        ```bibtex
+        @misc{foster2021shiftedode,
+            title={The shifted ODE method for underdamped Langevin MCMC},
+            author={James Foster and Terry Lyons and Harald Oberhauser},
+            year={2021},
+            eprint={2101.03446},
+            archivePrefix={arXiv},
+            primaryClass={math.NA},
+            url={https://arxiv.org/abs/2101.03446},
+        }
+        ```
+
+    """
+
+    interpolation_cls = LocalLinearInterpolation
+    minimal_levy_area = AbstractSpaceTimeLevyArea
+    taylor_threshold: RealScalarLike = eqx.field(static=True)
+    _is_fsal = True
+
+    def __init__(self, taylor_threshold: RealScalarLike = 0.1):
+        r"""**Arguments:**
+
+        - `taylor_threshold`: If the product `h*gamma` is less than this, then
+        the Taylor expansion will be used to compute the coefficients.
+        Otherwise they will be computed directly. When using float32, the
+        empirically optimal value is 0.1, and for float64 about 0.01.
+        """
+        self.taylor_threshold = taylor_threshold
+
+    def order(self, terms):
+        return 2
+
+    def strong_order(self, terms):
+        return 2.0
+
+    def _directly_compute_coeffs_leaf(
+        self, h: RealScalarLike, c: UnderdampedLangevinLeaf
+    ) -> _ALIGNCoeffs:
+        del self
+        # c is a leaf of gamma
+        # compute the coefficients directly (as opposed to via Taylor expansion)
+        al = c * h
+        beta = jnp.exp(-al)
+        a1 = (1 - beta) / c
+        b1 = (beta + al - 1) / (c * al)
+        aa = a1 / h
+
+        al2 = al**2
+        chh = 6 * (beta * (al + 2) + al - 2) / (al2 * c)
+
+        return _ALIGNCoeffs(
+            beta=beta,
+            a1=a1,
+            b1=b1,
+            aa=aa,
+            chh=chh,
+        )
+
+    def _tay_coeffs_single(self, c: UnderdampedLangevinLeaf) -> _ALIGNCoeffs:
+        del self
+        # c is a leaf of gamma
+        zero = jnp.zeros_like(c)
+        one = jnp.ones_like(c)
+        c2 = jnp.square(c)
+        c3 = c2 * c
+        c4 = c3 * c
+        c5 = c4 * c
+
+        # Coefficients of the Taylor expansion, starting from 5th power
+        # to 0th power. The descending power order is because of jnp.polyval
+        beta = jnp.stack([-c5 / 120, c4 / 24, -c3 / 6, c2 / 2, -c, one], axis=-1)
+        a1 = jnp.stack([c4 / 120, -c3 / 24, c2 / 6, -c / 2, one, zero], axis=-1)
+        b1 = jnp.stack([c4 / 720, -c3 / 120, c2 / 24, -c / 6, one / 2, zero], axis=-1)
+        aa = jnp.stack([-c5 / 720, c4 / 120, -c3 / 24, c2 / 6, -c / 2, one], axis=-1)
+        chh = jnp.stack([c4 / 168, -c3 / 30, 3 * c2 / 20, -c / 2, one, zero], axis=-1)
+
+        correct_shape = jnp.shape(c) + (6,)
+        assert (
+            beta.shape == a1.shape == b1.shape == aa.shape == chh.shape == correct_shape
+        )
+
+        return _ALIGNCoeffs(
+            beta=beta,
+            a1=a1,
+            b1=b1,
+            aa=aa,
+            chh=chh,
+        )
+
+    def _compute_step(
+        self,
+        h: RealScalarLike,
+        levy: AbstractSpaceTimeLevyArea,
+        x0: UnderdampedLangevinX,
+        v0: UnderdampedLangevinX,
+        underdamped_langevin_args: UnderdampedLangevinArgs,
+        coeffs: _ALIGNCoeffs,
+        rho: UnderdampedLangevinX,
+        prev_f: UnderdampedLangevinX,
+    ) -> tuple[
+        UnderdampedLangevinX,
+        UnderdampedLangevinX,
+        UnderdampedLangevinX,
+        UnderdampedLangevinTuple,
+    ]:
+        dtypes = jtu.tree_map(jnp.result_type, x0)
+        w: UnderdampedLangevinX = jtu.tree_map(jnp.asarray, levy.W, dtypes)
+        hh: UnderdampedLangevinX = jtu.tree_map(jnp.asarray, levy.H, dtypes)
+
+        gamma, u, f = underdamped_langevin_args
+
+        uh = (u**ω * h).ω
+        f0 = prev_f
+        x1 = (
+            x0**ω
+            + coeffs.a1**ω * v0**ω
+            - coeffs.b1**ω * uh**ω * f0**ω
+            + rho**ω * (coeffs.b1**ω * w**ω + coeffs.chh**ω * hh**ω)
+        ).ω
+        f1 = f(x1)
+        v1 = (
+            coeffs.beta**ω * v0**ω
+            - u**ω * ((coeffs.a1**ω - coeffs.b1**ω) * f0**ω + coeffs.b1**ω * f1**ω)
+            + rho**ω * (coeffs.aa**ω * w**ω - gamma**ω * coeffs.chh**ω * hh**ω)
+        ).ω
+
+        error_estimate = (
+            jtu.tree_map(jnp.zeros_like, x0),
+            (-(u**ω) * coeffs.b1**ω * (f1**ω - f0**ω)).ω,
+        )
+
+        return x1, v1, f1, error_estimate

diffrax/_solver/euler_heun.py

@@ -1,13 +1,13 @@
 from collections.abc import Callable
-from typing import ClassVar
+from typing import Any, ClassVar
 from typing_extensions import TypeAlias
 
 from equinox.internal import ω
 
 from .._custom_types import Args, BoolScalarLike, DenseInfo, RealScalarLike, VF, Y
 from .._local_interpolation import LocalLinearInterpolation
 from .._solution import RESULTS
-from .._term import AbstractTerm, MultiTerm, ODETerm
+from .._term import AbstractTerm, MultiTerm
 from .base import AbstractStratonovichSolver
 
 
@@ -26,7 +26,9 @@ class EulerHeun(AbstractStratonovichSolver):
     Used to solve SDEs, and converges to the Stratonovich solution.
     """
 
-    term_structure: ClassVar = MultiTerm[tuple[ODETerm, AbstractTerm]]
+    term_structure: ClassVar = MultiTerm[
+        tuple[AbstractTerm[Any, RealScalarLike], AbstractTerm]
+    ]
     interpolation_cls: ClassVar[
         Callable[..., LocalLinearInterpolation]
     ] = LocalLinearInterpolation
@@ -39,7 +41,7 @@ def strong_order(self, terms):
 
     def init(
         self,
-        terms: MultiTerm[tuple[ODETerm, AbstractTerm]],
+        terms: MultiTerm[tuple[AbstractTerm[Any, RealScalarLike], AbstractTerm]],
         t0: RealScalarLike,
         t1: RealScalarLike,
         y0: Y,
@@ -49,7 +51,7 @@ def init(
 
     def step(
         self,
-        terms: MultiTerm[tuple[ODETerm, AbstractTerm]],
+        terms: MultiTerm[tuple[AbstractTerm[Any, RealScalarLike], AbstractTerm]],
         t0: RealScalarLike,
         t1: RealScalarLike,
         y0: Y,