From 6088878e2602d194578c61e234041a3bc7fc7f94 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 1 Jul 2026 18:56:09 +0530 Subject: [PATCH] feat: set `Initial(x)` when setting `x` in `LinearizationProblem` --- Project.toml | 4 ++-- src/ModelingToolkit.jl | 1 + src/linearization.jl | 44 ++++++++++++++++++++++++++++-------------- test/linearize.jl | 24 +++++++++++++++++++++++ 4 files changed, 57 insertions(+), 16 deletions(-) diff --git a/Project.toml b/Project.toml index 69b096222b..1d7f57d7fa 100644 --- a/Project.toml +++ b/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"] diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl index 26005426bb..d13f850597 100644 --- a/src/ModelingToolkit.jl +++ b/src/ModelingToolkit.jl @@ -84,6 +84,7 @@ import PreallocationTools import PreallocationTools: DiffCache import FillArrays using BipartiteGraphs +import SciMLStructures @recompile_invalidations begin import StateSelection diff --git a/src/linearization.jl b/src/linearization.jl index 565c5298fa..4cb69d53ce 100644 --- a/src/linearization.jl +++ b/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,27 +250,27 @@ 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. """ @@ -274,35 +278,39 @@ struct LinearizationFunction{ """ 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) diff --git a/test/linearize.jl b/test/linearize.jl index e3b7973c2a..72ef60874c 100644 --- a/test/linearize.jl +++ b/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