Add the limited(actual, limiter) operator: symbolic PCNR iterate limiting for nonlinear solves#4774
Conversation
…ting Adds symbolic, component-level declaration of SPICE-style iterate limiting, lowered automatically to the Predictor/Corrector Newton-Raphson (PCNR) formulation of Aadithya, Keiter & Mei. Mirrors the architecture of the Modelica homotopy(actual, simplified) operator: limited(actual, limiter) stays opaque through System construction; mtkcompile lowers it for time-independent systems (auxiliary irreducible unknown per node, consistency equation, symbolic default/guess, limiter registry in LimitedCtx metadata) and strips it to actual for time-dependent systems so component libraries compile unchanged for transient simulation. NonlinearFunction construction compiles the limiter registry into the SciMLBase postcondition corrector hook, which NonlinearSolve.jl's native solvers apply at every accepted iterate. The limiter is written in terms of the reserved limitnew/limitold placeholders (fixed-name parameter sentinels, precompile-safe like HOMOTOPY_LAMBDA; namespaced occurrences from subsystems are recognized by name suffix and canonicalized during lowering). Because the auxiliary unknowns are irreducible, alias elimination keeps the limited quantity as the surviving representative of its alias class, i.e. the PCNR augmentation is reduced symbolically at no runtime cost. Measured on the classic Vs-R-diode DC circuit with pnjlim: plain NewtonRaphson 179 steps, with limited 11 steps, both as a flat system and through hierarchical namespaced components. Requires SciMLBase >= 3.37 (SciML/SciMLBase.jl#1449) for the postcondition field and SciML/NonlinearSolve.jl#1084 for solvers to apply it. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
|
Verification detail (local, Julia 1.12.4). The full What was run instead, in a coherent dev environment (SciMLBase#1449 + NonlinearSolve#1084 + this branch):
Also verified the operator's derivative rules ( 🤖 Generated with Claude Code |
|
CI triage (first full run): Runic Format Check and Spell Check pass. All 83 failing jobs (tests/sublibrary-ci/Downgrade/Documentation/Benchmarks/Catalyst downstream) fail at dependency resolution with Release order for green CI here: SciMLBase#1449 → tag 3.37.0 → re-run (tests pass with limiting inert) → NonlinearSolve#1084 + release → hook-gated solver-behavior assertions activate. 🤖 Generated with Claude Code |
| Return `true` iff `sys` contains a `limited(...)` node anywhere in its equations or its | ||
| observed equations. | ||
| """ | ||
| function has_any_limited(sys) |
There was a problem hiding this comment.
This should use https://github.com/JuliaSymbolics/SymbolicUtils.jl/blob/master/src/irstructure.jl#L659 to have any hope at scaling to Multibody models. Naive recursion on an expression is strictly off the table in such cases.
| Recursively replace every `limited(actual, limiter)` node in the unwrapped expression `x` | ||
| with `actual`, discarding the limiter (and with it any `limitnew`/`limitold` sentinels). | ||
| """ | ||
| function _strip_limited(x) |
There was a problem hiding this comment.
In a similar fashion, this should use IRStructure, though we don't have an inbuilt utility for this. An alternative formulation is to propagate a cache::Base.IdDict{SymbolicT, SymbolicT} throughout the recursion, populate it with the result of maketerm and check it at the beginning to avoid duplicating work.
| return filter(v -> _sentinel_kind(v) === :none, vars) | ||
| end | ||
|
|
||
| function _canonicalize_sentinels(x) |
There was a problem hiding this comment.
Similar cache required.
|
|
||
| # Replace every node in `x` that appears as a key of `repl` (an ordered | ||
| # node => variable map keyed by `isequal`) with its replacement. | ||
| function _replace_limited_nodes(x, repl) |
There was a problem hiding this comment.
Isn't this just substitute? It should use IRSubstituter.
| return maketerm(typeof(x), operation(x), new_args, metadata(x)) | ||
| end | ||
|
|
||
| function _collect_limited_nodes!(nodes, x) |
There was a problem hiding this comment.
This is just search_variables!. It should use the IRStructureSearchBuffer version.
| for v in Iterators.flatten((unknowns(sys), parameters(sys; initial_parameters = true))) | ||
| if hasname(v) && getname(v) in reserved | ||
| throw( | ||
| ArgumentError( |
There was a problem hiding this comment.
Can't the lowering use Unicode or a name with # in it?
| # parameters only; referencing other unknowns is not supported since the hook | ||
| # sees the limited entries, not the whole state. | ||
| bad = [ | ||
| v for v in unwrap.(Symbolics.get_variables(lexpr)) |
There was a problem hiding this comment.
Using IRStructureSearchBuffer and get_variables! here will be faster.
| lims = Tuple(lims) | ||
| if iip | ||
| return let lims = lims | ||
| function limited_postcondition_iip(up, uprev, p) |
There was a problem hiding this comment.
Can these closures be made callable structs so we don't accidentally forget to add something to the let block later?
| ) | ||
| ) | ||
| end | ||
| fn = generate_custom_function( |
There was a problem hiding this comment.
Do we want to FunctionWrapper these somehow to cut down on compile times?
Addresses @AayushSabharwal's review. All expression walking now goes through SymbolicUtils' IRStructure machinery instead of naive recursion, so the passes visit shared subexpressions once rather than once per occurrence (the Multibody scaling concern): - has_limited uses SU.query (DAG-memoized) with an IsLimitedNode callable predicate. - Node collection uses search_variables! with IRStructureSearchBuffer and is_atomic = IsLimitedNode(), which collects the limited nodes themselves; insertion order stays deterministic so the generated names remain a pure function of the system. - Node replacement, the strip pass, and sentinel canonicalization are now substitutions applied through SU.IRSubstituter rather than hand-rolled maketerm rebuilds. - The limiter free-variable check uses IRStructureSearchBuffer + search_variables!. - Auxiliary unknowns are named , which cannot be written as a Julia identifier, so the name-collision guard is gone entirely. - The postcondition closures became a LimitedPostcondition{iip} callable struct so captured state is explicit. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
|
Thanks @AayushSabharwal — all addressed in 204a5ce. Every traversal now goes through the IR machinery rather than naive recursion:
On the last one — FunctionWrapper for the compiled limiters: I left them unwrapped for now, because I don't think there's a signature that survives the use sites. The limiter is called as Re-verified after the rewrite: 🤖 Generated with Claude Code |
3 participants
Important
This PR should be ignored until reviewed by @ChrisRackauckas.
Note
Dependency chain: requires SciML/SciMLBase.jl#1449 (adds
NonlinearFunction.precondition/postcondition, v3.37.0) and SciML/NonlinearSolve.jl#1084 (solvers apply the hooks). CI here cannot resolve until SciMLBase 3.37 is released; solver-behavior test assertions are additionally gated onisdefined(NonlinearSolveBase, :apply_postcondition!!)so they activate when NonlinearSolve releases. Everything below was verified locally against the companion branches.Summary
Adds the
limited(actual, limiter)operator: symbolic, component-level declaration of SPICE-style iterate limiting, lowered automatically to the Predictor/Corrector Newton-Raphson (PCNR) formulation of Aadithya, Keiter & Mei. A device author annotates a sensitive quantity once —— and every nonlinear solve built from any model containing that component gets predictor/corrector limiting fully automatically:
mtkcompileperforms the PCNR augmentation andNonlinearProblemconstruction compiles the limiters into the newNonlinearFunction.postconditioncorrector hook.The operator follows the architecture of the Modelica
homotopy(actual, simplified)operator (#4601): the actual expression first, the helper second;limitnew/limitoldare reserved placeholder symbols (fixed sentinels, precompile-safe likeHOMOTOPY_LAMBDA) for the proposed and previously-accepted values inside the limiter expression.What the lowering does
For time-independent systems (in both the ModelingToolkitBase light pipeline and the ModelingToolkit structural pipeline, inserted post-
expand_connections):limited(actual, limiter)node gets an auxiliary irreducible unknownlimited_k; the node is replaced by it and the consistency equationlimited_k ~ actualis appended (EquationSourceInformationis padded accordingly forTearingState);limited_kis irreducible, alias elimination keeps the limited quantity as the surviving representative of its alias class — the PCNR augmentation is reduced symbolically, at no runtime cost;limited_kreceivesactualas both default and guess, so no user-provided initial value is needed;LimitedCtxsystem metadata;NonlinearFunctionconstruction compiles each limiter viagenerate_custom_function(parameters, including bound parameters, resolve through standard codegen) into an iip/ooppostconditionhook.For time-dependent systems the operator is stripped to
actualduringmtkcompile, so component libraries carrying limiters compile unchanged for transient simulation (limiting inside implicit steppers' nonlinear solves is future work). Un-loweredlimitednodes are rejected with clear errors atNonlinearFunction/ODEFunctionconstruction; nestedlimitedand limiters referencing other unknowns are rejected likewise.Changes
1.54.0 → 1.55.0): newsrc/systems/limited_operator.jl(operator, sentinels, detection/strip/lowering passes,generate_limited_postcondition); lowering invoked from__mtkcompile;postconditionwired intoNonlinearFunctionconstruction with guards;ODEFunctionguard; exportslimited,limitnew,limitold; SciMLBase compat3.37; tests intest/limited_operator.jl(registered in the InterfaceII group).11.34.0 → 11.35.0): lowering invoked from the structural__mtkcompile(with source-info padding); SciMLBase compat3.37; MTKBase compat1.55.docs/src/basics/Limiting.md(registered inpages.jl) with a component-based diode-circuit walkthrough, lowering semantics, and contracts.Verification (all local, Julia 1.12.4, dev'd SciMLBase#1449 + NonlinearSolve#1084)
limited_operator.jltestset: 28/28 in the MTKBase light pipeline.homotopy_lowering.jl: 43/43 (no regression from the shared pipeline edits).limited11 steps (identical to the hand-written PCNR/NonlinearSolve tutorial results), rootv = 0.6698509496766559exact to 1e-6;vcrit): 11 steps, correctv/irecovered through observed;System(eqs, t): operator stripped,ODEProblemcompiles, simulates, and the generated RHS matches the analytic value att=0.GROUP=InterfaceIIsuite (contains the new tests plus problem-construction/homotopy tests) run against the dev stack — result reported in a PR comment.Design notes
System(eqs, t)constructor's function-of-tvalidation accepts equations containing them; both lowering paths drop them from the discovered unknowns/parameters (namespaced occurrences —diode₊__limitnew_ₘₜₖ— are recognized by name suffix and canonicalized back to the toplevel sentinels).var_to_name, and crucially theirreduciblesfield, which structural simplification consults rather than per-variable metadata — found the hard way).∂₁ = 1,∂₂ = 0: Jacobians treat the operator asactual, matching the hook semantics (the limiter is a corrector, not part of the residual).Closes the ModelingToolkit third of the SciML/NonlinearSolve.jl#351 + PCNR feature arc.
🤖 Generated with Claude Code
https://claude.ai/code/session_01B9GHoo4pFa3DU9yhyTkevL