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,8 +92,8 @@ Libdl = "1"
 LinearAlgebra = "1"
 MLStyle = "0.4.17"
 NaNMath = "0.3, 1"
+NonlinearSolve = "3.12"
 OrderedCollections = "1"
-OrdinaryDiffEq = "6.73.0"
 PrecompileTools = "1"
 RecursiveArrayTools = "2.3, 3"
 Reexport = "0.2, 1"

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

@@ -1856,8 +1856,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,
@@ -1882,10 +1883,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)