Skip to content

Commit

Permalink
fix: u0 and p generation in remake
Browse files Browse the repository at this point in the history
  • Loading branch information
@AayushSabharwal
AayushSabharwal committed Jan 2, 2024
1 parent f4893f2 commit 98f357b
Showing 1 changed file with 48 additions and 81 deletions.
129 changes: 48 additions & 81 deletions src/remake.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,7 @@ for T in [
@eval remaker_of(::$T) = $T
end

"""
remake(thing; <keyword arguments>)
Re-construct `thing` with new field values specified by the keyword
arguments.
"""
function remake(thing; kwargs...)
function _remake_internal(thing; kwargs...)
T = remaker_of(thing)
if :kwargs ∈ fieldnames(typeof(thing))
if :kwargs ∉ keys(kwargs)
Expand All @@ -38,6 +32,21 @@ function remake(thing; kwargs...)
end
end

"""
remake(thing; <keyword arguments>)
Re-construct `thing` with new field values specified by the keyword
arguments.
"""
function remake(thing; kwargs...)
_remake_internal(thing; kwargs...)
end

function remake(prob::DiscreteProblem; u0 = missing, p = missing, kwargs...)
p, u0 = _remake_get_p_u0(prob; p, u0)
_remake_internal(prob; p, u0, kwargs...)
end

function isrecompile(prob::ODEProblem{iip}) where {iip}
(prob.f isa ODEFunction) ? !isfunctionwrapper(prob.f.f) : true
end
Expand All @@ -59,25 +68,7 @@ function remake(prob::ODEProblem; f = missing,
tspan = prob.tspan
end

if p === missing && u0 === missing
p, u0 = prob.p, prob.u0
else # at least one of them has a value
if p === missing
p = prob.p
end
if u0 === missing
u0 = prob.u0
end
if (eltype(p) <: Pair && !isempty(p)) || (eltype(u0) <: Pair && !isempty(u0)) # one is a non-empty symbolic map
hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) ||
throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
" Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order."))
hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) ||
throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
" Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order."))
p, u0 = process_p_u0_symbolic(prob, p, u0)
end
end
p, u0 = _remake_get_p_u0(prob; p, u0)

iip = isinplace(prob)

Expand Down Expand Up @@ -132,15 +123,11 @@ function remake(prob::BVProblem; f = missing, bc = missing, u0 = missing, tspan
tspan = prob.tspan
end

if p === missing && u0 === missing
p, u0 = prob.p, prob.u0
else # at least one of them has a value
if p === missing
p = prob.p
end
if u0 === missing
u0 = prob.u0
end
if p === missing
p = prob.p
end
if u0 === missing
u0 = prob.u0
end

iip = isinplace(prob)
Expand Down Expand Up @@ -202,13 +189,7 @@ function remake(prob::SDEProblem;
tspan = prob.tspan
end

if p === missing
p = prob.p
end

if u0 === missing
u0 = prob.u0
end
p, u0 = _remake_get_p_u0(prob; p, u0)

if noise === missing
noise = prob.noise
Expand Down Expand Up @@ -263,26 +244,7 @@ function remake(prob::OptimizationProblem;
sense = missing,
kwargs = missing,
_kwargs...)
if p === missing && u0 === missing
p, u0 = prob.p, prob.u0
else # at least one of them has a value
if p === missing
p = prob.p
end
if u0 === missing
u0 = prob.u0
end
if (eltype(p) <: Pair && !isempty(p)) || (eltype(u0) <: Pair && !isempty(u0)) # one is a non-empty symbolic map
hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) ||
throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
" Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order."))
hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) ||
throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
" Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order."))
p, u0 = process_p_u0_symbolic(prob, p, u0)
end
end

p, u0 = _remake_get_p_u0(prob; p, u0)
if f === missing
f = prob.f
end
Expand Down Expand Up @@ -332,25 +294,7 @@ function remake(prob::NonlinearProblem;
problem_type = missing,
kwargs = missing,
_kwargs...)
if p === missing && u0 === missing
p, u0 = prob.p, prob.u0
else # at least one of them has a value
if p === missing
p = prob.p
end
if u0 === missing
u0 = prob.u0
end
if (eltype(p) <: Pair && !isempty(p)) || (eltype(u0) <: Pair && !isempty(u0)) # one is a non-empty symbolic map
hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) ||
throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
" Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order."))
hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) ||
throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
" Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order."))
p, u0 = process_p_u0_symbolic(prob, p, u0)
end
end
p, u0 = _remake_get_p_u0(prob; p, u0)

if f === missing
f = prob.f
Expand Down Expand Up @@ -419,3 +363,26 @@ function remake(thing::AbstractEnsembleProblem; kwargs...)
en_kwargs = [k for k in kwargs if first(k) ∈ fieldnames(T)]
T(remake(thing.prob; setdiff(kwargs, en_kwargs)...); en_kwargs...)
end

function _remake_get_p_u0(prob; p = missing, u0 = missing)
if p === missing && u0 === missing
p, u0 = prob.p, prob.u0
else # at least one of them has a value
if p === missing
p = prob.p
end
if u0 === missing
u0 = prob.u0
end
if (eltype(p) <: Pair && !isempty(p)) || (eltype(u0) <: Pair && !isempty(u0)) # one is a non-empty symbolic map
hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) ||
throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
" Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order."))
hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) ||
throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." *
" Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order."))
p, u0 = process_p_u0_symbolic(prob, p, u0)
end
end
return p, u0
end

0 comments on commit 98f357b

Please sign in to comment.