Skip to content

NonlinearSolveAlg: forward precondition/postcondition into the implicit stage solve - #4158

Merged
ChrisRackauckas merged 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:nlsolvealg-nonlinear-preconditioning
Aug 8, 2026
Merged

ChrisRackauckas merged 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:nlsolvealg-nonlinear-preconditioning

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member

Please ignore this PR until it has been reviewed by @ChrisRackauckas. It is opened as
a draft and should not be merged before then.

Forwards NonlinearSolve.jl's nonlinear-preconditioning solve options — precondition (a
left preconditioner G on the residual) and postcondition (an iterate corrector H, the
corrector phase of PCNR that replaces SPICE-style limiting) — into the implicit stage solves
of NonlinearSolveAlg, so an implicit ODE method can carry a limiter or a residual scaling
down into its Newton iteration:

H(u, uprev, p, cache) = max(u, zero(u))          # keep a concentration positive
alg = ImplicitEuler(nlsolve = NonlinearSolveAlg(NewtonRaphson(); postcondition = H))

Both are released solve keywords; nothing here is gated on availability:

package version required by this PR previously
NonlinearSolveBase 2.41 2.39
NonlinearSolve 4.25 4.24
SciMLBase 3.43 3.39

SciMLBase = "3.43" is load-bearing rather than cosmetic: the inner stage init receives
these as ordinary solve keywords, so SciMLBase.allowedkeywords has to contain
:precondition/:postcondition or the unrecognized-keyword handling trips. 3.43 does.

The part that actually needs review: which variable does the corrector see?

An implicit stage does not solve for the ODE state. It solves for the increment z, and
the state at the stage is recovered as tmp + γ⋅z (DIRK) or as z itself (when
nlsolver.method is COEFFICIENT_MULTISTEP, i.e. inside BDF). A user writing
H(u, uprev, p, cache) = max(u, 0) for a physical variable means the state. Wiring the
option straight through to the stage NonlinearProblem would hand that rule an increment to
clip, which is silently and badly wrong — it is a footgun, not a rounding detail.

Measured on u' = -u, u0 = 1.0, ImplicitEuler, dt = 0.1, one step:

integrator.u (ODE state u_{n+1}) = 0.9090909090909091
inner NonlinearProblem iterate   = -0.09090909090909091     <- this is z
nlsolver.tmp = 1.0   nlsolver.γ = 1.0
tmp + γ*inner                    = 0.9090909090909091

A corrector passed straight through to that same stage problem, as NonlinearSolve's own
solve keyword, is handed [1.0, 1.0, -0.0909…, -0.0909…]: it converges on -0.0909…, the
increment, while the state there is 0.909…. A max(u, 0) written for a concentration would
be clipping the wrong quantity every iteration.

So this PR does not forward the options verbatim. Both are conjugated with the affine
stage map and are stated in ODE-state coordinates:

  • postcondition(u_stage, u_stage_prev, p, cache) — u_stage is uₙ₊₁ for ImplicitEuler,
    the stage value at t + cΔt for a DIRK stage, the state itself under
    COEFFICIENT_MULTISTEP. p is the ODE's parameter object, not the stage problem's
    internal parameter tuple.
  • precondition(fu, u_stage, p) — only the iterate and the parameters are remapped. fu
    stays the stage residual (Δt⋅f(u_stage, p, t + cΔt) - z)/(γΔt); it has no state-space
    reading to map onto.

Verified for a multi-stage method too. On u' = 1, u0 = 0, TRBDF2, dt = 0.5 (so the
exact state at time t is t), the values H receives are
[0.29289…, 0.5, 0.79289…, 1.0] — exactly t + γΔt and t + Δt for each of the two steps,
i.e. each stage's own time point. And under FBDF (method === COEFFICIENT_MULTISTEP
throughout) the values stay inside [exp(-1), 1], the state's range, never z's.

What a reviewer should decide is whether this remapping is the right call. The
alternative — forwarding verbatim and documenting "your corrector sees z" — is simpler and
composes more predictably with anything else that inspects the inner problem, but it makes
the natural way to write a limiter wrong by default. This PR takes the position that a
corrector is a statement about the model, so it should be stated in the model's variable;
the conjugation is confined to StageConditioner in utils.jl, and reverting to
pass-through would be a small deletion if the other call is preferred.

