Merge pull request #251 from patrick-kidger/implicit-euler-adaptive

Implicit Euler is now adaptive (basically necessary for it to be usef…

Author: patrick-kidger

diffrax/local_interpolation.py

@@ -41,18 +41,18 @@ def __init__(
         *,
         y0: PyTree[Array["dims":...]],  # noqa: F821
         y1: PyTree[Array["dims":...]],  # noqa: F821
-        f0: PyTree[Array["dims":...]],  # noqa: F821
-        f1: PyTree[Array["dims":...]],  # noqa: F821
+        k0: PyTree[Array["dims":...]],  # noqa: F821
+        k1: PyTree[Array["dims":...]],  # noqa: F821
         **kwargs
     ):
         super().__init__(**kwargs)
 
-        def _calculate(_y0, _y1, _f0, _f1):
-            _a = _f0 + _f1 + 2 * _y0 - 2 * _y1
-            _b = -2 * _f0 - _f1 - 3 * _y0 + 3 * _y1
-            return jnp.stack([_a, _b, _f0, _y0])
+        def _calculate(_y0, _y1, _k0, _k1):
+            _a = _k0 + _k1 + 2 * _y0 - 2 * _y1
+            _b = -2 * _k0 - _k1 - 3 * _y0 + 3 * _y1
+            return jnp.stack([_a, _b, _k0, _y0])
 
-        self.coeffs = jtu.tree_map(_calculate, y0, y1, f0, f1)
+        self.coeffs = jtu.tree_map(_calculate, y0, y1, k0, k1)
 
     @classmethod
     def from_k(
@@ -63,7 +63,7 @@ def from_k(
         k: PyTree[Array["order", "dims":...]],  # noqa: F821
         **kwargs
     ):
-        return cls(y0=y0, y1=y1, f0=ω(k)[0].ω, f1=ω(k)[-1].ω, **kwargs)
+        return cls(y0=y0, y1=y1, k0=ω(k)[0].ω, k1=ω(k)[-1].ω, **kwargs)
 
     def evaluate(
         self, t0: Scalar, t1: Optional[Scalar] = None, left: bool = True

diffrax/solver/base.py

@@ -85,6 +85,7 @@ def error_order(self, terms: PyTree[AbstractTerm]) -> Optional[Scalar]:
         else:
             return self.order(terms)
 
+    @abc.abstractmethod
     def init(
         self,
         terms: PyTree[AbstractTerm],
@@ -101,7 +102,6 @@ def init(
 
         The initial solver state, which should be used the first time `step` is called.
         """
-        return None
 
     @abc.abstractmethod
     def step(

diffrax/solver/euler.py

@@ -30,6 +30,16 @@ def order(self, terms):
     def strong_order(self, terms):
         return 0.5
 
+    def init(
+        self,
+        terms: AbstractTerm,
+        t0: Scalar,
+        t1: Scalar,
+        y0: PyTree,
+        args: PyTree,
+    ) -> _SolverState:
+        return None
+
     def step(
         self,
         terms: AbstractTerm,

diffrax/solver/euler_heun.py

@@ -28,6 +28,16 @@ def order(self, terms):
     def strong_order(self, terms):
         return 0.5
 
+    def init(
+        self,
+        terms: Tuple[ODETerm, AbstractTerm],
+        t0: Scalar,
+        t1: Scalar,
+        y0: PyTree,
+        args: PyTree,
+    ) -> _SolverState:
+        return None
+
     def step(
         self,
         terms: Tuple[ODETerm, AbstractTerm],

diffrax/solver/implicit_euler.py

@@ -3,13 +3,13 @@
 from equinox.internal import ω
 
 from ..custom_types import Bool, DenseInfo, PyTree, Scalar
+from ..heuristics import is_sde
 from ..local_interpolation import LocalLinearInterpolation
 from ..solution import RESULTS
 from ..term import AbstractTerm
 from .base import AbstractImplicitSolver
 
 
-_ErrorEstimate = None
 _SolverState = None
 
 
@@ -22,15 +22,36 @@ def _implicit_relation(z1, nonlinear_solve_args):
 class ImplicitEuler(AbstractImplicitSolver):
     r"""Implicit Euler method.
 
-    A-B-L stable 1st order SDIRK method. Does not support adaptive step sizing.
+    A-B-L stable 1st order SDIRK method. Has an embedded 2nd order method for adaptive
+    step sizing.
     """
 
     term_structure = AbstractTerm
+    # We actually have enough information to use 3rd order Hermite interpolation.
+    #
+    # We don't use it as this seems to be quite a bad choice for low-order solvers: it
+    # produces very oscillatory interpolations.
     interpolation_cls = LocalLinearInterpolation
 
     def order(self, terms):
         return 1
 
+    def error_order(self, terms):
+        if is_sde(terms):
+            return None
+        else:
+            return 2
+
+    def init(
+        self,
+        terms: AbstractTerm,
+        t0: Scalar,
+        t1: Scalar,
+        y0: PyTree,
+        args: PyTree,
+    ) -> _SolverState:
+        return None
+
     def step(
         self,
         terms: AbstractTerm,
@@ -40,20 +61,28 @@ def step(
         args: PyTree,
         solver_state: _SolverState,
         made_jump: Bool,
-    ) -> Tuple[PyTree, _ErrorEstimate, DenseInfo, _SolverState, RESULTS]:
-        del solver_state, made_jump
+    ) -> Tuple[PyTree, PyTree, DenseInfo, _SolverState, RESULTS]:
+        del made_jump
         control = terms.contr(t0, t1)
-        pred = terms.vf_prod(t0, y0, args, control)
+        # Could use FSAL here but that would mean we'd need to switch to working with
+        # `f0 = terms.vf(t0, y0, args)`, and that gets quite hairy quite quickly.
+        # (C.f. `AbstractRungeKutta.step`.)
+        # If we wanted FSAL then really the correct thing to do would just be to
+        # write out a `ButcherTableau` and use `AbstractSDIRK`.
+        k0 = terms.vf_prod(t0, y0, args, control)
         jac = self.nonlinear_solver.jac(
-            _implicit_relation, pred, (terms.vf_prod, t1, y0, args, control)
+            _implicit_relation, k0, (terms.vf_prod, t1, y0, args, control)
         )
         nonlinear_sol = self.nonlinear_solver(
-            _implicit_relation, pred, (terms.vf_prod, t1, y0, args, control), jac
+            _implicit_relation, k0, (terms.vf_prod, t1, y0, args, control), jac
         )
-        z1 = nonlinear_sol.root
-        y1 = (y0**ω + z1**ω).ω
+        k1 = nonlinear_sol.root
+        y1 = (y0**ω + k1**ω).ω
+        # Use the trapezoidal rule for adaptive step sizing.
+        y_error = (0.5 * (k1**ω - k0**ω)).ω
         dense_info = dict(y0=y0, y1=y1)
-        return y1, None, dense_info, None, nonlinear_sol.result
+        solver_state = None
+        return y1, y_error, dense_info, solver_state, nonlinear_sol.result
 
     def func(
         self,

diffrax/solver/milstein.py

@@ -45,9 +45,19 @@ def order(self, terms):
     def strong_order(self, terms):
         return 1  # assuming commutative noise
 
+    def init(
+        self,
+        terms: Tuple[ODETerm, AbstractTerm],
+        t0: Scalar,
+        t1: Scalar,
+        y0: PyTree,
+        args: PyTree,
+    ) -> _SolverState:
+        return None
+
     def step(
         self,
-        terms: Tuple[AbstractTerm, AbstractTerm],
+        terms: Tuple[ODETerm, AbstractTerm],
         t0: Scalar,
         t1: Scalar,
         y0: PyTree,
@@ -103,9 +113,19 @@ def order(self, terms):
     def strong_order(self, terms):
         return 1  # assuming commutative noise
 
+    def init(
+        self,
+        terms: Tuple[ODETerm, AbstractTerm],
+        t0: Scalar,
+        t1: Scalar,
+        y0: PyTree,
+        args: PyTree,
+    ) -> _SolverState:
+        return None
+
     def step(
         self,
-        terms: Tuple[AbstractTerm, AbstractTerm],
+        terms: Tuple[ODETerm, AbstractTerm],
         t0: Scalar,
         t1: Scalar,
         y0: PyTree,

diffrax/solver/semi_implicit_euler.py

@@ -25,6 +25,16 @@ class SemiImplicitEuler(AbstractSolver):
     def order(self, terms):
         return 1
 
+    def init(
+        self,
+        terms: Tuple[AbstractTerm, AbstractTerm],
+        t0: Scalar,
+        t1: Scalar,
+        y0: PyTree,
+        args: PyTree,
+    ) -> _SolverState:
+        return None
+
     def step(
         self,
         terms: Tuple[AbstractTerm, AbstractTerm],
@@ -35,7 +45,7 @@ def step(
         solver_state: _SolverState,
         made_jump: Bool,
     ) -> Tuple[Tuple[PyTree, PyTree], _ErrorEstimate, DenseInfo, _SolverState, RESULTS]:
-        del made_jump
+        del solver_state, made_jump
 
         term_1, term_2 = terms
         y0_1, y0_2 = y0

test/conftest.py

@@ -39,7 +39,8 @@ def clear_caches():
                         if hasattr(obj, "cache_clear"):
                             try:
                                 print(f"Clearing {obj}")
-                                obj.cache_clear()
+                                if "Weakref" not in type(obj).__name__:
+                                    obj.cache_clear()
                             except Exception:
                                 pass
         gc.collect()

test/test_local_interpolation.py

@@ -12,9 +12,7 @@ def test_local_linear_interpolation():
     t1_ = 2.9
     for y0 in (2.1, jnp.array(2.1), jnp.array([2.1, 3.1])):
         for y1 in (2.2, jnp.array(2.2), jnp.array([2.2, 3.2])):
-            interp = diffrax.local_interpolation.LocalLinearInterpolation(
-                t0=t0, t1=t1, y0=y0, y1=y1
-            )
+            interp = diffrax.LocalLinearInterpolation(t0=t0, t1=t1, y0=y0, y1=y1)
 
             # evaluate position
             pred = interp.evaluate(t0_)
@@ -36,12 +34,27 @@ def test_local_linear_interpolation():
             assert shaped_allclose(pred, jnp.zeros_like(pred))
 
             # evaluate over zero-length interval. Note t1=t0.
-            interp = diffrax.local_interpolation.LocalLinearInterpolation(
-                t0=t0, t1=t0, y0=y0, y1=y1
-            )
+            interp = diffrax.LocalLinearInterpolation(t0=t0, t1=t0, y0=y0, y1=y1)
             pred = interp.evaluate(t0)
             true, _ = jnp.broadcast_arrays(y0, y1)
             assert shaped_allclose(pred, true)
 
             _, pred = jax.jvp(interp.evaluate, (t0,), (jnp.ones_like(t0),))
             assert shaped_allclose(pred, jnp.zeros_like(pred))
+
+
+def test_third_order_hermite():
+    t0 = 2.0
+    t1 = 3.9
+
+    def y(t):
+        return 0.4 + 0.7 * t - 1.1 * t**2 + 0.4 * t**3
+
+    y0, f0 = jax.jvp(y, (t0,), (1.0,))
+    y1, f1 = jax.jvp(y, (t1,), (1.0,))
+    k0 = f0 * (t1 - t0)
+    k1 = f1 * (t1 - t0)
+    interp = diffrax.ThirdOrderHermitePolynomialInterpolation(
+        t0=t0, t1=t1, y0=y0, y1=y1, k0=k0, k1=k1
+    )
+    assert shaped_allclose(interp.evaluate(2.6), y(2.6))

test/test_solver.py

@@ -17,3 +17,29 @@ def test_half_solver():
 def test_instance_check():
     assert isinstance(diffrax.HalfSolver(diffrax.Euler()), diffrax.Euler)
     assert not isinstance(diffrax.HalfSolver(diffrax.Euler()), diffrax.Heun)
+
+
+def test_implicit_euler_adaptive():
+    term = diffrax.ODETerm(lambda t, y, args: -10 * y**3)
+    solver1 = diffrax.ImplicitEuler(
+        nonlinear_solver=diffrax.NewtonNonlinearSolver(rtol=1e-5, atol=1e-5)
+    )
+    solver2 = diffrax.ImplicitEuler()
+    t0 = 0
+    t1 = 1
+    dt0 = 1
+    y0 = 1.0
+    stepsize_controller = diffrax.PIDController(rtol=1e-5, atol=1e-5)
+    out1 = diffrax.diffeqsolve(term, solver1, t0, t1, dt0, y0, throw=False)
+    out2 = diffrax.diffeqsolve(
+        term,
+        solver2,
+        t0,
+        t1,
+        dt0,
+        y0,
+        stepsize_controller=stepsize_controller,
+        throw=False,
+    )
+    assert out1.result == diffrax.RESULTS.implicit_nonconvergence
+    assert out2.result == diffrax.RESULTS.successful