265 lines
9.6 KiB
Go
265 lines
9.6 KiB
Go
|
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
|||
|
|
// SPDX-License-Identifier: MIT
|
|||
|
|
|
|||
|
|
package integrate
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
|||
|
|
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"errors"
|
|||
|
|
"math"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// The Rosenbrock-Wanner workhorse: the four-stage L-stable scheme ROS4
|
|||
|
|
// of Hairer and Wanner, fourth order accurate with an embedded
|
|||
|
|
// third-order solution driving the step control. Where BDF solves each
|
|||
|
|
// step by Newton, a W method puts the Jacobian into the formula
|
|||
|
|
// itself: every stage is one linear solve against the frozen matrix
|
|||
|
|
// (1/(γh))I − J, so a step costs one numerical Jacobian, one LU
|
|||
|
|
// factorisation and four back-substitutions, and no Newton iteration
|
|||
|
|
// ever stalls. The scheme is L-stable: its stability function vanishes
|
|||
|
|
// at infinity, so a step far beyond the transient's time constant
|
|||
|
|
// damps the stiff mode instead of amplifying it.
|
|||
|
|
//
|
|||
|
|
// The stage form is the divided one the standard implementations use.
|
|||
|
|
// With A = (1/(γh))I − J factored once per step, stage one solves
|
|||
|
|
// A·k₁ = f(t, y) and stage i solves A·kᵢ = f(t + cᵢh, yᵢ) +
|
|||
|
|
// (Σⱼ<ᵢ cᵢⱼkⱼ)/h with the stage value yᵢ = y + Σⱼ<ᵢ aᵢⱼkⱼ; the step
|
|||
|
|
// advances by y + Σ bᵢkᵢ and the embedded estimate is Σ êᵢkᵢ. Stage
|
|||
|
|
// four shares its node and its stage value with stage three, so its f
|
|||
|
|
// evaluation carries over and a step costs three f calls. The digits
|
|||
|
|
// are the published ones, cross-checked against the standard
|
|||
|
|
// implementations of the scheme.
|
|||
|
|
//
|
|||
|
|
// Two honest limits of the tableau. The Jacobian enters at (t, y) and
|
|||
|
|
// the tableau's time-derivative weights are left out, because the
|
|||
|
|
// package's f contract carries no partial derivative in t: the fourth
|
|||
|
|
// order therefore holds for autonomous systems (every problem in this
|
|||
|
|
// package's stiff tests is one), while a genuinely time-dependent f
|
|||
|
|
// loses the second-order local terms those weights carry and can
|
|||
|
|
// degrade to first order globally; measured on y' = −y + t with
|
|||
|
|
// uniform steps the global ratios come out near 2, not near 16. And
|
|||
|
|
// the scheme is L-stable but not stiffly accurate: the fourth stage's
|
|||
|
|
// weights differ from the solution weights, the stiff damping coming
|
|||
|
|
// from the stability limit rather than from stage-end agreement.
|
|||
|
|
|
|||
|
|
var (
|
|||
|
|
// rowGamma is the abscissa of the stage solves and the source of
|
|||
|
|
// the L-stability: the stability function's denominator is
|
|||
|
|
// (1 − γz)⁴ and its numerator vanishes at infinity.
|
|||
|
|
rowGamma = 0.57282
|
|||
|
|
|
|||
|
|
// rowNodes are the stage times as multiples of h. The second node
|
|||
|
|
// is the published tableau's c2 = 0.114564 (Hairer and Wanner's
|
|||
|
|
// ROS4); its doubled form 2·gamma appears in some reprints, and the
|
|||
|
|
// digit never enters the arithmetic of an autonomous problem, where
|
|||
|
|
// stage times cancel, so either form integrates identically here.
|
|||
|
|
rowNodes = [4]float64{0, 0.114564, 0.65521686381559, 0.65521686381559}
|
|||
|
|
|
|||
|
|
// rowA holds the stage-value weights a[i][j], the coefficient of
|
|||
|
|
// k_j inside stage i's value.
|
|||
|
|
rowA = [4][3]float64{
|
|||
|
|
{},
|
|||
|
|
{2},
|
|||
|
|
{1.867943637803922, 0.2344449711399156},
|
|||
|
|
{1.867943637803922, 0.2344449711399156, 0},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// rowC holds the Jacobian-coupling weights c[i][j], the coefficient
|
|||
|
|
// of k_j inside stage i's right side, divided by h.
|
|||
|
|
rowC = [4][3]float64{
|
|||
|
|
{},
|
|||
|
|
{-7.13761503641231},
|
|||
|
|
{2.580708087951457, 0.6515950076447975},
|
|||
|
|
{-2.137148994382534, -0.3214669691237626, -0.6949742501781779},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// rowB advances the solution; rowE carries the embedded third-order
|
|||
|
|
// estimate (the difference between the fourth-order and the
|
|||
|
|
// embedded third-order weights).
|
|||
|
|
rowB = [4]float64{2.255570073418735, 0.2870493262186792, 0.435317943184018, 1.093502252409163}
|
|||
|
|
rowE = [4]float64{-0.2815431932141155, -0.0727619912493892, -0.1082196201495311, -1.093502252409163}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// IntegrateROS4 integrates y' = f(t, y) from t0 to t1 with the
|
|||
|
|
// four-stage L-stable Rosenbrock-Wanner scheme ROS4 and returns y(t1).
|
|||
|
|
// The fourth order holds for autonomous systems; a genuinely
|
|||
|
|
// time-dependent f loses the second-order local terms the tableau's
|
|||
|
|
// time-derivative weights would carry (the package's f contract has no
|
|||
|
|
// partial derivative in t), and can degrade to first order globally:
|
|||
|
|
// measured on y' = −y + t with uniform steps the ratios come out near
|
|||
|
|
// 2, not near 16. The adaptive controller still holds its tolerance
|
|||
|
|
// there, at the cost of more steps. The step size follows the embedded
|
|||
|
|
// error under the classic accept-or-shrink control, with the numerical
|
|||
|
|
// Jacobian taken once per step through the same central-difference
|
|||
|
|
// helper the Newton path uses; there is no user Jacobian parameter.
|
|||
|
|
// Backward integration works: a t1 < t0 simply integrates in the
|
|||
|
|
// negative direction. An exhausted step budget, a collapsed step size,
|
|||
|
|
// an f that returns a wrongly shaped state, or a stage that leaves the
|
|||
|
|
// finite range even as the step shrinks is an error, never a silently
|
|||
|
|
// truncated trajectory.
|
|||
|
|
func IntegrateROS4(f func(t float64, y *core.Array) (*core.Array, error),
|
|||
|
|
t0, t1 float64, y0 *core.Array, opts ODEOptions) (*core.Array, error) {
|
|||
|
|
const name = "IntegrateROS4"
|
|||
|
|
y, err := odeCheck(name, y0, &opts)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
h, err := bdf2InitialStep(name, f, t0, t1, y, &opts)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
budget := odeBudget{max: opts.MaxSteps}
|
|||
|
|
w := &odeWork{}
|
|||
|
|
t := t0
|
|||
|
|
for !odeArrived(t, t1) {
|
|||
|
|
if err := budget.spend(name, t, t1); err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
// Never step past t1; t1−t carries the integration direction.
|
|||
|
|
h = odeClampStep(h, t, t1)
|
|||
|
|
yEnd, errNorm, nerr := ros4Step(name, f, w, t, y, h, opts.AbsTol, opts.RelTol)
|
|||
|
|
if nerr != nil {
|
|||
|
|
if errors.Is(nerr, errNewtonStalled) {
|
|||
|
|
// The stage matrix closed onto singularity: halve the
|
|||
|
|
// step and retry the same interval, within the budget.
|
|||
|
|
h *= 0.5
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
return nil, nerr
|
|||
|
|
}
|
|||
|
|
if math.IsNaN(errNorm) || math.IsInf(errNorm, 0) {
|
|||
|
|
// A stage left the finite range: shrink hard and retry.
|
|||
|
|
h *= 0.25
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
factor := min(5, max(0.2, 0.9*math.Pow(1/errNorm, 1.0/4)))
|
|||
|
|
if errNorm <= 1 {
|
|||
|
|
copy(y, yEnd)
|
|||
|
|
prevT := t
|
|||
|
|
t += h
|
|||
|
|
h *= factor
|
|||
|
|
// Collapse is "t did not move": a span below the absolute
|
|||
|
|
// time scale is integrable, and an accepted step that
|
|||
|
|
// arrives at the end exactly is not a failure either.
|
|||
|
|
if t == prevT {
|
|||
|
|
return nil, base.Errf("%s: the step size shrank below the resolution of t at t=%g", name, prevT)
|
|||
|
|
}
|
|||
|
|
} else {
|
|||
|
|
// Rejected: retry the same interval with the smaller step.
|
|||
|
|
h *= max(0.2, factor)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return arrayFromVector(y), nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ros4Step attempts one ROS4 step of the given size from (t, y) and
|
|||
|
|
// returns the candidate end state, the embedded error norm against the
|
|||
|
|
// mixed absolute and relative tolerance, and a nil error when the
|
|||
|
|
// stages all solved. Every buffer belongs to the workspace, so a step
|
|||
|
|
// allocates nothing of its own: the caller's y is left untouched, and
|
|||
|
|
// the returned slice stays valid until the next step.
|
|||
|
|
func ros4Step(name string, f func(t float64, y *core.Array) (*core.Array, error),
|
|||
|
|
w *odeWork, t float64, y []float64, h, absTol, relTol float64) ([]float64, float64, error) {
|
|||
|
|
n := len(y)
|
|||
|
|
w.use(n)
|
|||
|
|
// One scratch buffer per role: the stage value, the stage right
|
|||
|
|
// side, f's result and the candidate end state are rebuilt at the
|
|||
|
|
// top of their use and read only by it.
|
|||
|
|
ks, stage, rhs, fy, yEnd := w.ks, w.stage, w.rhs, w.fy, w.yEnd
|
|||
|
|
// One numerical Jacobian per step, the house central-difference
|
|||
|
|
// helper the Newton path shares, and one factorisation of I − hγJ,
|
|||
|
|
// which is (1/(γh))I − J scaled by γh: the scale folds into the
|
|||
|
|
// stage right sides instead of the matrix.
|
|||
|
|
jac, jerr := odeJacobian(name, f, t, y, w)
|
|||
|
|
if jerr != nil {
|
|||
|
|
return nil, 0, jerr
|
|||
|
|
}
|
|||
|
|
lu := w.mat
|
|||
|
|
// The step-scaled matrix weight and the stage weight are one
|
|||
|
|
// product each, the same (−h·γ) and (h·γ) groupings the element
|
|||
|
|
// loops evaluated.
|
|||
|
|
hgr := -h * rowGamma
|
|||
|
|
hg := h * rowGamma
|
|||
|
|
for i := range n {
|
|||
|
|
row := lu[i]
|
|||
|
|
for j := range n {
|
|||
|
|
row[j] = hgr * jac[i*n+j]
|
|||
|
|
}
|
|||
|
|
row[i]++
|
|||
|
|
}
|
|||
|
|
w.perm, _ = base.Factor(lu)
|
|||
|
|
if err := base.CheckSingular(name, lu); err != nil {
|
|||
|
|
return nil, 0, base.Errf("%s: %w, singular stage matrix at t=%g", name, errNewtonStalled, t)
|
|||
|
|
}
|
|||
|
|
// Stage one.
|
|||
|
|
out, ferr := odeCall(name, f, t, y, n, &w.views)
|
|||
|
|
if ferr != nil {
|
|||
|
|
return nil, 0, ferr
|
|||
|
|
}
|
|||
|
|
readVector(fy, out)
|
|||
|
|
for i := range n {
|
|||
|
|
rhs[i] = hg * fy[i]
|
|||
|
|
}
|
|||
|
|
odePermuteColumn(rhs, w.perm, w.visited)
|
|||
|
|
base.SolveColumn(lu, rhs)
|
|||
|
|
copy(ks[0], rhs)
|
|||
|
|
// Stages two through four. Stage four repeats stage three's node
|
|||
|
|
// and stage value, so its derivative carries over.
|
|||
|
|
for s := 1; s < 4; s++ {
|
|||
|
|
if s < 3 {
|
|||
|
|
copy(stage, y)
|
|||
|
|
for j := range s {
|
|||
|
|
aj := rowA[s][j]
|
|||
|
|
if aj == 0 {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
kjs := ks[j]
|
|||
|
|
for i := range n {
|
|||
|
|
stage[i] += aj * kjs[i]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
out, ferr = odeCall(name, f, t+rowNodes[s]*h, stage, n, &w.views)
|
|||
|
|
if ferr != nil {
|
|||
|
|
return nil, 0, ferr
|
|||
|
|
}
|
|||
|
|
readVector(fy, out)
|
|||
|
|
}
|
|||
|
|
for i := range n {
|
|||
|
|
rhs[i] = hg * fy[i]
|
|||
|
|
}
|
|||
|
|
for j := range s {
|
|||
|
|
cj := rowC[s][j]
|
|||
|
|
if cj == 0 {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
gcj := rowGamma * cj
|
|||
|
|
kjs := ks[j]
|
|||
|
|
for i := range n {
|
|||
|
|
rhs[i] += gcj * kjs[i]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
odePermuteColumn(rhs, w.perm, w.visited)
|
|||
|
|
base.SolveColumn(lu, rhs)
|
|||
|
|
copy(ks[s], rhs)
|
|||
|
|
}
|
|||
|
|
// The embedded pair: the b weights advance, the gap to the
|
|||
|
|
// embedded third-order solution estimates the local error.
|
|||
|
|
errNorm := 0.0
|
|||
|
|
for i := range n {
|
|||
|
|
e, advance := 0.0, 0.0
|
|||
|
|
for s := range 4 {
|
|||
|
|
e += rowE[s] * ks[s][i]
|
|||
|
|
advance += rowB[s] * ks[s][i]
|
|||
|
|
}
|
|||
|
|
yEnd[i] = y[i] + advance
|
|||
|
|
scale := absTol + relTol*math.Max(math.Abs(y[i]), math.Abs(yEnd[i]))
|
|||
|
|
ratio := e / scale
|
|||
|
|
errNorm += ratio * ratio
|
|||
|
|
}
|
|||
|
|
return yEnd, math.Sqrt(errNorm/float64(n)) + 1e-10, nil
|
|||
|
|
}
|