feat: set Initial(x) when setting x in LinearizationProblem

Project.toml

@@ -33,6 +33,7 @@ RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SCCNonlinearSolve = "9dfe8606-65a1-4bb3-9748-cb89d1561431"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLPublic = "431bcebd-1456-4ced-9d72-93c2757fff0b"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
@@ -184,7 +185,6 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciCompDSL = "91a8cdf1-4ca6-467b-a780-87fda3fff15e"
-SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -197,4 +197,4 @@ TestEnv = "1e6cf692-eddd-4d53-88a5-2d735e33781b"
 URIs = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
 
 [targets]
-test = ["BenchmarkTools", "BoundaryValueDiffEqMIRK", "BoundaryValueDiffEqAscher", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqDefault", "REPL", "Random", "ReferenceTests", "SafeTestsets", "SciMLTesting", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "OrdinaryDiffEqNonlinearSolve", "Logging", "OptimizationBase", "LinearSolve", "Latexify", "Distributed", "DiffEqNoiseProcess", "DynamicQuantities", "DiffEqCallbacks", "URIs", "JumpProcesses", "RecursiveArrayTools", "SciMLStructures", "SpecialFunctions", "SciCompDSL", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqBDF", "OrdinaryDiffEqFunctionMap", "OrdinaryDiffEqFIRK"]
+test = ["BenchmarkTools", "BoundaryValueDiffEqMIRK", "BoundaryValueDiffEqAscher", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqDefault", "REPL", "Random", "ReferenceTests", "SafeTestsets", "SciMLTesting", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "OrdinaryDiffEqNonlinearSolve", "Logging", "OptimizationBase", "LinearSolve", "Latexify", "Distributed", "DiffEqNoiseProcess", "DynamicQuantities", "DiffEqCallbacks", "URIs", "JumpProcesses", "RecursiveArrayTools", "SpecialFunctions", "SciCompDSL", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqBDF", "OrdinaryDiffEqFunctionMap", "OrdinaryDiffEqFIRK"]

src/ModelingToolkit.jl

@@ -84,6 +84,7 @@ import PreallocationTools
 import PreallocationTools: DiffCache
 import FillArrays
 using BipartiteGraphs
+import SciMLStructures
 
 @recompile_invalidations begin
     import StateSelection

src/linearization.jl

@@ -94,6 +94,10 @@ function linearization_function(
         sys, merge(op, anydict(p)), (t, t); allow_incomplete = true,
         algebraic_only = true, guesses, missing_guess_value
     )
+    initial_idxs_for_unknowns = ParameterIndex{SciMLStructures.Initials, Int}[]
+    for v in unknowns(sys)
+        push!(initial_idxs_for_unknowns, parameter_index(sys, Initial(v))::eltype(initial_idxs_for_unknowns))
+    end
     u0 = state_values(prob)
 
     ps = parameters(sys)
