From 3582a2821547f27f551dd1b6ac0e006af4123207 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Thu, 5 Dec 2024 13:35:27 +0530
Subject: [PATCH] feat: add lazy initialization to new `remake` methods

---
 src/remake.jl | 54 +++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 50 insertions(+), 4 deletions(-)

diff --git a/src/remake.jl b/src/remake.jl
index 54165a1dd..d2b01a665 100644
--- a/src/remake.jl
+++ b/src/remake.jl
@@ -327,6 +327,7 @@ function remake(prob::SDEProblem;
         use_defaults = false,
         seed = missing,
         kwargs = missing,
+        lazy_initialization = nothing,
         build_initializeprob = true,
         _kwargs...)
     if tspan === missing
@@ -357,7 +358,7 @@ function remake(prob::SDEProblem;
     f = remake(prob.f; f, g, initialization_data)
     iip = isinplace(prob)
 
-    if kwargs === missing
+    prob = if kwargs === missing
         SDEProblem{iip}(f,
             newu0,
             tspan,
@@ -370,6 +371,20 @@ function remake(prob::SDEProblem;
     else
         SDEProblem{iip}(f, newu0, tspan, newp; noise, noise_rate_prototype, seed, kwargs...)
     end
+    if lazy_initialization === nothing
+        lazy_initialization = !is_trivial_initialization(initialization_data)
+    end
+    if !lazy_initialization
+        u0, p, _ = get_initial_values(
+            prob, prob, prob.f, OverrideInit(), Val(isinplace(prob)))
+        if u0 !== nothing && eltype(u0) == Any && isempty(u0)
+            u0 = nothing
+        end
+        @reset prob.u0 = u0
+        @reset prob.p = p
+    end
+
+    return prob
 end
 
 """
@@ -413,7 +428,8 @@ function remake(prob::DDEProblem; f = missing, h = missing, u0 = missing,
         tspan = missing, p = missing, constant_lags = missing,
         dependent_lags = missing, order_discontinuity_t0 = missing,
         neutral = missing, kwargs = missing, interpret_symbolicmap = true,
-        use_defaults = false, build_initializeprob = true, _kwargs...)
+        use_defaults = false, lazy_initialization = nothing, build_initializeprob = true,
+        _kwargs...)
     if tspan === missing
         tspan = prob.tspan
     end
@@ -438,7 +454,7 @@ function remake(prob::DDEProblem; f = missing, h = missing, u0 = missing,
 
     iip = isinplace(prob)
 
-    if kwargs === missing
+    prob = if kwargs === missing
         DDEProblem{iip}(f,
             newu0,
             h,
@@ -454,6 +470,20 @@ function remake(prob::DDEProblem; f = missing, h = missing, u0 = missing,
         DDEProblem{iip}(f, newu0, h, tspan, newp; constant_lags, dependent_lags,
             order_discontinuity_t0, neutral, kwargs...)
     end
+    if lazy_initialization === nothing
+        lazy_initialization = !is_trivial_initialization(initialization_data)
+    end
+    if !lazy_initialization
+        u0, p, _ = get_initial_values(
+            prob, prob, prob.f, OverrideInit(), Val(isinplace(prob)))
+        if u0 !== nothing && eltype(u0) == Any && isempty(u0)
+            u0 = nothing
+        end
+        @reset prob.u0 = u0
+        @reset prob.p = p
+    end
+
+    return prob
 end
 
 function remake(func::DDEFunction;
@@ -498,6 +528,7 @@ function remake(prob::SDDEProblem;
         use_defaults = false,
         seed = missing,
         kwargs = missing,
+        lazy_initialization = nothing,
         build_initializeprob = true,
         _kwargs...)
     if tspan === missing
@@ -535,7 +566,7 @@ function remake(prob::SDDEProblem;
     order_discontinuity_t0 = coalesce(order_discontinuity_t0, prob.order_discontinuity_t0)
     neutral = coalesce(neutral, prob.neutral)
 
-    if kwargs === missing
+    prob = if kwargs === missing
         SDDEProblem{iip}(f,
             g,
             newu0,
@@ -556,6 +587,21 @@ function remake(prob::SDDEProblem;
             f, newu0, tspan, newp; noise, noise_rate_prototype, seed, constant_lags,
             dependent_lags, order_discontinuity_t0, neutral, kwargs...)
     end
+
+    if lazy_initialization === nothing
+        lazy_initialization = !is_trivial_initialization(initialization_data)
+    end
+    if !lazy_initialization
+        u0, p, _ = get_initial_values(
+            prob, prob, prob.f, OverrideInit(), Val(isinplace(prob)))
+        if u0 !== nothing && eltype(u0) == Any && isempty(u0)
+            u0 = nothing
+        end
+        @reset prob.u0 = u0
+        @reset prob.p = p
+    end
+
+    return prob
 end
 
 """