diff --git a/Project.toml b/Project.toml index 9094e20..9db39d1 100644 --- a/Project.toml +++ b/Project.toml @@ -24,7 +24,7 @@ DiffEqCallbacks = "4.16.0" ForwardDiff = "1.0.1" LinearAlgebra = "1.10" LinearSolve = "5.4" -ModelingToolkit = "11.44" +ModelingToolkit = "11.45" NonlinearSolve = "4.12" NonlinearSolveBase = "2.40" OrdinaryDiffEq = "7" diff --git a/src/SteadyStateDiffEq.jl b/src/SteadyStateDiffEq.jl index e508efc..c2f0b8c 100644 --- a/src/SteadyStateDiffEq.jl +++ b/src/SteadyStateDiffEq.jl @@ -11,7 +11,7 @@ using SciMLPublic: @public using SciMLBase: SciMLBase, CallbackSet, LinearProblem, NonlinearProblem, ODEProblem, NonlinearSolution, ReturnCode, SteadyStateProblem, SteadyStateSolution, get_du, init, isinplace, remake, solve, successful_retcode -using SymbolicIndexingInterface: parameter_values +using SymbolicIndexingInterface: parameter_values, variable_symbols const infnorm = Base.Fix2(norm, Inf) diff --git a/src/solve.jl b/src/solve.jl index acacbee..beb2d6f 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -5,6 +5,40 @@ function __build_ssrootfind_solution(prob, nlsol) ) end +function __scc_original_residual(prob, u) + # The SCC lowering may drop observed/eliminated states from its residual, so + # `sol.resid` is not in the original state layout (and may even be empty). + # Recompute against the original steady-state residual at the restored `u`. + if isinplace(prob) + resid = similar(u) + prob.f(resid, u, prob.p, Inf) + return resid + else + return prob.f(u, prob.p, Inf) + end +end + +function __restore_scc_steady_state_solution( + prob, sol; save_idxs = nothing, alg = sol.alg, original = sol + ) + prob isa SteadyStateProblem || return nothing + prob.lowered_problem === nothing && return nothing + hasproperty(prob.f, :sys) || return nothing + + states = variable_symbols(prob) + length(states) == length(prob.u0) || return nothing + u = sol[states] + resid = __scc_original_residual(prob, u) + if save_idxs !== nothing + u = u[save_idxs] + resid = resid[save_idxs] + end + return SciMLBase.build_solution( + prob, alg, u, resid; + retcode = sol.retcode, stats = sol.stats, original + ) +end + function SciMLBase.__solve( prob::SciMLBase.AbstractSteadyStateProblem, alg::SSRootfind, args...; kwargs... @@ -18,9 +52,13 @@ function SciMLBase.__solve( fwd = (; (n => v for (n, v) in pairs(kwargs) if n !== :alias && n !== :verbose)...) end nlsol = solve(nlprob, alg.alg, args...; fwd...) - # 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. + restored = __restore_scc_steady_state_solution( + prob, nlsol; alg = SSRootfind(nlsol.alg), + save_idxs = get(kwargs, :save_idxs, nothing) + ) + restored !== nothing && return restored + # The fallback keeps the lowering as the solution problem when it has no + # symbolic map back to the original state variables. solprob = if prob isa SteadyStateProblem && prob.lowered_problem !== nothing nlprob else @@ -109,6 +147,10 @@ function __solve_scc_lowering( lp isa SciMLBase.AbstractSciMLProblem || (lp = lp(prob)) lp isa SciMLBase.SCCNonlinearProblem || return nothing sccsol = solve(lp, alg, args...; kwargs...) + restored = __restore_scc_steady_state_solution( + prob, sccsol; save_idxs, original = sccsol.original + ) + restored !== nothing && return restored save_idxs === nothing && return sccsol return SciMLBase.build_solution( lp, sccsol.alg, sccsol.u[save_idxs], sccsol.resid[save_idxs]; @@ -142,11 +184,18 @@ function SciMLBase.solve(prob::SteadyStateProblem, args...; kwargs...) # callable `lowered_problem` sees the updated operating point. _prob = NonlinearSolveBase.get_concrete_problem(prob; kwargs...) nlprob = NonlinearProblem(_prob) - nlprob isa SciMLBase.SCCNonlinearProblem && return solve(nlprob, args...; kwargs...) - return invoke( - solve, Tuple{SciMLBase.AbstractNonlinearProblem, Vararg{Any}}, - _prob, args...; kwargs... + sol = if nlprob isa SciMLBase.SCCNonlinearProblem + solve(nlprob, args...; kwargs...) + else + invoke( + solve, Tuple{SciMLBase.AbstractNonlinearProblem, Vararg{Any}}, + _prob, args...; kwargs... + ) + end + restored = __restore_scc_steady_state_solution( + _prob, sol; save_idxs = get(kwargs, :save_idxs, nothing) ) + return restored === nothing ? sol : restored end __get_tspan(u0, alg::Union{DynamicSS, SICNM}) = __get_tspan(u0, alg.tspan) diff --git a/test/scc.jl b/test/scc.jl index 4780d4f..16d69f9 100644 --- a/test/scc.jl +++ b/test/scc.jl @@ -391,14 +391,59 @@ end @test prob.lowered_problem !== nothing - # The `SSRootfind` solve happens on the `SCCNonlinearProblem` lowering, and - # `sol` is expressed on it. + # The solve uses the SCC lowering and returns values in the steady-state + # problem's state order. 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 + @test sol.prob isa SteadyStateProblem +end + +@testset "SteadyStateProblem SCC solutions use the original state order" begin + @parameters a = 1.0 b = 2.0 + @variables x(t) = 1.0 y(t) = 1.0 z(t) = 1.0 + @parameters α = 1.5 β = 1.0 γ = 3.0 δ = 1.0 + @variables u(t) = 1.0 v(t) = 1.0 + lv = mtkcompile( + System([D(u) ~ α * u - β * u * v, D(v) ~ -γ * v + δ * u * v], t; name = :lv) + ) + problems = ( + ( + mtkcompile( + System( + [D(x) ~ a - x, D(y) ~ x - b * y, D(z) ~ y - z], t; + name = :chain + ) + ), + [x => 0.1, y => 0.2, z => 0.3], + [x, y, z], + ), + (lv, [u => 2.9, v => 1.6], [u, v]), + ) + + for (sys, op, states) in problems + prob = SteadyStateProblem(sys, op) + sol = solve(prob, NewtonRaphson()) + reference = solve(remake(prob; lowered_problem = nothing), NewtonRaphson()) + + @test successful_retcode(sol) + @test sol.prob isa SteadyStateProblem + @test sol.u ≈ reference.u + @test sol[states] ≈ reference[states] + # Residual must match the original problem's state layout, not the + # (possibly shorter/empty) residual of the SCC lowering. + @test length(sol.resid) == length(sol.u) + @test sol.resid ≈ reference.resid atol = 1.0e-10 + + # `save_idxs` is applied after restoring the original state order. + saved = solve(prob, NewtonRaphson(); save_idxs = [2]) + @test length(saved.u) == 1 + @test saved.u ≈ sol.u[2:2] + @test length(saved.resid) == 1 + @test saved.resid ≈ sol.resid[2:2] atol = 1.0e-10 + end end @testset "DynamicSS on a SteadyStateProblem with an SCC lowering" begin @@ -415,7 +460,7 @@ end 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.prob isa SteadyStateProblem @test sol.original isa Tuple{SciMLBase.LinearSolution, NonlinearSolution} end @@ -432,6 +477,6 @@ end 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.prob isa SteadyStateProblem @test sol.original isa Tuple{SciMLBase.LinearSolution, NonlinearSolution} end