NonlinearIntegrators.jl

AbstractDenseNetBasis{T} <: NetworkBasis{T}

Abstract supertype for three-layer dense-network bases. The concrete implementation is DenseNetBasis{T}.

AbstractShallowNetBasis{T} <: NetworkBasis{T}

Abstract supertype for single-hidden-layer network bases. The concrete implementation is ShallowNetBasis{T}.

AngularGrid(; radii = (1.0,), amount = nothing)

Atoms placed on rays through the origin of (w, b) space: (w, b) = r · (cos θ, sin θ) for a uniform grid of amount + 1 angles over [0, 2π) and each radius in radii. amount = nothing uses the method's dict_amount.

This is the dictionary the underlying approximation theory is stated for — a grid on the unit sphere of ℝ^{d+1} — and it unifies the 1-D and 2-D cases. For a homogeneous activation the radius is redundant (it only rescales the atom), so a single radius suffices and the set covers the same ridge directions as {±1} × (bias grid), but sampled uniformly in atom space rather than uniformly in bias: the ±1 grid over [-π, π] clusters its resolution where |b| is large and the atom is nearly constant on t ∈ [0, 1], which is where it matters least. For a smooth activation the radius is the length scale, and log-spaced radii give the second degree of freedom without having to pick a weight interval.

The full circle is used rather than a half circle because ReLUᵏ(x) and ReLUᵏ(-x) are different functions — the sign of w is real shape information (only the scale is redundant).

AutodiffShallowNetCache{ST}(ics, nx, S, R, N; record_grid_points = 41)

Cache for the shallow-net integrators that differentiate a hand-written ansatz with ForwardDiff: ShallowNetAutodiff and ShallowNetAutodiffReversible.

nx is the number of unknowns per dimension.

BiasGrid1d()

The original dictionary: weights w = ±1 crossed with a uniform grid of dict_amount + 1 biases over the method's bias_interval, for 2·(dict_amount + 1) atoms.

Complete for ReLUᵏ (see the discussion at the top of this file) and the default. Atom order — the w = -1 block first, then w = +1 — is load-bearing: argmax breaks ties by first index, so reordering changes which atoms are selected.

DenseNet <: DenseNetMethod

Continuous Galerkin Variational Integrator using a DenseNetBasis basis (three hidden layers). The network ansatz q(t) = NN(t; θ) uses a deeper architecture than ShallowNet, at the cost of a larger parameter space per step.

Constructor

DenseNet(basis, quadrature; kwargs...)

Required:

• basis::AbstractDenseNetBasis{T} — e.g. DenseNetBasis{T}(activation, S₁, S)
• quadrature::QuadratureRule{T}

Keyword arguments:

• initial_trajectory_method — IntegratorExtrapolation() (default), HermiteExtrapolation(), or NoExtrapolation()
• initial_guess_method — LSGD() (default) or TrainingMethod()
• extrapolation_substep::Int = 10
• training_epochs::Int = 50000
• record_grid_points::Int = 41

Example

basis  = DenseNetBasis{Float64}(tanh, 8, 8)
quad   = GaussLegendreQuadrature(Float64, 8)
method = DenseNet(basis, quad; training_epochs = 1000)

DenseNetBasis{T} <: AbstractDenseNetBasis{T}

Three-layer dense-network basis for DenseNet. Architecture: Dense(1, S₁, σ) → Dense(S₁, S, σ) → Dense(S, 1). Symbolic derivatives are compiled at construction time.

Constructor

DenseNetBasis{T}(activation, S₁, S; backend = CPU(), cse = true, inplace = true)

• activation: elementwise activation (e.g. tanh)
• S₁::Int: first hidden layer width
• S::Int: second hidden layer width (= number of basis functions)
• cse, inplace: forwarded to SymbolicNeuralNetworks.build_nn_function, as for ShallowNetBasis. This is the basis where cse matters most: it is the extra layer that makes re-emitting the shared forward pass per gradient block expensive, and turning it off costs about four times the build (110 ms against 27 ms for tanh at S₁ = S = 3, Float64, once the code generation itself has been compiled).

There is no symbolic = false here: DenseNet has no ForwardDiff counterpart, so the compiled derivatives are always read.

Example

basis = DenseNetBasis{Float64}(tanh, 8, 8)

DenseNetMethod <: NetworkIntegratorMethod

Abstract supertype for integrators whose ansatz is a three-layer dense network: DenseNet. Accepts LSGD (default) or TrainingMethod as its initial_guess_method.

IncrementalQR()

Reuse the incrementally maintained QR factorisation (see IncrementalQRState): one triangular solve per greedy step instead of a fresh k × k factorisation.

Numerically equivalent to WeightedQR up to rounding, but O(k · nnodes) rather than O(k² · nnodes) per step, and it is the fit that shares its Q with OrthogonalProjection — pairing the two is the "textbook" efficient and stable OGA.

IncrementalQRState{T}(nnodes, maxcols)

Thin QR factorisation of the greedily selected design matrix, maintained by appending one column at a time.

OGA adds exactly one atom per iteration, so re-solving a k × k system from scratch at every step repeats work the previous step already did. Maintaining Q (orthonormal columns) and R (upper triangular) instead costs O(k · nnodes) per step, never forms a Gram matrix, and, as a side effect, yields two quantities the greedy loop wants anyway:

• the norm of the new column after deflation against the already-selected ones, which is the rank-gain the atom actually contributes (see OrthogonalProjection); and
• Q itself, which turns the orthogonal-greedy selection score into one matrix product against the dictionary.

Columns are stored √w-scaled, so the Euclidean inner products here are the quadrature-weighted ones and no re-scaling happens inside the loop.

InitialParametersMethod

Abstract supertype for the strategy used to initialise the network weights (parameter vector) at the start of each time step before the Newton solve.

IntegratorExtrapolation <: Extrapolation

Initial-trajectory method that seeds the per-step Newton solve by integrating a short sub-problem with ImplicitMidpoint over extrapolation_substep sub-steps. This is the default initial_trajectory_method for all network integrators and usually gives the best convergence.

See also: NoExtrapolation, HermiteExtrapolation (from GeometricIntegratorsBase).

