diff --git a/Project.toml b/Project.toml index 9094e20..584fb95 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "SteadyStateDiffEq" uuid = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f" -version = "2.18.0" +version = "2.19.0" [deps] ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" @@ -24,14 +24,14 @@ DiffEqCallbacks = "4.16.0" ForwardDiff = "1.0.1" LinearAlgebra = "1.10" LinearSolve = "5.4" -ModelingToolkit = "11.44" +ModelingToolkit = "11.46" NonlinearSolve = "4.12" NonlinearSolveBase = "2.40" OrdinaryDiffEq = "7" PrecompileTools = "1.2.1" SafeTestsets = "0.1" SCCNonlinearSolve = "1.13" -SciMLBase = "3.55" +SciMLBase = "3.59" SciMLPublic = "1" SciMLTesting = "2.4" Sundials = "6.4.1" diff --git a/docs/src/index.md b/docs/src/index.md index bae6ad1..0d79063 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -48,6 +48,40 @@ prob = SteadyStateProblem((u, p, t) -> 1 .- u, [0.0]) sol = solve(prob, SICNM(Rodas3d())) ``` +## SCC-split steady states + +`SCCSteadyStateProblem` is the SCC-split form of a steady-state problem: its +residual decomposes into a block-lower-triangular sequence of `LinearProblem`/ +`NonlinearProblem`/`HomotopyProblem` blocks. Build it from a block list with +`SCCSteadyStateProblem(probs, explicitfuns!, p)`, or from a model with +`ModelingToolkit.SCCSteadyStateProblem(sys, op)`. + +`DynamicSS` and `SICNM` solve the blocks sequentially in SCC order — each +nonlinear block is integrated to steady state on its own residual after +`explicitfuns!` updates its caches from the upstream solutions. `SSRootfind`, a +bare `solve`, or any nonlinear algorithm solve the problem through the SCC +solver (`SCCNonlinearSolve.jl`). The solution's `u` and `resid` are in the +compressed SCC state layout; symbolic indexing (`sol[x]`, `sol.ps[p]`) goes +through `prob.f.sys`, and `save_idxs` is a selection in that layout — integers, +or states resolved through `variable_index`. + +A solve with `save_idxs` returns a solution whose `prob` is a saved-state view +of the original problem: it shares the original's blocks and parameters, but +its states are the saved entries in selection order. Solving or remaking that +problem works in the saved layout — `solve(saved.prob, alg)` applies the same +selection and returns a solution over the saved states, and a `save_idxs` +passed there selects among them. `sol[s]` resolves saved states in any order; +states that were not saved and observed expressions (which need the full +state) throw. Parameters and independent variables index as usual, and +`sol.original` is the full solution over all states. + +Solve-time `u0` and `p` keywords are applied to the problem with `remake` +before dispatch: `u0` is given in the compressed state layout — in the saved +layout when solving a saved-state view — and `p` is subject to +`parameters_alias`, so an unaliased problem throws since its blocks hold +separate parameter objects. The returned solution's `prob` is the updated +problem. + ## API ```@docs diff --git a/src/SteadyStateDiffEq.jl b/src/SteadyStateDiffEq.jl index e508efc..0059871 100644 --- a/src/SteadyStateDiffEq.jl +++ b/src/SteadyStateDiffEq.jl @@ -9,9 +9,13 @@ using LinearAlgebra: Diagonal, norm using LinearSolve: LinearSolve using SciMLPublic: @public using SciMLBase: SciMLBase, CallbackSet, LinearProblem, NonlinearProblem, ODEProblem, - NonlinearSolution, ReturnCode, SteadyStateProblem, SteadyStateSolution, get_du, init, + NonlinearSolution, ReturnCode, SCCNonlinearProblem, SCCSteadyStateProblem, + SteadyStateProblem, SteadyStateSolution, get_du, init, isinplace, remake, solve, successful_retcode -using SymbolicIndexingInterface: parameter_values +using SymbolicIndexingInterface: + NotSymbolic, SymbolicIndexingInterface, finalize_parameters_hook!, + parameter_values, set_parameter!, set_state!, + state_values, symbolic_type, variable_index, variable_symbols const infnorm = Base.Fix2(norm, Inf) @@ -20,8 +24,8 @@ include("solve.jl") include("precompilation.jl") export SSRootfind, DynamicSS, SICNM -export NonlinearProblem, NonlinearSolution, ReturnCode, SteadyStateProblem, - SteadyStateSolution, remake, solve, successful_retcode +export NonlinearProblem, NonlinearSolution, ReturnCode, SCCSteadyStateProblem, + SteadyStateProblem, SteadyStateSolution, remake, solve, successful_retcode @public SteadyStateDiffEqAlgorithm end diff --git a/src/algorithms.jl b/src/algorithms.jl index 9e6e415..0760f69 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -92,6 +92,11 @@ requires the SCC solver implementation from `SCCNonlinearSolve.jl` to be loaded automatically), or an `SCCAlg` configured with separate linear and nonlinear block solvers. +An `SCCSteadyStateProblem` (for example the SCC decomposition built by +`ModelingToolkit.SCCSteadyStateProblem`) is solved through +`SCCNonlinearProblem(prob)`: it is forwarded to `alg` unchanged, and the +solution is expressed on the SCC problem's compressed state ordering. + A `SteadyStateProblem` that records a `lowered_problem` (for example the SCC decomposition stored by `ModelingToolkit`) is solved through that lowering; since the lowering defines its own state ordering, the returned solution is @@ -148,10 +153,14 @@ solutions through `explicitfuns!` as in the SCC solve. `LinearProblem` blocks are solved directly, and the remaining blocks are integrated to steady state by `DynamicSS` on the block residual — so convergence only requires each nonlinear block's residual to be attracting under its own pseudo-transient -dynamics, not a globally attracting concatenated field. A `SteadyStateProblem` -that records an `SCCNonlinearProblem` as its `lowered_problem` (for example one -built by `ModelingToolkit`) takes the same sequential solve, and its solution -is expressed on the lowered problem's state ordering. +dynamics, not a globally attracting concatenated field. An +`SCCSteadyStateProblem` (for example one built by +`ModelingToolkit.SCCSteadyStateProblem`) takes the same sequential solve, and +its solution is expressed on the compressed state ordering. A +`SteadyStateProblem` that records an `SCCNonlinearProblem` as its +`lowered_problem` (for example one built by `ModelingToolkit`) takes the same +sequential solve, and its solution is expressed on the lowered problem's state +ordering. # Arguments @@ -251,7 +260,10 @@ solutions through `explicitfuns!` as in the SCC solve. `LinearProblem` blocks are solved directly, and the remaining blocks each run their own SICNM continuous-Newton flow on the block residual — so convergence only requires each nonlinear block to be tractable for SICNM on its own, not the concatenated -system. A `SteadyStateProblem` that records an `SCCNonlinearProblem` as its +system. An `SCCSteadyStateProblem` (for example one built by +`ModelingToolkit.SCCSteadyStateProblem`) takes the same sequential solve, and +its solution is expressed on the compressed state ordering. A +`SteadyStateProblem` that records an `SCCNonlinearProblem` as its `lowered_problem` (for example one built by `ModelingToolkit`) takes the same sequential solve, and its solution is expressed on the lowered problem's state ordering. diff --git a/src/solve.jl b/src/solve.jl index acacbee..d6940eb 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -21,14 +21,20 @@ function SciMLBase.__solve( # A stored `lowered_problem` (e.g. an `SCCNonlinearProblem`) solves in the # lowering's own state ordering, so its solution is expressed on the # lowering rather than the steady-state problem. - solprob = if prob isa SteadyStateProblem && prob.lowered_problem !== nothing - nlprob - else - prob - end + solprob = __lowered_problem(prob) === nothing ? prob : nlprob return __build_ssrootfind_solution(solprob, nlsol) end + +# Solve-time `u0`/`p` go through `remake`, which chunks `u0` into the blocks and +# enforces `parameters_alias`; the solve then runs on the updated problem. +function __scc_apply_overrides(prob; u0 = nothing, p = nothing, kwargs...) + if u0 !== nothing || p !== nothing + prob = remake(prob; u0 = u0 === nothing ? missing : u0, p = p === nothing ? missing : p) + end + return prob, kwargs +end + # An SCCNonlinearProblem has no top-level `u0`/`kwargs` fields, so it cannot go # through the generic AbstractNonlinearProblem solve preprocessing. Forward it # directly to the wrapped algorithm instead, which dispatches to the SCC solver @@ -37,20 +43,40 @@ function SciMLBase.solve( prob::SciMLBase.SCCNonlinearProblem, alg::SSRootfind, args...; kwargs... ) + prob, kwargs = __scc_apply_overrides(prob; kwargs...) nlsol = solve(prob, alg.alg, args...; kwargs...) return __build_ssrootfind_solution(prob, nlsol) end -# `DynamicSS`/`SICNM` on an `SCCNonlinearProblem` solve the blocks sequentially +# Evaluate a stateless block's residual once: out-of-place blocks return it, +# in-place blocks write into a `resid_prototype`-sized buffer. `HomotopyProblem` +# residuals take the continuation parameter `λ`, evaluated at its end. +function __scc_stateless_resid( + block::Union{NonlinearProblem, SciMLBase.HomotopyProblem}, u + ) + p = parameter_values(block) + f = block.f + args = block isa SciMLBase.HomotopyProblem ? (u, p, last(block.λspan)) : (u, p) + if isinplace(block) + fbase = f isa SciMLBase.HomotopyNonlinearFunction ? f.f : f + proto = hasproperty(fbase, :resid_prototype) ? fbase.resid_prototype : nothing + resid = proto === nothing ? similar(u) : similar(proto) + f(resid, args...) + return resid + end + return f(args...) +end + +# `DynamicSS`/`SICNM` on an SCC problem solve the blocks sequentially # in SCC order: each block's parameter cache is updated from the upstream # solutions through `explicitfuns!`, `LinearProblem` blocks are solved directly, # and the remaining blocks are integrated to steady state on the block # residual. `explicitfuns!` mutate the block caches, matching the # `SCCNonlinearSolve` protocol; each upstream solution is stripped to a plain # solution so generated explicit functions can index `sols[j][k]`/`sols[j].u`. -function SciMLBase.solve( - prob::SciMLBase.SCCNonlinearProblem, alg::Union{DynamicSS, SICNM}, - args...; kwargs... +function __solve_scc_blocks( + prob::Union{SciMLBase.SCCNonlinearProblem, SciMLBase.SCCSteadyStateProblem}, + alg::Union{DynamicSS, SICNM}, args...; kwargs... ) sols = () for i in eachindex(prob.probs) @@ -68,6 +94,19 @@ function SciMLBase.solve( SciMLBase.build_linear_solution( nothing, linsol.u, resid, nothing; retcode = linsol.retcode ) + elseif state_values(block) !== nothing && isempty(state_values(block)) + # A stateless block carries no integration; its residual must be + # identically empty, which is verified by evaluating it once. + u = copy(state_values(block)) + resid = __scc_stateless_resid(block, u) + isempty(resid) || throw( + DimensionMismatch( + "SCC block $i has no states but a residual of length $(length(resid))" + ) + ) + SciMLBase.build_solution( + block, nothing, u, resid; retcode = ReturnCode.Success + ) else # `HomotopyProblem` has no `solve` preprocessing support # (`get_concrete_problem` is only defined for `NonlinearProblem` @@ -86,8 +125,8 @@ function SciMLBase.solve( sols = (sols..., _sol) end - u = vcat(map(sol -> sol.u, sols)...) - resid = vcat(map(sol -> sol.resid, sols)...) + u = isempty(sols) ? Float64[] : vcat(map(sol -> sol.u, sols)...) + resid = isempty(sols) ? Float64[] : vcat(map(sol -> sol.resid, sols)...) retcode = SciMLBase.ReturnCode.Success for s in sols if !SciMLBase.successful_retcode(s) @@ -98,13 +137,238 @@ function SciMLBase.solve( return SciMLBase.build_solution(prob, alg, u, resid; retcode, original = sols) end +function SciMLBase.solve( + prob::SciMLBase.SCCNonlinearProblem, alg::Union{DynamicSS, SICNM}, + args...; kwargs... + ) + prob, kwargs = __scc_apply_overrides(prob; kwargs...) + return __solve_scc_blocks(prob, alg, args...; kwargs...) +end + +function SciMLBase.solve( + prob::SciMLBase.SCCSteadyStateProblem, args...; + save_idxs = nothing, kwargs... + ) + prob, kwargs = __scc_apply_overrides(prob; kwargs...) + alg = isempty(args) ? nothing : first(args) + alg === nothing && (alg = get(kwargs, :alg, nothing)) + kwargs = (; (n => v for (n, v) in pairs(kwargs) if n !== :alg)...) + rest = isempty(args) ? () : Base.tail(args) + sol = __solve_scc_steady_state(prob, alg, rest...; kwargs...) + return __scc_save_idxs(prob, sol, save_idxs) +end + +function __solve_scc_steady_state( + prob::SciMLBase.SCCSteadyStateProblem, alg::Union{DynamicSS, SICNM}, + args...; kwargs... + ) + return __solve_scc_blocks(prob, alg, args...; kwargs...) +end + +function __solve_scc_steady_state( + prob::SciMLBase.SCCSteadyStateProblem, alg::SSRootfind, + args...; kwargs... + ) + nlsol = solve(SCCNonlinearProblem(prob), alg.alg, args...; kwargs...) + return SciMLBase.build_solution( + prob, SSRootfind(nlsol.alg), nlsol.u, nlsol.resid; + nlsol.retcode, nlsol.stats, nlsol.left, nlsol.right, original = nlsol + ) +end + +function __solve_scc_steady_state( + prob::SciMLBase.SCCSteadyStateProblem, ::Nothing, + args...; kwargs... + ) + nlsol = solve(SCCNonlinearProblem(prob), args...; kwargs...) + return SciMLBase.build_solution( + prob, nlsol.alg, nlsol.u, nlsol.resid; + nlsol.retcode, nlsol.stats, nlsol.left, nlsol.right, original = nlsol + ) +end + +function __solve_scc_steady_state( + prob::SciMLBase.SCCSteadyStateProblem, alg, + args...; kwargs... + ) + nlsol = solve(SCCNonlinearProblem(prob), alg, args...; kwargs...) + return SciMLBase.build_solution( + prob, nlsol.alg, nlsol.u, nlsol.resid; + nlsol.retcode, nlsol.stats, nlsol.left, nlsol.right, original = nlsol + ) +end + +# Index provider for the saved-state layout of an `SCCSteadyStateProblem`: +# `idxs` are the positions of the saved states in the full compressed layout. +# Parameters and independent variables forward to the wrapped `sys`; observed +# quantities are unavailable since they need the full state. +struct SCCSavedStates{S} + sys::S + idxs::Vector{Int} +end + +SymbolicIndexingInterface.symbolic_container(w::SCCSavedStates) = w.sys + +function SymbolicIndexingInterface.variable_index(w::SCCSavedStates, sym) + idx = variable_index(w.sys, sym) + idx === nothing && return nothing + if idx isa AbstractArray + out = map(i -> something(findfirst(==(i), w.idxs), 0), idx) + any(iszero, out) && return nothing + return out + end + return findfirst(==(idx), w.idxs) +end + +function SymbolicIndexingInterface.is_variable(w::SCCSavedStates, sym) + return variable_index(w, sym) !== nothing +end + +function SymbolicIndexingInterface.variable_symbols(w::SCCSavedStates) + return variable_symbols(w.sys)[w.idxs] +end + +SymbolicIndexingInterface.all_variable_symbols(w::SCCSavedStates) = + variable_symbols(w) +SymbolicIndexingInterface.is_observed(::SCCSavedStates, sym) = false + +function SymbolicIndexingInterface.all_symbols(w::SCCSavedStates) + return union( + variable_symbols(w), SymbolicIndexingInterface.parameter_symbols(w.sys), + SymbolicIndexingInterface.independent_variable_symbols(w.sys) + ) +end + +# A live view of an `SCCSteadyStateProblem` whose state layout is the +# saved-state selection: `idxs` are the saved positions in `prob`'s compressed +# layout. `state_values`, `solve` and `remake` all act in the saved layout; +# the view shares the wrapped problem's blocks, parameters and index provider, +# narrowed to the saved states through `SCCSavedStates`. +struct SCCSavedProblem{ + uType, P <: SciMLBase.SCCSteadyStateProblem, I, F, + } <: SciMLBase.AbstractSteadyStateProblem{uType, false} + prob::P + idxs::I + f::F +end + +function SCCSavedProblem(prob::SciMLBase.SCCSteadyStateProblem, idxs) + f = if prob.f.sys === nothing + SciMLBase.NonlinearFunction{false}(Returns(nothing)) + else + SciMLBase.NonlinearFunction{false}( + Returns(nothing); sys = SCCSavedStates(prob.f.sys, vec(collect(idxs))) + ) + end + uType = state_values(prob) === nothing ? Nothing : typeof(state_values(prob)[idxs]) + return SCCSavedProblem{uType, typeof(prob), typeof(idxs), typeof(f)}( + prob, idxs, f + ) +end + +__saved_positions(w::SCCSavedProblem) = w.idxs isa Integer ? (w.idxs,) : w.idxs + +function SymbolicIndexingInterface.state_values(w::SCCSavedProblem) + states = state_values(w.prob) + return states === nothing ? nothing : states[w.idxs] +end +SymbolicIndexingInterface.parameter_values(w::SCCSavedProblem) = + parameter_values(w.prob) +SymbolicIndexingInterface.symbolic_container(w::SCCSavedProblem) = w.f +function SymbolicIndexingInterface.set_state!(w::SCCSavedProblem, val, idx) + return set_state!(w.prob, val, __saved_positions(w)[idx]) +end +function SymbolicIndexingInterface.set_parameter!(w::SCCSavedProblem, val, idx) + return set_parameter!(w.prob, val, idx) +end +function SymbolicIndexingInterface.finalize_parameters_hook!(w::SCCSavedProblem, sym) + return finalize_parameters_hook!(w.prob, sym) +end + +function SciMLBase.solve( + w::SCCSavedProblem, args...; save_idxs = nothing, u0 = nothing, p = nothing, + kwargs... + ) + if u0 !== nothing || p !== nothing + w = remake(w; u0 = u0 === nothing ? missing : u0, p = p === nothing ? missing : p) + end + composed = if save_idxs === nothing + w.idxs + else + __saved_positions(w)[__resolve_save_idxs(w, save_idxs)] + end + return solve(w.prob, args...; save_idxs = composed, kwargs...) +end + +function SciMLBase.remake( + w::SCCSavedProblem; u0 = missing, p = missing, kwargs... + ) + if u0 isa Number || (u0 isa AbstractArray && eltype(u0) <: Number) + old = state_values(w.prob) + T = promote_type( + old === nothing ? Union{} : eltype(old), + u0 isa Number ? typeof(u0) : eltype(u0) + ) + full = if old === nothing + # A guess-less problem has no values at the unsaved positions; + # zeros match the guess `LinearSolve` uses when none is given. + zeros(T, sum(size(b.A, 2) for b in w.prob.probs)) + else + copyto!(Vector{T}(undef, length(old)), old) + end + positions = collect(__saved_positions(w)) + if u0 isa Number + full[positions] .= u0 + else + full[positions] = u0 + end + u0 = full + end + return SCCSavedProblem(remake(w.prob; u0, p, kwargs...), w.idxs) +end + +# `save_idxs` is a selection in the compressed state layout: integers pass +# through, symbolic entries are resolved through `prob`'s index provider and +# must be states of it. A non-array `save_idxs` is a single entry and keeps the +# scalar shape of `u`/`resid`, matching `DynamicSS` on a `SteadyStateProblem`. +function __resolve_save_idxs(prob, save_idxs) + function resolve(s) + symbolic_type(s) === NotSymbolic() && return s + idx = variable_index(prob, s) + idx === nothing && throw( + ArgumentError("`save_idxs` entry $s is not a state of the problem") + ) + return idx + end + return save_idxs isa AbstractArray ? map(resolve, save_idxs) : resolve(save_idxs) +end + +function __scc_save_idxs( + prob::SciMLBase.SCCSteadyStateProblem, sol, save_idxs + ) + save_idxs === nothing && return sol + idxs = __resolve_save_idxs(prob, save_idxs) + return SciMLBase.build_solution( + SCCSavedProblem(prob, idxs), sol.alg, sol.u[idxs], sol.resid[idxs]; + retcode = sol.retcode, stats = sol.stats, left = sol.left, + right = sol.right, original = sol + ) +end + +# `SteadyStateProblem.lowered_problem` is a SciMLBase 3.x field; reads go +# through `hasfield` so this package also loads where the field is gone. +function __lowered_problem(prob) + return prob isa SteadyStateProblem && hasfield(typeof(prob), :lowered_problem) ? + prob.lowered_problem : nothing +end + # A `SteadyStateProblem` that records an `SCCNonlinearProblem` lowering is # solved block-sequentially in the lowering's ordering instead of one # monolithic integration. Returns `nothing` when there is no SCC lowering. function __solve_scc_lowering( prob, alg, args...; save_idxs = nothing, kwargs... ) - lp = prob isa SteadyStateProblem ? prob.lowered_problem : nothing + lp = __lowered_problem(prob) lp === nothing && return nothing lp isa SciMLBase.AbstractSciMLProblem || (lp = lp(prob)) lp isa SciMLBase.SCCNonlinearProblem || return nothing diff --git a/test/qa/qa.jl b/test/qa/qa.jl index f5224f9..da958a2 100644 --- a/test/qa/qa.jl +++ b/test/qa/qa.jl @@ -2,8 +2,8 @@ using SciMLTesting, SteadyStateDiffEq, Test using JET const REEXPORTS = ( - :NonlinearProblem, :NonlinearSolution, :ReturnCode, :SteadyStateProblem, - :SteadyStateSolution, :remake, :solve, :successful_retcode, + :NonlinearProblem, :NonlinearSolution, :ReturnCode, :SCCSteadyStateProblem, + :SteadyStateProblem, :SteadyStateSolution, :remake, :solve, :successful_retcode, ) const QUALIFIED_PUBLIC_IGNORE = ( diff --git a/test/scc.jl b/test/scc.jl index 4780d4f..4270405 100644 --- a/test/scc.jl +++ b/test/scc.jl @@ -1,7 +1,12 @@ -using SteadyStateDiffEq, NonlinearSolve, OrdinaryDiffEq, Test +using SteadyStateDiffEq, NonlinearSolve, OrdinaryDiffEq, LinearAlgebra, Test +using ForwardDiff using ModelingToolkit using ModelingToolkit: t_nounits as t, D_nounits as D using SCCNonlinearSolve: SCCAlg +import SymbolicIndexingInterface +using SymbolicIndexingInterface: + parameter_symbols, parameter_values, setp, state_values, variable_index, + variable_symbols using SciMLBase: HomotopyProblem, LinearProblem, NonlinearProblem, SCCNonlinearProblem, SteadyStateSolution @@ -206,78 +211,593 @@ end end end -# A `SteadyStateProblem` recording an `SCCNonlinearProblem` lowering is solved -# through it: `DynamicSS` runs the sequential block solve and `SSRootfind` -# forwards the lowering to the nonlinear solver. The solution is expressed on -# the lowering (whose ordering need not match `prob.u0`), so `sol.prob` is the -# `SCCNonlinearProblem`. -@testset "SteadyStateProblem with an SCC lowering" begin - f_iip(du, u, p, t) = (du .= 1 .- u) - f_oop(u, p, t) = 1 .- u - @testset "iip=$iip builder=$builder alg=$alg" for iip in (false, true), - builder in (false, true), + +# An `SCCSteadyStateProblem` built from the same blocks solves through the same +# sequential/block or SCC-solver paths, with the solution expressed on the SCC +# problem itself. +@testset "SCCSteadyStateProblem from blocks" begin + @testset "iip=$iip vector=$use_vector alg=$alg" for iip in (false, true), + use_vector in (false, true), alg in ( DynamicSS(Tsit5()), SICNM(Rodas5P()), - SSRootfind(NewtonRaphson()), + SSRootfind(SCCAlg()), SSRootfind(SCCAlg(; nlalg = NewtonRaphson())), + NewtonRaphson(), + nothing, ) - sccprob = dynamicss_scc_problem(iip, false) - lowered = builder ? (prob -> sccprob) : sccprob - prob = SteadyStateProblem( - iip ? f_iip : f_oop, [0.0, 0.0]; lowered_problem = lowered - ) - sol = solve(prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) + prob = SciMLBase.SCCSteadyStateProblem(dynamicss_scc_problem(iip, use_vector)) + sol = alg === nothing ? solve(prob) : + solve(prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) @test successful_retcode(sol) + @test sol.prob === prob @test sol.u ≈ [1, 2, 1, 2] atol = 1.0e-8 - @test sol.prob === sccprob - @test sol.original !== nothing + @test sol.u isa Vector{Float64} + @test sol.resid isa Vector{Float64} + @test length(sol.resid) == length(sol.u) + @test sol.resid ≈ zero(sol.resid) atol = 1.0e-7 + if alg isa Union{DynamicSS, SICNM} + @test sol.original isa Tuple + @test sol.original[1] isa SciMLBase.LinearSolution + else + @test sol.original isa NonlinearSolution + @test sol.original.prob isa SCCNonlinearProblem + end end - # Bare and `AbstractNonlinearAlgorithm` solves have no `__solve` dispatch of - # their own here; they land on the generic nonlinear-solve path, which - # converts back through `prob.u0` — absent on `SCCNonlinearProblem`. `solve` - # forwards the lowering to its own `solve` dispatch (the SCC solver) - # instead. - @testset "bare and nonlinear-alg solve" for iip in (false, true), - builder in (false, true), - alg in (nothing, NewtonRaphson()) + # `save_idxs` selects entries of the compressed SCC state/resid layout and + # keeps the full solution in `original`. + @testset "save_idxs" for alg in (DynamicSS(Tsit5()), NewtonRaphson()) + prob = SciMLBase.SCCSteadyStateProblem(dynamicss_scc_problem(false, false)) + full = solve(prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) + saved = solve( + prob, alg; save_idxs = [2, 4], abstol = 1.0e-10, reltol = 1.0e-10 + ) + @test saved.u ≈ full.u[[2, 4]] + @test saved.resid ≈ full.resid[[2, 4]] + @test state_values(saved.prob) == state_values(prob)[[2, 4]] + @test solve( + saved.prob, alg; abstol = 1.0e-10, reltol = 1.0e-10 + ).u ≈ saved.u + @test saved.original.u ≈ full.u + end - sccprob = dynamicss_scc_problem(iip, false) - lowered = builder ? (prob -> sccprob) : sccprob - prob = SteadyStateProblem( - iip ? f_iip : f_oop, [0.0, 0.0]; lowered_problem = lowered + # A `tspan` too short to settle gives a failure retcode, like DynamicSS on + # a `SteadyStateProblem`. + @testset "short tspan" for alg in (DynamicSS(Tsit5(); tspan = 1.0e-3), SICNM(Rodas5P(); tspan = 1.0e-3)) + prob = SciMLBase.SCCSteadyStateProblem(dynamicss_scc_problem(false, false)) + sol = solve(prob, alg) + @test !successful_retcode(sol) + @test sol.prob === prob + end +end + +@testset "SSRootfind on a ModelingToolkit SCCSteadyStateProblem" begin + @variables a(t) b(t) x(t) [irreducible = true] + @named model = System( + [D(a) ~ 5 - 3a - b, D(b) ~ 5 - a - 2b, D(x) ~ a + b - x^3], t + ) + sys = mtkcompile(model) + prob = SciMLBase.SCCSteadyStateProblem(sys, [a => 0.8, b => 1.8, x => 0.8]) + + sol = solve(prob, SSRootfind(NewtonRaphson()); abstol = 1.0e-12, reltol = 1.0e-12) + @test successful_retcode(sol) + @test sol[[a, b, x]] ≈ [1, 2, cbrt(3)] atol = 1.0e-8 + @test sol.prob === prob +end + +@testset "DynamicSS on a ModelingToolkit SCCSteadyStateProblem" begin + # Scalar nonlinear block `x' = 3 - x^3` has a unique real root `∛3`, so the + # per-block pseudo-transient solve and the monolithic integration converge + # to the same steady state. + @variables a(t) b(t) x(t) [irreducible = true] + @named model = System( + [D(a) ~ 5 - 3a - b, D(b) ~ 5 - a - 2b, D(x) ~ a + b - x^3], t + ) + sys = mtkcompile(model) + prob = SciMLBase.SCCSteadyStateProblem(sys, [a => 0.8, b => 1.8, x => 0.8]) + + sol = solve(prob, DynamicSS(); abstol = 1.0e-10, reltol = 1.0e-10) + @test successful_retcode(sol) + @test sol[[a, b, x]] ≈ [1, 2, cbrt(3)] atol = 1.0e-8 + @test sol.prob === prob + @test sol.original isa Tuple{SciMLBase.LinearSolution, NonlinearSolution} +end + +@testset "SICNM on a ModelingToolkit SCCSteadyStateProblem" begin + # Scalar nonlinear block `3 - x^3` has a unique real root `∛3` that is + # attracting under SICNM's continuous-Newton flow. + @variables a(t) b(t) x(t) [irreducible = true] + @named model = System( + [D(a) ~ 5 - 3a - b, D(b) ~ 5 - a - 2b, D(x) ~ a + b - x^3], t + ) + sys = mtkcompile(model) + prob = SciMLBase.SCCSteadyStateProblem(sys, [a => 0.8, b => 1.8, x => 0.8]) + + sol = solve(prob, SICNM(Rodas5P()); abstol = 1.0e-10, reltol = 1.0e-10) + @test successful_retcode(sol) + @test sol[[a, b, x]] ≈ [1, 2, cbrt(3)] atol = 1.0e-8 + @test sol.prob === prob + @test sol.original isa Tuple{SciMLBase.LinearSolution, NonlinearSolution} +end + +# Systems shared by the two solution-layout testsets below: `op` is the +# operating point, `states`/`expected` the steady-state values, `observed` an +# `sym => value` pair of the `osys` system's observed `XY`. +function layout_cases() + @parameters ca = 1.0 cb = 2.0 + @variables x(t) y(t) z(t) + chain = mtkcompile( + System([D(x) ~ ca - x, D(y) ~ x - cb * y, D(z) ~ y - z], t; name = :chain) + ) + + @parameters α = 1.5 β = 1.0 γ = 3.0 δ = 1.0 + @variables u(t) v(t) + lv = mtkcompile( + System([D(u) ~ α * u - β * u * v, D(v) ~ -γ * v + δ * u * v], t; name = :lv) + ) + + @variables a(t) b(t) mx(t) [irreducible = true] + mix = mtkcompile( + System( + [D(a) ~ 5 - 3a - b, D(b) ~ 5 - a - 2b, D(mx) ~ a + b - mx^3], t; + name = :mix ) - sol = alg === nothing ? solve(prob) : solve(prob, alg) + ) + + @parameters kp kd k1 k2 + @variables X(t) Y(t) XY(t) + osys = complete( + System( + [D(X) ~ kp - kd * X - k1 * X + k2 * Y, D(Y) ~ 1 + k1 * X - k2 * Y - Y], t; + observed = [XY ~ X + Y], name = :osys + ) + ) + Xss, Yss = [0.35 -0.5; -0.25 1.5] \ [1.0, 1.0] + + # `DynamicSS` integrates each block's residual to steady state. With the + # explicit `Tsit5` the step size is pinned at the stability limit near the + # fixed point, so the residual plateaus at a tolerance-sized noise floor + # right at the termination threshold and flips between `Success` and + # `Failure` under small perturbations. `Rodas5P` is not step-size limited at + # the fixed point, so the residual keeps decaying until the terminator + # fires (here at t ≈ 14). + allalgs = ( + NewtonRaphson(), SSRootfind(NewtonRaphson()), DynamicSS(Rodas5P()), + SICNM(Rodas5P()), + ) + nlalgs = (NewtonRaphson(), SSRootfind(NewtonRaphson())) + return ( + (; + sys = chain, op = [x => 0.1, y => 0.2, z => 0.3], states = [x, y, z], + expected = [1.0, 0.5, 0.5], algs = allalgs, observed = nothing, + ), + # Lotka–Volterra's nonzero steady state is a center, so the + # pseudo-transient methods do not apply. + (; + sys = lv, op = [u => 2.9, v => 1.6], states = [u, v], + expected = [3.0, 1.5], algs = nlalgs, observed = nothing, + ), + (; + sys = mix, op = [a => 0.8, b => 1.8, mx => 0.8], states = [a, b, mx], + expected = [1.0, 2.0, cbrt(3)], algs = allalgs, observed = nothing, + ), + (; + sys = osys, + op = [X => 4.0, Y => 5.0, kp => 1.0, kd => 0.1, k1 => 0.25, k2 => 0.5], + states = [X, Y], expected = [Xss, Yss], algs = allalgs, + observed = XY => Xss + Yss, + ), + ) +end + +# https://github.com/SciML/SteadyStateDiffEq.jl/issues/171 +# A `SteadyStateProblem` solves in its own layout: `sol.u`/`sol.resid` have the +# length and ordering of `prob.u0`. Nonlinear algorithms return their solution +# on the converted `NonlinearProblem`; the steady-state algorithms on `prob`. +@testset "SteadyStateProblem solutions are in the problem's layout" begin + @testset "case=$(nameof(case.sys)) alg=$(nameof(typeof(alg)))" for + case in layout_cases(), alg in case.algs + + (; sys, op, states, expected, observed) = case + prob = SteadyStateProblem(sys, op) + sol = solve(prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) @test successful_retcode(sol) - @test sol.u ≈ [1, 2, 1, 2] atol = 1.0e-9 + @test sol.prob isa ( + alg isa SciMLBase.AbstractNonlinearAlgorithm ? + NonlinearProblem : SteadyStateProblem + ) + @test length(sol.u) == length(prob.u0) + @test sol.u[variable_index.((sys,), states)] ≈ expected atol = 1.0e-7 + @test sol[states] ≈ expected atol = 1.0e-7 + observed === nothing || @test sol[first(observed)] ≈ last(observed) atol = 1.0e-7 + @test length(sol.resid) == length(prob.u0) + @test sol.resid ≈ zero(sol.resid) atol = 1.0e-7 + + if alg isa Union{DynamicSS, SICNM} + saved = solve( + prob, alg; save_idxs = [2], abstol = 1.0e-10, reltol = 1.0e-10 + ) + @test saved.u ≈ sol.u[2:2] atol = 1.0e-7 + @test length(saved.resid) == 1 + end end - # Non-SCC lowerings and problems without a lowering keep the upstream - # nonlinear-solve path. - @testset "non-SCC lowering" begin - nlprob = NonlinearProblem((u, p) -> [u[1]^2 - 4], [1.0]) - prob = SteadyStateProblem(f_oop, [0.0]; lowered_problem = nlprob) - sol = solve(prob) + # The residual solves `x` to an integer literal; the solution must keep + # `prob.u0`'s type. + @variables ex(t) + csys = complete(System([D(ex) ~ 1 - ex], t; name = :csys)) + @testset "eltype, alg=$(nameof(typeof(alg)))" for alg in ( + SSRootfind(Broyden()), DynamicSS(Tsit5()), SICNM(Rodas5P()), + ) + prob = SteadyStateProblem(csys, [ex => 0.2]) + sol = solve(prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) @test successful_retcode(sol) - @test sol.u ≈ [2.0] atol = 1.0e-9 + @test sol.u isa typeof(prob.u0) + @test sol.u ≈ [1.0] + end +end - plain = SteadyStateProblem(f_oop, [0.0]) - psol = solve(plain) - @test successful_retcode(psol) - @test psol.u ≈ [1.0] atol = 1.0e-9 +# The same systems via `SCCSteadyStateProblem`: the solution is expressed on the +# compressed block layout, symbolic indexing resolves eliminated variables as +# observed, and `save_idxs` indexes that layout. +@testset "SCCSteadyStateProblem from ModelingToolkit" begin + @testset "case=$(nameof(case.sys)) alg=$(nameof(typeof(alg)))" for + case in layout_cases(), alg in case.algs + + (; sys, op, states, expected, observed) = case + prob = SciMLBase.SCCSteadyStateProblem(sys, op) + sol = solve(prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) + @test successful_retcode(sol) + @test sol.prob === prob + @test length(sol.u) == length(state_values(prob)) + @test sol[states] ≈ expected atol = 1.0e-7 + observed === nothing || @test sol[first(observed)] ≈ last(observed) atol = 1.0e-7 + for s in variable_symbols(prob) + @test sol.u[variable_index(prob, s)] == sol[s] + end + @test length(sol.resid) == length(sol.u) + @test sol.resid ≈ zero(sol.resid) atol = 1.0e-7 + + isempty(state_values(prob)) && continue + saved = solve( + prob, alg; save_idxs = [1], abstol = 1.0e-10, reltol = 1.0e-10 + ) + @test saved.u ≈ sol.u[1:1] atol = 1.0e-7 + s1 = first(variable_symbols(prob)) + saved = solve( + prob, alg; save_idxs = [s1], abstol = 1.0e-10, reltol = 1.0e-10 + ) + @test saved.u ≈ [sol[s1]] atol = 1.0e-7 + @test saved.original.u ≈ sol.u atol = 1.0e-7 + # States not in `save_idxs` and observed expressions are unresolvable on + # the saved-state layout; parameters still index through it. + for s in Iterators.drop(variable_symbols(prob), 1) + @test_throws ErrorException saved[s] + end + observed === nothing || @test_throws ErrorException saved[first(observed)] + if !isempty(parameter_symbols(prob)) + p1 = first(parameter_symbols(prob)) + @test saved.ps[p1] == sol.ps[p1] + end + + # A symbolic `save_idxs` selection is applied in selection order. + rev = reverse(variable_symbols(prob)) + saved = solve( + prob, alg; save_idxs = rev, abstol = 1.0e-10, reltol = 1.0e-10 + ) + for s in rev + @test saved[s] ≈ sol[s] atol = 1.0e-7 + end + re = solve(saved.prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) + for s in rev + @test re[s] ≈ sol[s] atol = 1.0e-7 + end + @test re.u ≈ saved.u end +end - # `remake`d values reach a callable lowering through the materialized - # problem, so the block solve sees the new operating point. - sccprob = dynamicss_scc_problem(false, false) - prob = SteadyStateProblem( - f_oop, [0.0, 0.0]; lowered_problem = prob -> sccprob +@testset "SCCSteadyStateProblem save_idxs symbolic indexing" begin + # A single 2-state linear block with solution x = 1, y = 2; `sys` provides + # symbolic names so `saved[s]` resolves against the saved-state layout. + sysc = SymbolicIndexingInterface.SymbolCache([:x, :y], [], nothing) + linblock() = LinearProblem(Diagonal(ones(2)), [1.0, 2.0]; u0 = [10.0, 20.0]) + + @testset "alg=$(nameof(typeof(alg)))" for alg in ( + NewtonRaphson(), SSRootfind(NewtonRaphson()), + DynamicSS(Rodas5P()), SICNM(Rodas5P()), + ) + prob = SCCSteadyStateProblem( + (linblock(),), (Returns(nothing),), nothing, true; sys = sysc + ) + saved = solve(prob, alg; save_idxs = [2, 1]) + @test saved.u == [2, 1] + @test saved[:x] == 1 + @test saved[:y] == 2 + @test saved[[:x, :y]] == [1, 2] + @test saved.original[:x] == 1 + @test saved.original[:y] == 2 + + # Re-solving the saved problem works in the saved layout: the same + # selection is applied again, and `u0`/`save_idxs` are interpreted in + # saved positions. + re2 = solve(saved.prob, alg) + @test re2.u == [2, 1] + @test re2[:x] == 1 + @test re2[:y] == 2 + @test solve(saved.prob, alg; save_idxs = [:y]).u == [2] + @test state_values(saved.prob) == state_values(prob)[[2, 1]] + @test SymbolicIndexingInterface.getu(saved.prob, :y)(saved.prob) == + SymbolicIndexingInterface.getu(prob, :y)(prob) + @test SymbolicIndexingInterface.getu(saved.prob, :x)(saved.prob) == + SymbolicIndexingInterface.getu(prob, :x)(prob) + + pre_remake = copy(state_values(prob)) + r = remake(saved.prob; u0 = [30.0, 40.0]) + @test state_values(r) == [30.0, 40.0] + @test state_values(prob) == pre_remake + @test solve(r, alg).u == [2, 1] + + saved = solve(prob, alg; save_idxs = [:y]) + @test saved.u == [2] + @test saved[:y] == 2 + @test_throws ErrorException saved[:x] + @test saved.original[:x] == 1 + + re = solve(saved.prob, alg) + @test re.u == [2] + @test re[:y] == 2 + @test_throws ErrorException re[:x] + @test SymbolicIndexingInterface.getu(saved.prob, :y)(saved.prob) == + SymbolicIndexingInterface.getu(prob, :y)(prob) + end +end + +@testset "saved-problem parameter updates reach the blocks" begin + blk(c, s) = NonlinearProblem( + NonlinearFunction((u, p) -> [c * p[1] - u[1]]; sys = SymbolicIndexingInterface.SymbolCache([s], [:k])), + [0.0], [1.0] ) - @test NonlinearProblem(prob) === sccprob - prob2 = remake(prob; u0 = [5.0, 6.0]) - @test NonlinearProblem(prob2) === sccprob + mkprob() = SCCSteadyStateProblem( + (blk(1, :x), blk(2, :y)), (Returns(nothing), Returns(nothing)), + [1.0], false; sys = SymbolicIndexingInterface.SymbolCache([:x, :y], [:k]) + ) + @testset "alg=$(nameof(typeof(alg)))" for alg in ( + NewtonRaphson(), SSRootfind(NewtonRaphson()), + DynamicSS(Rodas5P()), SICNM(Rodas5P()), + ) + prob = mkprob() + setp(prob, :k)(prob, 3.0) + @test solve(prob, alg; abstol = 1.0e-10, reltol = 1.0e-10).u ≈ [3, 6] atol = 1.0e-6 + + prob = mkprob() + saved = solve(prob, alg; save_idxs = [:y], abstol = 1.0e-10, reltol = 1.0e-10) + @test saved.u ≈ [2] atol = 1.0e-6 + setp(saved.prob, :k)(saved.prob, 3.0) + @test all(b -> parameter_values(b) == [3.0], saved.prob.prob.probs) + re = solve(saved.prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) + @test re.ps[:k] == 3 + @test re.u ≈ [6.0] atol = 1.0e-6 + @test re[:y] ≈ 6 atol = 1.0e-6 + end +end + +@testset "guess-less blocks support save_idxs" begin + @testset "alg=$(nameof(typeof(alg)))" for alg in ( + NewtonRaphson(), SSRootfind(NewtonRaphson()), + DynamicSS(Rodas5P()), SICNM(Rodas5P()), + ) + prob = SCCSteadyStateProblem( + (LinearProblem(Diagonal(ones(2)), [1.0, 2.0]),), + (Returns(nothing),); sys = SymbolicIndexingInterface.SymbolCache([:x, :y], []) + ) + @test state_values(prob) === nothing + @test solve(prob, alg).u == [1, 2] + + for (sel, expected) in ( + ([2], [2]), ([:y], [2]), ([2, 1], [2, 1]), + ([:y, :x], [2, 1]), (2, 2), (:y, 2), + ) + saved = solve(prob, alg; save_idxs = sel) + @test saved.u == expected + @test saved[:y] == 2 + @test state_values(saved.prob) === nothing + @test solve(saved.prob, alg).u == expected + if sel isa AbstractArray + @test solve(saved.prob, alg; save_idxs = [:y]).u == [2] + end + end + + saved = solve(prob, alg; save_idxs = [2, 1]) + r = remake(saved.prob; u0 = [5.0, 6.0]) + @test state_values(r) == [5.0, 6.0] + @test solve(r, alg).u == [2, 1] + end +end + +@testset "saved remake preserves replacement numeric type" begin + linblock() = LinearProblem(Diagonal(ones(2)), [1.0, 2.0]; u0 = [10.0, 20.0]) + @testset "alg=$(nameof(typeof(alg)))" for alg in ( + NewtonRaphson(), SSRootfind(NewtonRaphson()), + DynamicSS(Rodas5P()), SICNM(Rodas5P()), + ) + prob = SCCSteadyStateProblem( + (linblock(),), (Returns(nothing),), nothing, true; + sys = SymbolicIndexingInterface.SymbolCache([:x, :y], [], nothing) + ) + saved = solve(prob, alg; save_idxs = [2]) + unsaved = state_values(prob)[1] + b = big"1.2345678901234567890123456789" + r = remake(saved.prob; u0 = [b]) + @test state_values(r) == [b] + @test eltype(state_values(r)) == BigFloat + @test state_values(r.prob)[1] == unsaved + @test state_values(remake(prob; u0 = [10.0, b]))[2] == b + + d = ForwardDiff.Dual(1.5, 1.0) + r = remake(saved.prob; u0 = [d]) + @test state_values(r) == [d] + @test ForwardDiff.partials(only(state_values(r))) == ForwardDiff.partials(d) + @test state_values(r.prob) == [unsaved, d] + + saved = solve(prob, alg; save_idxs = 2) + r = remake(saved.prob; u0 = b) + @test state_values(r) == b + end +end + +@testset "solve-time u0 and p overrides" begin + tol = (; abstol = 1.0e-10, reltol = 1.0e-10) + kblock(p) = NonlinearProblem( + NonlinearFunction( + (u, p) -> [p[1] - u[1]]; + sys = SymbolicIndexingInterface.SymbolCache([:x], [:k]) + ), [0.0], p + ) + function xblock(u0) + return NonlinearProblem((u, p) -> [u[1] - u[1]^3], [u0]) + end + function yblock(u0) + return NonlinearProblem((u, p) -> [4u[1] - u[1]^3], [u0]) + end + cubic() = NonlinearProblem( + (u, p) -> [u[1] - u[1]^3, 4u[2] - u[2]^3], [0.8, 1.5] + ) + function mkparam() + p = [1.0] + return SCCSteadyStateProblem( + (kblock(p),), (Returns(nothing),), p, true; + sys = SymbolicIndexingInterface.SymbolCache([:x], [:k]) + ) + end + mkblocks() = ( + SCCSteadyStateProblem( + (xblock(0.8), yblock(1.5)), (Returns(nothing), Returns(nothing)), + nothing, true; sys = SymbolicIndexingInterface.SymbolCache([:x, :y], []) + ), + SCCSteadyStateProblem( + (cubic(),), (Returns(nothing),), + nothing, true; sys = SymbolicIndexingInterface.SymbolCache([:x, :y], []) + ), + ) + @testset "alg=$(nameof(typeof(alg)))" for alg in ( + NewtonRaphson(), SSRootfind(NewtonRaphson()), + DynamicSS(Rodas5P()), SICNM(Rodas5P()), + ) + prob = mkparam() + sol = solve(prob, alg; p = [3.0], tol...) + @test sol.u ≈ [3] atol = 1.0e-6 + @test sol.ps[:k] == 3 + @test parameter_values(sol.prob) == [3.0] + @test sol.resid ≈ [parameter_values(sol.prob)[1] - sol.u[1]] atol = 1.0e-8 + @test prob.ps[:k] == 1 + + saved = solve(prob, alg; save_idxs = [:x], tol...) + re = solve(saved.prob, alg; p = [3.0], tol...) + @test re.u ≈ [3] atol = 1.0e-6 + @test re.ps[:k] == 3 + + unaliased = SCCSteadyStateProblem( + (kblock([1.0]),), (Returns(nothing),), [1.0], false; + sys = SymbolicIndexingInterface.SymbolCache([:x], [:k]) + ) + @test_throws ArgumentError solve(unaliased, alg; p = [3.0]) + + if !(alg isa typeof(NewtonRaphson())) + p = [1.0] + nlprob = SCCNonlinearProblem( + (kblock(p),), (Returns(nothing),), p, true; + sys = SymbolicIndexingInterface.SymbolCache([:x], [:k]) + ) + nlsol = solve(nlprob, alg; p = [3.0], tol...) + @test nlsol.u ≈ [3] atol = 1.0e-6 + @test parameter_values(nlsol.prob) == [3.0] + end + + for prob in mkblocks() + sol = solve(prob, alg; u0 = [-0.8, -1.5], tol...) + @test sol.u ≈ [-1, -2] atol = 1.0e-6 + @test state_values(prob) == [0.8, 1.5] + + saved = solve(prob, alg; save_idxs = [:y, :x], tol...) + @test saved.u ≈ [2, 1] atol = 1.0e-6 + re = solve(saved.prob, alg; u0 = [-1.5, 0.8], tol...) + @test re.u ≈ [-2, 1] atol = 1.0e-6 + @test re.u ≈ solve(remake(saved.prob; u0 = [-1.5, 0.8]), alg; tol...).u atol = 1.0e-6 + + saved = solve(prob, alg; save_idxs = [:y], tol...) + r = solve(saved.prob, alg; u0 = [-1.5], tol...) + @test r.u ≈ [-2] atol = 1.0e-6 + @test r.original.u ≈ [1, -2] atol = 1.0e-6 + + saved = solve(prob, alg; save_idxs = :y, tol...) + @test solve(saved.prob, alg; u0 = -1.5, tol...).u ≈ -2 atol = 1.0e-6 + end + end +end + +@testset "stateless and empty SCC blocks" begin + @testset "alg=$(nameof(typeof(alg)))" for alg in ( + DynamicSS(Rodas5P()), SICNM(Rodas5P()), + ) + # A stateless block with a nonempty residual is rejected — the residual + # is still evaluated once. + for iip in (false, true) + counter = Ref(0) + blk = if iip + NonlinearProblem( + NonlinearFunction( + (res, u, p) -> (counter[] += 1; res .= [1.0]); + resid_prototype = [0.0] + ), Float64[] + ) + else + NonlinearProblem( + NonlinearFunction( + (u, p) -> (counter[] += 1; [1.0]); resid_prototype = [0.0] + ), Float64[] + ) + end + for T in (SCCSteadyStateProblem, SCCNonlinearProblem) + prob = T((blk,), (Returns(nothing),), nothing, true) + counter[] = 0 + @test_throws DimensionMismatch solve(prob, alg) + @test counter[] > 0 + end + end + + # A stateless block whose residual is empty returns Success; the + # residual is evaluated rather than assumed zero. + counter = Ref(0) + blk = NonlinearProblem( + NonlinearFunction( + (u, p) -> (counter[] += 1; Float64[]); resid_prototype = Float64[] + ), Float64[] + ) + for T in (SCCSteadyStateProblem, SCCNonlinearProblem) + prob = T((blk,), (Returns(nothing),), nothing, true) + counter[] = 0 + sol = solve(prob, alg) + @test successful_retcode(sol) + @test counter[] > 0 + end + + # An empty linear block is a no-op, alone or before another block. + lblk() = LinearProblem(zeros(0, 0), Float64[]; u0 = Float64[]) + for T in (SCCSteadyStateProblem, SCCNonlinearProblem) + sol = solve(T((lblk(),), (Returns(nothing),), nothing, true), alg) + @test successful_retcode(sol) + @test sol.u == Float64[] + prob = T( + (lblk(), LinearProblem([2.0;;], [4.0]; u0 = [0.0])), + (Returns(nothing), Returns(nothing)), nothing, true + ) + sol = solve(prob, alg) + @test successful_retcode(sol) + @test sol.u ≈ [2.0] + end + end end @testset "SSRootfind on a ModelingToolkit SCC decomposition" begin @@ -368,70 +888,93 @@ end @test sol.original[2].original.original isa SciMLBase.AbstractODESolution end -@testset "SteadyStateProblem solves through an SCC lowering" begin - @variables a(t) b(t) x(t) y(t) c(t) d(t) [irreducible = true] - @named model = System( - [ - D(a) ~ 5 - 3a - b, - D(b) ~ 5 - a - 2b, - D(x) ~ a + b - x^2 - y, - D(y) ~ 3a + b - x - y^2, - D(c) ~ x + y + 7 - 2c - d, - D(d) ~ 2x + y + 11 - c - 3d, - ], t - ) - sys = mtkcompile(model) - op = [a => 0.8, b => 1.8, x => 0.8, y => 1.8, c => 2.8, d => 3.8] - prob = SteadyStateProblem(sys, op) +@testset "SteadyStateProblem with an SCC lowering" begin + f_iip(du, u, p, t) = (du .= 1 .- u) + f_oop(u, p, t) = 1 .- u + @testset "iip=$iip builder=$builder alg=$alg" for iip in (false, true), + builder in (false, true), + alg in ( + DynamicSS(Tsit5()), + SICNM(Rodas5P()), + SSRootfind(NewtonRaphson()), + SSRootfind(SCCAlg(; nlalg = NewtonRaphson())), + ) + + sccprob = dynamicss_scc_problem(iip, false) + lowered = builder ? (prob -> sccprob) : sccprob + prob = SteadyStateProblem( + iip ? f_iip : f_oop, [0.0, 0.0]; lowered_problem = lowered + ) + sol = solve(prob, alg; abstol = 1.0e-10, reltol = 1.0e-10) + @test successful_retcode(sol) + @test sol.u ≈ [1, 2, 1, 2] atol = 1.0e-8 + @test sol.prob === sccprob + @test sol.original !== nothing + end - @test prob.f.initialization_data.initializeprob isa SCCNonlinearProblem + # Bare and `AbstractNonlinearAlgorithm` solves have no `__solve` dispatch of + # their own here; they land on the generic nonlinear-solve path, which + # converts back through `prob.u0` — absent on `SCCNonlinearProblem`. `solve` + # forwards the lowering to its own `solve` dispatch (the SCC solver) + # instead. + @testset "bare and nonlinear-alg solve" for iip in (false, true), + builder in (false, true), + alg in (nothing, NewtonRaphson()) - states = [a, b, x, y, c, d] - expected = [1, 2, 1, 2, 3, 4] + sccprob = dynamicss_scc_problem(iip, false) + lowered = builder ? (prob -> sccprob) : sccprob + prob = SteadyStateProblem( + iip ? f_iip : f_oop, [0.0, 0.0]; lowered_problem = lowered + ) + sol = alg === nothing ? solve(prob) : solve(prob, alg) + @test successful_retcode(sol) + @test sol.u ≈ [1, 2, 1, 2] atol = 1.0e-9 + end - @test prob.lowered_problem !== nothing + # Non-SCC lowerings and problems without a lowering keep the upstream + # nonlinear-solve path. + @testset "non-SCC lowering" begin + nlprob = NonlinearProblem((u, p) -> [u[1]^2 - 4], [1.0]) + prob = SteadyStateProblem(f_oop, [0.0]; lowered_problem = nlprob) + sol = solve(prob) + @test successful_retcode(sol) + @test sol.u ≈ [2.0] atol = 1.0e-9 - # The `SSRootfind` solve happens on the `SCCNonlinearProblem` lowering, and - # `sol` is expressed on it. - sol = solve( - prob, SSRootfind(NewtonRaphson()); abstol = 1.0e-12, reltol = 1.0e-12 - ) - @test successful_retcode(sol) - @test sol[states] ≈ expected atol = 1.0e-9 - @test sol.prob isa SCCNonlinearProblem -end + plain = SteadyStateProblem(f_oop, [0.0]) + psol = solve(plain) + @test successful_retcode(psol) + @test psol.u ≈ [1.0] atol = 1.0e-9 + end -@testset "DynamicSS on a SteadyStateProblem with an SCC lowering" begin - # Scalar nonlinear block `x' = 3 - x^3` has a unique real root `∛3`, so the - # per-block pseudo-transient solve and the monolithic integration converge - # to the same steady state. - @variables a(t) b(t) x(t) [irreducible = true] - @named model = System( - [D(a) ~ 5 - 3a - b, D(b) ~ 5 - a - 2b, D(x) ~ a + b - x^3], t + # `remake`d values reach a callable lowering through the materialized + # problem, so the block solve sees the new operating point. + sccprob = dynamicss_scc_problem(false, false) + prob = SteadyStateProblem( + f_oop, [0.0, 0.0]; lowered_problem = prob -> sccprob ) - sys = mtkcompile(model) - prob = SteadyStateProblem(sys, [a => 0.8, b => 1.8, x => 0.8]) - - sol = solve(prob, DynamicSS(); abstol = 1.0e-10, reltol = 1.0e-10) - @test successful_retcode(sol) - @test sol[[a, b, x]] ≈ [1, 2, cbrt(3)] atol = 1.0e-8 - @test sol.prob isa SCCNonlinearProblem - @test sol.original isa Tuple{SciMLBase.LinearSolution, NonlinearSolution} + @test NonlinearProblem(prob) === sccprob + prob2 = remake(prob; u0 = [5.0, 6.0]) + @test NonlinearProblem(prob2) === sccprob end -@testset "SICNM on a SteadyStateProblem with an SCC lowering" begin - # Scalar nonlinear block `3 - x^3` has a unique real root `∛3` that is - # attracting under SICNM's continuous-Newton flow. - @variables a(t) b(t) x(t) [irreducible = true] - @named model = System( - [D(a) ~ 5 - 3a - b, D(b) ~ 5 - a - 2b, D(x) ~ a + b - x^3], t +# A `SteadyStateProblem` may record a non-SCC `lowered_problem`; it solves +# through the lowering and the solution is expressed on the lowered problem's +# own coordinates. +@testset "SteadyStateProblem with a generic lowering" begin + ssys = SymbolicIndexingInterface.SymbolCache([:x, :y, :z], [], nothing) + lowsys = SymbolicIndexingInterface.SymbolCache([:y], [], nothing) + low = NonlinearProblem( + SciMLBase.NonlinearFunction((u, p) -> [2 - u[1]]; sys = lowsys), [0.0] ) - sys = mtkcompile(model) - prob = SteadyStateProblem(sys, [a => 0.8, b => 1.8, x => 0.8]) - - sol = solve(prob, SICNM(Rodas5P()); abstol = 1.0e-10, reltol = 1.0e-10) - @test successful_retcode(sol) - @test sol[[a, b, x]] ≈ [1, 2, cbrt(3)] atol = 1.0e-8 - @test sol.prob isa SCCNonlinearProblem - @test sol.original isa Tuple{SciMLBase.LinearSolution, NonlinearSolution} + for lp in (low, prob -> low), alg in (SSRootfind(NewtonRaphson()), SSRootfind()) + prob = SteadyStateProblem( + SciMLBase.ODEFunction((u, p, t) -> [1 - u[1], 2 - u[2], 3 - u[3]]; sys = ssys), + zeros(3); lowered_problem = lp + ) + sol = solve(prob, alg) + @test sol.u ≈ [2] + @test variable_symbols(sol) == [:y] + @test sol[:y] ≈ 2 + @test sol.prob === low + end end