fix: properly initialize initialization problem in linearization_function

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -1774,18 +1774,87 @@ function linearization_function(sys::AbstractSystem, inputs,
     sys = ssys
     initsys = complete(generate_initializesystem(
         sys, guesses = guesses(sys), algebraic_only = true))
+    if p isa SciMLBase.NullParameters
+        p = Dict()
+    else
+        p = todict(p)
+    end
+    p[get_iv(sys)] = 0.0
+    if has_index_cache(initsys) && get_index_cache(initsys) !== nothing
+        oldps = MTKParameters(initsys, p, merge(guesses(sys), defaults(sys), op))
+        initsys_ps = parameters(initsys)
+        initsys_idxs = [parameter_index(initsys, param) for param in initsys_ps]
+        tunable_ps = [initsys_ps[i]
+                      for i in eachindex(initsys_ps)
+                      if initsys_idxs[i].portion == SciMLStructures.Tunable()]
+        tunable_getter = isempty(tunable_ps) ? nothing : getu(sys, tunable_ps)
+        discrete_ps = [initsys_ps[i]
+                       for i in eachindex(initsys_ps)
+                       if initsys_idxs[i].portion == SciMLStructures.Discrete()]
+        disc_getter = isempty(discrete_ps) ? nothing : getu(sys, discrete_ps)
+        constant_ps = [initsys_ps[i]
+                       for i in eachindex(initsys_ps)
+                       if initsys_idxs[i].portion == SciMLStructures.Constants()]
+        const_getter = isempty(constant_ps) ? nothing : getu(sys, constant_ps)
+        nonnum_ps = [initsys_ps[i]
+                     for i in eachindex(initsys_ps)
+                     if initsys_idxs[i].portion == NONNUMERIC_PORTION]
+        nonnum_getter = isempty(nonnum_ps) ? nothing : getu(sys, nonnum_ps)
+        u_getter = isempty(unknowns(initsys)) ? (_...) -> nothing :
+                   getu(sys, unknowns(initsys))
+        get_initprob_u_p = let tunable_getter = tunable_getter,
+            disc_getter = disc_getter,
+            const_getter = const_getter,
+            nonnum_getter = nonnum_getter,
+            oldps = oldps,
+            u_getter = u_getter
+
+            function (u, p, t)
+                state = ProblemState(; u, p, t)
+                if tunable_getter !== nothing
+                    oldps = SciMLStructures.replace!(
+                        SciMLStructures.Tunable(), oldps, tunable_getter(state))
+                end
+                if disc_getter !== nothing
+                    oldps = SciMLStructures.replace!(
+                        SciMLStructures.Discrete(), oldps, disc_getter(state))
+                end
+                if const_getter !== nothing
+                    oldps = SciMLStructures.replace!(
+                        SciMLStructures.Constants(), oldps, const_getter(state))
+                end
+                if nonnum_getter !== nothing
+                    oldps = SciMLStructures.replace!(
+                        NONNUMERIC_PORTION, oldps, nonnum_getter(state))
+                end
+                newu = u_getter(state)
+                return newu, oldps
+            end
+        end
+    else
+        get_initprob_u_p = let p_getter = getu(sys, parameters(initsys)),
+            u_getter = getu(sys, unknowns(initsys))
+
+            function (u, p, t)
+                state = ProblemState(; u, p, t)
+                return u_getter(state), p_getter(state)
+            end
+        end
+    end
     initfn = NonlinearFunction(initsys)
     initprobmap = getu(initsys, unknowns(sys))
     ps = parameters(sys)
     lin_fun = let diff_idxs = diff_idxs,
         alge_idxs = alge_idxs,
         input_idxs = input_idxs,
         sts = unknowns(sys),
+        get_initprob_u_p = get_initprob_u_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)
+        chunk = ForwardDiff.Chunk(input_idxs),
+        initialize = initialize
 
         function (u, p, t)
             if u !== nothing # Handle systems without unknowns
