from ode.core import (
Dense, OdeCfg, OdeState, DOPRI5, MAX_FACTOR, MIN_FACTOR, NOT_FINITE, SAFETY, STEP_TOO_SMALL,
all_finite, min_step_at,
)
# the nodes
C2 = 1.0 / 5.0
C3 = 3.0 / 10.0
C4 = 4.0 / 5.0
C5 = 8.0 / 9.0
# the stage weights a_ij
A21 = 1.0 / 5.0
A31 = 3.0 / 40.0
A32 = 9.0 / 40.0
A41 = 44.0 / 45.0
A42 = -56.0 / 15.0
A43 = 32.0 / 9.0
A51 = 19372.0 / 6561.0
A52 = -25360.0 / 2187.0
A53 = 64448.0 / 6561.0
A54 = -212.0 / 729.0
A61 = 9017.0 / 3168.0
A62 = -355.0 / 33.0
A63 = 46732.0 / 5247.0
A64 = 49.0 / 176.0
A65 = -5103.0 / 18656.0
# the fifth-order weights (also row 7 of the tableau)
B1 = 35.0 / 384.0
B3 = 500.0 / 1113.0
B4 = 125.0 / 192.0
B5 = -2187.0 / 6784.0
B6 = 11.0 / 84.0
# fifth-order minus fourth-order weights: the error estimate
E1 = -71.0 / 57600.0
E3 = 71.0 / 16695.0
E4 = -71.0 / 1920.0
E5 = 17253.0 / 339200.0
E6 = -22.0 / 525.0
E7 = 1.0 / 40.0
# the dense output's fifth-order correction (Hairer's d_i)
D1 = -12715105075.0 / 11282082432.0
D3 = 87487479700.0 / 32700410799.0
D4 = -10690763975.0 / 1880347072.0
D5 = 701980252875.0 / 199316789632.0
D6 = -1453857185.0 / 822651844.0
D7 = 69997945.0 / 29380423.0
# the controller's exponent, -1 / (order of the error estimate + 1)
ERR_EXPONENT = -1.0 / 5.0
# rhs evaluations per step (the seventh is reused)
STAGE_EVALS: Int = 6
# the dense output's coefficient rows
DENSE_ROWS = 5
# the stiffness test: h |lambda| above this bound, this many steps in a row
STIFF_BOUND = 3.25
STIFF_STEPS: Int = 15
# this many non-stiff steps in a row forget a stiff run
NONSTIFF_RESET: Int = 6
#| One accepted Dormand-Prince step from s, retrying with smaller steps
#| until the error test passes. Returns the state at the new point, with its
#| dense output; status STEP_TOO_SMALL when the step underflows.
def dopri_step(f, s0: OdeState, c: OdeCfg, p: Arr) -> OdeState {
s = s0
t = s.t
y = s.y
k1 = s.fy
n = len(y)
min_step = min_step_at(t)
h_abs = min(s.h_abs, c.max_step)
if h_abs < min_step {
h_abs = min_step
}
rejected = false
accepted = false
h_next = h_abs
while not accepted {
if h_abs < min_step {
s.status = STEP_TOO_SMALL
return s
}
t_new = t + h_abs * s.dir
if s.dir * (t_new - c.t_bound) > 0.0 {
t_new = c.t_bound
}
h = t_new - t
h_abs = abs(h)
#: the six new stages (element loops: no temporary arrays in the hot path)
yt = zeros(n)
for i to n {
yt[i] = y[i] + h * (A21 * k1[i])
}
k2 = f(t + C2 * h, yt, p)
for i to n {
yt[i] = y[i] + h * (A31 * k1[i] + A32 * k2[i])
}
k3 = f(t + C3 * h, yt, p)
for i to n {
yt[i] = y[i] + h * (A41 * k1[i] + A42 * k2[i] + A43 * k3[i])
}
k4 = f(t + C4 * h, yt, p)
for i to n {
yt[i] = y[i] + h * (A51 * k1[i] + A52 * k2[i] + A53 * k3[i] + A54 * k4[i])
}
k5 = f(t + C5 * h, yt, p)
y6 = zeros(n)
for i to n {
y6[i] = y[i] + h * (A61 * k1[i] + A62 * k2[i] + A63 * k3[i] + A64 * k4[i] + A65 * k5[i])
}
k6 = f(t + h, y6, p)
y_new = zeros(n)
for i to n {
y_new[i] = y[i] + h * (B1 * k1[i] + B3 * k3[i] + B4 * k4[i] + B5 * k5[i] + B6 * k6[i])
}
k7 = f(t + h, y_new, p)
s.nfev += STAGE_EVALS
if not all_finite(y_new) {
s.status = NOT_FINITE
return s
}
#: the error test, in the weighted RMS norm
acc = 0.0
for i to n {
e = h * (E1 * k1[i] + E3 * k3[i] + E4 * k4[i] + E5 * k5[i] + E6 * k6[i] + E7 * k7[i])
w = e / (c.atol[i] + max(abs(y[i]), abs(y_new[i])) * c.rtol)
acc += w * w
}
err_norm = sqrt(acc / n)
if err_norm < 1.0 {
factor = MAX_FACTOR
if err_norm > 0.0 {
factor = min(MAX_FACTOR, SAFETY * pow(err_norm, ERR_EXPONENT))
}
if rejected {
factor = min(1.0, factor)
}
h_next = h_abs * factor
accepted = true
#: Hairer's stiffness estimate: h |lambda| ~ h |k7 - k6| / |y_new - y6|
num2 = 0.0
den2 = 0.0
for i to n {
num2 += (k7[i] - k6[i]) * (k7[i] - k6[i])
den2 += (y_new[i] - y6[i]) * (y_new[i] - y6[i])
}
if den2 > 0.0 {
hlamb = abs(h) * sqrt(num2 / den2)
if hlamb > STIFF_BOUND {
s.nonstiff_count = 0
s.stiff_count += 1
if s.stiff_count >= STIFF_STEPS {
s.stiff = true
}
}
else {
s.nonstiff_count += 1
if s.nonstiff_count >= NONSTIFF_RESET {
s.stiff_count = 0
}
}
}
#: the dense output's coefficients
Q = zeros(DENSE_ROWS, n)
for i to n {
ydiff = y_new[i] - y[i]
bspl = h * k1[i] - ydiff
Q[0, i] = y[i]
Q[1, i] = ydiff
Q[2, i] = bspl
Q[3, i] = ydiff - h * k7[i] - bspl
Q[4, i] = h * (D1 * k1[i] + D3 * k3[i] + D4 * k4[i] + D5 * k5[i] + D6 * k6[i] + D7 * k7[i])
}
s.dense = Dense { kind: DOPRI5, t_old: t, t_new: t_new, h: h, Q: Q, y_old: y, order: 0 }
s.has_dense = true
s.t = t_new
s.y = y_new
s.fy = k7
}
else {
h_abs = h_abs * max(MIN_FACTOR, SAFETY * pow(err_norm, ERR_EXPONENT))
rejected = true
s.nreject += 1
}
}
s.h_abs = h_next
s.naccept += 1
return s
}
test "the tableau is consistent: rows sum to the nodes, the weights to 1" {
assert_near(A21, C2, 1.0e-16)
assert_near(A31 + A32, C3, 1.0e-16)
assert_near(A41 + A42 + A43, C4, 1.0e-15)
assert_near(A51 + A52 + A53 + A54, C5, 1.0e-14)
assert_near(A61 + A62 + A63 + A64 + A65, 1.0, 1.0e-14)
assert_near(B1 + B3 + B4 + B5 + B6, 1.0, 1.0e-15)
# the embedded fourth-order weights also sum to 1, so the error weights sum to 0
assert_near(E1 + E3 + E4 + E5 + E6 + E7, 0.0, 1.0e-16)
# the fifth-order weights integrate t^4 exactly: sum b_i c_i^4 = 1/5
assert_near(B3 * pow(C3, 4.0) + B4 * pow(C4, 4.0) + B5 * pow(C5, 4.0) + B6, 0.2, 1.0e-15)
}
# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors