fixup! feat: support parameter updates in initialize_dae!

SciML · Aug 6, 2024 · 2064680 · 2064680
1 parent 60a45fa
commit 2064680
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -5,6 +5,7 @@ version = "6.87.0"
 
 [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"
@@ -46,6 +47,7 @@ TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 
 [compat]
 ADTypes = "0.2, 1"
+Accessors = "0.1.36"
 Adapt = "3.0, 4"
 ArrayInterface = "7"
 DataStructures = "0.18"

diff --git a/src/OrdinaryDiffEq.jl b/src/OrdinaryDiffEq.jl
@@ -14,6 +14,8 @@ using Reexport
 
 using Logging
 
+using Accessors: @reset
+
 using MuladdMacro, SparseArrays, FastClosures
 
 using LinearAlgebra

diff --git a/src/initialize_dae.jl b/src/initialize_dae.jl
@@ -164,8 +164,11 @@ function _initialize_dae!(integrator, prob::Union{ODEProblem, DAEProblem},
     else
         error("Unreachable reached. Report this error.")
     end
-    if isdefined(prob.f, :initializeprobpmap) && prob.f.initializeprobpmap !== nothing
+    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