@@ -1794,7 +1863,8 @@ function linearization_function(sys::AbstractSystem, inputs,
                 if initialize && !isempty(alge_idxs) # This is expensive and can be omitted if the user knows that the system is already initialized
                     residual = fun(u, p, t)
                     if norm(residual[alge_idxs]) > √(eps(eltype(residual)))
-                        initprob = NonlinearLeastSquaresProblem(initfn, u, p)
+                        initu0, initp = get_initprob_u_p(u, p, t)
+                        initprob = NonlinearLeastSquaresProblem(initfn, initu0, initp)
                         @set! fun.initializeprob = initprob
                         prob = ODEProblem(fun, u, (t, t + 1), p)
                         integ = init(prob, OrdinaryDiffEq.Rodas5P())

test/downstream/linearization_dd.jl

@@ -0,0 +1,62 @@
+## Test that dummy_derivatives can be set to zero
+# The call to Link(; m = 0.2, l = 10, I = 1, g = -9.807) hangs forever on Julia v1.6
+using LinearAlgebra
+using ModelingToolkit
+using ModelingToolkitStandardLibrary
+using ModelingToolkitStandardLibrary.Blocks
+using ModelingToolkitStandardLibrary.Mechanical.MultiBody2D
+using ModelingToolkitStandardLibrary.Mechanical.TranslationalPosition
+using Test
+
+using ControlSystemsMTK
+using ControlSystemsMTK.ControlSystemsBase: sminreal, minreal, poles
+connect = ModelingToolkit.connect
+
+@parameters t
+D = Differential(t)
+
+@named link1 = Link(; m = 0.2, l = 10, I = 1, g = -9.807)
+@named cart = TranslationalPosition.Mass(; m = 1, s = 0)
+@named fixed = Fixed()
+@named force = Force(use_support = false)
+
+eqs = [connect(link1.TX1, cart.flange)
+       connect(cart.flange, force.flange)
+       connect(link1.TY1, fixed.flange)]
+
+@named model = ODESystem(eqs, t, [], []; systems = [link1, cart, force, fixed])
+def = ModelingToolkit.defaults(model)
+def[cart.s] = 10
+def[cart.v] = 0
+def[link1.A] = -pi / 2
+def[link1.dA] = 0
+lin_outputs = [cart.s, cart.v, link1.A, link1.dA]
+lin_inputs = [force.f.u]
+
+@info "named_ss"
+G = named_ss(model, lin_inputs, lin_outputs, allow_symbolic = true, op = def,
+    allow_input_derivatives = true, zero_dummy_der = true)
+G = sminreal(G)
+@info "minreal"
+G = minreal(G)
+@info "poles"
+ps = poles(G)
+
+@test minimum(abs, ps) < 1e-6
+@test minimum(abs, complex(0, 1.3777260367206716) .- ps) < 1e-10
+
+lsys, syss = linearize(model, lin_inputs, lin_outputs, allow_symbolic = true, op = def,
+    allow_input_derivatives = true, zero_dummy_der = true)
+lsyss, sysss = ModelingToolkit.linearize_symbolic(model, lin_inputs, lin_outputs;
+    allow_input_derivatives = true)
+
+dummyder = setdiff(unknowns(sysss), unknowns(model))
+def = merge(ModelingToolkit.guesses(model), def, Dict(x => 0.0 for x in dummyder))
+def[link1.fy1] = -def[link1.g] * def[link1.m]
+
+@test substitute(lsyss.A, def) ≈ lsys.A
+# We cannot pivot symbolically, so the part where a linear solve is required
+# is not reliable.
+@test substitute(lsyss.B, def)[1:6, 1] ≈ lsys.B[1:6, 1]
+@test substitute(lsyss.C, def) ≈ lsys.C
+@test substitute(lsyss.D, def) ≈ lsys.D

test/downstream/linearize.jl

@@ -195,67 +195,3 @@ lsys, ssys = linearize(sat, [u], [y]; op = Dict(u => 2))
 @test isempty(lsys.B)
 @test isempty(lsys.C)
 @test lsys.D[] == 0
-
-## Test that dummy_derivatives can be set to zero
-if VERSION >= v"1.8"
-    # The call to Link(; m = 0.2, l = 10, I = 1, g = -9.807) hangs forever on Julia v1.6
-    using LinearAlgebra
-    using ModelingToolkit
-    using ModelingToolkitStandardLibrary
-    using ModelingToolkitStandardLibrary.Blocks
-    using ModelingToolkitStandardLibrary.Mechanical.MultiBody2D
-    using ModelingToolkitStandardLibrary.Mechanical.TranslationalPosition
-
-    using ControlSystemsMTK
-    using ControlSystemsMTK.ControlSystemsBase: sminreal, minreal, poles
-    connect = ModelingToolkit.connect
-
-    @parameters t
-    D = Differential(t)
-
-    @named link1 = Link(; m = 0.2, l = 10, I = 1, g = -9.807)
-    @named cart = TranslationalPosition.Mass(; m = 1, s = 0)
-    @named fixed = Fixed()
-    @named force = Force(use_support = false)
-
-    eqs = [connect(link1.TX1, cart.flange)
-           connect(cart.flange, force.flange)
-           connect(link1.TY1, fixed.flange)]
-
-    @named model = ODESystem(eqs, t, [], []; systems = [link1, cart, force, fixed])
-    def = ModelingToolkit.defaults(model)
-    def[cart.s] = 10
-    def[cart.v] = 0
-    def[link1.A] = -pi / 2
-    def[link1.dA] = 0
-    lin_outputs = [cart.s, cart.v, link1.A, link1.dA]
-    lin_inputs = [force.f.u]
-
-    @info "named_ss"
-    G = named_ss(model, lin_inputs, lin_outputs, allow_symbolic = true, op = def,
-        allow_input_derivatives = true, zero_dummy_der = true)
-    G = sminreal(G)
-    @info "minreal"
-    G = minreal(G)
-    @info "poles"
-    ps = poles(G)
-
-    @test minimum(abs, ps) < 1e-6
-    @test minimum(abs, complex(0, 1.3777260367206716) .- ps) < 1e-10
-
-    lsys, syss = linearize(model, lin_inputs, lin_outputs, allow_symbolic = true, op = def,
-        allow_input_derivatives = true, zero_dummy_der = true)
-    lsyss, sysss = ModelingToolkit.linearize_symbolic(model, lin_inputs, lin_outputs;
-        allow_input_derivatives = true)
-
-    dummyder = setdiff(unknowns(sysss), unknowns(model))
-    def = merge(def, Dict(x => 0.0 for x in dummyder))
-    def[link1.fy1] = -def[link1.g] * def[link1.m]
-
-    @test substitute(lsyss.A, def) ≈ lsys.A
-    # We cannot pivot symbolically, so the part where a linear solve is required
-    # is not reliable.
-    @test substitute(lsyss.B, def)[1:6, 1] ≈ lsys.B[1:6, 1]
-    @test substitute(lsyss.C, def) ≈ lsys.C
-    @test substitute(lsyss.D, def) ≈ lsys.D
-end

test/runtests.jl

@@ -92,6 +92,7 @@ end
     if GROUP == "All" || GROUP == "Downstream"
         activate_downstream_env()
         @safetestset "Linearization Tests" include("downstream/linearize.jl")
+        @safetestset "Linearization Dummy Derivative Tests" include("downstream/linearization_dd.jl")
         @safetestset "Inverse Models Test" include("downstream/inversemodel.jl")
     end