Consequences of the choice, all documented in the NonlinearSolveAlg docstring:

  • The inverse map is applied to the change the corrector made, z + (H(u) - u)/γ, not by
    recomputing (H(u) - tmp)/γ. Only the former returns z bit-for-bit when the corrector
    leaves the state alone, and that matters: a limiter goes inactive exactly when the
    iteration is converging, and a one-ulp perturbation there is enough to stop the stage
    residual from being exactly zero, which stops the (deliberately zero-tolerance) inner cache
    from terminating, which nlsolve! reads as a rejected trial step rather than convergence.
    There is a test that an identity corrector reproduces the untouched solve exactly.
  • H is also applied once to the stage predictor, the stage analogue of the H(u0, u0, p, nothing) correction init performs. reinit! does not redo it, so without this the first
    residual of every stage would be evaluated at an uncorrected iterate.
  • u_stage_prev on a stage's first in-loop correction is that stage's corrected predictor.
    reinit! refreshes the inner cache's iterate but not its previous-iterate buffer, so the
    raw value there is the previous stage's last z read through this stage's tmp — not a
    state at all.
  • A corrector that is still clamping when the iteration ends makes the displacement exceed
    the raw Newton step, so the stage reads as non-converged and the integrator rejects the
    step and retries with a smaller Δt. That is the intended response, and there is a test
    that a limiter tight enough to clip repeatedly still converges to the unlimited solution.
  • precondition changes the residual the inner solver differentiates, so W is no longer
    its Jacobian: W reuse (and with it the W-based smoothed error estimate) is disabled for
    the stage solve when a precondition is supplied. A postcondition does not change the
    residual, so W reuse stays on.

Second finding, also worth a look: in-place precondition cannot use ForwardDiff

Dropping W reuse has a consequence I did not expect. #4137's own test states the premise:

W reuse stays enabled for no-init caches in place: […] the residual writes through
preallocated Float64 buffers that AD cannot go through.

That is exactly right — odenlf writes the stage state and f's output into the Float64
buffers carried in the parameter tuple, and an AutoSpecialize f is a FunctionWrapper
accepting only Float64 and OrdinaryDiffEq's own one-chunk duals. Nothing normally notices,
because the reused W is handed to the inner solver as its Jacobian so the residual is never
differentiated. A precondition removes that reuse, and the inner solver then differentiates
the residual itself. Measured, u' = -u in place, ImplicitEuler, dt = 0.1:

inner algorithm with precondition without
NewtonRaphson() (autodiff unset → ForwardDiff) throws MethodError ok
NewtonRaphson(autodiff = AutoForwardDiff()) throws MethodError ok
NewtonRaphson(autodiff = AutoFiniteDiff()) ok ok
SimpleNewtonRaphson() throws MethodError ok
SimpleNewtonRaphson(autodiff = AutoFiniteDiff()) ok ok
TrustRegion() throws MethodError ok

Out-of-place is unrestricted (its residual allocates; NewtonRaphson(), AutoForwardDiff
and SimpleNewtonRaphson() all work), and postcondition is unaffected in both forms.

NonlinearSolveAlg's default alg is already NewtonRaphson(autodiff = AutoFiniteDiff()),
so the default path works — but NonlinearSolveAlg(NewtonRaphson(); precondition = G!), the
most natural thing to write, is a 200-line MethodError from inside ForwardDiff. This PR
therefore rejects an in-place precondition combined with a dual-based (or unset) inner AD
backend, up front, with a message naming the fix. This is a limitation of the surrounding
in-place stage-residual design rather than of the options themselves; making it work would
mean giving _compute_rhs! dual-capable buffers and getting past the AutoSpecialize
function wrapper, which is well outside this PR.

What accepts the options

algorithm behaviour
NonlinearSolveAlg both options honoured, in stage-state coordinates
NLNewton, NLFunctional, NLAnderson accept the keywords only to throw
HomotopyNonlinearSolveAlg accepts the keywords only to throw

The hand-rolled loops in NLNewton/NLFunctional/NLAnderson do not go through
NonlinearSolve's iterate-commit points, so there is no place a corrector could be applied;
HomotopyNonlinearSolveAlg additionally solves the λ-embedding rather than the stage
residual. Rather than ignore the option, their constructors take it and raise an
ArgumentError naming NonlinearSolveAlg as the alternative.

This is deliberately stricter than NonlinearSolve's own convention, which reports an
unsupported corrector through the SciMLLogging unsupported_postcondition toggle (raising at
ErrorLevel by default, but turn-downable). The toggle makes sense there because
postcondition is a solve keyword that a user may hold fixed while sweeping algorithms;
here it is a constructor keyword of one specific nonlinear solver, so you have to have named
NLNewton explicitly to pass it, and there is no verbosity object in scope at construction
time anyway. Happy to switch to the toggle if consistency is preferred.

