refactor: improve remake_initializeprob

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -178,10 +178,13 @@ end
 function SciMLBase.remake_initializeprob(sys::ODESystem, odefn, u0, t0, p)
     if (u0 === missing || !(eltype(u0) <: Pair) || isempty(u0)) &&
        (p === missing || !(eltype(p) <: Pair) || isempty(p))
-        return odefn.initializeprob, odefn.initializeprobmap, odefn.initializeprobpmap
+        return odefn.initializeprob, odefn.update_initializeprob!, odefn.initializeprobmap,
+        odefn.initializeprobpmap
     end
-    if u0 === missing
+    if u0 === missing || isempty(u0)
         u0 = Dict()
+    elseif !(eltype(u0) <: Pair)
+        u0 = Dict(unknowns(sys) .=> u0)
     end
     if p === missing
         p = Dict()
@@ -190,15 +193,33 @@ function SciMLBase.remake_initializeprob(sys::ODESystem, odefn, u0, t0, p)
         t0 = 0.0
     end
     u0 = todict(u0)
+    defs = defaults(sys)
+    varmap = merge(defs, u0)
+    varmap = canonicalize_varmap(varmap)
+    missingvars = setdiff(unknowns(sys), collect(keys(varmap)))
+    setobserved = filter(keys(varmap)) do var
+        has_observed_with_lhs(sys, var) || has_observed_with_lhs(sys, default_toterm(var))
+    end
     p = todict(p)
-    initprob = InitializationProblem(sys, t0, u0, p)
-    initprobmap = getu(initprob, unknowns(sys))
-    punknowns = [p for p in all_variable_symbols(initprob) if is_parameter(sys, p)]
-    getpunknowns = getu(initprob, punknowns)
-    setpunknowns = setp(sys, punknowns)
-    initprobpmap = GetUpdatedMTKParameters(getpunknowns, setpunknowns)
-    reqd_syms = vcat(variable_symbols(initprob), parameter_symbols(initprob))
-    update_initializeprob! = UpdateInitializeprob(
-        getu(sys, reqd_syms), setu(initprob, reqd_syms))
-    return initprob, update_initializeprob!, initprobmap, initprobpmap
+    guesses = ModelingToolkit.guesses(sys)
+    solvablepars = [par
+                    for par in parameters(sys)
+                    if is_parameter_solvable(par, p, defs, guesses)]
+    if (((!isempty(missingvars) || !isempty(solvablepars) ||
+          !isempty(setobserved)) &&
+         ModelingToolkit.get_tearing_state(sys) !== nothing) ||
+        !isempty(initialization_equations(sys)))
+        initprob = InitializationProblem(sys, t0, u0, p)
+        initprobmap = getu(initprob, unknowns(sys))
+        punknowns = [p for p in all_variable_symbols(initprob) if is_parameter(sys, p)]
+        getpunknowns = getu(initprob, punknowns)
+        setpunknowns = setp(sys, punknowns)
+        initprobpmap = GetUpdatedMTKParameters(getpunknowns, setpunknowns)
+        reqd_syms = parameter_symbols(initprob)
+        update_initializeprob! = UpdateInitializeprob(
+            getu(sys, reqd_syms), setu(initprob, reqd_syms))
+        return initprob, update_initializeprob!, initprobmap, initprobpmap
+    else
+        return nothing, nothing, nothing, nothing
+    end
 end