LSGD <: InitialParametersMethod

Initialise dense-network parameters with a Lipschitz-SGD step (LSGD). Used as the default initial_guess_method for DenseNet.

MirrorPairs()

One atom (w, b) → two neurons (w, b) and (-w, w + b) → two independent design columns, so each member of the pair gets its own output weight.

The map (w, b) ↦ (-w, w + b) sends σ(w t + b) to σ(w (1 - t) + b), i.e. reflects the neuron about the midpoint of the unit time interval. Used by ShallowNetReversible.

Since neurons come in pairs, the requested neuron count must be even; oga_fit rejects an odd one rather than quietly placing one neuron fewer.

NetworkBasis{T} <: Basis{T}

Abstract supertype for all neural-network basis types in this package. The type parameter T is the floating-point element type (e.g. Float64, Float32). Common fields (NN, activation, SNN, dqdθ, V_func, dvdθ) live in a NetworkBasisCore sub-struct and are forwarded through getproperty, so call sites can write basis.NN, basis.activation, etc.

NetworkBasisCore{AT,NT,BT,SNNT,QWFT,VT,VWFT}

Common sub-struct shared by all NetworkBasis concrete types. Bundles the neural network together with its symbolically-compiled derivatives, which are built once at construction time and reused at every integration step.

SNN, dqdθ, V_func and dvdθ are all nothing for a ShallowNetBasis built with symbolic = false — the form the ForwardDiff-based integrators want, since they differentiate their ansatz at run time and never read these fields. Test for it with has_symbolic_derivatives rather than by comparing a field against nothing; the four type parameters are Nothing in that case, so the distinction is visible to dispatch.

Fields

• activation: activation function used in the hidden layer.
• NN: neural-network model (the forward map q(t; θ)).
• backend: computation backend (e.g. CPU array backend).
• SNN: SymbolicNeuralNetwork wrapping NN, used to derive symbolic expressions for all required derivatives.
• dqdθ: compiled function returning ∂q/∂θ — the Jacobian of the network output (position) with respect to the network parameters θ.
• V_func: compiled function returning the velocity v = dq/dt — the time-derivative of the network output, obtained from the symbolic Jacobian with respect to the time input.
• dvdθ: compiled function returning ∂v/∂θ — the Jacobian of the velocity with respect to the network parameters θ.

NetworkIntegratorMethod <: LODEMethod

Abstract supertype for all neural-network-based variational integrators in this package. Every concrete subtype wraps a NetworkIntegratorCore under the field common and exposes its fields (basis, quadrature, extrapolation settings, …) via getproperty forwarding, so call sites can write method.basis, method.record_grid_points, etc.

NoSymmetry()

One atom → one neuron → one design-matrix column.

NormalEquationsFit(; ridge = true, island = false)

Solve the normal equations G x = Φ diag(w) y with G = Φ diag(w) Φᵀ, optionally with the precision-scaled Tikhonov ridge of oga_tikhonov (ridge = true) and optionally in Float64 regardless of the working precision (island = true).

This is the baseline, not a recommendation: forming G squares the condition number. It exists so that "island vs. working precision" and "ridge vs. no ridge" are two knobs on one code path that can be ablated in a benchmark, rather than differences buried in forked implementations. NormalEquationsFit(ridge = false, island = true) reproduces the arithmetic of OGA1dNormalEquations, the original-paper reference.

Note island = true deliberately violates the package's precision discipline; it is the thing being measured against.

NormalizedProjection()

Score by the cosine-like ratio |⟨r, g⟩_w| / ‖g‖_w.

This is the textbook greedy criterion: it measures how much of the residual the atom explains, independently of the atom's amplitude, which the output weight absorbs anyway. Scale invariance makes it mandatory for WeightBiasGrid2d and AngularGrid, where atoms differ in norm by orders of magnitude.

OGA(dictionary = BiasGrid1d(), selection = RawProjection(), fit = WeightedQR();
    coherence = true, norm_guard = true, fill_unused = true)

Orthogonal Greedy Algorithm initial guess for the network integrators.

At every time step the integrator must solve a nonlinear system for the parameters of a shallow network u(x) = Σₖ cₖ σ(wₖ x + bₖ). The OGA produces the starting point: it repeatedly picks, from a fixed dictionary of candidate neurons (w, b), the atom most correlated with the current fit residual, then refits all output weights c by a quadrature-weighted least-squares solve — the "orthogonal" part, as opposed to plain matching pursuit, which would only fit the new atom.

The three axes:

• dictionary::OGADictionary — the candidate neuron set: the original {±1} × (bias grid) (BiasGrid1d), a genuine 2-D (w, b) grid (WeightBiasGrid2d), an angular grid on the atom sphere (AngularGrid), or any of them with off-grid polish (Refined).
• selection::OGASelection — how candidates are ranked: RawProjection, NormalizedProjection or OrthogonalProjection.
• fit::OGAFit — how the output weights are refit: WeightedQR, IncrementalQR, PivotedQR, TruncatedSVD or NormalEquationsFit.

Guard rails, all scaled to eps(T) rather than to absolute constants:

• coherence — after an atom is selected, block dictionary atoms whose weighted L² coherence with it exceeds 1 - sqrt(eps(T)). Inert at Float64/Float32; it only bites where distinct atoms have rounded together.
• norm_guard — treat atoms whose weighted norm falls below oga_norm_floor as unusable rather than normalising by noise. Atoms with a non-finite norm (reachable at Float16 with a high ReLUᵏ power over a wide bias interval, where σ(b)ᵏ overflows) are always excluded, guard or not.
• fill_unused — when the greedy loop runs out of usable atoms before all S neurons are placed, give the remaining neurons distinct, well-separated (w, b) with zero output weight. Without this they would all keep (0, 0) and become identical rows of the Newton Jacobian — trading a rank-deficient seed for a rank-deficient solve.

Everything runs at the solver's working precision T; see the precision note in src/oga/numerics.jl.

Examples

ShallowNet(basis, quadrature)                                # OGA1d(), the default
ShallowNet(basis, quadrature; initial_guess_method = OGA2d())
ShallowNet(basis, quadrature;
    initial_guess_method = OGA(BiasGrid1d(), OrthogonalProjection(), TruncatedSVD()))

