From bcd324a7f025e0a2d4c4ad42a27616bf8844fe7b Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Tue, 6 Aug 2024 12:46:30 +0530
Subject: [PATCH 1/2] feat: support parameter updates in `initialize_dae!`

---
 lib/OrdinaryDiffEqCore/Project.toml              |  2 ++
 lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl |  1 +
 lib/OrdinaryDiffEqCore/src/initialize_dae.jl     | 10 ++++++++++
 3 files changed, 13 insertions(+)

diff --git a/lib/OrdinaryDiffEqCore/Project.toml b/lib/OrdinaryDiffEqCore/Project.toml
index 274351cdde..44ad3c80d3 100644
--- a/lib/OrdinaryDiffEqCore/Project.toml
+++ b/lib/OrdinaryDiffEqCore/Project.toml
@@ -5,6 +5,7 @@ version = "1.6.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -42,6 +43,7 @@ OrdinaryDiffEqCoreEnzymeCoreExt = "EnzymeCore"
 
 [compat]
 ADTypes = "0.2, 1"
+Accessors = "0.1.36"
 Adapt = "3.0, 4"
 ArrayInterface = "7"
 DataStructures = "0.18"
diff --git a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl
index d3ae689ec0..9276a467f9 100644
--- a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl
+++ b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl
@@ -70,6 +70,7 @@ import DiffEqBase: calculate_residuals,
 import Polyester
 using MacroTools, Adapt
 import ADTypes: AutoFiniteDiff, AutoForwardDiff
+import Accessors: @reset
 
 using SciMLStructures: canonicalize, Tunable, isscimlstructure
 
diff --git a/lib/OrdinaryDiffEqCore/src/initialize_dae.jl b/lib/OrdinaryDiffEqCore/src/initialize_dae.jl
index c352cb3ddf..e88deeef1a 100644
--- a/lib/OrdinaryDiffEqCore/src/initialize_dae.jl
+++ b/lib/OrdinaryDiffEqCore/src/initialize_dae.jl
@@ -153,6 +153,10 @@ function _initialize_dae!(integrator, prob::Union{ODEProblem, DAEProblem},
         alg::OverrideInit, isinplace::Union{Val{true}, Val{false}})
     initializeprob = prob.f.initializeprob
 
+    if SciMLBase.has_update_initializeprob!(prob.f)
+        prob.f.update_initializeprob!(initializeprob, prob)
+    end
+
     # If it doesn't have autodiff, assume it comes from symbolic system like ModelingToolkit
     # Since then it's the case of not a DAE but has initializeprob
     # In which case, it should be differentiable
@@ -173,6 +177,12 @@ function _initialize_dae!(integrator, prob::Union{ODEProblem, DAEProblem},
     else
         error("Unreachable reached. Report this error.")
     end
+    if SciMLBase.has_initializeprobpmap(prob.f)
+        integrator.p = prob.f.initializeprobpmap(prob, nlsol)
+        sol = integrator.sol
+        @reset sol.prob.p = integrator.p
+        integrator.sol = sol
+    end
 
     if nlsol.retcode != ReturnCode.Success
         integrator.sol = SciMLBase.solution_new_retcode(integrator.sol,

From fef8e0fab3a1a46f1c68aa0bee89b8b4314b123e Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 7 Oct 2024 20:11:30 +0530
Subject: [PATCH 2/2] build: bump SciMLBase and DiffEqBase compat

---
 lib/OrdinaryDiffEqCore/Project.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/OrdinaryDiffEqCore/Project.toml b/lib/OrdinaryDiffEqCore/Project.toml
index 44ad3c80d3..396ea028ca 100644
--- a/lib/OrdinaryDiffEqCore/Project.toml
+++ b/lib/OrdinaryDiffEqCore/Project.toml
@@ -47,7 +47,7 @@ Accessors = "0.1.36"
 Adapt = "3.0, 4"
 ArrayInterface = "7"
 DataStructures = "0.18"
-DiffEqBase = "6.147"
+DiffEqBase = "6.157"
 DiffEqDevTools = "2.44.4"
 DocStringExtensions = "0.9"
 EnumX = "1"
@@ -67,7 +67,7 @@ Random = "<0.0.1, 1"
 RecursiveArrayTools = "2.36, 3"
 Reexport = "1.0"
 SafeTestsets = "0.1.0"
-SciMLBase = "2.50.4"
+SciMLBase = "2.56"
 SciMLOperators = "0.3"
 SciMLStructures = "1"
 SimpleUnPack = "1"