feat: preserve the u0map passed to ODEProblem and use it in remake_initializeprob

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -170,16 +170,23 @@ function generate_initializesystem(sys::ODESystem;
     nleqs = Symbolics.substitute.(nleqs, (paramsubs,))
     unks = [full_states; collect(values(paramsubs))]
     u0 = Dict(k => substitute(v, paramsubs) for (k, v) in u0)
+    meta = InitializationSystemMetadata(Dict{Any, Any}(u0map), Dict{Any, Any}(pmap))
     sys_nl = NonlinearSystem(nleqs,
         unks,
         pars;
         defaults = merge(ModelingToolkit.defaults(sys), todict(u0), dd_guess, pmap),
         checks = check_units,
         name,
+        metadata = meta,
         kwargs...)
     return sys_nl
 end
 
+struct InitializationSystemMetadata
+    u0map::Dict{Any, Any}
+    pmap::Dict{Any, Any}
+end
+
 function is_parameter_solvable(p, pmap, defs, guesses)
     _val1 = pmap isa AbstractDict ? get(pmap, p, nothing) : nothing
     _val2 = get(defs, p, nothing)
@@ -267,6 +274,15 @@ function SciMLBase.remake_initializeprob(sys::ODESystem, odefn, u0, t0, p)
           !isempty(setobserved) || !isempty(setparobserved)) &&
          ModelingToolkit.get_tearing_state(sys) !== nothing) ||
         !isempty(initialization_equations(sys)))
+        if SciMLBase.has_initializeprob(odefn)
+            oldsys = odefn.initializeprob.f.sys
+            meta = get_metadata(oldsys)
+            if meta isa InitializationSystemMetadata
+                u0 = merge(meta.u0map, u0)
+                p = merge(meta.pmap, p)
+            end
+        end
+
         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)]

test/initializationsystem.jl

@@ -770,3 +770,19 @@ end
     @test eltype(prob2.f.initializeprob.p.tunable) <: ForwardDiff.Dual
     @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
 end
+
+@testset "`remake` preserves old u0map and pmap" begin
+    @variables x(t) y(t)
+    @parameters p
+    @mtkbuild sys = ODESystem(
+        [D(x) ~ x + p * y, y^2 + 4y * p^2 ~ x], t; guesses = [y => 1.0, p => 1.0])
+    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0])
+    @test is_variable(prob.f.initializeprob, y)
+    prob2 = @test_nowarn remake(prob; p = [p => 3.0]) # ensure no over/under-determined warning
+    @test is_variable(prob.f.initializeprob, y)
+
+    prob = ODEProblem(sys, [y => 1.0, x => 2.0], (0.0, 1.0), [p => missing])
+    @test is_variable(prob.f.initializeprob, p)
+    prob2 = @test_nowarn remake(prob; u0 = [y => 0.5])
+    @test is_variable(prob.f.initializeprob, p)
+end