Implementing a corrector inside the hand-rolled NLNewton/NLFunctional/NLAnderson
loops is the other option, and it is not attempted here — those loops would need commit
points and an equivalent of the predictor correction, which is a separate change with its own
convergence-test consequences. Rejecting is the deliberate choice for now, not an oversight.

Also rejected, at build_nlsolver time:

  • DAEProblems — the residual is the implicit DAE residual and its unknown is the state
    increment against uprev, so the state-space contract does not apply.
  • ModelingToolkit nlstep_data stage systems — the stage system is code-generated and
    OrdinaryDiffEq has no map from its unknowns back to the ODE state to state the option in.

An inner NonlinearSolve.jl algorithm that does not itself apply iterate corrections (a
SimpleNonlinearSolve one, say) is caught by NonlinearSolveBase's own
unsupported_postcondition report when the stage cache is built; there is a test for that
too.

Compatible with SciML/NonlinearSolve.jl#1148 (PostconditionSpecifier(H; space = ...)) if it
lands: the specifier is callable with the same 4-argument signature, so it passes through
StageConditioner unchanged, and postcondition_space on the adapter falls back to
:original, which is right because the stage NonlinearProblem carries no lb/ub.

Tests

New file lib/OrdinaryDiffEqNonlinearSolve/test/nsa_conditioning_tests.jl, wired into
runtests.jl. It covers the coordinate contract in and out of place, the multistep form, the
predictor correction, the previous-iterate identity, an identity corrector being bit-for-bit
inert, a limiting corrector converging to the unlimited solution, precondition composing
into the residual, W reuse being disabled by precondition but not by postcondition,
survival across resize!, and every rejection path.

Run against the rebased branch with the registry versions above (NonlinearSolveBase 2.41.0,
NonlinearSolve 4.25.0, SciMLBase 3.43.0):

developer_api_tests         |   21 passed | 18.5s
newton_tests                |    6 passed | 11.0s
nsa_jacobian_reuse_tests    |   21 passed | 1m34.1s
nsa_sparse_tests            |   10 passed | 2m01.2s
nsa_matrixfree_tests        |   14 passed | 48.8s
nsa_smooth_est_tests        |   10 passed | 27.2s
nsa_stats_tests             |   24 passed | 39.6s
nsa_polyalg_tests           |    8 passed | 41.8s
nsa_noinit_tests            |  102 passed | 4m05.4s
nested_ad_nlsolvealg_tests  |    5 passed | 1m01.9s
nsa_conditioning_tests      |   70 passed | 2m20.8s
                            |  291 total  |

Runic reports no changes on all touched files, and
Test.detect_ambiguities/detect_unbound_args on OrdinaryDiffEqNonlinearSolve both report
0.

Not verified locally, left to CI: the Downgrade job (whether the new 2.41/4.25/3.43
lower bounds resolve against each other and against the sublibraries), and the QA job.


🤖 Generated with Claude Code

https://claude.ai/code/session_01B9GHoo4pFa3DU9yhyTkevL

`NonlinearSolveAlg` gains the `precondition` (left preconditioner `G` on the
residual) and `postcondition` (iterate corrector `H`, the PCNR corrector phase)
solve options of NonlinearSolve.jl, so an implicit ODE method can carry a
limiter or a residual scaling into its stage solves.

The stage unknown is the increment `z`, not the ODE state: the state at a stage
is `tmp + γ⋅z` for a DIRK stage and `z` itself when `nlsolver.method` is
`COEFFICIENT_MULTISTEP`. Forwarding a user's corrector verbatim would hand a
limiter written for a physical variable an increment to clip, so both options
are conjugated with that affine map and are stated in ODE-state coordinates:
`H(u_stage, u_stage_prev, p, cache)` and `G(fu, u_stage, p)`, with `p` the ODE's
parameters rather than the stage problem's internal tuple.

`H` is applied to the stage predictor as well as at every inner-solver commit
point, since `reinit!` does not redo the initial-guess correction that `init`
performs. `W` reuse is disabled under a `precondition` because `W` is the
Jacobian of the raw stage residual, not of the composed one — which in turn
leaves the inner solver differentiating the in-place stage residual, something
its preallocated `Float64` buffers and `AutoSpecialize` function wrapper cannot
support, so an in-place `precondition` combined with a dual-based inner AD
backend is rejected up front rather than left to fail inside ForwardDiff.

`NLNewton`, `NLFunctional`, `NLAnderson` and `HomotopyNonlinearSolveAlg` run
their own iteration with no commit point a corrector could be applied at; they
accept the keywords only to reject them, so supplying one is an error rather
than a silent no-op. The options are likewise rejected for `DAEProblem`s and
for ModelingToolkit `nlstep_data` stage systems, whose unknowns have no
available map back to the ODE state.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants