Skip to content

Commit

Permalink
feat: add lazy initialization to new remake methods
Browse files Browse the repository at this point in the history
  • Loading branch information
AayushSabharwal committed Dec 5, 2024
1 parent db56a4f commit 3582a28
Showing 1 changed file with 50 additions and 4 deletions.
54 changes: 50 additions & 4 deletions src/remake.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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

"""
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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

"""
Expand Down

0 comments on commit 3582a28

Please sign in to comment.