Add nlstep_data to DAEFunction (SciML#1509)

`ODENLStepData` lets a symbolic front end hand an implicit stage solve a
pre-built `NonlinearProblem` instead of a stage-equation closure.
`ODEFunction` and `SplitFunction` carry it as `nlstep_data`; `DAEFunction`
did not, so the fully implicit `0 = F(du, u, p, t)` form — which is what
modified nodal analysis produces for circuit models — had no way to reach
that path.

`ODENLStepData` is reused unchanged rather than adding a DAE-specific type.
For a fully implicit DAE both arguments of `F` are affine in the stage
unknown, so the same six hooks parametrize the stage system:

    g(z, p') = F(gamma1 * z + outer_tmp, gamma2 * z + inner_tmp, p, t_c)

`gamma2`/`inner_tmp` build the state argument exactly as in the mass-matrix
form and `gamma1`/`outer_tmp` build the derivative argument, so with
`gamma2 = 1`, `inner_tmp = 0` and a BDF-type `du ≈ (u - tmp) / (γ * dt)`,
`gamma1` is the `gamma` of `jac(J, du, u, p, gamma, t)`. The `ODENLStepData`
docstring now states this form alongside the mass-matrix one.

`NLP <: Union{Nothing, ODENLStepData}` is bounded as it is on `ODEFunction`,
so `widen_bounded_type_params` erases it for `AutoSpecialize` and `remake`
does not narrow it back.

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>

Author: ChrisRackauckas-Claude

src/odenlstep.jl

@@ -1,34 +1,50 @@
 """
     ODENLStepData(nlprob, u0perm, set_gamma_c, set_outer_tmp, set_inner_tmp, nlprobmap)
 
-A collection of hooks for custom nonlinear stage solves in implicit ODE
+A collection of hooks for custom nonlinear stage solves in implicit ODE and DAE
 algorithms.
 
-`ODENLStepData` lets an `ODEFunction` provide a structured
+`ODENLStepData` lets an `ODEFunction`, `SplitFunction` or `DAEFunction` provide a structured
 `AbstractNonlinearProblem` template for solver packages that form implicit stage
 equations. Before each nonlinear solve, the algorithm updates the stage guess,
 scaling factors, time information, and temporary vectors through the stored
 setter callables. After the nonlinear solve, `nlprobmap` converts the nonlinear
-unknown back to the state vector used by the original ODE problem.
+unknown back to the state vector used by the original problem.
 
-The nonlinear problem should represent a stage equation of the form
-`M * z = outer_tmp + gamma1 * f(gamma2 * z + inner_tmp, p, t_c)`, equivalently
+## Mass-matrix form
+
+For `M * du/dt = f(u, p, t)` the nonlinear problem should represent a stage equation of the
+form `M * z = outer_tmp + gamma1 * f(gamma2 * z + inner_tmp, p, t_c)`, equivalently
 `g(z, p') = gamma1 * f(gamma2 * z + inner_tmp, p, t_c) + outer_tmp - M * z`.
 Here `z` is the nonlinear stage unknown, `p` is the ODE parameter object, `t_c`
 is the stage evaluation time, and `gamma1`, `gamma2`, `outer_tmp`, and
 `inner_tmp` are supplied by the ODE algorithm.
 
+## Fully implicit form
+
+For `0 = F(du, u, p, t)` (a `DAEFunction`) the stage equation has the same shape, with both
+arguments of `F` affine in the stage unknown:
+`g(z, p') = F(gamma1 * z + outer_tmp, gamma2 * z + inner_tmp, p, t_c)`.
+`gamma2` and `inner_tmp` build the state argument from the stage unknown exactly as in the
+mass-matrix form, while `gamma1` and `outer_tmp` build the derivative argument. Taking the
+stage unknown to be the stage state (`gamma2 = 1`, `inner_tmp = 0`), a BDF-type step with
+`du ≈ (u - tmp) / (γ * dt)` gives `gamma1 = inv(γ * dt)` and `outer_tmp = -tmp / (γ * dt)`.
+With that convention `gamma1` is the `gamma` of the `DAEFunction` Jacobian signature
+`jac(J, du, u, p, gamma, t)`: the Jacobian of the stage residual with respect to `z` is
+`gamma1 * dF/d(du) + dF/du`.
+
 # Fields
 
 $(TYPEDFIELDS)
 
 # Extension Rules
 
-Symbolic-system packages construct this value and store it as an `ODEFunction`'s
-nonlinear-stage metadata. Solver packages may consume the six fields through their
-callable contracts, but must not assume concrete callable types or mutate the container.
-Each setter must update the object it closes over consistently with `nlprob`, and
-`nlprobmap` must map a completed nonlinear solution back to the ODE stage representation.
+Symbolic-system packages construct this value and store it as the `nlstep_data` of an
+`ODEFunction`, `SplitFunction` or `DAEFunction`. Solver packages may consume the six fields
+through their callable contracts, but must not assume concrete callable types or mutate the
+container. Each setter must update the object it closes over consistently with `nlprob`, and
+`nlprobmap` must map a completed nonlinear solution back to the stage representation of the
+original problem.
 """
 struct ODENLStepData{NLProb, SetU0, SetGammaC, SetOuterTmp, SetInnerTmp, NLProbMap}
     """

src/scimlfunctions.jl

@@ -1586,7 +1586,8 @@ DAEFunction{iip,specialize}(f;
                            jac_prototype = __has_jac_prototype(f) ? f.jac_prototype : nothing,
                            sparsity = __has_sparsity(f) ? f.sparsity : jac_prototype,
                            colorvec = __has_colorvec(f) ? f.colorvec : nothing,
-                           sys = __has_sys(f) ? f.sys : nothing)
+                           sys = __has_sys(f) ? f.sys : nothing,
+                           nlstep_data = __has_nlstep_data(f) ? f.nlstep_data : nothing)
 ```
 
 Note that only the function `f` itself is required. This function should
@@ -1622,6 +1623,11 @@ the usage of `f`. These include:
   based on the sparsity pattern. Defaults to `nothing`, which means a color vector will be
   internally computed on demand when required. The cost of this operation is highly dependent
   on the sparsity pattern.
+- `nlstep_data`: an [`ODENLStepData`](@ref SciMLBase.ODENLStepData) holding a structured
+  nonlinear problem for the implicit stage solve, or `nothing`. Implicit DAE integrators
+  which support it solve this problem in place of building a stage-equation closure. See the
+  `ODENLStepData` documentation for the stage equation the nonlinear problem must represent
+  in the fully implicit case.
 
 ## iip: In-Place vs Out-Of-Place
 
@@ -1685,7 +1691,7 @@ numerically-defined functions.
 struct DAEFunction{
         iip, specialize, F, Ta, Tt, TJ, TJU, TJD, JVP, VJP, JP, SP, TW, TWt, TPJ, O,
         TCV,
-        SYS, ID,
+        SYS, ID, NLP <: Union{Nothing, ODENLStepData},
     } <:
     AbstractDAEFunction{iip}
     f::F
@@ -1705,6 +1711,7 @@ struct DAEFunction{
     colorvec::TCV
     sys::SYS
     initialization_data::ID
+    nlstep_data::NLP
 end
 
 """
@@ -4143,7 +4150,8 @@ function DAEFunction{iip, specialize}(
         initializeprobmap = __has_initializeprobmap(f) ? f.initializeprobmap : nothing,
         initializeprobpmap = __has_initializeprobpmap(f) ? f.initializeprobpmap : nothing,
         initialization_data = __has_initialization_data(f) ? f.initialization_data :
-            nothing
+            nothing,
+        nlstep_data = __has_nlstep_data(f) ? f.nlstep_data : nothing
     ) where {
         iip,
         specialize,
@@ -4188,12 +4196,12 @@ function DAEFunction{iip, specialize}(
             iip, specialize, Any, Any, Any,
             Any, Any, Any, Any, Any, Any, Any,
             Any, Any, Any,
-            Any, typeof(_colorvec), Any, Any,
+            Any, typeof(_colorvec), Any, Any, Union{Nothing, ODENLStepData},
         }(
             _f, analytic, tgrad, jac, jac_u, jac_du, jvp,
             vjp, jac_prototype, sparsity,
             Wfact, Wfact_t, paramjac, observed,
-            _colorvec, sys, initdata
+            _colorvec, sys, initdata, nlstep_data
         )
     else
         DAEFunction{
@@ -4203,12 +4211,12 @@ function DAEFunction{iip, specialize}(
             typeof(sparsity), typeof(Wfact), typeof(Wfact_t),
             typeof(paramjac),
             typeof(observed), typeof(_colorvec),
-            typeof(sys), typeof(initdata),
+            typeof(sys), typeof(initdata), typeof(nlstep_data),
         }(
             _f, analytic, tgrad, jac, jac_u, jac_du, jvp, vjp,
             jac_prototype, sparsity, Wfact, Wfact_t,
             paramjac, observed,
-            _colorvec, sys, initdata
+            _colorvec, sys, initdata, nlstep_data
         )
     end
 end

test/function_building_error_messages.jl

@@ -387,6 +387,60 @@ DAEFunction(dfiip, vjp = dvjp)
 DAEFunction(dfoop, vjp = dvjp)
 DAEFunction{true, SciMLBase.NoSpecialize}(dfiip, observed = 1)
 
+@testset "DAEFunction nlstep_data" begin
+    nlstep = SciMLBase.ODENLStepData(
+        NonlinearProblem((z, p) -> z, [1.0]), identity,
+        (gamma1, gamma2, c) -> nothing, identity, identity, identity
+    )
+
+    @testset "construction" begin
+        @test DAEFunction(dfiip).nlstep_data === nothing
+        @test typeof(DAEFunction(dfiip)).parameters[end] === Nothing
+        @test DAEFunction(dfiip; nlstep_data = nlstep).nlstep_data === nlstep
+        @test DAEFunction(dfoop; nlstep_data = nlstep).nlstep_data === nlstep
+        # a `DAEFunction` passed as the function carries its `nlstep_data` over
+        wrapped = DAEFunction{true, SciMLBase.FullSpecialize}(
+            DAEFunction(dfiip; nlstep_data = nlstep)
+        )
+        @test wrapped.nlstep_data === nlstep
+    end
+
+    @testset "specialization $spec" for spec in (
+            SciMLBase.FullSpecialize, SciMLBase.NoSpecialize, SciMLBase.AutoSpecialize,
+        )
+        f = DAEFunction{true, spec}(dfiip; nlstep_data = nlstep)
+        @test f.nlstep_data === nlstep
+
+        # `remake` reconstructs the struct through the keyword constructor, which builds
+        # it positionally with an explicit type-parameter list. `unwrapped_f` has no
+        # `DAEFunction` method and so is the identity here; assert it anyway so a future
+        # method has to keep the field.
+        @test SciMLBase.unwrapped_f(f).nlstep_data === nlstep
+        @test SciMLBase.remake(f).nlstep_data === nlstep
+        @test SciMLBase.remake(f; jac_prototype = zeros(1, 1)).nlstep_data === nlstep
+        # explicit overrides still win, matching `ODEFunction`
+        @test SciMLBase.remake(f; nlstep_data = nothing).nlstep_data === nothing
+
+        # AutoSpecialize erases the bounded parameters to keep the function type
+        # model-independent; `remake` must not narrow them back.
+        widened = SciMLBase.widen_bounded_type_params(f)
+        @test typeof(widened).parameters[end] === Union{Nothing, SciMLBase.ODENLStepData}
+        @test widened.nlstep_data === nlstep
+        rewidened = SciMLBase.remake(widened; jac_prototype = zeros(1, 1))
+        @test typeof(rewidened).parameters[end] ===
+            Union{Nothing, SciMLBase.ODENLStepData}
+        @test rewidened.nlstep_data === nlstep
+    end
+
+    @testset "problem construction" begin
+        prob = DAEProblem(
+            DAEFunction(dfiip; nlstep_data = nlstep), [0.0], [1.0], (0.0, 1.0)
+        )
+        @test prob.f.nlstep_data === nlstep
+        @test SciMLBase.remake(prob; u0 = [2.0]).f.nlstep_data === nlstep
+    end
+end
+
 # DDEFunction
 
 ddefoop(u, h, p, t) = u