OGA1dNormalEquations()

The reference implementation from the original paper, kept as a selectable baseline for comparison.

The dictionary and the greedy least-squares fit are assembled in Float64 — a "double-precision island" — the output weights come from the normal equations Gₖ xₖ = bₖ, and the result is rounded into the working-precision cache. It carries none of the precision-scaled guard rails: no norm floor, no coherence guard, no ridge, and no rank detection.

That combination is why it is the baseline rather than the default. Forming Gₖ squares the condition number, so the fit needs roughly twice the digits the problem does, which is what the Float64 island supplies; and at 16 bits the atom selection degrades until the third or fourth selected neuron is linearly dependent on its predecessors, at which point the Gram solve raises SingularException before the Newton iteration has begun. See the "Orthogonal Greedy Algorithm" section of the documentation for the full analysis, and NormalEquationsFit for the same arithmetic available as a fit strategy inside the modern composable OGA, where the island and the ridge can be toggled independently.

Select it with ShallowNet(...; initial_guess_method = OGA1dNormalEquations()).

OGADictionary

The candidate neuron set the greedy step selects from. One of BiasGrid1d, WeightBiasGrid2d, AngularGrid, or any of those wrapped in Refined.

OGAFit

How the OGA refits the output weights of the selected atoms. One of WeightedQR, IncrementalQR, PivotedQR, TruncatedSVD or NormalEquationsFit; see each for the trade-off.

OGAResult{T}

What oga_fit returns.

• W, b, c — hidden weights, hidden biases and output weights, one entry per neuron.
• atoms — the dictionary indices selected, in order.
• neurons — how many neurons the greedy loop actually placed; a smaller number than requested means it ran out of atoms that add a new direction, and the remainder were filled with zero-weight placeholders (see fill_unused in OGA).
• residual — the weighted L² norm of the final fit residual.
• gains — the rank gain ‖g⊥‖ of each accepted atom: how much genuinely new direction it contributed. A gain collapsing towards zero across the sequence is the fingerprint of the reduced-precision failure this subsystem exists to remove.
• rejected — how many candidate atoms were skipped for adding no new direction.

OGASelection

How the greedy step scores candidate atoms against the current residual. One of RawProjection, NormalizedProjection or OrthogonalProjection.

OGASymmetry

How one selected dictionary atom turns into network neurons. NoSymmetry is the plain case; the two mirror variants exist for the time-reversible integrators, whose ansatz requires neurons to come in pairs related by t ↦ 1 - t.

OrthogonalProjection(; min_gain = nothing)

Score by the projection onto the part of the atom orthogonal to the already selected ones, |⟨r, g⟩_w| / ‖g⊥‖_w, and refuse any atom whose orthogonal part has collapsed (‖g⊥‖ < min_gain · ‖g‖); min_gain = nothing uses sqrt(eps(T)).

This is what makes the algorithm orthogonal greedy rather than matching pursuit, and it is the direct fix for the observed reduced-precision failure. The residual is already orthogonal to the selected span, so the numerator is unchanged from NormalizedProjection — but the denominator penalises an atom that mostly duplicates what is already there, and the min_gain floor rules it out entirely. An atom that adds no new direction is therefore never selected, which is the condition that otherwise surfaces downstream as SingularException: zero pivot found at index 3 out of four neurons.

Costs one dictionary-sized matrix product against the maintained Q per step, the same order as the score itself.

PivotedQR(; rtol = nothing)

Rank-revealing Householder QR with column pivoting, truncated below rtol · (largest pivot); rtol = nothing uses eps(T) · max(4, k) with k = min(nnodes, ncols) — see the note above _rtol for why not sqrt(eps(T)).

Unlike WeightedQR, a numerically dependent selected atom is detected and given a zero coefficient rather than solved through. Hand-rolled because qr(Â, ColumnNorm()) is LAPACK-only and so does not exist at Float16 — the precision that needs it.

RawProjection()

Score by the bare weighted inner product |⟨r, g⟩_w|.

The default, and what the Float64/Float32 regression tests pin. Note it is not scale-invariant — an atom with a large norm outranks a better-aligned small one — which is harmless for the {±1} × (bias grid) dictionary, whose atoms all have comparable norms, and wrong for a 2-D (w, b) dictionary, whose atoms do not.

Refined(inner; iterations = 3, shrink = 0.5)

Wrap any dictionary so that, after the greedy argmax picks a grid atom, the atom's (w, b) are polished off the grid by locally maximising the same selection score.

The grid then only has to get the neighbourhood right, which decouples accuracy from dictionary size: a few dozen atoms plus refinement can match a dictionary of hundreds of thousands, and the greedy step is linear in the dictionary size. This is the standard "OGA with inner optimisation".

The local search is a derivative-free compass search — evaluate the score at (w ± h, b) and (w, b ± h), step to the best improvement, shrink h by shrink when none improves — repeated iterations times. Derivative-free on purpose: the score is only piecewise smooth for ReLUᵏ (the kink crosses a quadrature node), and it keeps the activation off the ForwardDiff path entirely, so no Dual tag can leak into the working precision.

The polished objective is always the normalised score, even when the selection rule is RawProjection — see the note in _candidate_score. Maximising the raw inner product continuously over (w, b) would reward growing the atom rather than fitting the residual.

ShallowNet <: ShallowNetMethod

Continuous Galerkin Variational Integrator using a ShallowNetBasis basis. The network ansatz q(t) = NN(t; θ) is a single-hidden-layer network; the optimal parameters θ are found by Newton's method applied to the discrete Euler-Lagrange equations at each time step.

Constructor

ShallowNet(basis, quadrature; kwargs...)

Required:

• basis::AbstractShallowNetBasis{T} — e.g. ShallowNetBasis{T}(activation, S)
• quadrature::QuadratureRule{T} — e.g. GaussLegendreQuadrature(T, R)

Keyword arguments:

• initial_trajectory_method — IntegratorExtrapolation() (default), HermiteExtrapolation(), or NoExtrapolation()
• initial_guess_method — an OGA seed: OGA1d() (default), OGA1dNormalized(), OGA1dStable(), OGA2d(), OGASphere(), or a hand-built OGA(dictionary, selection, fit). Also OGA1dNormalEquations() (the original-paper reference path) and TrainingMethod().
• extrapolation_substep::Int = 10 — sub-steps for the IntegratorExtrapolation warm start
• training_epochs::Int = 50000 — gradient-descent epochs when initial_guess_method = TrainingMethod()
• bias_interval — bias search range for OGA dictionary, default [-π, π]
• dict_amount::Int = 50000 — number of atoms in the OGA dictionary
• record_grid_points::Int = 41 — number of grid points per step stored in stage_values

Example

using NonlinearIntegrators, QuadratureRules
basis = ShallowNetBasis{Float64}(tanh, 8)
quad  = GaussLegendreQuadrature(Float64, 8)
method = ShallowNet(basis, quad; bias_interval = [-π, π], dict_amount = 400)

ShallowNetAutodiff <: ShallowNetMethod

Shallow-net variational integrator that computes network derivatives with ForwardDiff instead of the pre-compiled symbolic derivatives used by ShallowNet. The ansatz and optimisation are otherwise identical to ShallowNet.

Constructor

ShallowNetAutodiff(basis, quadrature; kwargs...)

Keyword arguments are the same as ShallowNet: initial_trajectory_method, initial_guess_method, extrapolation_substep, training_epochs, show_status, bias_interval, dict_amount, record_grid_points.

ShallowNetAutodiffReversible <: ShallowNetMethod

Time-symmetric variant of ShallowNetAutodiff: computes network derivatives with ForwardDiff (no symbolic pre-compilation) and enforces the palindromic time-reversal symmetry (issymmetric(method) == true). Combines the forward-differentiation approach of ShallowNetAutodiff with the symmetry structure of ShallowNetReversible.

Constructor

ShallowNetAutodiffReversible(basis, quadrature; kwargs...)

Keyword arguments are the same as ShallowNetAutodiff. Only the OGA seeds are supported as initial_guess_method. The basis must have an even number of neurons — they come in mirrored pairs sharing one output weight.

ShallowNetBasis{T} <: AbstractShallowNetBasis{T}

Single-hidden-layer network basis, built with AbstractNeuralNetworks. The network maps a scalar time input to a scalar position: Dense(1, S, σ) → Dense(S, 1). Symbolic derivatives (dqdθ, V_func, dvdθ) are compiled once at construction time via SymbolicNeuralNetworks.jl, unless symbolic = false.

Constructor

ShallowNetBasis{T}(activation, S; backend = CPU(), symbolic = true,
                   cse = true, inplace = true)

• activation: any elementwise activation function (e.g. tanh, relu_k(3))
• S::Int: number of hidden neurons (= number of basis functions)
• symbolic: compile the symbolic derivatives (default true). Pass false to build the network only and leave SNN, dqdθ, V_func and dvdθ as nothing.
• cse, inplace: forwarded to SymbolicNeuralNetworks.build_nn_function; see below.

symbolic = false exists for ShallowNetAutodiff and ShallowNetAutodiffReversible: those two differentiate their ansatz with ForwardDiff at run time and never read the compiled derivatives, so building them is pure overhead — 15 ms against 29 ns for tanh at S = 8, Float64, once the code generation itself has been compiled; the basis is then just a Chain. The integrators that do read them (ShallowNet, ShallowNetReversible, DenseNet) reject such a basis in their constructor; see has_symbolic_derivatives.

cse (common-subexpression elimination during code generation) and inplace (evaluate a batch through a kernel writing into one preallocated array) both default to true, which is also what SymbolicNeuralNetworks 0.6 uses. They are pinned here rather than left to the upstream default so that a change there cannot silently change this package's code generation. They are exposed to be turned off: cse = false, inplace = false emits the whole shared forward pass once per gradient block and evaluates a batch out of place, one allocation per sample, which is what benchmark/compare_derivative_backends.jl measures the two settings against each other for. Note that inplace = true mutates its output and so cannot be differentiated with Zygote; nothing in this package does that, but a caller who wants to needs inplace = false.

Example

basis = ShallowNetBasis{Float64}(tanh, 8)
autodiff_basis = ShallowNetBasis{Float64}(tanh, 8; symbolic = false)
plain_codegen = ShallowNetBasis{Float64}(tanh, 8; cse = false, inplace = false)

ShallowNetMethod <: NetworkIntegratorMethod

Abstract supertype for integrators whose ansatz is a single-hidden-layer network: ShallowNet, ShallowNetAutodiff, ShallowNetReversible, and ShallowNetAutodiffReversible. All take an OGA seed as their initial_guess_method — OGA1d() by default, OGA1dNormalized() for ShallowNetAutodiff.

ShallowNetReversible <: ShallowNetMethod

Time-symmetric variant of ShallowNet. The variational integrator is constructed so that reversing the time direction recovers the original trajectory, giving issymmetric(method) == true. Uses a symmetric quadrature / ansatz structure while keeping the same ShallowNetBasis basis.

Constructor

ShallowNetReversible(basis, quadrature; kwargs...)

Keyword arguments are the same as ShallowNet. Only the OGA seeds are meaningful as initial_guess_method; TrainingMethod is not specialised for this integrator. The basis must have an even number of neurons — they come in mirrored pairs.

SharedMirrorPairs()

One atom → two mirrored neurons → one design column, their sum, so the pair shares a single output weight.

Sharing the weight is what actually enforces time-reversal symmetry of the ansatz (with independent weights the pair can drift apart). Used by ShallowNetAutodiffReversible.

As for MirrorPairs, the requested neuron count must be even.

SymbolicShallowNetCache{ST}(ics, nx, S, R, N; record_grid_points = 41)

Cache for the shallow-net integrators that evaluate the symbolically compiled derivatives of their basis: ShallowNet and ShallowNetReversible.

nx is the number of unknowns per dimension.

TrainingMethod <: InitialParametersMethod

Initialise network parameters by gradient descent (GeometricOptimizers.Adam with a DecayingStatic line search) against a mean-absolute-error target built from the extrapolated trajectory. Applies to both ShallowNet and DenseNet; controlled by the training_epochs constructor kwarg.

