feat: support update_initializeprob!

Author: AayushSabharwal

src/systems/diffeqs/abstractodesystem.jl

@@ -331,6 +331,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         analytic = nothing,
         split_idxs = nothing,
         initializeprob = nothing,
+        update_initializeprob! = nothing,
         initializeprobmap = nothing,
         initializeprobpmap = nothing,
         kwargs...) where {iip, specialize}
@@ -434,6 +435,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
         analytic = analytic,
         initializeprob = initializeprob,
+        update_initializeprob! = update_initializeprob!,
         initializeprobmap = initializeprobmap,
         initializeprobpmap = initializeprobpmap)
 end
@@ -778,6 +780,17 @@ function (f::GetUpdatedMTKParameters)(prob, initializesol)
     mtkp
 end
 
+struct UpdateInitializeprob{G, S}
+    # `getu` functor which gets all values from prob
+    getvals::G
+    # `setu` functor which updates initializeprob with values
+    setvals::S
+end
+
+function (f::UpdateInitializeprob)(initializeprob, prob)
+    f.setvals(initializeprob, f.getvals(prob))
+end
+
 function get_temporary_value(p)
     stype = symtype(unwrap(p))
     return if stype == Real
@@ -865,6 +878,10 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         getpunknowns = getu(initializeprob, punknowns)
         setpunknowns = setp(sys, punknowns)
         initializeprobpmap = GetUpdatedMTKParameters(getpunknowns, setpunknowns)
+        reqd_syms = vcat(
+            variable_symbols(initializeprob), parameter_symbols(initializeprob))
+        update_initializeprob! = UpdateInitializeprob(
+            getu(sys, reqd_syms), setu(initializeprob, reqd_syms))
 
         zerovars = Dict(setdiff(unknowns(sys), keys(defaults(sys))) .=> 0.0)
         if parammap isa SciMLBase.NullParameters
@@ -880,6 +897,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
             (trueinit = SVector{length(trueinit)}(trueinit))
     else
         initializeprob = nothing
+        update_initializeprob! = nothing
         initializeprobmap = nothing
         initializeprobpmap = nothing
         trueinit = u0map
@@ -929,6 +947,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         sparse = sparse, eval_expression = eval_expression,
         eval_module = eval_module,
         initializeprob = initializeprob,
+        update_initializeprob! = update_initializeprob!,
         initializeprobmap = initializeprobmap,
         initializeprobpmap = initializeprobpmap,
         kwargs...)

src/systems/nonlinear/initializesystem.jl

@@ -201,5 +201,8 @@ function SciMLBase.remake_initializeprob(sys::ODESystem, odefn, u0, t0, p)
     getpunknowns = getu(initprob, punknowns)
     setpunknowns = setp(sys, punknowns)
     initprobpmap = GetUpdatedMTKParameters(getpunknowns, setpunknowns)
-    return initprob, initprobmap, initprobpmap
+    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
 end

test/initializationsystem.jl

@@ -690,3 +690,22 @@ end
     prob5 = remake(prob)
     @test init(prob, Tsit5()).ps[p] ≈ 2.0
 end
+
+@testset "Update initializeprob parameters" begin
+    @variables x(t) y(t)
+    @parameters p q
+    @mtkbuild sys = ODESystem(
+        [D(x) ~ x, p ~ x + y], t; guesses = [x => 0.0, p => 0.0])
+    prob = ODEProblem(sys, [y => 1.0], (0.0, 1.0), [p => 3.0])
+    @test prob.f.initializeprob.ps[p] ≈ 3.0
+    @test init(prob, Tsit5())[x] ≈ 2.0
+    prob.ps[p] = 2.0
+    @test prob.f.initializeprob.ps[p] ≈ 3.0
+    @test init(prob, Tsit5())[x] ≈ 1.0
+    ModelingToolkit.defaults(prob.f.sys)[p] = missing
+    prob2 = remake(prob; u0 = [y => 1.0], p = [p => 3x])
+    @test !is_variable(prob2.f.initializeprob, p) &&
+          !is_parameter(prob2.f.initializeprob, p)
+    @test init(prob2, Tsit5())[x] ≈ 0.5
+    @test_nowarn solve(prob2, Tsit5())
+end