From edd1c635bd791ea41e93a25e5c1e8a622755d465 Mon Sep 17 00:00:00 2001 From: Philipp Jurasic Date: Wed, 7 Oct 2026 11:57:34 +0000 Subject: [PATCH 1/2] vern9 integrator added --- diffrax/__init__.py | 1 + diffrax/_autocitation.py | 7 +- diffrax/_solver/__init__.py | 1 + diffrax/_solver/vern9.py | 305 +++++++++++++++++++++++++++ docs/api/solvers/ode_solvers.md | 5 + docs/usage/how-to-choose-a-solver.md | 2 +- test/helpers.py | 1 + test/test_integrate.py | 7 +- test/test_solver.py | 44 ++++ 9 files changed, 368 insertions(+), 5 deletions(-) create mode 100644 diffrax/_solver/vern9.py diff --git a/diffrax/__init__.py b/diffrax/__init__.py index d35a7fac..1a2b91c6 100644 --- a/diffrax/__init__.py +++ b/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, diff --git a/diffrax/_autocitation.py b/diffrax/_autocitation.py index b176bd75..a948e98d 100644 --- a/diffrax/_autocitation.py +++ b/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, diff --git a/diffrax/_solver/__init__.py b/diffrax/_solver/__init__.py index 0a840413..42753f11 100644 --- a/diffrax/_solver/__init__.py +++ b/diffrax/_solver/__init__.py @@ -55,3 +55,4 @@ StochasticButcherTableau as StochasticButcherTableau, ) from .tsit5 import Tsit5 as Tsit5 +from .vern9 import Vern9 as Vern9 diff --git a/diffrax/_solver/vern9.py b/diffrax/_solver/vern9.py new file mode 100644 index 00000000..c6275176 --- /dev/null +++ b/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 . + """ + + 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""" diff --git a/docs/api/solvers/ode_solvers.md b/docs/api/solvers/ode_solvers.md index 10547bff..9492a7b9 100644 --- a/docs/api/solvers/ode_solvers.md +++ b/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 diff --git a/docs/usage/how-to-choose-a-solver.md b/docs/usage/how-to-choose-a-solver.md index 2dce12d1..41280bf0 100644 --- a/docs/usage/how-to-choose-a-solver.md +++ b/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`][]. diff --git a/test/helpers.py b/test/helpers.py index 85f8e8ce..8ddf9d8d 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -34,6 +34,7 @@ diffrax.LeapfrogMidpoint(), diffrax.ReversibleHeun(), diffrax.Tsit5(), + diffrax.Vern9(), diffrax.ImplicitEuler(), diffrax.Kvaerno3(), diffrax.Kvaerno4(), diff --git a/test/test_integrate.py b/test/test_integrate.py index 9918aa20..beac7d41 100644 --- a/test/test_integrate.py +++ b/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) diff --git a/test/test_solver.py b/test/test_solver.py index dfff3f6b..037ad454 100644 --- a/test/test_solver.py +++ b/test/test_solver.py @@ -6,8 +6,10 @@ import jax.numpy as jnp import jax.random as jr import jax.tree_util as jtu +import numpy as np import optimistix as optx import pytest +import scipy.stats from .helpers import implicit_tol, tree_allclose @@ -502,3 +504,45 @@ def test_adaptive_dt0_milstein(getkey): diffrax.diffeqsolve( terms, solver, 0, 1, None, 1, stepsize_controller=stepsize_controller ) + + +def test_vern9_tableau(): + tableau = diffrax.Vern9.tableau + # The stage times are `(0, c)`; row `i` of `a_lower` is stage `i+1`. + c = np.concatenate([[0.0], tableau.c]) + for i, row in enumerate(tableau.a_lower): + np.testing.assert_allclose(np.sum(row), c[i + 1], atol=1e-15) + # The solution is 9th order, and the embedded solution (the difference of which is + # `b_error`) is 8th order. Written as quadrature conditions `sum_i b_i c_i^(q-1)`. + for q in range(1, 10): + np.testing.assert_allclose( + np.sum(tableau.b_sol * c ** (q - 1)), 1 / q, atol=1e-15 + ) + for q in range(1, 9): + np.testing.assert_allclose( + np.sum(tableau.b_error * c ** (q - 1)), 0, atol=1e-15 + ) + assert abs(np.sum(tableau.b_error * c**8)) > 1e-8 + + +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_vern9_order(dtype): + akey, ykey = jr.split(jr.PRNGKey(5678)) + A = jr.normal(akey, (10, 10), dtype=dtype) * 0.5 + y0 = jr.normal(ykey, (10,), dtype=dtype) + true_y1 = jax.scipy.linalg.expm(4 * A) @ y0 + exponents = [0, -1, -2, -3] # Smaller steps are dominated by round-off. + errors = [] + for exponent in exponents: + sol = diffrax.diffeqsolve( + diffrax.ODETerm(lambda t, y, args: A @ y), + diffrax.Vern9(), + 0, + 4, + 2.0**exponent, + y0, + ) + errors.append(jnp.log2(jnp.sum(jnp.abs(sol.ys[-1] - true_y1)))) + order = scipy.stats.linregress(exponents, errors).slope + # Four points are not asymptotic: we measure 8.1 (real) and 10.2 (complex). + assert abs(order - 9) < 1.3 From a39b644b86dd6fedcab105539e0e2458c6b41731 Mon Sep 17 00:00:00 2001 From: Philipp Jurasic Date: Wed, 7 Oct 2026 14:46:41 +0000 Subject: [PATCH 2/2] Pyright fix --- test/test_solver.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_solver.py b/test/test_solver.py index 037ad454..d994ae17 100644 --- a/test/test_solver.py +++ b/test/test_solver.py @@ -9,7 +9,6 @@ import numpy as np import optimistix as optx import pytest -import scipy.stats from .helpers import implicit_tol, tree_allclose @@ -542,7 +541,8 @@ def test_vern9_order(dtype): 2.0**exponent, y0, ) + assert sol.ys is not None errors.append(jnp.log2(jnp.sum(jnp.abs(sol.ys[-1] - true_y1)))) - order = scipy.stats.linregress(exponents, errors).slope - # Four points are not asymptotic: we measure 8.1 (real) and 10.2 (complex). + order, _ = np.polyfit(exponents, errors, 1) + # Few points poorly estimate the asymptotic convergence for high order methods. assert abs(order - 9) < 1.3