Merge pull request #2809 from AayushSabharwal/as/linearization-initsy…

…s-missing-vars

fix: fix unknowns not present in initialization system in linearization

src/systems/abstractsystem.jl

@@ -1896,6 +1896,7 @@ function linearization_function(sys::AbstractSystem, inputs,
         zero_dummy_der = false,
         initialization_solver_alg = TrustRegion(),
         eval_expression = false, eval_module = @__MODULE__,
+        warn_initialize_determined = true,
         kwargs...)
     inputs isa AbstractVector || (inputs = [inputs])
     outputs isa AbstractVector || (outputs = [outputs])
@@ -1909,10 +1910,26 @@ function linearization_function(sys::AbstractSystem, inputs,
         op = merge(defs, op)
     end
     sys = ssys
+    u0map = Dict(k => v for (k, v) in op if is_variable(ssys, k))
     initsys = structural_simplify(
         generate_initializesystem(
-            sys, guesses = guesses(sys), algebraic_only = true),
+            sys, u0map = u0map, guesses = guesses(sys), algebraic_only = true),
         fully_determined = false)
+
+    # HACK: some unknowns may not be involved in any initialization equations, and are
+    # thus removed from the system during `structural_simplify`.
+    # This causes `getu(initsys, unknowns(sys))` to fail, so we add them back as parameters
+    # for now.
+    missing_unknowns = setdiff(unknowns(sys), all_symbols(initsys))
+    if !isempty(missing_unknowns)
+        if warn_initialize_determined
+            @warn "Initialization system is underdetermined. No equations for $(missing_unknowns). Initialization will default to using least squares. To suppress this warning pass warn_initialize_determined = false."
+        end
+        new_parameters = [parameters(initsys); missing_unknowns]
+        @set! initsys.ps = new_parameters
+        initsys = complete(initsys)
+    end
+
     if p isa SciMLBase.NullParameters
         p = Dict()
     else
@@ -1983,7 +2000,9 @@ function linearization_function(sys::AbstractSystem, inputs,
                 p = todict(p)
                 newps = deepcopy(sys_ps)
                 for (k, v) in p
-                    setp(sys, k)(newps, v)
+                    if is_parameter(sys, k)
+                        setp(sys, k)(newps, v)
+                    end
                 end
                 p = newps
             end

src/systems/nonlinear/initializesystem.jl

@@ -70,25 +70,23 @@ function generate_initializesystem(sys::ODESystem;
     defs = merge(defaults(sys), filtered_u0)
     guesses = merge(get_guesses(sys), todict(guesses), dd_guess)
 
-    if !algebraic_only
-        for st in full_states
-            if st ∈ keys(defs)
-                def = defs[st]
+    for st in full_states
+        if st ∈ keys(defs)
+            def = defs[st]
 
-                if def isa Equation
-                    st ∉ keys(guesses) && check_defguess &&
-                        error("Invalid setup: unknown $(st) has an initial condition equation with no guess.")
-                    push!(eqs_ics, def)
-                    push!(u0, st => guesses[st])
-                else
-                    push!(eqs_ics, st ~ def)
-                    push!(u0, st => def)
-                end
-            elseif st ∈ keys(guesses)
+            if def isa Equation
+                st ∉ keys(guesses) && check_defguess &&
+                    error("Invalid setup: unknown $(st) has an initial condition equation with no guess.")
+                push!(eqs_ics, def)
                 push!(u0, st => guesses[st])
-            elseif check_defguess
-                error("Invalid setup: unknown $(st) has no default value or initial guess")
+            else
+                push!(eqs_ics, st ~ def)
+                push!(u0, st => def)
             end
+        elseif st ∈ keys(guesses)
+            push!(u0, st => guesses[st])
+        elseif check_defguess
+            error("Invalid setup: unknown $(st) has no default value or initial guess")
         end
     end
 

test/downstream/linearize.jl

@@ -195,3 +195,113 @@ lsys, ssys = linearize(sat, [u], [y]; op = Dict(u => 2))
 @test isempty(lsys.B)
 @test isempty(lsys.C)
 @test lsys.D[] == 0
+
+# Test case when unknowns in system do not have equations in initialization system
+using ModelingToolkit, OrdinaryDiffEq, LinearAlgebra
+using ModelingToolkitStandardLibrary.Mechanical.Rotational
+using ModelingToolkitStandardLibrary.Blocks: Add, Sine, PID, SecondOrder, Step, RealOutput
+using ModelingToolkit: connect
+
+# Parameters
+m1 = 1
+m2 = 1
+k = 1000 # Spring stiffness
+c = 10   # Damping coefficient
+@named inertia1 = Inertia(; J = m1)
+@named inertia2 = Inertia(; J = m2)
+@named spring = Spring(; c = k)
+@named damper = Damper(; d = c)
+@named torque = Torque()
+
+function SystemModel(u = nothing; name = :model)
+    eqs = [connect(torque.flange, inertia1.flange_a)
+           connect(inertia1.flange_b, spring.flange_a, damper.flange_a)
+           connect(inertia2.flange_a, spring.flange_b, damper.flange_b)]
+    if u !== nothing
+        push!(eqs, connect(torque.tau, u.output))
+        return ODESystem(eqs, t;
+            systems = [
+                torque,
+                inertia1,
+                inertia2,
+                spring,
+                damper,
+                u
+            ],
+            name)
+    end
+    ODESystem(eqs, t; systems = [torque, inertia1, inertia2, spring, damper], name)
+end
+
+@named r = Step(start_time = 0)
+model = SystemModel()
+@named pid = PID(k = 100, Ti = 0.5, Td = 1)
+@named filt = SecondOrder(d = 0.9, w = 10)
+@named sensor = AngleSensor()
+@named er = Add(k2 = -1)
+
+connections = [connect(r.output, :r, filt.input)
+               connect(filt.output, er.input1)
+               connect(pid.ctr_output, :u, model.torque.tau)
+               connect(model.inertia2.flange_b, sensor.flange)
+               connect(sensor.phi, :y, er.input2)
+               connect(er.output, :e, pid.err_input)]
+
+closed_loop = ODESystem(connections, t, systems = [model, pid, filt, sensor, r, er],
+    name = :closed_loop, defaults = [
+        model.inertia1.phi => 0.0,
+        model.inertia2.phi => 0.0,
+        model.inertia1.w => 0.0,
+        model.inertia2.w => 0.0,
+        filt.x => 0.0,
+        filt.xd => 0.0
+    ])
+
+@test_nowarn linearize(closed_loop, :r, :y)
+
+# https://discourse.julialang.org/t/mtk-change-in-linearize/115760/3
+@mtkmodel Tank_noi begin
+    # Model parameters
+    @parameters begin
+        ρ = 1, [description = "Liquid density"]
+        A = 5, [description = "Cross sectional tank area"]
+        K = 5, [description = "Effluent valve constant"]
+        h_ς = 3, [description = "Scaling level in valve model"]
+    end
+    # Model variables, with initial values needed
+    @variables begin
+        m(t) = 1.5 * ρ * A, [description = "Liquid mass"]
+        md_i(t), [description = "Influent mass flow rate"]
+        md_e(t), [description = "Effluent mass flow rate"]
+        V(t), [description = "Liquid volume"]
+        h(t), [description = "level"]
+    end
+    # Providing model equations
+    @equations begin
+        D(m) ~ md_i - md_e
+        m ~ ρ * V
+        V ~ A * h
+        md_e ~ K * sqrt(h / h_ς)
+    end
+end
+
+@named tank_noi = Tank_noi()
+@unpack md_i, h, m = tank_noi
+m_ss = 2.4000000003229878
+@test_nowarn linearize(tank_noi, [md_i], [h]; op = Dict(m => m_ss, md_i => 2))
+
+# Test initialization
+@variables x(t) y(t) u(t)=1.0
+@parameters p = 1.0
+eqs = [D(x) ~ p * u, x ~ y]
+@named sys = ODESystem(eqs, t)
+
+matrices1, _ = linearize(sys, [u], []; op = Dict(x => 2.0))
+matrices2, _ = linearize(sys, [u], []; op = Dict(y => 2.0))
+@test matrices1 == matrices2
+
+# Ensure parameter values passed as `Dict` are respected
+linfun, _ = linearization_function(sys, [u], []; op = Dict(x => 2.0))
+matrices = linfun([1.0], Dict(p => 3.0), 1.0)
+# this would be 1 if the parameter value isn't respected
+@test matrices.f_u[] == 3.0