fix: create and solve initialization problem in linearization_function

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -1756,24 +1756,18 @@ function linearization_function(sys::AbstractSystem, inputs,
         op = merge(defs, op)
     end
     sys = ssys
-    x0 = merge(defaults(sys), Dict(missing_variable_defaults(sys)), op)
-    u0, _p, _ = get_u0_p(sys, x0, p; use_union = false, tofloat = true)
+    initsys = complete(generate_initializesystem(
+        sys, guesses = guesses(sys), algebraic_only = true))
+    initfn = NonlinearFunction(initsys)
+    initprobmap = getu(initsys, unknowns(sys))
     ps = parameters(sys)
-    if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        p = MTKParameters(sys, p, u0)
-    else
-        p = _p
-        p, split_idxs = split_parameters_by_type(p)
-        if p isa Tuple
-            ps = Base.Fix1(getindex, ps).(split_idxs)
-            ps = (ps...,) #if p is Tuple, ps should be Tuple
-        end
-    end
     lin_fun = let diff_idxs = diff_idxs,
         alge_idxs = alge_idxs,
         input_idxs = input_idxs,
         sts = unknowns(sys),
-        fun = ODEFunction{true, SciMLBase.FullSpecialize}(sys, unknowns(sys), ps; p = p),
+        fun = ODEFunction{true, SciMLBase.FullSpecialize}(
+            sys, unknowns(sys), ps; initializeprobmap = initprobmap),
+        initfn = initfn,
         h = build_explicit_observed_function(sys, outputs),
         chunk = ForwardDiff.Chunk(input_idxs)
 
@@ -1782,6 +1776,8 @@ function linearization_function(sys::AbstractSystem, inputs,
                 length(sts) == length(u) ||
                     error("Number of unknown variables ($(length(sts))) does not match the number of input unknowns ($(length(u)))")
                 if initialize && !isempty(alge_idxs) # This is expensive and can be omitted if the user knows that the system is already initialized
+                    initprob = NonlinearLeastSquaresProblem(initfn, u, p)
+                    @set! fun.initializeprob = initprob
                     residual = fun(u, p, t)
                     if norm(residual[alge_idxs]) > √(eps(eltype(residual)))
                         prob = ODEProblem(fun, u, (t, t + 1), p)
@@ -2051,8 +2047,8 @@ lsys_sym, _ = ModelingToolkit.linearize_symbolic(cl, [f.u], [p.x])
 """
 function linearize(sys, lin_fun; t = 0.0, op = Dict(), allow_input_derivatives = false,
         p = DiffEqBase.NullParameters())
-    x0 = merge(defaults(sys), op)
-    u0, p2, _ = get_u0_p(sys, x0, p; use_union = false, tofloat = true)
+    x0 = merge(defaults(sys), Dict(missing_variable_defaults(sys)), op)
+    u0, defs = get_u0(sys, x0, p)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         if p isa SciMLBase.NullParameters
             p = op
@@ -2063,7 +2059,7 @@ function linearize(sys, lin_fun; t = 0.0, op = Dict(), allow_input_derivatives =
         elseif p isa Vector
             p = merge(Dict(parameters(sys) .=> p), op)
         end
-        p2 = MTKParameters(sys, p, Dict(unknowns(sys) .=> u0))
+        p2 = MTKParameters(sys, p, merge(Dict(unknowns(sys) .=> u0), x0, guesses(sys)))
     end
     linres = lin_fun(u0, p2, t)
     f_x, f_z, g_x, g_z, f_u, g_u, h_x, h_z, h_u = linres

src/systems/nonlinear/initializesystem.jl

@@ -8,6 +8,7 @@ function generate_initializesystem(sys::ODESystem;
         name = nameof(sys),
         guesses = Dict(), check_defguess = false,
         default_dd_value = 0.0,
+        algebraic_only = false,
         kwargs...)
     sts, eqs = unknowns(sys), equations(sys)
     idxs_diff = isdiffeq.(eqs)
@@ -68,28 +69,34 @@ function generate_initializesystem(sys::ODESystem;
     defs = merge(defaults(sys), filtered_u0)
     guesses = merge(get_guesses(sys), todict(guesses), dd_guess)
 
-    for st in full_states
-        if st ∈ keys(defs)
-            def = defs[st]
+    if !algebraic_only
+        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)
+                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)
                 push!(u0, st => guesses[st])
-            else
-                push!(eqs_ics, st ~ def)
-                push!(u0, st => def)
+            elseif check_defguess
+                error("Invalid setup: unknown $(st) has no default value or initial guess")
             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
 
     pars = [parameters(sys); get_iv(sys)]
-    nleqs = [eqs_ics; get_initialization_eqs(sys); observed(sys)]
+    nleqs = if algebraic_only
+        [eqs_ics; observed(sys)]
+    else
+        [eqs_ics; get_initialization_eqs(sys); observed(sys)]
+    end
 
     sys_nl = NonlinearSystem(nleqs,
         full_states,

test/downstream/linearize.jl

@@ -120,10 +120,14 @@ lsys = ModelingToolkit.reorder_unknowns(lsys0, unknowns(ssys), desired_order)
 lsyss, _ = ModelingToolkit.linearize_symbolic(pid, [reference.u, measurement.u],
     [ctr_output.u])
 
-@test substitute(lsyss.A, ModelingToolkit.defaults(pid)) == lsys.A
-@test substitute(lsyss.B, ModelingToolkit.defaults(pid)) == lsys.B
-@test substitute(lsyss.C, ModelingToolkit.defaults(pid)) == lsys.C
-@test substitute(lsyss.D, ModelingToolkit.defaults(pid)) == lsys.D
+@test substitute(
+    lsyss.A, merge(ModelingToolkit.defaults(pid), ModelingToolkit.guesses(pid))) == lsys.A
+@test substitute(
+    lsyss.B, merge(ModelingToolkit.defaults(pid), ModelingToolkit.guesses(pid))) == lsys.B
+@test substitute(
+    lsyss.C, merge(ModelingToolkit.defaults(pid), ModelingToolkit.guesses(pid))) == lsys.C
+@test substitute(
+    lsyss.D, merge(ModelingToolkit.defaults(pid), ModelingToolkit.guesses(pid))) == lsys.D
 
 # Test with the reverse desired unknown order as well to verify that similarity transform and reoreder_unknowns really works
 lsys = ModelingToolkit.reorder_unknowns(lsys, unknowns(ssys), reverse(desired_order))

test/input_output_handling.jl

@@ -144,7 +144,9 @@ if VERSION >= v"1.8" # :opaque_closure not supported before
             drop_expr = identity)
         x = randn(size(A, 1))
         u = randn(size(B, 2))
-        p = getindex.(Ref(ModelingToolkit.defaults(ssys)), parameters(ssys))
+        p = getindex.(
+            Ref(merge(ModelingToolkit.defaults(ssys), ModelingToolkit.guesses(ssys))),
+            parameters(ssys))
         y1 = obsf(x, u, p, 0)
         y2 = C * x + D * u
         @test y1[] ≈ y2[]