Skip to content

feat: save discrete variables in callback init #3007

New issue

Have a question about this project? No Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

No Sign up for GitHub

By clicking “No Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? No Sign in to your account

Merged
merged 1 commit into from
Sep 2, 2024
Merged

feat: save discrete variables in callback init #3007

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 60 additions & 5 deletions src/systems/callbacks.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -565,8 +565,22 @@ function generate_single_rootfinding_callback(
rf_oop(u, parameter_values(integ), t)
end
end

if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing &&
(save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
initfn = let save_idxs = save_idxs
function (cb, u, t, integrator)
for idx in save_idxs
SciMLBase.save_discretes!(integrator, idx)
end
end
end
else
initfn = SciMLBase.INITIALIZE_DEFAULT
end
return ContinuousCallback(
cond, affect_function.affect, affect_function.affect_neg, rootfind = cb.rootfind)
cond, affect_function.affect, affect_function.affect_neg,
rootfind = cb.rootfind, initialize = initfn)
end

function generate_vector_rootfinding_callback(
Expand Down Expand Up @@ -618,8 +632,25 @@ function generate_vector_rootfinding_callback(
affect_neg(integ)
end
end
if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
save_idxs = mapreduce(
cb -> get(ic.callback_to_clocks, cb, Int[]), vcat, cbs; init = Int[])
initfn = if isempty(save_idxs)
SciMLBase.INITIALIZE_DEFAULT
else
let save_idxs = save_idxs
function (cb, u, t, integrator)
for idx in save_idxs
SciMLBase.save_discretes!(integrator, idx)
end
end
end
end
else
initfn = SciMLBase.INITIALIZE_DEFAULT
end
return VectorContinuousCallback(
cond, affect, affect_neg, length(eqs), rootfind = rootfind)
cond, affect, affect_neg, length(eqs), rootfind = rootfind, initialize = initfn)
end

"""
Expand Down Expand Up @@ -727,12 +758,24 @@ function generate_timed_callback(cb, sys, dvs, ps; postprocess_affect_expr! = no
cond = condition(cb)
as = compile_affect(affects(cb), cb, sys, dvs, ps; expression = Val{false},
postprocess_affect_expr!, kwargs...)
if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing &&
(save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
initfn = let save_idxs = save_idxs
function (cb, u, t, integrator)
for idx in save_idxs
SciMLBase.save_discretes!(integrator, idx)
end
end
end
else
initfn = SciMLBase.INITIALIZE_DEFAULT
end
if cond isa AbstractVector
# Preset Time
return PresetTimeCallback(cond, as)
return PresetTimeCallback(cond, as; initialize = initfn)
else
# Periodic
return PeriodicCallback(as, cond)
return PeriodicCallback(as, cond; initialize = initfn)
end
end

Expand All @@ -745,7 +788,19 @@ function generate_discrete_callback(cb, sys, dvs, ps; postprocess_affect_expr! =
c = compile_condition(cb, sys, dvs, ps; expression = Val{false}, kwargs...)
as = compile_affect(affects(cb), cb, sys, dvs, ps; expression = Val{false},
postprocess_affect_expr!, kwargs...)
return DiscreteCallback(c, as)
if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing &&
(save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
initfn = let save_idxs = save_idxs
function (cb, u, t, integrator)
for idx in save_idxs
SciMLBase.save_discretes!(integrator, idx)
end
end
end
else
initfn = SciMLBase.INITIALIZE_DEFAULT
end
return DiscreteCallback(c, as; initialize = initfn)
end
end

Expand Down
6 changes: 3 additions & 3 deletions test/symbolic_events.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -882,7 +882,7 @@ end
@test sort(canonicalize(Discrete(), prob.p)[1]) == [0.0, 1.0, 2.0]
sol = solve(prob, Tsit5())

@test sol[a] == [-1.0]
@test sol[b] == [5.0, 5.0]
@test sol[c] == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
@test sol[a] == [1.0, -1.0]
@test sol[b] == [2.0, 5.0, 5.0]
@test sol[c] == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
end
Loading