TruncatedSVD(; rtol = nothing)

Minimum-norm solve through a truncated pseudo-inverse, dropping singular directions with σ < rtol · σ_max; rtol = nothing uses eps(T) · max(4, k) with k = min(nnodes, ncols) — see the note above _rtol for why not sqrt(eps(T)).

The most robust of the fits — a rank-deficient selected set gives a bounded solution instead of amplified rounding noise — and the one to reach for if a single variant must work unchanged across every precision. Uses the generic one-sided Jacobi jacobi_svd, since svd is likewise LAPACK-only.

VISE(basis::VISEBasis{T}, quadrature, init_w; extrapolation_substep = 10,
     record_grid_points = 41)

Variational integrator on a symbolic ansatz: one closed-form expression per degree of freedom, whose free weights are solved for at every step.

Where the network integrators fit a shallow or dense network, VISE takes an ansatz the caller writes down — W₁·cos(W₂·t + W₃), say — and lets VISEBasis differentiate it symbolically with respect to t and to each weight. The discrete Euler–Lagrange equations are then solved for the weights directly. An ansatz that spans the exact solution therefore reproduces it to Newton's residual floor rather than to a discretisation order.

Arguments

• basis: a VISEBasis carrying the compiled ansatz and its derivatives.
• quadrature: the quadrature rule for the action integral.
• init_w: one initial weight vector per degree of freedom. Also the fallback restart point: initial_guess! reuses the previous step's weights unless they have drifted from init_w by more than 1 in norm.

Keywords

• extrapolation_substep = 10: sub-steps of the warm-start extrapolation.
• record_grid_points = 41: rows of the per-step stage_values record — the continuous solution between two discrete steps, returned as the second element of integrate's tuple. Same keyword as on the network integrators.

Note

Unlike the network integrators, integrate returns a three-element tuple here: (sol, internal_values, each_step_solution), the last being the converged weight vector of every step.

quadrature is a type parameter rather than an untyped field: it was the one untyped field of this struct, and method.quadrature.nodes is read at the top of components!, so quad_nodes came back Any and poisoned every expression it appeared in — including the arguments of the compiled basis functions.

VISEBasis{T}(q_expr, W, t, D)

Basis for VISE: a symbolic ansatz q_expr[d](W[d], t) per degree of freedom d, together with its time derivative and its derivatives with respect to every weight.

q_expr is a vector of D Symbolics expressions, W a vector of D symbolic weight arrays, and t the symbolic time variable. Everything is compiled at construction.

Implementation

The compiled callables come from Symbolics.build_function(…; expression = Val(false)), which returns a RuntimeGeneratedFunction. The previous form was Symbolics.eval(Symbolics.build_function(…)), which had two costs:

• World age. eval adds methods to the table in a newer world than the one the calling code is running in, so building a basis and evaluating it within the same function body raised MethodError: … The applicable method may be too new. It happened to work when the basis was built at top level, because top-level statements advance the world age between them — which is why the only test of this integrator built it that way, and why a test that used a helper function to build it failed.
• Typing. eval returns Any, so q_expr, dqdW, v_expr and dvdW could not be given concrete field types no matter how the struct was declared, and every call through them in components! — R × W_size × D of them per residual evaluation — was a dynamic dispatch. RuntimeGeneratedFunction is a concrete type, so the fields below are parametrised on it and the calls resolve statically.

WeightBiasGrid2d(; octaves = (-3.0, 3.0), weight_amount = 6, signed = true,
                   bias_amount = nothing)

A genuine 2-D grid over (w, b): weight_amount + 1 weight magnitudes spaced logarithmically over 2^octaves[1] … 2^octaves[2], optionally sign-symmetric, crossed with the bias grid.

The weight axis spans length scales by ratio, so the default covers a factor of 64 between the sharpest and gentlest transition. bias_amount overrides the method's dict_amount on the bias axis (nothing keeps it), which is how the total dictionary size is held roughly constant while the weight axis is added — the greedy step is linear in the dictionary size.

Set octaves = (0.0, 0.0), weight_amount = 0 to recover BiasGrid1d exactly: the extra weight degrees of freedom are redundant for ReLUᵏ, so this is a strict generalisation — neutral for the homogeneous activations and enabling for the smooth ones. Pair it with NormalizedProjection: atoms of very different |w| have very different raw norms, so RawProjection would rank them by amplitude rather than by fit.

WeightedQR()

QR least squares on the √w-scaled design matrix (see weighted_lstsq): conditioned on κ(Φ) rather than κ(Φ)², with no Gram matrix and no ridge unless the plain solve comes back non-finite.

The default, and the fit whose behaviour the Float64/Float32 regression tests pin. It re-solves from scratch at each greedy step.

NN_ansatz(ps, S, activation, t, q̄, q)

The trajectory ansatz of ShallowNetAutodiff and ShallowNetAutodiffReversible, on the unit interval t ∈ [0, 1]:

q_h(t) = (1-t)·q̄ + t·q + t(1-t)·N(t),    N(t) = Σᵢ W2ᵢ·σ(W1ᵢ·t + b1ᵢ)

ps is the flat [W2 | W1 | b1] parameter vector of 3S entries; q̄ and q are the two endpoints. The t(1-t) factor makes the network vanish at both ends, so the ansatz interpolates the endpoints exactly — q_h(0) = q̄ and q_h(1) = q by construction, whatever the parameters. That is why components! needs no boundary-point parameter gradients.

Its time derivative is VNN_ansatz.

OGA1d(; kwargs...)

The default seed: the original {±1} × (bias grid) dictionary, raw-projection selection, and a QR fit of the √w-scaled design matrix.

Its atom choices are pinned by the regression tests: normalising before selection steers the Newton solve into a different and empirically worse basin.

OGA1dNormalized(; kwargs...)

OGA1d, but selecting on the normalized inner product.

ShallowNetAutodiff's constructor default: alone among the four integrators it ranks candidates by |⟨r, g⟩_w| / ‖g‖_w rather than by the raw projection. Which of the two an integrator uses changes which neurons get picked and therefore which basin the Newton solve lands in, so each keeps the rule it was tuned with rather than inheriting a single default.

