Add reinitializationmaps

Companion PR to SciML/SciMLBase.jl#641, but after implementing it I'm not so sure it's a good idea since you may not want to reinit with the system conditions.
SciML · Mar 1, 2024 · c979fd7 · c979fd7
1 parent ef8171a
commit c979fd7
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -323,6 +323,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         split_idxs = nothing,
         initializeprob = nothing,
         initializeprobmap = nothing,
+        reinitializemap = nothing,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEFunction`")
@@ -505,7 +506,8 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
         analytic = analytic,
         initializeprob = initializeprob,
-        initializeprobmap = initializeprobmap)
+        initializeprobmap = initializeprobmap
+        reinitializemap = reinitializemap)
 end
 
 """
@@ -537,6 +539,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         checkbounds = false,
         initializeprob = nothing,
         initializeprobmap = nothing,
+        reinitializemap = nothing,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DAEFunction`")
@@ -610,7 +613,8 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         jac_prototype = jac_prototype,
         observed = observedfun,
         initializeprob = initializeprob,
-        initializeprobmap = initializeprobmap)
+        initializeprobmap = initializeprobmap,
+        reinitializemap = reinitializemap)
 end
 
 function DiffEqBase.DDEFunction(sys::AbstractODESystem, args...; kwargs...)
@@ -878,12 +882,14 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         initializeprob = ModelingToolkit.InitializationProblem(
             sys, t, u0map, parammap; guesses, warn_initialize_determined)
         initializeprobmap = getu(initializeprob, unknowns(sys))
+        reinitializemap = setp(initializeprob, get_iv(sys))
 
         zerovars = setdiff(unknowns(sys), keys(defaults(sys))) .=> 0.0
         trueinit = identity.([zerovars; u0map])
     else
         initializeprob = nothing
         initializeprobmap = nothing
+        reinitializemap = nothing
         trueinit = u0map
     end
 
@@ -925,6 +931,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         sparse = sparse, eval_expression = eval_expression,
         initializeprob = initializeprob,
         initializeprobmap = initializeprobmap,
+        reinitializemap = reinitializemap,
         kwargs...)
     implicit_dae ? (f, du0, u0, p) : (f, u0, p)
 end