Module Jera.System
Zeros of systems.
A system is a function f from values of a structure 'x to values of the same structure, dtypes and shapes: its zero is the x with f x = 0. The float tensors of a value are its vector, of one dtype; its other tensors are carried from the guess. A system is one problem, with one status; Rune.vmap gives each lane its own.
(* x² + y² = 4 and x = y, from (1, 2) *) let f v = let x = Nx.get [ 0 ] v and y = Nx.get [ 1 ] v in Nx.stack [ Nx.sub_s (Nx.add (Nx.square x) (Nx.square y)) 4.; Nx.sub x y ] in System.solve Nx.Ptree.tensor (System.newton ~derivative:(fun v dv -> snd (Rune.jvp' f v dv))) ~linear:Linear.dense ~tol:(Tol.rel 1e-12) ~budget:20 f (Nx.create Nx.float64 [| 2 |] [| 1.; 2. |])
Error. Every method takes undamped steps δ: Newton's or Broyden's direction before the line search scales it, Anderson's fixed-point residual f x. A lane's error is e = |δ| / (1 − q) per component, at the answer x + δ, where f is not evaluated; y is the answer. q is the contraction of the undamped map N x = x + δ: the larger of its two newest readings |N x − N x'| / |x − x'|, or the one while only one exists, with lengths in tol's scaled norm, the root mean square of the components over their scales. A reading counts once the rounding of N's two values settles whether it is below 1; until then x' stays and the next reading spans more steps, and one whose least value exceeds q drops both. e is unbounded while q ≥ 1 or before a reading. A component of δ below the least normal float counts the subnormals' spacing as well, the floats' resolution there. One that is zero counts zero unless f may have underflowed short of the zero: where the answer there is a subnormal, or for Anderson's residual, which pairs equation i with unknown i, where f's component is not exactly zero, or for Newton's and Broyden's solved steps, each of whose components answers to every equation, where f is neither zero nor of largest magnitude at least tiny / ε. A shortened or mixed step moves the estimate but its length never enters e, so a search cannot converge by shrinking its steps.
If the step from x reduces the distance to the zero x* by a factor L in that norm, |x − x*| ≤ |δ| / (1 − L) and the answer is within L |δ| / (1 − L) of x*: within e where L ≤ q, and within tol wherever e met it and that last step at least halved the distance, as Newton's and Broyden's steps do near a simple zero. Otherwise e is an estimate, and the answer is within (1 − q) / (1 − L) times tol, which exceeds it in these cases. The map's rate rises near the zero beyond what its readings saw. The derivative is wrong at the answer by orders of magnitude, so that the last step is tiny and L near 1. f's evaluation error exceeds its rounding. Or one component contracts much more slowly than another: the readings see the fast one while the error lies in the slow one, as for (−x₀, (z − x₁) / 100) from (1, z + 300 ulps), which converges at its first reading about 50 times past Tol.ulps 4.
A lane whose step cannot move the estimate, one within a unit in the last place of every component, or whose line search finds no decrease, is at the floats' limit, where no further reading comes: its q is the newest settled reading, and it converges if e meets tol and ends Stalled otherwise. A tolerance finer than the steps a method still takes within an ulp of the answer cannot be met, and such a lane ends Stalled even where its estimate holds the zero: up to κ ulps at a minimum of condition κ, where a quasi-Newton model's steps and the gradient's bound both reach that far. A failed linear solve ends a lane Stalled; a non-finite f at the guess or at an estimate ends it Not_finite; budget iterations end it Budget_spent.
Derivative. The answer is stated as the zero of f through Rune.root, so its derivative is the implicit one, −J⁻¹ ∂f/∂θ through every tracked value f reads, with J rune's derivative of f at the answer, solved by linear; zero at a lane that did not converge. At a zero where J is singular the derivative does not exist and is not finite.
type 'x tThe type for methods for systems on values of type 'x.
newton ~derivative is Newton's method: derivative x dx is f's Jacobian at x applied to dx, such as snd (Rune.jvp' f x dx). It steers the steps only: one too steep by a modest factor slows a solve, one too shallow by 2 or more need not converge, and one wrong by orders of magnitude at the answer can end it early, as the error above says.
Method. Each step solves J δ = −f x with linear, then backtracks from x + δ to the sufficient decrease of |f|² / 2, by safeguarded quadratic steps. Stability. The line search makes every step decrease |f|, so the search can end at a local minimum of |f| that is no zero, where it stalls. Cost. Per iteration, one solve of linear on derivative and one evaluation of f per trial of the line search; quadratic convergence near a simple zero, and with a Krylov solver it is Newton–Krylov.
broyden is Broyden's method, for an f with no usable derivative.
Method. The Jacobian's estimate B starts as forward differences and takes Broyden's rank-one update after each step; each step solves B δ = −f x with linear and backtracks by halving to Li and Fukushima's derivative-free decrease test, |f (x + α δ)| ≤ (1 + η_k) |f x| − σ |α δ|² with η_k = 1 / (k + 1)² and σ = 10⁻⁴, which short steps always meet. Cost. n evaluations for the differences, n the size of the vector, then one per trial; it keeps the n × n matrix B, so it suits up to some hundreds of unknowns, and converges superlinearly near a simple zero.
anderson ~memory is for a fixed point x = g x stated as the residual f x = g x − x.
Method. Iterates g with Anderson mixing of the last memory steps: the mixed point is g x − ΔG γ, with ΔF and ΔG the changes of f and g over those steps and γ least |f x − ΔF γ| by QR, dropping the oldest steps while the ratio of R's diagonal entries exceeds ε^(−1/2). A mixed point that does not reduce |f| is replaced by the Picard point g x. memory = 0 is Picard iteration. Cost. One evaluation per iteration, two when the Picard point replaces the mixed one, and a QR of n × memory.
Raises Invalid_argument if memory < 0.
val solve :
('x, 'd) Nx.Ptree.t ->
'x t ->
linear:'x Linear.t ->
tol:Tol.t ->
budget:int ->
('x -> 'x) ->
'x ->
('x, 'd) Solution.tsolve x m ~linear ~tol ~budget f guess is a zero of f near guess by m. linear solves the method's linear systems and the derivative's. Its evaluations count the calls of f, budget its iterations.
Raises Invalid_argument if budget < 1, if the float tensors of guess differ in dtype, or if f returns a value of another structure, dtype or shape than its argument.
val lanes :
tol:Tol.t ->
budget:int ->
jacobian:((float, 'b, 'd) Nx.t -> (float, 'b, 'd) Nx.t) ->
((float, 'b, 'd) Nx.t -> (float, 'b, 'd) Nx.t) ->
(float, 'b, 'd) Nx.t ->
((float, 'b, 'd) Nx.t, 'd) Solution.tlanes ~tol ~budget ~jacobian f guess is a zero of f in each lane: one problem per lane, the last axis holding its k coordinates, as Root.newton is elementwise for one. f computes each lane's result from that lane alone; jacobian x, of shape lanes @ [k; k], is f's Jacobian per lane, and steers the steps only, as newton's derivative does.
Method. newton's search, with each lane's step a k × k direct solve of jacobian, checked against the backward error of an LU factorisation as Linear.dense checks, and its own line search and status. Error. As solve's, per lane. Cost. One evaluation of f and jacobian and a batch of k × k factorisations per iteration, plus the line search's trials. Derivative. Each lane's implicit one, −J⁻¹ ∂f/∂θ with J rune's derivative of f, whose blocks it probes with k products in every lane at once; it checks its solution with one more product and raises Invalid_argument where f read another lane.
(* u + d u = x for each point of [x], shape [n; 2] *) System.lanes ~tol:(Tol.v ~rel:1e-12 ~abs:1e-14) ~budget:20 ~jacobian:(fun u -> Nx.add (Nx.eye Nx.float64 2) (jacobian_of_d u)) (fun u -> Nx.sub (Nx.add u (d u)) x) x
Raises Invalid_argument if budget < 1, if guess is a scalar, or if f or jacobian returns another shape.