Restore original state order for SCC steady-state solutions

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"

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)
 

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)

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