From 76d3bded416b5d707621cc37d2013f22d295870d Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Tue, 6 Aug 2024 15:31:27 +0530
Subject: [PATCH] feat: perform limited DAE initialization for null
 integrators/solutions

---
 src/solve.jl | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/solve.jl b/src/solve.jl
index 1b0a7c54e..6de8273f4 100644
--- a/src/solve.jl
+++ b/src/solve.jl
@@ -626,6 +626,9 @@ end
 function build_null_integrator(prob::AbstractDEProblem, args...;
         kwargs...)
     sol = solve(prob, args...; kwargs...)
+    # The DAE initialization in `build_null_solution` may change the parameter
+    # object `prob.p` via `@set!`, hence use the "new" prob instead of the "old" one.
+    prob = sol.prob
     return NullODEIntegrator{
         isinplace(prob), typeof(prob), eltype(prob.tspan), typeof(sol),
         typeof(prob.f), typeof(prob.p)
@@ -675,7 +678,12 @@ function build_null_solution(prob::AbstractDEProblem, args...;
     end
 
     timeseries = [Float64[] for i in 1:length(ts)]
-
+    
+    if SciMLBase.has_initializeprob(prob.f) && SciMLBase.has_initializeprobpmap(prob.f)
+        initializeprob = prob.f.initializeprob
+        nlsol = solve(initializeprob)
+        @set! prob.p = prob.f.initializeprobpmap(prob, nlsol)
+    end
     build_solution(prob, nothing, ts, timeseries, retcode = ReturnCode.Success)
 end