OGA1dStable(; kwargs...)

The same 1-D dictionary made robust at reduced precision: orthogonal-greedy selection with a rank-gain floor, on top of the incrementally maintained QR.

The combination aimed squarely at the 16-bit failure mode: an atom that adds no new direction is never selected, so the selected design matrix cannot go rank-deficient regardless of precision.

OGA2d(; dictionary = WeightBiasGrid2d(), kwargs...)

A 2-D (w, b) dictionary with normalised selection and the incremental QR fit: the variant for activations that are not positively homogeneous (ELU, GELU, tanh), where |w| is a genuine length-scale degree of freedom rather than redundant with b and c.

OGASphere(; dictionary = AngularGrid(), kwargs...)

Atoms sampled uniformly on the sphere of (w, b) space rather than uniformly in bias — the dictionary the underlying approximation theory is stated for. See AngularGrid.

VNN_ansatz(ps, S, activation, t, q̄, q)

d/dt of NN_ansatz, in closed form.

components! used to get this from Zygote.gradient — reverse mode, for a scalar ℝ→ℝ derivative, once per quadrature node per dimension per Newton residual and per Jacobian column. The obvious replacement, ForwardDiff.derivative(tt -> NN_ansatz(…, tt, …), t), does not work at that call site: SimpleSolvers builds its Jacobian with untagged (Tag = Nothing) Duals, so ps and q arrive untagged, and ForwardDiff cannot order Nothing against the tag its own inner derivative introduces. That is what kept the Zygote call alive.

Differentiating by hand sidesteps the nesting. With

q_h(t) = (1-t)·q̄ + t·q + t(1-t)·N(t),    N(t) = Σᵢ W2ᵢ·σ(W1ᵢ·t + b1ᵢ)

we have

q_h'(t) = q - q̄ + (1-2t)·N(t) + t(1-t)·N'(t),    N'(t) = Σᵢ W2ᵢ·W1ᵢ·σ'(W1ᵢ·t + b1ᵢ)

so only the scalar activation still needs differentiating, and that stays within a single tag level whatever ps is. No allocation, no reverse-mode tape, and one fewer dependency.

_param_arrays(params) -> Tuple

Every layer parameter array of params, vec'd, in layer-then-field order. The backing tuple for flatten_params! and flatten_params; it aliases params rather than copying.

@generated because the layer/field structure is in the type, so the sequence can be emitted once at compile time. Walking it at run time is what made the old flatten_params type-unstable — values(params) iterates a heterogeneous NamedTuple, so fieldnames(typeof(layer)) inside the loop cannot be folded.

This is not NeuralNetworkParameters.flatten!, which does the same walk. That one is allocation-free only when it is handed a ParameterLayout, and building the layout per call is what has to be avoided here: the four call sites in DenseNet's components! flatten a freshly built gradient set, and there is nowhere on the cache to keep a layout for it today. A @generated walk needs no layout at all. The training loops, which flatten one long-lived parameter set, do use the upstream pair — see initial_params! in shallownet.jl.

bias_grid(lo, hi, n, ::Type{T}) -> Vector{T}

Uniform grid of n + 1 bias values from lo to hi in precision T.

The coordinates are generated from an integer-indexed range in Float64 and then cast to T, so a large n cannot overflow the step to zero in reduced precision: at Float16, T(n) overflows to Inf for n > 65504 and (hi - lo)/n evaluates to zero, which a range built in T rejects with ArgumentError: range step cannot be zero. Only the grid coordinates touch Float64; the seed's dictionary and solve run entirely in T.

box_init_plain(input_dim, output_dim, ::Type{T}; rng = Random.default_rng())

Draw a "box" initialisation of a output_dim × input_dim weight matrix and its bias at element type T.

T is mandatory. It used to default to Float32, and every call site omitted it, so a Float64 network was initialised at single precision and then converted on assignment — invisible to a test that checks the eltype of the result. The package's central invariant is that a run started at T stays at T, so the initialiser has to be told which T.

rng is drawn from, not re-seeded. It used to be a keyword defaulting to the expression Random.seed!(1), which is evaluated afresh on every call that omits it: each call silently reseeded Julia's global RNG and then drew from it, so consecutive calls returned correlated draws and any seeding the caller had done was discarded. Seed at the call site instead.

build_network_derivatives(NN; cse = true, inplace = true) -> (SNN, dqdθ, V_func, dvdθ)

Compile the four symbolic slots of a NetworkBasisCore for the network NN: the gradient of the output with respect to the parameters, the time derivative of the output, and the gradient of that with respect to the parameters. cse and inplace go straight to SymbolicNeuralNetworks.build_nn_function; see ShallowNetBasis.

Shared by ShallowNetBasis and DenseNetBasis, which differ only in the NN they hand over — keeping the two in one place rather than in two copies that drifted apart once already. It also lets ShallowNetBasis's symbolic = false branch read as the single expression it is.

create_internal_stage_vector(DT, D, S) -> Vector{Vector{DT}}

S zero vectors of length D, one per internal stage of a D-dimensional problem. Used by the integrator caches to hold the stage values Q, P, V and F.

flatten_params!(dest, params) -> dest

Copy every layer parameter array of params, in layer-then-field order, into the flat vector dest.

Replaces an allocating flatten_params that built a flat_list = [] — a Vector{Any} — and returned vcat(flat_list...), which infers as Any when splatted. DenseNet's components! calls this 2 + 2R times per dimension per residual evaluation, i.e. per Newton iteration and per Jacobian column, so it is worth having a form that writes into a caller-supplied buffer.

flatten_params(params) -> Vector

Allocating form of flatten_params!. Not on any hot path; kept because it reads better in one-off code and in the benchmarks. reduce(vcat, ...) over the concretely typed tuple keeps the element type generic, where the old splatted vcat over a Vector{Any} inferred as Any.

has_symbolic_derivatives(basis) -> Bool

Whether basis carries the derivatives (dqdθ, V_func, dvdθ) that SymbolicNeuralNetworks.jl compiles at construction time.

false only for a ShallowNetBasis built with symbolic = false, which is the form the ForwardDiff-based integrators want — they differentiate their ansatz at run time and never read these fields. Every other basis builds them unconditionally.

