879 lines
29 KiB
Go
879 lines
29 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"
|
||
)
|
||
|
||
// Ordinary differential equation solvers for initial value problems
|
||
// y' = f(t, y). The state y is a rank-1 vector of length n; a system
|
||
// of higher rank flattens to its leading-axis vector first.
|
||
//
|
||
// Three schemes cover the standard regimes. IntegrateODE is the
|
||
// workhorse: an adaptive embedded Runge-Kutta pair (Dormand-Prince
|
||
// 4(5)) that controls the local error against a mixed absolute and
|
||
// relative tolerance. IntegrateRK4 is the classical fixed-step
|
||
// fourth-order scheme, useful when a uniform step or simple
|
||
// reproducibility per step matters. IntegrateBackwardEuler is the
|
||
// entry-level stiff scheme: fully implicit, with each step's
|
||
// nonlinear equation solved by Newton over a numerical Jacobian and
|
||
// the library's LU solver.
|
||
|
||
// ODEOptions tunes the adaptive integrator. RelTol ≤ 0 means 1e-6,
|
||
// AbsTol ≤ 0 means 1e-9, MaxSteps ≤ 0 means 100000.
|
||
type ODEOptions struct {
|
||
RelTol float64
|
||
AbsTol float64
|
||
MaxSteps int
|
||
}
|
||
|
||
// Dormand-Prince 4(5): node offsets, stage coefficients, and the
|
||
// 5th- and 4th-order solution weights. Stage 7 shares the 5th-order
|
||
// weights (the FSAL property), which is why it needs no separate row.
|
||
var (
|
||
odeC = [7]float64{0, 1.0 / 5, 3.0 / 10, 4.0 / 5, 8.0 / 9, 1, 1}
|
||
odeA = [][]float64{
|
||
{},
|
||
{1.0 / 5},
|
||
{3.0 / 40, 9.0 / 40},
|
||
{44.0 / 45, -56.0 / 15, 32.0 / 9},
|
||
{19372.0 / 6561, -25360.0 / 2187, 64448.0 / 6561, -212.0 / 729},
|
||
{9017.0 / 3168, -355.0 / 33, 46732.0 / 5247, 49.0 / 176, -5103.0 / 18656},
|
||
{35.0 / 384, 0, 500.0 / 1113, 125.0 / 192, -2187.0 / 6784, 11.0 / 84},
|
||
}
|
||
odeB5 = [7]float64{35.0 / 384, 0, 500.0 / 1113, 125.0 / 192, -2187.0 / 6784, 11.0 / 84, 0}
|
||
odeB4 = [7]float64{5179.0 / 57600, 0, 7571.0 / 16695, 393.0 / 640, -92097.0 / 339200, 187.0 / 2100, 1.0 / 40}
|
||
)
|
||
|
||
// IntegrateODE integrates y' = f(t, y) from t0 to t1 with the adaptive
|
||
// Dormand-Prince 4(5) pair and returns y(t1). Backward integration
|
||
// works: a t1 < t0 simply integrates in the negative direction. An
|
||
// exhausted step budget, a collapsed step size or an f that returns a
|
||
// wrongly shaped state is an error, never a silently truncated
|
||
// trajectory.
|
||
func IntegrateODE(f func(t float64, y *core.Array) (*core.Array, error),
|
||
t0, t1 float64, y0 *core.Array, opts ODEOptions) (*core.Array, error) {
|
||
return odeRun(f, t0, t1, y0, opts, nil)
|
||
}
|
||
|
||
// readVector copies a's elements into dst, sweeping the raw float64
|
||
// payload when a is a dense float64 array and falling back to the
|
||
// widening accessor for views and other dtypes. The values written are
|
||
// identical either way.
|
||
func readVector(dst []float64, a *core.Array) {
|
||
if !a.Strided() && a.Dtype() == core.Float {
|
||
copy(dst, a.RawFloats())
|
||
return
|
||
}
|
||
for i := range dst {
|
||
dst[i] = a.FloatAt(i)
|
||
}
|
||
}
|
||
|
||
// denseFloats returns a's elements as a plain float64 slice, sharing
|
||
// the payload when a is a dense float64 array and copying the widened
|
||
// values otherwise. The values read are the ones the accessor
|
||
// returned; a shared slice is read-only, and only an array the caller
|
||
// owns may be written through it. A caller sweeping the elements of a
|
||
// solver's result uses this instead of one accessor call per element.
|
||
func denseFloats(a *core.Array) []float64 {
|
||
if !a.Strided() && a.Dtype() == core.Float {
|
||
return a.RawFloats()
|
||
}
|
||
out := make([]float64, a.Len())
|
||
for i := range out {
|
||
out[i] = a.FloatAt(i)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// odeRun drives the adaptive Dormand-Prince loop over the whole span.
|
||
// When watch is not nil it is called after every accepted step with
|
||
// the interval just integrated and clones of the states at both ends;
|
||
// a true return stops the integration there, and the watch's error
|
||
// aborts it. Everything else behaves exactly like IntegrateODE.
|
||
func odeRun(f func(t float64, y *core.Array) (*core.Array, error),
|
||
t0, t1 float64, y0 *core.Array, opts ODEOptions,
|
||
watch func(tPrev, tNow float64, yPrev, yNow []float64) (bool, error)) (*core.Array, error) {
|
||
const name = "IntegrateODE"
|
||
y, err := odeCheck(name, y0, &opts)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
n := len(y)
|
||
w := &odeWork{}
|
||
w.useStage(n)
|
||
k := make([][]float64, 8) // k[1..7] are the stages; k[0] unused
|
||
for i := 1; i <= 7; i++ {
|
||
k[i] = make([]float64, n)
|
||
}
|
||
// One scratch accumulator serves every stage of every step: it is
|
||
// rebuilt from y at the top of each stage call and read only by
|
||
// that call's f evaluation, the same transient view the package's
|
||
// fixed-step solvers hand out.
|
||
acc := w.stage
|
||
stage := func(i int, t float64, h float64) error {
|
||
copy(acc, y)
|
||
row := odeA[i-1]
|
||
for j := 1; j < i; j++ {
|
||
if row[j-1] == 0 {
|
||
continue
|
||
}
|
||
// The step-scaled weight is one product, the same (h·a)·k
|
||
// grouping the plain component loop evaluated.
|
||
hj := h * row[j-1]
|
||
kj := k[j]
|
||
for m := range n {
|
||
acc[m] += hj * kj[m]
|
||
}
|
||
}
|
||
out, err := odeCall(name, f, t+odeC[i-1]*h, acc, n, &w.views)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
readVector(k[i], out)
|
||
return nil
|
||
}
|
||
|
||
t := t0
|
||
h := odeInitialStep(t0, t1, y)
|
||
budget := odeBudget{max: opts.MaxSteps}
|
||
yEnd := w.yEnd
|
||
// The solution weights scaled by the step size: one product each,
|
||
// the same (h·b)·k grouping the component loop evaluated.
|
||
var hb5, hb4 [7]float64
|
||
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)
|
||
for j := 1; j <= 7; j++ {
|
||
hb5[j-1] = h * odeB5[j-1]
|
||
hb4[j-1] = h * odeB4[j-1]
|
||
}
|
||
for i := 1; i <= 7; i++ {
|
||
if err := stage(i, t, h); err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
// The embedded pair: the 5th-order solution advances, the gap
|
||
// to the 4th-order one estimates the local error.
|
||
errNorm := 0.0
|
||
for m := range n {
|
||
y5, y4 := y[m], y[m]
|
||
for j := 1; j <= 7; j++ {
|
||
y5 += hb5[j-1] * k[j][m]
|
||
y4 += hb4[j-1] * k[j][m]
|
||
}
|
||
yEnd[m] = y5
|
||
scale := opts.AbsTol + opts.RelTol*math.Max(math.Abs(y[m]), math.Abs(y5))
|
||
ratio := (y5 - y4) / scale
|
||
errNorm += ratio * ratio
|
||
}
|
||
errNorm = math.Sqrt(errNorm/float64(n)) + 1e-10
|
||
|
||
factor := math.Min(5, math.Max(0.2, 0.9*math.Pow(1/errNorm, 1.0/5)))
|
||
if errNorm <= 1 {
|
||
if watch != nil {
|
||
stop, werr := watch(t, t+h, cloneDenseSlice(y), cloneDenseSlice(yEnd))
|
||
if werr != nil {
|
||
return nil, werr
|
||
}
|
||
if stop {
|
||
return arrayFromVector(yEnd), nil
|
||
}
|
||
}
|
||
copy(y, yEnd)
|
||
prevT := t
|
||
t += h
|
||
h *= factor
|
||
// Collapse is "t did not move", not "h is small": a span
|
||
// far below the absolute time scale is perfectly
|
||
// integrable, and the old absolute floor refused it.
|
||
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 *= math.Max(0.2, factor)
|
||
}
|
||
}
|
||
return arrayFromVector(y), nil
|
||
}
|
||
|
||
// IntegrateODEPath integrates y' = f(t, y) from t0 to t1 and returns
|
||
// the trajectory sampled at nSamples evenly spaced points, endpoints
|
||
// included: times[i] is the sample time and states[i] the state there,
|
||
// so states[0] is the initial state and states[nSamples−1] the answer
|
||
// IntegrateODE would return. Every interval between neighbouring
|
||
// samples is integrated on its own, so the adaptive step control never
|
||
// has to align with the sampling grid. Backward integration
|
||
// (t1 < t0) works, and the error contract of IntegrateODE applies
|
||
// per interval.
|
||
func IntegrateODEPath(f func(t float64, y *core.Array) (*core.Array, error),
|
||
t0, t1 float64, y0 *core.Array, nSamples int, opts ODEOptions) ([]float64, []*core.Array, error) {
|
||
if nSamples < 2 {
|
||
return nil, nil, base.Errf("IntegrateODEPath: nSamples must be ≥ 2, got %d", nSamples)
|
||
}
|
||
y, err := odeCheck("IntegrateODEPath", y0, &opts)
|
||
if err != nil {
|
||
return nil, nil, base.Errf("IntegrateODEPath: %w", err)
|
||
}
|
||
times := make([]float64, nSamples)
|
||
states := make([]*core.Array, nSamples)
|
||
times[0] = t0
|
||
states[0] = wrapVector(y)
|
||
// The last sample is pinned to t1 exactly; the intermediate ones
|
||
// are the evenly spaced grid.
|
||
for i := 1; i < nSamples; i++ {
|
||
times[i] = t0 + float64(i)*(t1-t0)/float64(nSamples-1)
|
||
}
|
||
times[nSamples-1] = t1
|
||
for i := 1; i < nSamples; i++ {
|
||
states[i], err = IntegrateODE(f, times[i-1], times[i], states[i-1], opts)
|
||
if err != nil {
|
||
return nil, nil, base.Errf("IntegrateODEPath: %w", err)
|
||
}
|
||
}
|
||
return times, states, nil
|
||
}
|
||
|
||
// IntegrateODESteps integrates y' = f(t, y) from t0 to t1 and returns
|
||
// the trajectory as recorded at every accepted solver step: times[i]
|
||
// carries states[i] = y(times[i]), starting with (t0, y0) and ending
|
||
// with (t1, y(t1)). The accepted steps are where the adaptive control
|
||
// judged the local error within tolerance, so they are the natural
|
||
// interpolation nodes for post-processing, sensitivity analysis and
|
||
// adjoint passes. Backward integration records descending times; the
|
||
// error contract of IntegrateODE applies.
|
||
func IntegrateODESteps(f func(t float64, y *core.Array) (*core.Array, error),
|
||
t0, t1 float64, y0 *core.Array, opts ODEOptions) ([]float64, []*core.Array, error) {
|
||
y, err := odeCheck("IntegrateODESteps", y0, &opts)
|
||
if err != nil {
|
||
return nil, nil, base.Errf("IntegrateODESteps: %w", err)
|
||
}
|
||
times := []float64{t0}
|
||
states := []*core.Array{wrapVector(y)}
|
||
watch := func(tPrev, tNow float64, yPrev, yNow []float64) (bool, error) {
|
||
times = append(times, tNow)
|
||
// yNow is the run's own per-call clone; nothing aliases it.
|
||
states = append(states, wrapVector(yNow))
|
||
return false, nil
|
||
}
|
||
if _, err := odeRun(f, t0, t1, y0, opts, watch); err != nil {
|
||
return nil, nil, base.Errf("IntegrateODESteps: %w", err)
|
||
}
|
||
// The run's last accepted boundary is t+h with h = t1−t, which
|
||
// rounds a few ulps off t1 whenever the magnitudes demand it; the
|
||
// documented endpoint is t1 exactly, and the recorded state there
|
||
// is already the run's answer y(t1).
|
||
times[len(times)-1] = t1
|
||
return times, states, nil
|
||
}
|
||
|
||
// IntegrateRK4 integrates y' = f(t, y) with the classical fixed-step
|
||
// fourth-order Runge-Kutta scheme over the given number of equal
|
||
// steps, returning y(t1).
|
||
func IntegrateRK4(f func(t float64, y *core.Array) (*core.Array, error),
|
||
t0, t1 float64, y0 *core.Array, steps int) (*core.Array, error) {
|
||
if steps <= 0 {
|
||
return nil, base.Errf("IntegrateRK4: steps must be ≥ 1, got %d", steps)
|
||
}
|
||
y, err := odeCheck("IntegrateRK4", y0, nil)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
n := len(y)
|
||
h := (t1 - t0) / float64(steps)
|
||
k1 := make([]float64, n)
|
||
k2 := make([]float64, n)
|
||
k3 := make([]float64, n)
|
||
k4 := make([]float64, n)
|
||
tmp := make([]float64, n)
|
||
views := &odeViews{}
|
||
call := func(t float64, v []float64, out []float64) error {
|
||
o, err := odeCall("IntegrateRK4", f, t, v, n, views)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
readVector(out, o)
|
||
// A non-finite stage flows straight into the state with no
|
||
// rejection mechanism to catch it, and the fixed-step run
|
||
// would publish NaN with a nil error; the adaptive drivers
|
||
// reject it, this one has to refuse it.
|
||
for i := range n {
|
||
if math.IsNaN(out[i]) || math.IsInf(out[i], 0) {
|
||
return base.Errf("IntegrateRK4: f returned the non-finite value %g at coordinate %d, t=%g", out[i], i, t)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
// The stage times come from the exact grid t0 + i·h, never from an
|
||
// accumulated t += h: the addition's rounding walks over a long run
|
||
// (measured on y' = cos t from t0 = 1e6 the walk contributes an
|
||
// error of 1.1e-8 that no step count refines away), while each grid
|
||
// point carries a single rounding that stays put.
|
||
for i := range steps {
|
||
t := t0 + float64(i)*h
|
||
if err := call(t, y, k1); err != nil {
|
||
return nil, err
|
||
}
|
||
for i := range n {
|
||
tmp[i] = y[i] + h*k1[i]/2
|
||
}
|
||
if err := call(t+h/2, tmp, k2); err != nil {
|
||
return nil, err
|
||
}
|
||
for i := range n {
|
||
tmp[i] = y[i] + h*k2[i]/2
|
||
}
|
||
if err := call(t+h/2, tmp, k3); err != nil {
|
||
return nil, err
|
||
}
|
||
for i := range n {
|
||
tmp[i] = y[i] + h*k3[i]
|
||
}
|
||
if err := call(t+h, tmp, k4); err != nil {
|
||
return nil, err
|
||
}
|
||
for i := range n {
|
||
y[i] += h * (k1[i] + 2*k2[i] + 2*k3[i] + k4[i]) / 6
|
||
}
|
||
}
|
||
return arrayFromVector(y), nil
|
||
}
|
||
|
||
// IntegrateBackwardEuler integrates y' = f(t, y) with the fully
|
||
// implicit Euler scheme y_{n+1} = y_n + h·f(t_{n+1}, y_{n+1}), solving
|
||
// each step by Newton over a numerical Jacobian and the library's LU
|
||
// solver. The extra work per step is what buys stability on stiff
|
||
// systems, where the explicit schemes need step sizes far below what
|
||
// accuracy alone would ask for.
|
||
func IntegrateBackwardEuler(f func(t float64, y *core.Array) (*core.Array, error),
|
||
t0, t1 float64, y0 *core.Array, steps int, opts ODEOptions) (*core.Array, error) {
|
||
if steps <= 0 {
|
||
return nil, base.Errf("IntegrateBackwardEuler: steps must be ≥ 1, got %d", steps)
|
||
}
|
||
y, err := odeCheck("IntegrateBackwardEuler", y0, &opts)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
n := len(y)
|
||
h := (t1 - t0) / float64(steps)
|
||
yn := cloneDenseSlice(y)
|
||
seed := make([]float64, n)
|
||
fy := make([]float64, n)
|
||
// One Newton result buffer serves every step: it aliases neither the
|
||
// state nor the seed, and each step overwrites it fully.
|
||
zbuf := make([]float64, n)
|
||
w := &odeWork{}
|
||
// The step times come from the exact grid t0 + i·h, never from an
|
||
// accumulated t += h: the addition's rounding walks over a long
|
||
// run, while each grid point carries a single rounding that stays
|
||
// put.
|
||
for i := range steps {
|
||
tNext := t0 + float64(i+1)*h
|
||
// Newton on G(z) = z − y_n − h·f(t_{n+1}, z) = 0, seeded with
|
||
// the semi-implicit Euler prediction.
|
||
out, ferr := odeCall("IntegrateBackwardEuler", f, tNext, yn, n, &w.views)
|
||
if ferr != nil {
|
||
return nil, ferr
|
||
}
|
||
readVector(fy, out)
|
||
for i := range n {
|
||
seed[i] = yn[i] + h*fy[i]
|
||
}
|
||
if nerr := odeNewton("IntegrateBackwardEuler", f, w, tNext, 1, h, yn, seed,
|
||
zbuf, opts.AbsTol, opts.RelTol); nerr != nil {
|
||
return nil, nerr
|
||
}
|
||
copy(yn, zbuf)
|
||
}
|
||
return arrayFromVector(yn), nil
|
||
}
|
||
|
||
// errNewtonStalled marks an implicit solve whose Newton iteration ran
|
||
// out of budget or hit a singular matrix without converging. A driver
|
||
// that can shrink the step retries on it; a failed f evaluation is a
|
||
// different, fatal error.
|
||
var errNewtonStalled = errors.New("the Newton iteration did not converge")
|
||
|
||
// odeCall evaluates f at (t, v) and validates that the result is a
|
||
// vector of the expected length n, returning it unchanged. views,
|
||
// when not nil, caches the read-only wrapper handed to f, which a
|
||
// driver calling f repeatedly wants; a cold call site passes nil.
|
||
// Callers that need plain floats follow up with odeEval; callers that
|
||
// want to place the values themselves read them straight off the
|
||
// array.
|
||
func odeCall(name string, f func(t float64, y *core.Array) (*core.Array, error),
|
||
t float64, v []float64, n int, views *odeViews) (*core.Array, error) {
|
||
out, err := f(t, views.of(v))
|
||
if err != nil {
|
||
return nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
if out.NDim() != 1 || out.Len() != n {
|
||
return nil, base.Errf("%s: f returned shape %s, want a vector of length %d",
|
||
name, base.ShapeText(out.Shape()), n)
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// odeEval calls f at (t, v) and returns the derivative as a plain
|
||
// float64 slice.
|
||
func odeEval(name string, f func(t float64, y *core.Array) (*core.Array, error),
|
||
t float64, v []float64, n int, views *odeViews) ([]float64, error) {
|
||
out, err := odeCall(name, f, t, v, n, views)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
r := make([]float64, n)
|
||
readVector(r, out)
|
||
return r, nil
|
||
}
|
||
|
||
// odeBudget counts attempted solver steps against the MaxSteps option:
|
||
// the budget is spent before the step's stages are evaluated, so a
|
||
// rejected step consumes it like an accepted one.
|
||
type odeBudget struct {
|
||
used int
|
||
max int
|
||
}
|
||
|
||
// spend spends one step of the budget, failing once it is exhausted.
|
||
func (b *odeBudget) spend(name string, t, t1 float64) error {
|
||
if b.used >= b.max {
|
||
return base.Errf("%s: reached MaxSteps=%d at t=%g before t1=%g", name, b.max, t, t1)
|
||
}
|
||
b.used++
|
||
return nil
|
||
}
|
||
|
||
// odeClampStep caps h so a step never overshoots t1; t1−t carries the
|
||
// integration direction.
|
||
func odeClampStep(h, t, t1 float64) float64 {
|
||
if math.Abs(h) > math.Abs(t1-t) {
|
||
return t1 - t
|
||
}
|
||
return h
|
||
}
|
||
|
||
// odeArrived reports whether t sits within a few ulps of t1. The
|
||
// accumulating t += h can miss the exact endpoint by rounding once
|
||
// t and t1 differ in magnitude beyond Sterbenz territory, and the
|
||
// residual distance is indistinguishable from zero at working
|
||
// precision, so the solvers treat it as arrived rather than report a
|
||
// collapsed step over it.
|
||
func odeArrived(t, t1 float64) bool {
|
||
if t == t1 {
|
||
return true
|
||
}
|
||
return math.Abs(t1-t) <= 8*base.EpsF*math.Max(math.Abs(t), math.Abs(t1))
|
||
}
|
||
|
||
// odeWork holds the scratch the implicit schemes reuse across the
|
||
// steps of one run: the Newton vectors, one flat numerical Jacobian
|
||
// with its difference stencils, the factored matrix's rows, the ROS4
|
||
// stage buffers and the read-only views handed to f. Every buffer is
|
||
// overwritten before it is read, so a run builds the workspace once
|
||
// and no step allocates scratch of its own.
|
||
type odeWork struct {
|
||
// The Newton iteration's residual, step and derivative buffers,
|
||
// plus the DAE mass-matrix product M·z.
|
||
g, col, fzs []float64
|
||
mz []float64
|
||
// jac is the numerical Jacobian as one flat n×n buffer, row-major:
|
||
// jac[i*n+j] is ∂f_i/∂z_j.
|
||
jac []float64
|
||
// The Jacobian's central-difference stencils: the two perturbed
|
||
// states and their two results.
|
||
zp, zm, fp, fm []float64
|
||
// mat is the implicit relation's matrix in row-major rows, ROS4's
|
||
// (1/(γh))I − J or Newton's α·I − h·J depending on the caller, with
|
||
// perm the row permutation its factorisation produced. Factor
|
||
// shuffles the rows in place, and every rebuild rewrites the lot.
|
||
mat [][]float64
|
||
perm []int
|
||
// The ROS4 stage buffers: the four divided differences, the stage
|
||
// value, its right side, f's result and the candidate end state.
|
||
ks [4][]float64
|
||
stage, rhs, fy []float64
|
||
yEnd []float64
|
||
// The Newton iterate the implicit schemes iterate in, and the visit
|
||
// bitmap the stage permutation walks: both reused across every step
|
||
// and every attempt of one run.
|
||
zwork []float64
|
||
visited []bool
|
||
// views caches the wrapper handed to f per scratch slice.
|
||
views odeViews
|
||
}
|
||
|
||
// use returns the workspace's buffers at the state length n, growing
|
||
// them on first use. The content is left as the previous step wrote
|
||
// it: every consumer overwrites its buffer before reading it.
|
||
func (w *odeWork) use(n int) {
|
||
w.g = sizedBuf(w.g, n)
|
||
w.col = sizedBuf(w.col, n)
|
||
w.fzs = sizedBuf(w.fzs, n)
|
||
w.mz = sizedBuf(w.mz, n)
|
||
w.jac = sizedBuf(w.jac, n*n)
|
||
w.zp = sizedBuf(w.zp, n)
|
||
w.zm = sizedBuf(w.zm, n)
|
||
w.fp = sizedBuf(w.fp, n)
|
||
w.fm = sizedBuf(w.fm, n)
|
||
w.mat = sizedRows(w.mat, n)
|
||
for s := range w.ks {
|
||
w.ks[s] = sizedBuf(w.ks[s], n)
|
||
}
|
||
w.stage = sizedBuf(w.stage, n)
|
||
w.rhs = sizedBuf(w.rhs, n)
|
||
w.fy = sizedBuf(w.fy, n)
|
||
w.yEnd = sizedBuf(w.yEnd, n)
|
||
w.zwork = sizedBuf(w.zwork, n)
|
||
w.visited = sizedBools(w.visited, n)
|
||
}
|
||
|
||
// useStage returns the explicit step loop's buffers at the state
|
||
// length n, growing them on first use: the stage accumulator and the
|
||
// candidate end state are all an explicit pair needs, and sizing the
|
||
// implicit buffers here would allocate the Jacobian and the factored
|
||
// matrix for a loop that never takes a derivative.
|
||
func (w *odeWork) useStage(n int) {
|
||
w.stage = sizedBuf(w.stage, n)
|
||
w.yEnd = sizedBuf(w.yEnd, n)
|
||
}
|
||
|
||
// sizedBuf returns b cut to length n, reusing its storage when it is
|
||
// large enough.
|
||
func sizedBuf(b []float64, n int) []float64 {
|
||
if cap(b) < n {
|
||
return make([]float64, n)
|
||
}
|
||
return b[:n]
|
||
}
|
||
|
||
// sizedBools returns b cut to length n, reusing its storage when it is
|
||
// large enough.
|
||
func sizedBools(b []bool, n int) []bool {
|
||
if cap(b) < n {
|
||
return make([]bool, n)
|
||
}
|
||
return b[:n]
|
||
}
|
||
|
||
// odePermuteColumn reorders col in place so that col[i] takes the
|
||
// value that sat at perm[i], the permutation the workspace's LU
|
||
// factorisation produced. The walk carries each displaced value around
|
||
// its cycle exactly as the library's PermuteColumn does, plain
|
||
// assignments moving each value once, so the column ends up bit-
|
||
// identical; the difference is that the visit bitmap is the caller's
|
||
// reused scratch rather than a fresh allocation per call. The bitmap
|
||
// is cleared on entry, so a dirty buffer behaves exactly like a fresh
|
||
// one.
|
||
func odePermuteColumn(col []float64, perm []int, visited []bool) {
|
||
vis := visited[:len(col)]
|
||
clear(vis)
|
||
for i := range col {
|
||
if vis[i] || perm[i] == i {
|
||
vis[i] = true
|
||
continue
|
||
}
|
||
// Carry the displaced value around the cycle.
|
||
tmp := col[i]
|
||
j := i
|
||
for {
|
||
vis[j] = true
|
||
k := perm[j]
|
||
if k == i {
|
||
break
|
||
}
|
||
col[j] = col[k]
|
||
j = k
|
||
}
|
||
col[j] = tmp
|
||
}
|
||
}
|
||
|
||
// arrayFromVector copies a float64 slice into a fresh rank-1 float64
|
||
// array: the trajectory endpoint contract, without the intermediate
|
||
// wrapper a cloneArray(wrapVector(...)) pair built. The result never
|
||
// aliases the input.
|
||
func arrayFromVector(v []float64) *core.Array {
|
||
a := core.New(core.Float, len(v))
|
||
copy(a.RawFloats(), v)
|
||
return a
|
||
}
|
||
|
||
// sizedRows returns m as n row slices of length n, reusing the rows it
|
||
// already holds. The rows' content is the caller's to overwrite.
|
||
func sizedRows(m [][]float64, n int) [][]float64 {
|
||
if cap(m) < n {
|
||
m = make([][]float64, n)
|
||
}
|
||
m = m[:n]
|
||
for i := range m {
|
||
m[i] = sizedBuf(m[i], n)
|
||
}
|
||
return m
|
||
}
|
||
|
||
// odeViews caches the read-only wrapper handed to f for one scratch
|
||
// slice, so a driver that calls f thousands of times builds the
|
||
// wrapper once per slice instead of once per call. The values behind
|
||
// the wrapper are the driver's own scratch and keep changing exactly
|
||
// as they did; only the Array header is reused. The slices a driver
|
||
// hands in are few, so a linear scan beats a map; a driver that
|
||
// presents a fresh slice every call cannot grow the cache without
|
||
// bound, because the oldest entry makes way.
|
||
type odeViews struct {
|
||
entries []odeView
|
||
}
|
||
|
||
type odeView struct {
|
||
vals []float64
|
||
arr *core.Array
|
||
}
|
||
|
||
// odeViewSlots bounds the cache. A driver holds a handful of scratch
|
||
// slices at once, and each slot pins one sliced buffer, so the bound
|
||
// keeps both the scan and the retention small.
|
||
const odeViewSlots = 8
|
||
|
||
// of returns a read-only view of s, reusing the one already built for
|
||
// that slice. A nil cache builds a fresh view, which is what a cold
|
||
// call site wants.
|
||
func (v *odeViews) of(s []float64) *core.Array {
|
||
if v == nil || len(s) == 0 {
|
||
return wrapVector(s)
|
||
}
|
||
for i := range v.entries {
|
||
e := &v.entries[i]
|
||
if len(e.vals) == len(s) && &e.vals[0] == &s[0] {
|
||
return e.arr
|
||
}
|
||
}
|
||
arr := wrapVector(s)
|
||
if len(v.entries) == odeViewSlots {
|
||
copy(v.entries, v.entries[1:])
|
||
v.entries = v.entries[:odeViewSlots-1]
|
||
}
|
||
v.entries = append(v.entries, odeView{vals: s, arr: arr})
|
||
return arr
|
||
}
|
||
|
||
// odeJacobian fills the workspace's flat Jacobian with the central
|
||
// differences of f at (t, z), one column per state component:
|
||
// entry i*n+j is ∂f_i/∂z_j. The returned slice is the workspace's, so
|
||
// it stays valid until the next Jacobian. The two perturbed stencils
|
||
// and their two result buffers are reused across columns: each round
|
||
// rebuilds the stencils from z and overwrites both results before
|
||
// reading them.
|
||
func odeJacobian(name string, f func(t float64, y *core.Array) (*core.Array, error),
|
||
t float64, z []float64, w *odeWork) ([]float64, error) {
|
||
n := len(z)
|
||
w.use(n)
|
||
jac := w.jac
|
||
for j := range n {
|
||
eps := math.Sqrt(base.EpsF) * math.Max(1, math.Abs(z[j]))
|
||
copy(w.zp, z)
|
||
copy(w.zm, z)
|
||
w.zp[j] += eps
|
||
w.zm[j] -= eps
|
||
out, e1 := odeCall(name, f, t, w.zp, n, &w.views)
|
||
if e1 != nil {
|
||
return nil, e1
|
||
}
|
||
readVector(w.fp, out)
|
||
out, e2 := odeCall(name, f, t, w.zm, n, &w.views)
|
||
if e2 != nil {
|
||
return nil, e2
|
||
}
|
||
readVector(w.fm, out)
|
||
for i := range n {
|
||
jac[i*n+j] = (w.fp[i] - w.fm[i]) / (2 * eps)
|
||
}
|
||
}
|
||
return jac, nil
|
||
}
|
||
|
||
// odeNewton solves the implicit step equation α·z − h·f(tNext, z) = β
|
||
// for z by Newton over a numerical Jacobian and the library's LU
|
||
// solver, writing the converged state into dst and returning an error
|
||
// otherwise. The Jacobian is frozen from the seed and rebuilt twice
|
||
// when convergence drags, so a converged step costs one Jacobian and a
|
||
// handful of f evaluations. The iteration runs in the workspace's own
|
||
// buffer and the column permutation walks the workspace's bitmap, so
|
||
// the only allocation a converged step costs is the caller's dst: a
|
||
// driver whose dst recycles through a ring allocates its result
|
||
// buffers once per solve, not once per step, and a rejected or stalled
|
||
// attempt allocates nothing. dst must not alias seed or beta; the
|
||
// workspace overwrites it fully at convergence. Convergence is
|
||
// measured on the residual against the error scale the caller
|
||
// integrates to, an order of magnitude below it, but never below the
|
||
// floating-point floor of the residual's own terms, which would
|
||
// otherwise be unreachable at the tiny steps a stiff start begins
|
||
// with. An iteration that outlives twenty rounds, or a singular Newton
|
||
// matrix, surfaces as errNewtonStalled so a stepping driver can retry
|
||
// with a smaller step; an f that fails is the fatal error it is.
|
||
func odeNewton(name string, f func(t float64, y *core.Array) (*core.Array, error),
|
||
w *odeWork, tNext float64, alpha, h float64, beta, seed, dst []float64,
|
||
absTol, relTol float64) error {
|
||
n := len(seed)
|
||
w.use(n)
|
||
// The iterate starts as the seed and stays the workspace's buffer:
|
||
// the views cache then hands f one stable wrapper for every Newton
|
||
// call of the run.
|
||
z := sizedBuf(w.zwork, n)
|
||
copy(z, seed)
|
||
for iteration := range 20 {
|
||
out, err := odeCall(name, f, tNext, z, n, &w.views)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
readVector(w.fzs, out)
|
||
worst, terms := 0.0, 0.0
|
||
for i := range n {
|
||
w.g[i] = alpha*z[i] - h*w.fzs[i] - beta[i]
|
||
worst = math.Max(worst, math.Abs(w.g[i]))
|
||
terms = math.Max(terms, math.Abs(alpha*z[i])+math.Abs(h*w.fzs[i])+math.Abs(beta[i]))
|
||
}
|
||
limit := math.Max(0.1*(absTol+relTol*normInfOfStep(z)), 8*base.EpsF*terms)
|
||
if worst <= limit {
|
||
copy(dst, z)
|
||
return nil
|
||
}
|
||
if iteration == 0 || iteration == 4 || iteration == 10 {
|
||
jac, jerr := odeJacobian(name, f, tNext, z, w)
|
||
if jerr != nil {
|
||
return jerr
|
||
}
|
||
// Newton matrix α·I − h·J, a fresh LU for the frozen
|
||
// Jacobian; the iterations that follow only substitute.
|
||
// Every row is rebuilt entry by entry before the
|
||
// factorisation reads it.
|
||
for i := range n {
|
||
row := w.mat[i]
|
||
for j := range n {
|
||
row[j] = -h * jac[i*n+j]
|
||
}
|
||
row[i] += alpha
|
||
}
|
||
w.perm, _ = base.Factor(w.mat)
|
||
if err := base.CheckSingular(name, w.mat); err != nil {
|
||
return base.Errf("%s: %w, singular Newton matrix at t=%g",
|
||
name, errNewtonStalled, tNext)
|
||
}
|
||
}
|
||
for i := range n {
|
||
w.col[i] = -w.g[i]
|
||
}
|
||
odePermuteColumn(w.col, w.perm, w.visited)
|
||
base.SolveColumn(w.mat, w.col)
|
||
for i := range n {
|
||
z[i] += w.col[i]
|
||
}
|
||
}
|
||
return base.Errf("%s: %w at t=%g", name, errNewtonStalled, tNext)
|
||
}
|
||
|
||
// odeCheck validates the initial state, applies option defaults and
|
||
// returns the flat float64 working state.
|
||
func odeCheck(name string, y0 *core.Array, opts *ODEOptions) ([]float64, error) {
|
||
if y0.NDim() != 1 {
|
||
return nil, base.Errf("%s: the state must be a vector, got shape %s", name, base.ShapeText(y0.Shape()))
|
||
}
|
||
if y0.Len() == 0 {
|
||
return nil, base.Errf("%s: the state must not be empty", name)
|
||
}
|
||
if y0.Dtype() == core.Complex {
|
||
return nil, base.Errf("%s: complex states are not supported", name)
|
||
}
|
||
if opts != nil {
|
||
if opts.RelTol <= 0 {
|
||
opts.RelTol = 1e-6
|
||
}
|
||
if opts.AbsTol <= 0 {
|
||
opts.AbsTol = 1e-9
|
||
}
|
||
if opts.MaxSteps <= 0 {
|
||
opts.MaxSteps = 100000
|
||
}
|
||
}
|
||
return cloneDense(y0), nil
|
||
}
|
||
|
||
// odeInitialStep guesses the first step size as a small fraction of
|
||
// the integration span, carrying the direction in its sign.
|
||
func odeInitialStep(t0, t1 float64, y []float64) float64 {
|
||
h := 0.01 * math.Abs(t1-t0)
|
||
if h == 0 {
|
||
h = 1e-6
|
||
}
|
||
// The step sign carries the integration direction: a t1 < t0
|
||
// integrates backwards.
|
||
if t1 < t0 {
|
||
h = -h
|
||
}
|
||
return h
|
||
}
|
||
|
||
// cloneArray copies an array element by element, so the integration
|
||
// steps' results never alias a buffer already handed out.
|
||
func cloneArray(a *core.Array) *core.Array {
|
||
out := core.New(a.Dtype(), a.Shape()...)
|
||
switch a.Dtype() {
|
||
case core.Float:
|
||
copy(out.RawFloats(), a.RawFloats())
|
||
case core.Float32:
|
||
copy(out.RawFloat32s(), a.RawFloat32s())
|
||
case core.Int:
|
||
copy(out.RawInts(), a.RawInts())
|
||
default:
|
||
copy(out.RawComplexes(), a.RawComplexes())
|
||
}
|
||
return out
|
||
}
|
||
|
||
// normInfOfStep returns the infinity norm of a step vector.
|
||
func normInfOfStep(step []float64) float64 {
|
||
worst := 0.0
|
||
for _, v := range step {
|
||
if a := math.Abs(v); a > worst {
|
||
worst = a
|
||
}
|
||
}
|
||
return worst
|
||
}
|
||
|
||
// wrapVector views a float64 slice as a rank-1 Array without copying.
|
||
// The caller must treat the result as read-only.
|
||
func wrapVector(v []float64) *core.Array {
|
||
a, _ := core.FloatsFromArray(v, len(v))
|
||
return a
|
||
}
|
||
|
||
// cloneDense copies an array's elements into a plain float64 slice,
|
||
// widening int and float32 elements exactly.
|
||
func cloneDense(y0 *core.Array) []float64 {
|
||
vals := make([]float64, y0.Len())
|
||
for i := range vals {
|
||
vals[i] = y0.FloatAt(i)
|
||
}
|
||
return vals
|
||
}
|
||
|
||
// cloneDenseSlice copies a float64 slice.
|
||
func cloneDenseSlice(v []float64) []float64 {
|
||
out := make([]float64, len(v))
|
||
copy(out, v)
|
||
return out
|
||
}
|