From c979fd70122ebfdc1435e6cb526d5da3d9b60ad5 Mon Sep 17 00:00:00 2001
From: Christopher Rackauckas <accounts@chrisrackauckas.com>
Date: Fri, 1 Mar 2024 02:33:55 -0600
Subject: [PATCH] Add reinitializationmaps

Companion PR to https://github.com/SciML/SciMLBase.jl/pull/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.
---
 src/systems/diffeqs/abstractodesystem.jl | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 3ad6bff0ff..811cc27e0d 100644
--- 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