vern9 integrator

diffrax/__init__.py

@@ -117,6 +117,7 @@
     StochasticButcherTableau as StochasticButcherTableau,
     StratonovichMilstein as StratonovichMilstein,
     Tsit5 as Tsit5,
+    Vern9 as Vern9,
 )
 from ._step_size_controller import (
     AbstractAdaptiveStepSizeController as AbstractAdaptiveStepSizeController,

diffrax/_autocitation.py

@@ -34,6 +34,7 @@
     SPaRK,
     SRA1,
     Tsit5,
+    Vern9,
 )
 from ._step_size_controller import ClipStepSizeController, PIDController
 
@@ -361,9 +362,8 @@ def _virtual_brownian_tree(terms):
 
 @citation_rules.append
 def _space_time_levy_area(terms):
-    has_levy_area = (
-        lambda x: isinstance(x, AbstractBrownianPath)
-        and x.levy_area != BrownianIncrement
+    has_levy_area = lambda x: (
+        isinstance(x, AbstractBrownianPath) and x.levy_area != BrownianIncrement
     )
     leaves = jtu.tree_leaves(terms, is_leaf=has_levy_area)
     if any(has_levy_area(leaf) for leaf in leaves):
@@ -487,6 +487,7 @@ def _sde(terms):
 def _solvers(solver, saveat=None):
     if type(solver) in (
         Tsit5,
+        Vern9,
         Kvaerno3,
         Kvaerno4,
         Kvaerno5,

diffrax/_solver/__init__.py

@@ -55,3 +55,4 @@
     StochasticButcherTableau as StochasticButcherTableau,
 )
 from .tsit5 import Tsit5 as Tsit5
+from .vern9 import Vern9 as Vern9

diffrax/_solver/vern9.py

@@ -0,0 +1,305 @@
+from collections.abc import Callable
+from typing import ClassVar
+
+import equinox.internal as eqxi
+import numpy as np
+
+from .._local_interpolation import ThirdOrderHermitePolynomialInterpolation
+from .runge_kutta import AbstractERK, ButcherTableau
+
+
+# Coefficients of Verner's "most efficient" 9(8) pair. The stage times are
+# (0, c2..c16); stages 15 and 16 are evaluated at the end of the step. Stage 16 only
+# enters the error estimate, and so `b_sol` has a trailing zero.
+_vern9_tableau = ButcherTableau(
+    a_lower=(
+        np.array(
+            [
+                0.03462000000000000000000000000000000000000,
+            ]
+        ),
+        np.array(
+            [
+                -0.03893354388572873270170426872292844785319,
+                0.1359578945245091786499878854939346230295,
+            ]
+        ),
+        np.array(
+            [
+                0.03638413148954266723060635628912731569112,
+                0.0,
+                0.1091523944686280016918190688673819470733,
+            ]
+        ),
+        np.array(
+            [
+                2.025763914393969636805657604282571047511,
+                0.0,
+                -7.638023836496292020387602153091964592952,
+                6.173259922102322383581944548809393545441,
+            ]
+        ),
+        np.array(
+            [
+                0.05112275589406060872792270881648288397197,
+                0.0,
+                0.0,
+                0.1770823794555021537929910813839068684086,
+                0.0008027762409222501453613869810802528375789,
+            ]
+        ),
+        np.array(
+            [
+                0.1316006357975216279279871693164256985334,
+                0.0,
+                0.0,
+                -0.2957276252669636417685183174672273730698,
+                0.08781378035642952374211247040538866670829,
+                0.6213052975225274774321435005639430026100,
+            ]
+        ),
+        np.array(
+            [
+                0.07166666666666666666666666666666666666667,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                0.3305533578915319409260346730051472207727,
+                0.2427799754418013924072986603281861125606,
+            ]
+        ),
+        np.array(
+            [
+                0.07180664062500000000000000000000000000000,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                0.3294380283228177160744825466257672816401,
+                0.1165190029271822839255174533742327183599,
+                -0.03401367187500000000000000000000000000000,
+            ]
+        ),
+        np.array(
+            [
+                0.04836757646340646986611287718844085773549,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                0.03928989925676163974333190042057047002854,
+                0.1054740945890344608263649267140088017604,
+                -0.02143865284648312665982642293830533996214,
+                -0.1041229174627194437759832813847147895623,
+            ]
+        ),
+        np.array(
+            [
+                -0.02664561487201478635337289243849737340534,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                0.03333333333333333333333333333333333333333,
+                -0.1631072244872467239162704487554706387141,
+                0.03396081684127761199487954930015522928243,
+                0.1572319413814626097110769806810024118077,
+                0.2152267478031879552303534778794770376960,
+            ]
+        ),
+        np.array(
+            [
+                0.03689009248708622334786359863227633989719,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                -0.1465181576725542928653609891758501156785,
+                0.2242577768172024345345469822625833796001,
+                0.02294405717066072637090897902753790803033,
+                -0.003585005290572876135739442488933033433436,
+                0.08669223316444385506869203619044453906051,
+                0.4383840651968337846196219974168630120572,
+            ]
+        ),
+        np.array(
+            [
+                -0.4866012215113340846662212357570395295089,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                -6.304602650282852990657772792012007122988,
+                -0.2812456182894725647782841837901184181118,
+                -2.679019236219849057687906597489223155566,
+                0.5188156639241575115653111646150125220241,
+                1.365353187603341710683633635235238678626,
+                5.885091088503946585721274891680604830712,
+                2.802808786272062889819965117517532194812,
+            ]
+        ),
+        np.array(
+            [
+                0.4185367457753471441471025246471931649630,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                6.724547581906459363100870806514855026677,
+                -0.4254442801646117906069834096971130646156,
+                3.343279153001265577811816947557982637750,
+                0.6170816631175377595284211175077097847373,
+                -0.9299661239399328339377495239888008520137,
+                -6.099948804751010722472962837945508844846,
+                -3.002206187889399044804158084895173690015,
+                0.2553202529443445472336424602988558373637,
+            ]
+        ),
+        np.array(
+            [
+                -0.7793740861228846646446230408438405063434,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                -13.93734253810777678786523664804936051203,
+                1.252048853379357320949735183924200895140,
+                -14.69150040801686878191527989293072091588,
+                -0.4947050585331416856551919921369628735756,
+                2.242974909146236657906984549543692874754,
+                13.36789380382864375813864978592679139881,
+                14.39665048665068644512236935340272139005,
+                -0.7975813331776800379127866056663258667437,
+                0.4409353709534277758753793068298041158235,
+            ]
+        ),
+        np.array(
+            [
+                2.058051337466886442151242368989994043993,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                22.35793772796803295519317565842520212899,
+                0.9094981099755633274500919813797189078288,
+                35.89110098240264104710550686568482456493,
+                -3.442515027624453437985000403608480262215,
+                -4.865481358036368826566013387928704014496,
+                -18.90980381354342625688427480879773032857,
+                -34.26354448030451782929251177395134170515,
+                1.264756521695642578827783499806516664686,
+                0.0,
+                0.0,
+            ]
+        ),
+    ),
+    b_sol=np.array(
+        [
+            0.01461197685842315252051541915018784713459,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            -0.3915211862331339089410228267288242030810,
+            0.2310932500289506415909675644868993669908,
+            0.1274766769992852382560589467488989175618,
+            0.2246434176204157731566981937082069688984,
+            0.5684352689748512932705226972873692126743,
+            0.05825871557215827200814768021863420902155,
+            0.1364317403482215641609022744494239843327,
+            0.03057013983082797397721005067920369646664,
+            0.0,
+        ]
+    ),
+    b_error=np.array(
+        [
+            -0.005357988290444578334669665030800840430051,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            -2.583020491182463963471769926039661727945,
+            0.1425225315468662612826303441690357783613,
+            0.01342065351268867600756328623975571429445,
+            -0.02867296291409493339975957974875822883630,
+            2.624999655215792304429522291883350513167,
+            -0.2825509643291537215078618092038201722614,
+            0.1364317403482215641609022744494239843327,
+            0.03057013983082797397721005067920369646664,
+            -0.04834231373823958314376726739772871714902,
+        ]
+    ),
+    c=np.array(
+        [
+            0.03462000000000000000000000000000000000000,
+            0.09702435063878044594828361677100617517631,
+            0.1455365259581706689224254251565092627645,
+            0.5610000000000000000000000000000000000000,
+            0.2290079115904850126662751771814700052182,
+            0.5449920884095149873337248228185299947818,
+            0.6450000000000000000000000000000000000000,
+            0.4837500000000000000000000000000000000000,
+            0.06757000000000000000000000000000000000000,
+            0.2500000000000000000000000000000000000000,
+            0.6590650618730998549405331618649220295334,
+            0.8206000000000000000000000000000000000000,
+            0.9012000000000000000000000000000000000000,
+            1.000000000000000000000000000000000000000,
+            1.000000000000000000000000000000000000000,
+        ]
+    ),
+)
+
+
+class Vern9(AbstractERK):
+    r"""Verner's "most efficient" 9(8) method.
+
+    9th order explicit Runge--Kutta method. Has an embedded 8th order method for
+    adaptive step sizing. Uses 16 stages, and is not FSAL. Uses 3rd order Hermite
+    interpolation for dense/ts output.
+
+    This is a good choice for smooth problems that must be solved to tight tolerances
+    (around `rtol=1e-9` or smaller), where it typically needs far fewer function
+    evaluations than lower-order methods.
+
+    !!! warning
+
+        The interpolation is only 3rd order accurate. When using `SaveAt(ts=...)` or
+        `SaveAt(dense=True)` the saved values between steps may therefore be less
+        accurate than the solution at the end of each step.
+
+    ??? cite "Reference"
+
+        ```bibtex
+        @article{verner2010numerically,
+            title={Numerically optimal {R}unge--{K}utta pairs with interpolants},
+            author={Verner, James H.},
+            journal={Numerical Algorithms},
+            volume={53},
+            number={2},
+            pages={383--396},
+            year={2010},
+            publisher={Springer}
+        }
+        ```
+
+        Coefficients from <https://www.sfu.ca/~jverner/>.
+    """
+
+    tableau: ClassVar[ButcherTableau] = _vern9_tableau
+    interpolation_cls: ClassVar[
+        Callable[..., ThirdOrderHermitePolynomialInterpolation]
+    ] = ThirdOrderHermitePolynomialInterpolation.from_k
+
+    def order(self, terms):
+        del terms
+        return 9
+
+
+eqxi.doc_remove_args("scan_kind")(Vern9.__init__)
+Vern9.__init__.__doc__ = """**Arguments:** None"""

docs/api/solvers/ode_solvers.md

@@ -54,6 +54,11 @@ These methods are suitable for most problems.
         members:
             - __init__
 
+::: diffrax.Vern9
+    options:
+        members:
+            - __init__
+
 ---
 
 ### Implicit Runge--Kutta (IRK) methods

docs/usage/how-to-choose-a-solver.md

@@ -16,7 +16,7 @@ For non-stiff problems then [`diffrax.Tsit5`][] is a good general-purpose solver
 
     For a long time the recommend default solver for many problems was [`diffrax.Dopri5`][]. This is the default solver used in [`torchdiffeq`](https://github.com/rtqichen/torchdiffeq/), and is the solver used in MATLAB's `ode45`. However `Tsit5` is now reckoned on being slightly more efficient overall. (Try both if you wish.)
 
-If you need accurate solutions at tight tolerances then try [`diffrax.Dopri8`][].
+If you need accurate solutions at tight tolerances then try [`diffrax.Dopri8`][] or [`diffrax.Vern9`][]. For smooth problems at very tight tolerances (around `rtol=1e-9` or smaller) `Vern9` is often cheaper still, but its dense output is only third-order accurate.
 
 If you are solving a neural differential equation, and training via discretise-then-optimise (corresponding to `diffeqsolve(..., adjoint=RecursiveCheckpointAdjoint())`, which is the default), then accurate solutions are often not needed and a low-order solver will be most efficient. For example something like [`diffrax.Heun`][].
 

test/helpers.py

@@ -34,6 +34,7 @@
     diffrax.LeapfrogMidpoint(),
     diffrax.ReversibleHeun(),
     diffrax.Tsit5(),
+    diffrax.Vern9(),
     diffrax.ImplicitEuler(),
     diffrax.Kvaerno3(),
     diffrax.Kvaerno4(),

test/test_integrate.py

@@ -141,7 +141,12 @@ def f(t, y, args):
         assert tree_allclose(y1, true_y1, atol=1e-2, rtol=1e-2)
 
 
-@pytest.mark.parametrize("solver", all_ode_solvers + all_split_solvers)
+# `Vern9` is covered by `test_vern9_order`.
+@pytest.mark.parametrize(
+    "solver",
+    [s for s in all_ode_solvers if not isinstance(s, diffrax.Vern9)]
+    + list(all_split_solvers),
+)
 @pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128])
 def test_ode_order(solver, dtype):
     solver = implicit_tol(solver)