feat: initial release

Assisted-by: GLM 5.3 Flash

optim/rootsystem.go

@@ -0,0 +1,437 @@
+// 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
+}