refactor: directly solve initialization problem in linearization_function

Author: AayushSabharwal

Project.toml

@@ -32,8 +32,8 @@ Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -92,6 +92,7 @@ Libdl = "1"
 LinearAlgebra = "1"
 MLStyle = "0.4.17"
 NaNMath = "0.3, 1"
+NonlinearSolve = "3.12"
 OrderedCollections = "1"
 OrdinaryDiffEq = "6.82.0"
 PrecompileTools = "1"
@@ -129,6 +130,7 @@ NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationMOI = "fd9f6733-72f4-499f-8506-86b2bdd0dea1"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
@@ -142,4 +144,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET"]
+test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET"]

src/ModelingToolkit.jl

@@ -46,7 +46,7 @@ using SciMLBase: StandardODEProblem, StandardNonlinearProblem, handle_varmap
 using Distributed
 import JuliaFormatter
 using MLStyle
-import OrdinaryDiffEq
+using NonlinearSolve
 using Reexport
 using RecursiveArrayTools
 import Graphs: SimpleDiGraph, add_edge!, incidence_matrix

src/systems/abstractsystem.jl

@@ -1813,8 +1813,10 @@ function linearization_function(sys::AbstractSystem, inputs,
         op = merge(defs, op)
     end
     sys = ssys
-    initsys = complete(generate_initializesystem(
-        sys, guesses = guesses(sys), algebraic_only = true))
+    initsys = structural_simplify(
+        generate_initializesystem(
+            sys, guesses = guesses(sys), algebraic_only = true),
+        fully_determined = false)
     if p isa SciMLBase.NullParameters
         p = Dict()
     else
@@ -1897,8 +1899,9 @@ function linearization_function(sys::AbstractSystem, inputs,
         sts = unknowns(sys),
         get_initprob_u_p = get_initprob_u_p,
         fun = ODEFunction{true, SciMLBase.FullSpecialize}(
-            sys, unknowns(sys), ps; initializeprobmap = initprobmap),
+            sys, unknowns(sys), ps),
         initfn = initfn,
+        initprobmap = initprobmap,
         h = build_explicit_observed_function(sys, outputs),
         chunk = ForwardDiff.Chunk(input_idxs),
         sys_ps = sys_ps,
@@ -1923,10 +1926,8 @@ function linearization_function(sys::AbstractSystem, inputs,
                     if norm(residual[alge_idxs]) > √(eps(eltype(residual)))
                         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())
-                        u = integ.u
+                        nlsol = solve(initprob)
+                        u = initprobmap(nlsol)
                     end
                 end
                 uf = SciMLBase.UJacobianWrapper(fun, t, p)