feat: propagate ODEProblem guesses to remake

Author: AayushSabharwal

src/systems/diffeqs/abstractodesystem.jl

@@ -1310,11 +1310,11 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
     elseif isempty(u0map) && get_initializesystem(sys) === nothing
         isys = structural_simplify(
             generate_initializesystem(
-                sys; initialization_eqs, check_units, pmap = parammap); fully_determined)
+                sys; initialization_eqs, check_units, pmap = parammap, guesses); fully_determined)
     else
         isys = structural_simplify(
             generate_initializesystem(
-                sys; u0map, initialization_eqs, check_units, pmap = parammap); fully_determined)
+                sys; u0map, initialization_eqs, check_units, pmap = parammap, guesses); fully_determined)
     end
 
     ts = get_tearing_state(isys)

src/systems/nonlinear/initializesystem.jl

@@ -30,7 +30,8 @@ function generate_initializesystem(sys::ODESystem;
     # 1) process dummy derivatives and u0map into initialization system
     eqs_ics = eqs[idxs_alge] # start equation list with algebraic equations
     defs = copy(defaults(sys)) # copy so we don't modify sys.defaults
-    guesses = merge(get_guesses(sys), todict(guesses))
+    additional_guesses = anydict(guesses)
+    guesses = merge(get_guesses(sys), additional_guesses)
     schedule = getfield(sys, :schedule)
     if !isnothing(schedule)
         for x in filter(x -> !isnothing(x[1]), schedule.dummy_sub)
@@ -178,7 +179,7 @@ function generate_initializesystem(sys::ODESystem;
     for k in keys(defs)
         defs[k] = substitute(defs[k], paramsubs)
     end
-    meta = InitializationSystemMetadata(Dict{Any, Any}(u0map), Dict{Any, Any}(pmap))
+    meta = InitializationSystemMetadata(anydict(u0map), anydict(pmap), additional_guesses)
     return NonlinearSystem(eqs_ics,
         vars,
         pars;
@@ -193,6 +194,7 @@ end
 struct InitializationSystemMetadata
     u0map::Dict{Any, Any}
     pmap::Dict{Any, Any}
+    additional_guesses::Dict{Any, Any}
 end
 
 function is_parameter_solvable(p, pmap, defs, guesses)
@@ -263,75 +265,29 @@ function SciMLBase.remake_initializeprob(sys::ODESystem, odefn, u0, t0, p)
         return initprob, odefn.update_initializeprob!, odefn.initializeprobmap,
         odefn.initializeprobpmap
     end
-    if u0 === missing || isempty(u0)
-        u0 = Dict()
-    elseif !(eltype(u0) <: Pair)
-        u0 = Dict(unknowns(sys) .=> u0)
-    end
-    if p === missing
-        p = Dict()
+    dvs = unknowns(sys)
+    ps = parameters(sys)
+    u0map = to_varmap(u0, dvs)
+    pmap = to_varmap(p, ps)
+    guesses = Dict()
+    if SciMLBase.has_initializeprob(odefn)
+        oldsys = odefn.initializeprob.f.sys
+        meta = get_metadata(oldsys)
+        if meta isa InitializationSystemMetadata
+            u0map = merge(meta.u0map, u0map)
+            pmap = merge(meta.pmap, pmap)
+            merge!(guesses, meta.additional_guesses)
+        end
     end
     if t0 === nothing
         t0 = 0.0
     end
-    u0 = todict(u0)
-    defs = defaults(sys)
-    varmap = merge(defs, u0)
-    for k in collect(keys(varmap))
-        if varmap[k] === nothing
-            delete!(varmap, k)
-        end
-    end
-    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)
-    guesses = ModelingToolkit.guesses(sys)
-    solvablepars = [par
-                    for par in parameters(sys)
-                    if is_parameter_solvable(par, p, defs, guesses)]
-    pvarmap = merge(defs, p)
-    setparobserved = filter(keys(pvarmap)) do var
-        has_parameter_dependency_with_lhs(sys, var)
-    end
-    if (((!isempty(missingvars) || !isempty(solvablepars) ||
-          !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
-        for k in collect(keys(u0))
-            if u0[k] === nothing
-                delete!(u0, k)
-            end
-        end
-        for k in collect(keys(p))
-            if p[k] === nothing
-                delete!(p, k)
-            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)]
-        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
+    filter_missing_values!(u0map)
+    filter_missing_values!(pmap)
+    f, _ = process_SciMLProblem(EmptySciMLFunction, sys, u0map, pmap; guesses, t = t0)
+    kws = f.kwargs
+    return get(kws, :initializeprob, nothing), get(kws, :update_initializeprob!, nothing), get(kws, :initializeprobmap, nothing),
+        get(kws, :initializeprobpmap, nothing)
 end
 
 """

src/systems/problem_utils.jl

@@ -4,12 +4,13 @@ const AnyDict = Dict{Any, Any}
     $(TYPEDSIGNATURES)
 
 If called without arguments, return `Dict{Any, Any}`. Otherwise, interpret the input
-as a symbolic map and turn it into a `Dict{Any, Any}`. Handles `SciMLBase.NullParameters`
-and `nothing`.
+as a symbolic map and turn it into a `Dict{Any, Any}`. Handles `SciMLBase.NullParameters`,
+`missing` and `nothing`.
 """
 anydict() = AnyDict()
 anydict(::SciMLBase.NullParameters) = AnyDict()
 anydict(::Nothing) = AnyDict()
+anydict(::Missing) = AnyDict()
 anydict(x::AnyDict) = x
 anydict(x) = AnyDict(x)
 
@@ -388,6 +389,15 @@ function evaluate_varmap!(varmap::AbstractDict, vars; limit = 100)
     end
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Remove keys in `varmap` whose values are `nothing`.
+"""
+function filter_missing_values!(varmap::AbstractDict)
+    filter!(kvp -> kvp[2] !== nothing, varmap)
+end
+
 struct GetUpdatedMTKParameters{G, S}
     # `getu` functor which gets parameters that are unknowns during initialization
     getpunknowns::G

test/initializationsystem.jl

@@ -975,3 +975,26 @@ end
     @test integ.ps[p] ≈ 1.0
     @test integ.ps[q]≈cbrt(2) rtol=1e-6
 end
+
+@testset "Guesses provided to `ODEProblem` are used in `remake`" begin
+    @variables x(t) y(t)=2x
+    @parameters p q=3x
+    @mtkbuild sys = ODESystem([D(x) ~ x * p + q, x^3 + y^3 ~ 3], t)
+    prob = ODEProblem(
+        sys, [], (0.0, 1.0), [p => 1.0]; guesses = [x => 1.0, y => 1.0, q => 1.0])
+    @test prob[x] == 0.0
+    @test prob[y] == 0.0
+    @test prob.ps[p] == 1.0
+    @test prob.ps[q] == 0.0
+    integ = init(prob)
+    @test integ[x] ≈ 1 / cbrt(3)
+    @test integ[y] ≈ 2 / cbrt(3)
+    @test integ.ps[p] == 1.0
+    @test integ.ps[q] ≈ 3 / cbrt(3)
+    prob2 = remake(prob; u0 = [y => 3x], p = [q => 2x])
+    integ2 = init(prob2)
+    @test integ2[x] ≈ cbrt(3 / 28)
+    @test integ2[y] ≈ 3cbrt(3 / 28)
+    @test integ2.ps[p] == 1.0
+    @test integ2.ps[q] ≈ 2cbrt(3 / 28)
+end