438 lines
16 KiB
Go
438 lines
16 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package optim
|
||
|
||
import (
|
||
"fmt"
|
||
"math"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
"sourcedock.dev/petrbalvin/tensor/linalg"
|
||
)
|
||
|
||
// Root finding for systems of nonlinear equations r(x) = 0 with as
|
||
// many equations as unknowns. FindRoot and FindRootNewton cover the
|
||
// one-dimensional case; implicit solvers, equilibrium conditions and
|
||
// closure relations live in many dimensions, where a residual vector
|
||
// replaces the single equation.
|
||
//
|
||
// The method is the damped Newton iteration: each step solves the
|
||
// linearised system J·δ = −r with a central-difference Jacobian and
|
||
// the library's LU solver, then backtracks along δ until the residual
|
||
// norm actually falls, which keeps the iteration from leaping out of
|
||
// the basin of attraction when the linear model overshoots.
|
||
//
|
||
// RootSystemOptions.UseBroyden switches how the Jacobian is
|
||
// maintained: it is built by central differences once at the start
|
||
// and carried between steps by the rank-one Broyden update on
|
||
// its inverse, with a numerical rebuild whenever the update degrades
|
||
// (broydenMaintain documents the triggers). The default false keeps
|
||
// the per-step Jacobian and the iteration exactly as described here.
|
||
|
||
// RootSystemOptions tunes FindRootSystem. Tolerance ≤ 0 means 1e-10
|
||
// (an infinity-norm threshold on both the residual and the scaled
|
||
// step), MaxIterations ≤ 0 means 100.
|
||
type RootSystemOptions struct {
|
||
Tolerance float64
|
||
MaxIterations int
|
||
// UseBroyden maintains the Jacobian across steps instead of
|
||
// rebuilding it every round: one central-difference Jacobian at
|
||
// the start, then the rank-one Broyden update on its inverse
|
||
// after every accepted step, with a numerical rebuild whenever the
|
||
// update degrades (the triggers are documented on
|
||
// broydenMaintain). The default false leaves the per-step
|
||
// Jacobian and the iteration's results exactly as they are
|
||
// without the option.
|
||
UseBroyden bool
|
||
// ParallelJacobian lets the central-difference Jacobian sweep its
|
||
// columns on several goroutines, the build at the start and every
|
||
// Broyden restart included. Setting it is the caller's consent
|
||
// that the residual callback may run concurrently from more than
|
||
// one goroutine: the default false keeps every evaluation on the
|
||
// caller's goroutine, and the run is bit for bit the same either
|
||
// way, because the columns are independent and each one is
|
||
// differenced by the same stencil.
|
||
ParallelJacobian bool
|
||
}
|
||
|
||
// FindRootSystem solves r(x) = 0 for a vector function r of an n-vector,
|
||
// returning the solution and the residual infinity norm at it. The
|
||
// function must return a vector of the same length as x0. A singular
|
||
// Jacobian, a step that collapses before the tolerance is met or an
|
||
// exhausted iteration budget are errors, never silent answers.
|
||
//
|
||
// The iteration is a local solver: it follows the residual norm
|
||
// downhill, so a start whose basin contains no root, or one sitting
|
||
// in a parasitic minimum of the residual norm, reports the unfinished
|
||
// residual instead of pretending to converge. When the Newton system
|
||
// is singular the step falls back to the steepest-descent direction
|
||
// of the residual norm, which lets the iteration slide off the
|
||
// singular locus rather than dying there.
|
||
func FindRootSystem(f func(x *core.Array) (*core.Array, error), x0 *core.Array,
|
||
opts RootSystemOptions) (*core.Array, float64, error) {
|
||
if opts.Tolerance <= 0 {
|
||
opts.Tolerance = 1e-10
|
||
}
|
||
if opts.MaxIterations <= 0 {
|
||
opts.MaxIterations = 100
|
||
}
|
||
if x0.Dtype() == core.Complex {
|
||
return nil, 0, base.Errf("FindRootSystem: complex unknowns are not supported")
|
||
}
|
||
n := x0.Len()
|
||
if n == 0 {
|
||
return nil, 0, base.Errf("FindRootSystem: the initial guess must not be empty")
|
||
}
|
||
// eval reads the residual at x. The finiteness gate is strict for
|
||
// the states the iteration adopts (the start point and accepted
|
||
// iterates): a non-finite residual there is invisible to the
|
||
// norm's strict comparisons and would read as a converged root,
|
||
// the same refusal LevenbergMarquardt makes. Backtracking trials
|
||
// take the lenient variant: a trial that stepped into overflow is
|
||
// a rejected candidate to halve past, not a dead run, exactly the
|
||
// recovery the damping exists for.
|
||
eval := func(x []float64, strict bool, dst []float64) ([]float64, error) {
|
||
// The callback receives a copy, not an aliasing view of the
|
||
// reused iterate or stencil buffers: arrays are contractually
|
||
// immutable, and a callback that retains its argument must not
|
||
// observe the later writes. Minimise, L-BFGS and Levenberg-
|
||
// Marquardt copy for the same reason.
|
||
xa, aerr := core.FromFloats(x, len(x))
|
||
if aerr != nil {
|
||
return nil, aerr
|
||
}
|
||
out, err := f(xa)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if out.NDim() != 1 || out.Len() != n {
|
||
return nil, base.Errf("residual has shape %s, want a vector of length %d",
|
||
base.ShapeText(out.Shape()), n)
|
||
}
|
||
if err := requireReal("FindRootSystem", "residuals", out); err != nil {
|
||
return nil, err
|
||
}
|
||
// dst carries a buffer the Jacobian's stencil reuses across
|
||
// columns; the states the iteration keeps come back freshly
|
||
// allocated. Every entry of the buffer is written before it is
|
||
// read.
|
||
res := dst
|
||
if cap(res) < n {
|
||
res = make([]float64, n)
|
||
}
|
||
r := res[:n]
|
||
for i := range n {
|
||
r[i] = out.FloatAt(i)
|
||
if strict && (math.IsNaN(r[i]) || math.IsInf(r[i], 0)) {
|
||
// Bare, as Minimise's eval explains: every call site
|
||
// wraps once with the entry-point name, so the surfaced
|
||
// message carries exactly one prefix.
|
||
return nil, fmt.Errorf("the residual returned the non-finite value %g at coordinate %d", r[i], i)
|
||
}
|
||
}
|
||
return r, nil
|
||
}
|
||
|
||
// cloneDense promotes through FloatAt, so Int and Float32 starting
|
||
// vectors behave exactly like Float64 ones (a RawFloats fast path
|
||
// would hand the iteration a nil slice for those dtypes).
|
||
x := cloneDense(x0)
|
||
r, err := eval(x, true, nil)
|
||
if err != nil {
|
||
return nil, 0, base.Errf("FindRootSystem: %w", err)
|
||
}
|
||
res := normInfOfStep(r)
|
||
jac := make([][]float64, n)
|
||
for i := range n {
|
||
jac[i] = make([]float64, n)
|
||
}
|
||
// The Newton solve runs on a working copy of the Jacobian, refilled
|
||
// from jac every round: base.Factor consumes its argument in place
|
||
// before it can fail, and the fallback below needs the pristine
|
||
// central-difference matrix.
|
||
jacWork := make([][]float64, n)
|
||
for i := range n {
|
||
jacWork[i] = make([]float64, n)
|
||
}
|
||
// Buffers reused across iterations: the two perturbed stencils, the
|
||
// backtracking candidate, the Newton column and the step. Each is
|
||
// fully rewritten before it is read, and only transiently wrapped
|
||
// views of them reach the callback. The two residual vectors of the
|
||
// Jacobian's stencil are carried the same way: a column is
|
||
// differenced from them and they are dead once it is written, so
|
||
// one pair per round replaces one pair per column.
|
||
xp := make([]float64, n)
|
||
xm := make([]float64, n)
|
||
candidate := make([]float64, n)
|
||
col := make([]float64, n)
|
||
colWrap := [][]float64{nil}
|
||
step := make([]float64, n)
|
||
var resPlus, resMinus []float64
|
||
// Broyden state, allocated only under UseBroyden: the maintained
|
||
// inverse of the Jacobian and the flag saying it is fit to use,
|
||
// false until the first Jacobian is inverted and again after any
|
||
// restart, which sends the next round down the same
|
||
// build-and-invert path as the first.
|
||
useBroyden := opts.UseBroyden
|
||
var invJac [][]float64
|
||
var sVec, yVec []float64
|
||
haveInv := false
|
||
stalled := 0
|
||
if useBroyden {
|
||
invJac = make([][]float64, n)
|
||
for i := range n {
|
||
invJac[i] = make([]float64, n)
|
||
}
|
||
sVec = make([]float64, n)
|
||
yVec = make([]float64, n)
|
||
}
|
||
for range opts.MaxIterations {
|
||
if res <= opts.Tolerance {
|
||
return linalg.ArrayFromFloatsSafe(x, len(x)), res, nil
|
||
}
|
||
var derr error
|
||
if useBroyden && haveInv {
|
||
// The maintained inverse turns the linearised solve into a
|
||
// matrix-vector product, δ = −H·r: no Jacobian build and
|
||
// no factorisation this round.
|
||
broydenStep(invJac, r, step)
|
||
} else {
|
||
// Central-difference Jacobian, one column per unknown. The
|
||
// stencil carries the offset on one unknown at a time,
|
||
// restored as soon as the column is differenced, so the
|
||
// whole point is copied once per round rather than once per
|
||
// column. column walks one unknown's stencil and writes
|
||
// that column of jac, and nothing else, so the bits it
|
||
// produces do not depend on which driver walks the columns.
|
||
column := func(j int, sp, sm, rp, rm []float64) ([]float64, []float64, error) {
|
||
eps := math.Sqrt(base.EpsF) * math.Max(1, math.Abs(x[j]))
|
||
sp[j] += eps
|
||
sm[j] -= eps
|
||
rp, e1 := eval(sp, true, rp)
|
||
rm, e2 := eval(sm, true, rm)
|
||
sp[j], sm[j] = x[j], x[j]
|
||
if e1 != nil || e2 != nil {
|
||
return rp, rm, firstError(e1, e2)
|
||
}
|
||
for i := range n {
|
||
jac[i][j] = (rp[i] - rm[i]) / (2 * eps)
|
||
}
|
||
return rp, rm, nil
|
||
}
|
||
if opts.ParallelJacobian {
|
||
// The consent the option records lets the columns go to
|
||
// the engine's workers: each goroutine owns a disjoint
|
||
// run of columns, writes only into those columns of jac
|
||
// and reads only x, so no two workers write the same
|
||
// address and the sweep needs no locks. The build runs
|
||
// once at the start and again at every Broyden restart,
|
||
// both through this sweep. A failing column records its
|
||
// error and abandons the rest of its run; the reported
|
||
// one is the lowest failing column, the one the serial
|
||
// walk would hit first.
|
||
colErrs := make([]error, n)
|
||
engine.ParallelMin(n, 1, func(start, end int) {
|
||
sp, sm := make([]float64, n), make([]float64, n)
|
||
copy(sp, x)
|
||
copy(sm, x)
|
||
var rp, rm []float64
|
||
for j := start; j < end; j++ {
|
||
var err error
|
||
rp, rm, err = column(j, sp, sm, rp, rm)
|
||
if err != nil {
|
||
colErrs[j] = err
|
||
return
|
||
}
|
||
}
|
||
})
|
||
for _, err := range colErrs {
|
||
if err != nil {
|
||
return nil, 0, base.Errf("FindRootSystem: %w", err)
|
||
}
|
||
}
|
||
} else {
|
||
copy(xp, x)
|
||
copy(xm, x)
|
||
for j := range n {
|
||
var err error
|
||
resPlus, resMinus, err = column(j, xp, xm, resPlus, resMinus)
|
||
if err != nil {
|
||
return nil, 0, base.Errf("FindRootSystem: %w", err)
|
||
}
|
||
}
|
||
}
|
||
for i := range n {
|
||
col[i] = -r[i]
|
||
}
|
||
// Factor consumes the matrix in place: the pivot search swaps
|
||
// the rows and the elimination overwrites the subdiagonal with
|
||
// the multipliers, and the mutation lands in base.go before
|
||
// CheckSingular can reject the matrix. The solve therefore runs
|
||
// on the copy, so the steepest-descent fallback reads the
|
||
// pristine Jacobian: a direction formed from the factored
|
||
// matrix applies the permuted LU factors to r, which is not
|
||
// −Jᵀr and does not descend.
|
||
for i := range n {
|
||
copy(jacWork[i], jac[i])
|
||
}
|
||
colWrap[0] = col
|
||
if useBroyden {
|
||
// First round or restart after a degraded update:
|
||
// rebuild the inverse and step with it at once, so a
|
||
// fresh Jacobian is never paid for without moving. A
|
||
// singular Jacobian takes the same steepest-descent
|
||
// escape the Newton path uses, with the inverse left
|
||
// unfit so the next round rebuilds.
|
||
if ierr := broydenInvert(jac, jacWork, invJac); ierr != nil {
|
||
derr = ierr
|
||
steepestDescentStep(jac, r, step)
|
||
} else {
|
||
haveInv = true
|
||
broydenStep(invJac, r, step)
|
||
}
|
||
} else {
|
||
sol, serr := base.SolveSystem("FindRootSystem", jacWork, colWrap)
|
||
derr = serr
|
||
if serr != nil {
|
||
// A singular Jacobian has no Newton direction, but the
|
||
// steepest-descent direction of the merit function ‖r‖²
|
||
// always exists when r ≠ 0 and always decreases it, so
|
||
// the iteration escapes the singular locus instead of
|
||
// dying there.
|
||
steepestDescentStep(jac, r, step)
|
||
} else {
|
||
for i := range n {
|
||
step[i] = sol[0][i]
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// Backtrack along the step until the residual norm falls, at
|
||
// the Armijo fraction of the predicted linear decrease.
|
||
resSq := 0.0
|
||
for i := range n {
|
||
resSq += r[i] * r[i]
|
||
}
|
||
// Scale the descent fallback down to Newton's magnitude so the
|
||
// first trial is comparable.
|
||
if derr != nil {
|
||
stepNorm := normInfOfStep(step)
|
||
if stepNorm > 0 {
|
||
scale := normInfOfStep(col) / stepNorm
|
||
if scale > 0 && !math.IsInf(scale, 0) {
|
||
for i := range step {
|
||
step[i] *= scale
|
||
}
|
||
}
|
||
}
|
||
}
|
||
alpha := 1.0
|
||
improved := false
|
||
for range 60 {
|
||
for j := range n {
|
||
candidate[j] = x[j] + alpha*step[j]
|
||
}
|
||
rc, cerr := eval(candidate, false, nil)
|
||
if cerr != nil {
|
||
return nil, 0, base.Errf("FindRootSystem: %w", cerr)
|
||
}
|
||
candSq := 0.0
|
||
for i := range n {
|
||
candSq += rc[i] * rc[i]
|
||
}
|
||
// A trial whose residual overflowed is a rejected candidate:
|
||
// its norm is Inf or NaN, compares false below, and the
|
||
// halving continues, which is what the damping exists for.
|
||
// A finite sum of squares admits only finite components, so
|
||
// whatever passes both comparisons is adoptable.
|
||
if candSq <= (1-1e-4*alpha)*resSq || candSq < opts.Tolerance*opts.Tolerance {
|
||
// The candidate buffer is reused by the next round's
|
||
// backtracking, so the accepted point is copied into
|
||
// the iteration's own state.
|
||
if useBroyden {
|
||
// The update's step and residual differences come
|
||
// from the state the accepted step leaves, so they
|
||
// are read before the copy overwrites it.
|
||
for j := range n {
|
||
sVec[j] = candidate[j] - x[j]
|
||
yVec[j] = rc[j] - r[j]
|
||
}
|
||
}
|
||
copy(x, candidate)
|
||
r = rc
|
||
resNew := normInfOfStep(r)
|
||
// The update corrects a maintained inverse, so it only
|
||
// runs when one is fit to be corrected: after a
|
||
// singular round the rebuild next loop takes over.
|
||
if useBroyden && haveInv {
|
||
// A false report leaves the inverse untouched and
|
||
// sends the next round back to a fresh Jacobian.
|
||
haveInv, stalled = broydenMaintain(invJac, sVec, yVec, r, resNew, res, stalled)
|
||
}
|
||
res = resNew
|
||
improved = true
|
||
break
|
||
}
|
||
alpha /= 2
|
||
}
|
||
if !improved {
|
||
return nil, 0, base.Errf("FindRootSystem: the residual cannot be reduced below %g by damping", res)
|
||
}
|
||
// The documented success criterion is both a small step and a
|
||
// residual under the tolerance: a vanished step alone can come
|
||
// from a singular Jacobian whose steepest-descent fallback is
|
||
// numerically zero, and that point is a stall to keep working
|
||
// on, never a root to publish.
|
||
if res <= opts.Tolerance && normInfOfStep(step) <= opts.Tolerance*(1+normInfOfStep(x)) {
|
||
return linalg.ArrayFromFloatsSafe(x, len(x)), res, nil
|
||
}
|
||
}
|
||
// The last accepted step updated r and res after the loop-top test,
|
||
// so a run that met the tolerance exactly on the final iteration
|
||
// must re-test before the budget refusal reports it.
|
||
if res <= opts.Tolerance {
|
||
return linalg.ArrayFromFloatsSafe(x, len(x)), res, nil
|
||
}
|
||
return nil, 0, base.Errf("FindRootSystem: reached MaxIterations=%d with residual %g",
|
||
opts.MaxIterations, res)
|
||
}
|
||
|
||
// 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
|
||
}
|
||
|
||
// steepestDescentStep writes the negative gradient of the merit
|
||
// function ‖r‖², −Jᵀr, into step: entry i sums column i of the
|
||
// Jacobian against the residual, so entry j of the Jacobian's row k
|
||
// carries ∂r_k/∂x_j. jac must be the pristine central-difference
|
||
// matrix, never a factored one: Factor overwrites the subdiagonal
|
||
// with the multipliers and swaps the rows in place, so a factored
|
||
// matrix yields a direction that is no descent direction at all.
|
||
func steepestDescentStep(jac [][]float64, r, step []float64) {
|
||
for i := range step {
|
||
s := 0.0
|
||
for k := range r {
|
||
s += jac[k][i] * r[k]
|
||
}
|
||
step[i] = -s
|
||
}
|
||
}
|
||
|
||
// firstError returns the first non-nil of two errors.
|
||
func firstError(e1, e2 error) error {
|
||
if e1 != nil {
|
||
return e1
|
||
}
|
||
return e2
|
||
}
|