refactor: improve type-stability of maybe_build_initialization_problem

Author: AayushSabharwal

lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl

@@ -633,8 +633,8 @@ function SciMLBase.remake_initialization_data(
     u0_constructor = get_u0_constructor(identity, typeof(newu0), floatT, false)
     p_constructor = get_p_constructor(identity, typeof(newu0), floatT)
     kws = maybe_build_initialization_problem(
-        sys, SciMLBase.isinplace(odefn), op, t0, guesses;
-        time_dependent_init, use_scc, initialization_eqs, floatT, fast_path = true,
+        sys, Val{SciMLBase.isinplace(odefn)}(), op, t0, guesses, floatT;
+        time_dependent_init, use_scc, initialization_eqs, fast_path = true,
         u0_constructor, p_constructor, allow_incomplete = true, check_units = false,
         missing_guess_value = meta.missing_guess_value
     )

lib/ModelingToolkitBase/src/systems/problem_utils.jl

@@ -1088,16 +1088,15 @@ struct GetUpdatedU0{GG, GIU}
     get_initial_unknowns::GIU
 end
 
-function GetUpdatedU0(sys::AbstractSystem, initprob::SciMLBase.AbstractNonlinearProblem, op::AbstractDict)
-    @nospecialize initprob
+function GetUpdatedU0(sys::AbstractSystem, initsys::AbstractSystem, op::AbstractDict)
     dvs = unknowns(sys)
     eqs = equations(sys)
     guessvars = trues(length(dvs))
     for (i, var) in enumerate(dvs)
         varval = get(op, var, COMMON_NOTHING)
         guessvars[i] = varval === COMMON_NOTHING || !SU.isconst(varval)
     end
-    get_guessvars = getu(initprob, dvs[guessvars])
+    get_guessvars = getu(initsys, dvs[guessvars])
     get_initial_unknowns = getu(sys, Initial.(dvs))
     return GetUpdatedU0(guessvars, get_guessvars, get_initial_unknowns)
 end
@@ -1170,14 +1169,14 @@ constructed is in implicit DAE form (`DAEProblem`). All other keyword arguments
 to `InitializationProblem`.
 """
 function maybe_build_initialization_problem(
-        sys::AbstractSystem, iip, op::AbstractDict, t, guesses;
+        sys::AbstractSystem, ::Val{iip}, op::SymmapT, t, guesses, ::Type{floatT};
         time_dependent_init = is_time_dependent(sys), u0_constructor = identity,
-        p_constructor = identity, floatT = Float64, initialization_eqs = [],
+        p_constructor = identity, initialization_eqs = [],
         use_scc = true, eval_expression = false, eval_module = @__MODULE__,
         missing_guess_value = default_missing_guess_value(),
         # Intercept `expression` because we don't support it here yet
         implicit_dae = false, is_steadystateprob = false, expression = Val{false}, kwargs...
-    )
+    ) where {iip, floatT}
     guesses = merge(ModelingToolkitBase.guesses(sys), todict(guesses))
 
     if t === nothing && is_time_dependent(sys)
@@ -1190,6 +1189,7 @@ function maybe_build_initialization_problem(
         use_scc, u0_constructor, p_constructor, eval_expression, eval_module,
         missing_guess_value, is_steadystateprob, kwargs...
     )
+    initsys = initializeprob.f.sys::System
     needs_remake = false
     _u0 = state_values(initializeprob)
      if _u0 !== nothing
@@ -1236,7 +1236,7 @@ function maybe_build_initialization_problem(
     end
 
     get_initial_unknowns = if time_dependent_init
-        GetUpdatedU0(sys, initializeprob, op)
+        GetUpdatedU0(sys, initsys, op)
     else
         nothing
     end
@@ -1246,7 +1246,7 @@ function maybe_build_initialization_problem(
         Vector{Equation}(initialization_eqs),
         use_scc, time_dependent_init,
         ReconstructInitializeprob(
-            sys, initializeprob.f.sys; u0_constructor,
+            sys, initsys; u0_constructor,
             p_constructor, eval_expression, eval_module, is_steadystateprob, kwargs...
         ),
         get_initial_unknowns, SetInitialUnknowns(sys), missing_guess_value
@@ -1274,7 +1274,7 @@ function maybe_build_initialization_problem(
         initializeprobpmap = nothing
     else
         initializeprobpmap = construct_initializeprobpmap(
-            sys, initializeprob.f.sys; p_constructor, eval_expression, eval_module, kwargs...
+            sys, initsys; p_constructor, eval_expression, eval_module, kwargs...
         )
     end
 
@@ -1302,7 +1302,6 @@ function maybe_build_initialization_problem(
             end
         end
         if implicit_dae
-            initsys = initializeprob.f.sys
             for v in unknowns(sys)
                 v = Differential(get_iv(sys))(v)
                 ttv = default_toterm(v)
@@ -1328,10 +1327,10 @@ function maybe_build_initialization_problem(
     end
     missingvars = collect(missingvars)
 
-    for (i, v) in enumerate(unknowns(initializeprob.f.sys))
+    for (i, v) in enumerate(unknowns(initsys))
         write_possibly_indexed_array!(temp_op, v, SConst(_u0[i]), COMMON_NOTHING)
     end
-    add_observed!(initializeprob.f.sys, temp_op)
+    add_observed!(initsys, temp_op)
     left_merge!(temp_op, ModelingToolkitBase.guesses(sys))
     subber = Symbolics.FixpointSubstituter{true}(AADSubWrapper(temp_op))
     for p in missingvars