Files
sim-numerics/odestep/driver.py
T

115 lines
4.2 KiB
Python

"""Integration driver whose only original part is the step-size controller.
The tableau, the stage combination and the embedded error estimate are
SciPy's - `scipy.integrate._ivp.rk.rk_step` with RK45's A, B, C and E - so
nothing about the method itself is a transcription of ours. An earlier
version of this file carried a hand-typed Dormand-Prince tableau; that was an
unnecessary risk and an unnecessary thing for a reader to have to check.
What remains ours is the loop that decides the next step size, which is the
object under study. Everything else is shared between controllers, so a
difference between two runs can only come from the controller.
The driver is validated against SciPy's own `solve_ivp(method="RK45")`: with
the classical controller - SciPy's own safety factor and clamps - the two
must agree closely on the number of right-hand-side evaluations.
"""
from __future__ import annotations
import math
import numpy as np
from scipy.integrate._ivp.rk import RK45, rk_step
ORDER = RK45.error_estimator_order # 4; the step exponent is ORDER + 1
class Counter:
"""Counts evaluations of the right-hand side - the cost model."""
__slots__ = ("f", "n")
def __init__(self, f):
self.f = f
self.n = 0
def __call__(self, t, y):
self.n += 1
return np.asarray(self.f(t, y), dtype=float)
def _norm(v, y, y_new, atol, rtol):
scale = atol + rtol * np.maximum(np.abs(y), np.abs(y_new))
return float(np.sqrt(np.mean((v / scale) ** 2)))
def integrate(f, t0, t1, y0, controller, atol=1e-8, rtol=1e-8,
max_steps=200000):
"""Integrate with the given controller. Returns what the comparison
needs. The initial step is chosen the same way for every controller, so
none can win by guessing it better."""
g = Counter(f)
t = float(t0)
y = np.asarray(y0, dtype=float)
fv = g(t, y)
K = np.empty((RK45.n_stages + 1, y.size), dtype=float)
d0 = _norm(y, y, y, atol, rtol)
d1 = _norm(fv, y, y, atol, rtol)
h = 1e-6 if (d0 < 1e-5 or d1 < 1e-5) else 0.01 * d0 / d1
h = min(h, abs(t1 - t0))
controller.reset()
accepted = rejected = 0
while t < t1 and accepted + rejected < max_steps:
h = min(h, t1 - t)
if h <= 0.0:
break
y_new, f_new = rk_step(g, t, y, fv, h, RK45.A, RK45.B, RK45.C, K)
err = np.dot(K.T, RK45.E) * h # RK45's own error estimate
e = _norm(err, y, y_new, atol, rtol)
ok = e <= 1.0
h_next = controller.update(h, e, ok, ORDER)
if ok:
accepted += 1
t, y, fv = t + h, y_new, f_new
else:
rejected += 1
if not math.isfinite(h_next) or h_next <= 0.0:
break
h = h_next
return {"t": t, "y": list(y), "nfev": g.n, "accepted": accepted,
"rejected": rejected,
"reached_end": abs(t - t1) < 1e-12 * max(1.0, abs(t1))}
def validate_against_scipy(atol=1e-8, rtol=1e-8):
"""The driver against SciPy's own solver, with SciPy's own controller.
This replaces the hand-written order check that the previous version
needed: if the driver reproduces solve_ivp's evaluation count with the
same control law, the tableau, the stage combination, the error estimate
and the loop are all the library's behaviour."""
from scipy.integrate import solve_ivp
from .controllers import Classical
out = []
probs = [
("decay", lambda t, y: [-y[0]], 0.0, 10.0, [1.0]),
("forced", lambda t, y: [-y[0] + math.sin(t)], 0.0, 10.0, [0.0]),
("lotka", lambda t, y: [1.5 * y[0] - y[0] * y[1],
-3.0 * y[1] + y[0] * y[1]], 0.0, 15.0,
[10.0, 5.0]),
("vdp", lambda t, y: [y[1], 5.0 * (1 - y[0] ** 2) * y[1] - y[0]],
0.0, 25.0, [2.0, 0.0]),
]
for name, f, t0, t1, y0 in probs:
mine = integrate(f, t0, t1, y0, Classical(), atol=atol, rtol=rtol)
ref = solve_ivp(f, (t0, t1), y0, method="RK45", atol=atol,
rtol=rtol)
out.append((name, mine["nfev"], int(ref.nfev),
max(abs(a - b) for a, b in zip(mine["y"],
ref.y[:, -1]))))
return out