initial_params!(int, ::OGA1dNormalEquations, sol)

Reference OGA initial guess for ShallowNet: the implementation from the original paper, kept as a selectable baseline.

The dictionary and the greedy least-squares fit are assembled in Float64 (a "double-precision island"), the output weights come from the normal equations Gk \ rhs, and the result is rounded into the working-precision cache. There is no norm floor, no coherence guard, no ridge and no rank detection. See the "Orthogonal Greedy Algorithm" section of the documentation for why the working-precision QR fit of OGA1d replaced it as the default, and OGA1dNormalEquations for the failure mode this variant exhibits at 16 bits.

Defined only for ShallowNet — it is a comparison baseline, not a production seed.

jacobi_svd(Â; sweeps = 30) -> (σ, U, V)

Singular values and factors of Â (nodes × atoms, tall or square) by the one-sided Jacobi method: orthogonalise the columns pairwise by plane rotations until they are mutually orthogonal, at which point the column norms are the singular values, the normalised columns are U, and the accumulated rotations are V.

One-sided Jacobi is chosen over bidiagonalisation because it is short, generic in eltype(Â) (LinearAlgebra's svd is LAPACK-only), and has high relative accuracy on the small singular values — exactly the ones that decide whether the greedily selected atoms are still independent at reduced precision. At k ≤ 8 columns a handful of sweeps converges and the cost is irrelevant.

σ is returned in the column order of V, not sorted.

mae_loss(x, y, NN, ps; λ = 0)

Mean absolute error of NN(x, ps) against target y, with optional boundary penalty λ * |NN(x[1], ps) - y[1]|². Used as the training objective for TrainingMethod.

Named mse_loss until the audit: the name said squared error, the docstring said absolute error, and the body computed absolute error. The body is authoritative — renaming leaves the numerics of every TrainingMethod seed exactly as they were, where switching to a squared error would have changed them silently. A dead μ = 0.00001 keyword was also dropped.

λ defaults to an untyped 0, not 0.0: a Float64 literal here promoted the whole loss to Float64 for a Float32 or Float16 network.

oga_atoms(dict, bias_interval, dict_amount, ::Type{T}) -> Matrix{T}

Build the candidate atom matrix: one row per atom, column 1 the hidden weight w and column 2 the bias b, in precision T.

oga_check_neuron_count(nneurons, symmetry)

Assert that nneurons is a multiple of neurons_per_atom(symmetry).

The loop places whole atoms, so under a mirrored symmetry an odd nneurons runs nneurons ÷ 2 steps and leaves the last neuron at (0, 0); _fill_unused! fills pairs too, so it cannot repair it. That is the duplicated-neuron state the fill exists to avoid, where a rank-deficient seed becomes a rank-deficient Newton Jacobian.

oga_check_precision(σ, ::Type{T})

Assert that the activation evaluates at the working precision, i.e. that σ(::T) is a T.

This is the one trap that silently invalidates a reduced-precision run: an activation written max(0.0, x)^k instead of max(zero(x), x)^k promotes every evaluation to Float64, so the seed is computed in double precision and the measurement says nothing about T. It costs one scalar call per fit to rule out, and the failure it catches is otherwise visible only as suspiciously good half-precision accuracy.

oga_fit(oga, σ, nodes, w, y, nneurons; bias_interval, dict_amount,
        modulation = nothing, symmetry = NoSymmetry()) -> OGAResult

Greedily fit nneurons neurons of a shallow (single-hidden-layer) network to the target y sampled at nodes, under the quadrature weights w.

• oga::OGA — the dictionary, selection rule, fit and guard rails.
• σ — the activation; must evaluate at eltype(nodes) (checked, see oga_check_precision).
• modulation — optional per-node factor multiplying every atom, for the boundary ansatz q(t) = (1-t) q̄ + t q + t(1-t) u(t), where the dictionary is t(1-t) σ(w t + b). Pass the t(1-t) vector; nothing means no modulation.
• symmetry::OGASymmetry — how atoms map to neurons.

nneurons must be a multiple of neurons_per_atom(symmetry) — even, for the two mirrored symmetries. An odd count is an ArgumentError rather than a silently short fit; see oga_check_neuron_count.

Runs entirely at T = eltype(nodes).

oga_norm_floor(::Type{T}, ref) -> T

Amplitude-scale floor below which a dictionary-atom norm is treated as numerically zero (so normalization is skipped rather than dividing by noise). Scales with sqrt(eps(T)) * ref, i.e. relative to the largest atom ref and to the working precision: a norm smaller than this cannot be inverted reliably in T.

Replaces the hard-coded absolute 1e-12 guard, which sat below eps(Float32) and so never fired in reduced precision.

oga_qr_append!(qr, a; min_gain) -> ρ

Append column a to the factorisation and return the deflated norm ρ = ‖a⊥‖, the part of a orthogonal to the columns already present. The column is rejected (nothing is appended, ρ still returned) when ρ ≤ min_gain · ‖a‖, i.e. when a lies in the existing span to within the requested tolerance — the rank drop that otherwise surfaces downstream as SingularException: zero pivot found at index 3.

Orthogonalisation is modified Gram–Schmidt with one reorthogonalisation pass. That second pass is not optional bookkeeping: plain MGS loses orthogonality in proportion to the condition number, which at Float16 is the whole problem, whereas reorthogonalising once restores it to O(eps(T)) for any conditioning that has not already collapsed.

oga_qr_reset!(qr)

Drop all columns, reusing the storage. Called once per dimension d of the fit.

oga_qr_solve(qr, ŷ) -> Vector

Least-squares coefficients from the maintained factorisation: one triangular solve R z = Qᵀ ŷ. Falls back to ridged_lstsq on the reconstructed design matrix if the triangular solve is not finite, so this can never return a non-finite seed.

oga_refine(dict, score, w, b) -> (w, b)

Polish a selected atom off the grid. score(w, b) returns the selection score of the candidate atom (larger is better; -Inf/NaN for an invalid one). The default is a no-op, so only Refined does any work.

oga_scores!(score, rule, Ψ, rownorms, nfloor, r̂, qr, proj)

Fill score with the selection score of every dictionary atom.

• Ψ — (natoms × nnodes) √w-scaled dictionary, one atom per row.
• rownorms — the √w-weighted L² norm of each atom.
• nfloor — norm floor below which an atom counts as numerically zero.
• r̂ — the √w-scaled current residual.
• qr — the factorisation of the selected columns (used by OrthogonalProjection).
• proj — (natoms × maxcols) scratch for Ψ Q.

Returns score. Unusable atoms are set to -one(T).

oga_seed(int, oga, symmetry, targets, modulation) -> Vector{OGAResult}

Run the greedy fit once per solution component, sharing the dictionary configuration and quadrature taken from the integrator's method.

targets[d] is the fit target for component d at the network's input nodes, and modulation is the optional per-node ansatz factor (see oga_fit).

oga_solve(fit, Â, ŷ, qr) -> Vector

Refit the output weights of the currently selected atoms. Â (nnodes × k) is the √w-scaled design matrix, ŷ the √w-scaled target, and qr the incrementally maintained factorisation of Â (used only by IncrementalQR, but kept in the signature so every fit is interchangeable).

Guarantees, for every fit, at every precision: the result has one entry per column of Â and every entry is finite. That is enforced here rather than in each fit, so it holds for a new fit by construction. Where a factorisation throws on a rank-deficient design, or returns Inf/NaN from a division by a pivot that survived truncation, the fit falls back to the ridged solve of ridged_lstsq. A seed the Newton solve can start from is worth more than a faithful report that the fit was impossible — and the rank-deficiency itself is already reported, through OGAResult's gains and rejected.

oga_tikhonov(G; C = 100) -> eltype(G)

Scale- and precision-relative Tikhonov floor C * eps(T) * tr(G)/n for stabilizing a Gram / normal-equations solve G \ rhs at working precision T = eltype(G). Relative to the mean diagonal tr(G)/n, so the effective condition-number cap tracks the precision; C is a modest safety factor.

Replaces the hard-coded absolute 1e-12 / 1e-14 ridges, which round away entirely below eps(Float32). Used by NormalEquationsFit as its ridge; the default OGA fit is weighted_lstsq, which avoids forming G at all.

pivoted_qr_lstsq(Â, ŷ, rtol) -> Vector

Rank-revealing least squares by Householder QR with column pivoting, truncated where the pivot norm drops below rtol times the first (largest) pivot.

The pivoting is what makes it rank-revealing: at each step the column with the largest remaining norm is brought forward, so a numerically dependent column is pushed to the end and detected by its collapsed pivot rather than being solved through. Columns past the detected rank get a zero coefficient — the returned vector always has one entry per input column, so callers need not know the rank.

Generic in eltype(Â); see the note above on why qr(Â, ColumnNorm()) cannot be used.

require_symbolic_derivatives(basis, method_name)

Throw an ArgumentError unless basis carries compiled symbolic derivatives.

Called from the constructors of the integrators whose components! evaluates them, so a basis built with symbolic = false is rejected where the mistake was made rather than several call levels down as a nothing being called on the first Newton iteration.

ridged_lstsq(Â, ŷ; C = 100) -> Vector

Tikhonov-ridged least squares for the √w-scaled design matrix Â, solved as the augmented QR problem [Â; √λ I] x ≈ [ŷ; 0] rather than on the normal equations, so the ridge stabilises the solve without squaring the condition number. The ridge λ = C · eps(T) · tr(ÂᵀÂ)/natoms is the precision-scaled floor of oga_tikhonov.

scaled_lstsq(Â, ŷ; C = 100) -> Vector

Plain QR least squares  \ ŷ with the precision-scaled ridged retry, for an already √w-scaled design matrix Â (nodes × atoms) and target ŷ.

This is the body of weighted_lstsq split out so that the greedy loop, which works entirely in the √w-scaled space, does not re-scale on every iteration. The ridged retry is shared by every OGAFit as a last resort, so no variant can return a non-finite seed.

truncated_svd_lstsq(Â, ŷ, rtol) -> Vector

Minimum-norm least squares via the truncated pseudo-inverse: singular directions with σ < rtol · σ_max are dropped instead of divided by, so a rank-deficient selected set yields a bounded solution rather than amplified rounding noise. The most robust of the fits, and the one to reach for when a variant must never fail at any precision.

Uses jacobi_svd, so it is generic in eltype(Â).

weight_grid(lo_octave, hi_octave, n, ::Type{T}) -> Vector{T}

Logarithmically spaced grid of n + 1 positive weight magnitudes 2^lo_octave … 2^hi_octave in precision T, for the weight axis of a 2-D (w, b) dictionary.

Spacing is logarithmic because w sets the length scale of a non-homogeneous activation (how sharply the unit transitions), and length scales are naturally compared by ratio rather than by difference. The exponents are generated by the same integer-indexed Float64 range as bias_grid, so a large n cannot collapse the step in reduced precision, and the exponentiation is done in Float64 before the single cast to T so that an octave outside the range of T saturates predictably instead of overflowing mid-computation.

weighted_lstsq(Φ, w, y; C = 100) -> Vector

Solve the (ridged) quadrature-weighted least-squares fit for the output weights x,

minₓ  Σⱼ wⱼ (Σᵢ xᵢ Φ[i,j] − yⱼ)²  +  λ‖x‖² ,

where each row of Φ is a dictionary atom sampled at the quadrature nodes, w are the (positive) quadrature weights, and y is the fit target at those nodes.

Solved by QR on the √w-scaled design matrix rather than the normal equations Φ diag(w) Φᵀ, so accuracy is governed by κ(Φ) instead of κ(Φ)² — the difference that lets the fit run in reduced precision without a rank-deficient Gram matrix.

The plain QR solve is used whenever it is finite (so the Float64/Float32 atom choice matches the Gram solution bit for bit and the greedy residual is unperturbed). Only if it returns a non-finite result — a genuinely rank-deficient design matrix, i.e. the Float16 case — is the fit retried with a √λ·I augmentation, where the ridge λ = C · eps(T) · tr(ÂᵀÂ)/natoms is the precision-scaled Tikhonov floor (see oga_tikhonov) that keeps the solution bounded and the seed finite.