@@ -170,7 +174,7 @@ function linearization_function(
     lin_fun = LinearizationFunction(
         diff_idxs, alge_idxs, inputs, length(unknowns(sys)),
         prob, h, u0 === nothing ? nothing : similar(u0, T), uf_jac, h_jac, pf_jac,
-        hp_jac, initializealg, initialization_kwargs
+        hp_jac, initializealg, initialization_kwargs, initial_idxs_for_unknowns
     )
     return lin_fun, sys
 end
@@ -246,63 +250,67 @@ A callable struct which linearizes a system.
 
 $(TYPEDFIELDS)
 """
-struct LinearizationFunction{
+mutable struct LinearizationFunction{
         DI <: AbstractVector{Int}, AI <: AbstractVector{Int}, I, P <: ODEProblem,
         H, C, J1, J2, J3, J4, IA <: SciMLBase.DAEInitializationAlgorithm, IK,
     }
     """
     The indexes of differential equations in the linearized system.
     """
-    diff_idxs::DI
+    const diff_idxs::DI
     """
     The indexes of algebraic equations in the linearized system.
     """
-    alge_idxs::AI
+    const alge_idxs::AI
     """
     The indexes of parameters in the linearized system which represent
     input variables.
     """
-    inputs::I
+    const inputs::I
     """
     The number of unknowns in the linearized system.
     """
-    num_states::Int
+    const num_states::Int
     """
     The `ODEProblem` of the linearized system.
     """
     prob::P
     """
     A function which takes `(u, p, t)` and returns the outputs of the linearized system.
     """
-    h::H
+    const h::H
     """
     Any required cache buffers.
     """
-    caches::C
+    const caches::C
     """
     `PreparedJacobian` for calculating jacobian of `prob.f` w.r.t. `u`
     """
-    uf_jac::J1
+    const uf_jac::J1
     """
     `PreparedJacobian` for calculating jacobian of `h` w.r.t. `u`
     """
-    h_jac::J2
+    const h_jac::J2
     """
     `PreparedJacobian` for calculating jacobian of `prob.f` w.r.t. `p`
     """
-    pf_jac::J3
+    const pf_jac::J3
     """
     `PreparedJacobian` for calculating jacobian of `h` w.r.t. `p`
     """
-    hp_jac::J4
+    const hp_jac::J4
     """
     The initialization algorithm to use.
     """
-    initializealg::IA
+    const initializealg::IA
     """
     Keyword arguments to be passed to `SciMLBase.get_initial_values`.
     """
-    initialize_kwargs::IK
+    const initialize_kwargs::IK
+    """
+    Index of `Initial(x)` for every `x` in unknowns.
+    """
+    const initial_idxs_for_unknowns::Vector{ParameterIndex{SciMLStructures.Initials, Int}}
 end
 
 SymbolicIndexingInterface.symbolic_container(f::LinearizationFunction) = f.prob
@@ -311,6 +319,10 @@ function SymbolicIndexingInterface.parameter_values(f::LinearizationFunction)
     return parameter_values(f.prob)
 end
 SymbolicIndexingInterface.current_time(f::LinearizationFunction) = current_time(f.prob)
+function SymbolicIndexingInterface.set_state!(f::LinearizationFunction, val, idx)
+    f.prob.u0[idx] = val
+    return f.prob.p[f.initial_idxs_for_unknowns[idx]] = val
+end
 
 function Base.show(io::IO, mime::MIME"text/plain", lf::LinearizationFunction)
     printstyled(io, "LinearizationFunction"; bold = true, color = :blue)
@@ -473,6 +485,10 @@ SymbolicIndexingInterface.symbolic_container(p::LinearizationProblem) = p.f
 SymbolicIndexingInterface.state_values(p::LinearizationProblem) = state_values(p.f)
 SymbolicIndexingInterface.parameter_values(p::LinearizationProblem) = parameter_values(p.f)
 SymbolicIndexingInterface.current_time(p::LinearizationProblem) = p.t
+function SymbolicIndexingInterface.set_state!(p::LinearizationProblem, val, idx)
+    return SymbolicIndexingInterface.set_state!(p.f, val, idx)
+end
+
 
 function Base.getindex(prob::LinearizationProblem, idx)
     return getu(prob, idx)(prob)

test/linearize.jl

@@ -2,6 +2,7 @@ using ModelingToolkit, ADTypes, Test
 using NonlinearSolve
 using Symbolics: value
 using CommonSolve: solve
+using SymbolicIndexingInterface
 
 # Test reorder_unknowns
 # sys = ssrand(1,1,4);
@@ -399,3 +400,26 @@ end
         tank_noi, [md_i], [h]; p = [md_i => 1.0, m => m_ss]
     )
 end
+
+@testset "Issue#4694: `setsym` on a variable `x` also updates `Initial(x)`" begin
+    @parameters a = 1.0
+    @variables x(t) = 1.0 u(t) = 0.0 [input = true]
+    @named sys = System([D(x) ~ -a * x^3 + u], t, [x, u], [a])
+    sys = complete(sys)
+    # Reference: a full rebuild at (x=2, a=3) gives A = -3*3*2^2 = -36.
+    Afresh(xv, av) = linearize(sys, [sys.u], [sys.x]; op = Dict(sys.x => xv, sys.a => av))[1].A[1]
+    @test Afresh(2.0, 3.0) ≈ -36.0
+    build() = LinearizationProblem(
+        linearization_function(sys, [sys.u], [sys.x]; op = Dict(sys.x => 1.0, sys.a => 1.0))[1], 0.0
+    )
+    A(lp) = solve(lp)[1].A[1]
+    ix = Initial(sys.x)
+    lp = build(); setsym(lp, [sys.x, sys.a])(lp, [2.0, 3.0])
+    @test A(lp) ≈ -36.0
+    lp = build(); setu(lp, sys.x)(lp, 2.0); setp(lp, sys.a)(lp, 3.0)
+    @test A(lp) ≈ -36.0
+    lp = build(); setsym(lp, [ix, sys.a])(lp, [2.0, 3.0])
+    @test A(lp) ≈ -36.0
+    lp = build(); setp(lp, [ix, sys.a])(lp, [2.0, 3.0])
+    @test A(lp) ≈ -36.0
+end