Skip to content

Commit e0e0e46

Browse files
BenChungBenChung
committed
Format
1 parent 6424ca9 commit e0e0e46

File tree

2 files changed

+121
-88
lines changed

2 files changed

+121
-88
lines changed

‎src/systems/callbacks.jl‎

Lines changed: 88 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ struct SymbolicContinuousCallback
131131
finalize = NULL_AFFECT,
132132
rootfind = SciMLBase.LeftRootFind,
133133
reinitializealg = SciMLBase.CheckInit())
134-
new(eqs, initialize, finalize, make_affect(affect),
134+
new(eqs, initialize, finalize, make_affect(affect),
135135
make_affect(affect_neg), rootfind, reinitializealg)
136136
end # Default affect to nothing
137137
end
@@ -227,18 +227,19 @@ function SymbolicContinuousCallback(args...)
227227
end # wrap eq in vector
228228
SymbolicContinuousCallback(p::Pair) = SymbolicContinuousCallback(p[1], p[2])
229229
SymbolicContinuousCallback(cb::SymbolicContinuousCallback) = cb # passthrough
230-
function SymbolicContinuousCallback(eqs::Equation, affect = NULL_AFFECT;
231-
initialize=NULL_AFFECT, finalize=NULL_AFFECT,
232-
affect_neg = affect, rootfind = SciMLBase.LeftRootFind)
230+
function SymbolicContinuousCallback(eqs::Equation, affect = NULL_AFFECT;
231+
initialize = NULL_AFFECT, finalize = NULL_AFFECT,
232+
affect_neg = affect, rootfind = SciMLBase.LeftRootFind)
233233
SymbolicContinuousCallback(
234-
eqs = [eqs], affect = affect, affect_neg = affect_neg,
235-
initialize=initialize, finalize=finalize, rootfind = rootfind)
234+
eqs = [eqs], affect = affect, affect_neg = affect_neg,
235+
initialize = initialize, finalize = finalize, rootfind = rootfind)
236236
end
237237
function SymbolicContinuousCallback(eqs::Vector{Equation}, affect = NULL_AFFECT;
238-
affect_neg = affect, initialize=NULL_AFFECT, finalize=NULL_AFFECT,
238+
affect_neg = affect, initialize = NULL_AFFECT, finalize = NULL_AFFECT,
239239
rootfind = SciMLBase.LeftRootFind)
240240
SymbolicContinuousCallback(
241-
eqs = eqs, affect = affect, affect_neg = affect_neg, initialize=initialize, finalize=finalize, rootfind = rootfind)
241+
eqs = eqs, affect = affect, affect_neg = affect_neg,
242+
initialize = initialize, finalize = finalize, rootfind = rootfind)
242243
end
243244

244245
SymbolicContinuousCallbacks(cb::SymbolicContinuousCallback) = [cb]
@@ -334,10 +335,11 @@ struct SymbolicDiscreteCallback
334335

335336
function SymbolicDiscreteCallback(
336337
condition, affects = NULL_AFFECT; reinitializealg = SciMLBase.CheckInit(),
337-
initialize=NULL_AFFECT, finalize=NULL_AFFECT)
338+
initialize = NULL_AFFECT, finalize = NULL_AFFECT)
338339
c = scalarize_condition(condition)
339340
a = scalarize_affects(affects)
340-
new(c, a, scalarize_affects(initialize), scalarize_affects(finalize), reinitializealg)
341+
new(c, a, scalarize_affects(initialize),
342+
scalarize_affects(finalize), reinitializealg)
341343
end # Default affect to nothing
342344
end
343345

@@ -376,14 +378,17 @@ function Base.show(io::IO, db::SymbolicDiscreteCallback)
376378
end
377379

378380
function Base.:(==)(e1::SymbolicDiscreteCallback, e2::SymbolicDiscreteCallback)
379-
isequal(e1.condition, e2.condition) && isequal(e1.affects, e2.affects) &&
380-
isequal(e1.initialize, e2.initialize) && isequal(e1.finalize, e2.finalize)
381+
isequal(e1.condition, e2.condition) && isequal(e1.affects, e2.affects) &&
382+
isequal(e1.initialize, e2.initialize) && isequal(e1.finalize, e2.finalize)
381383
end
382384
function Base.hash(cb::SymbolicDiscreteCallback, s::UInt)
383385
s = hash(cb.condition, s)
384-
s = cb.affects isa AbstractVector ? foldr(hash, cb.affects, init = s) : hash(cb.affects, s)
385-
s = cb.initialize isa AbstractVector ? foldr(hash, cb.initialize, init = s) : hash(cb.initialize, s)
386-
s = cb.finalize isa AbstractVector ? foldr(hash, cb.finalize, init = s) : hash(cb.finalize, s)
386+
s = cb.affects isa AbstractVector ? foldr(hash, cb.affects, init = s) :
387+
hash(cb.affects, s)
388+
s = cb.initialize isa AbstractVector ? foldr(hash, cb.initialize, init = s) :
389+
hash(cb.initialize, s)
390+
s = cb.finalize isa AbstractVector ? foldr(hash, cb.finalize, init = s) :
391+
hash(cb.finalize, s)
387392
s = hash(cb.reinitializealg, s)
388393
return s
389394
end
@@ -405,7 +410,6 @@ function reinitialization_algs(cbs::Vector{SymbolicDiscreteCallback})
405410
reinitialization_alg, vcat, cbs, init = SciMLBase.DAEInitializationAlgorithm[])
406411
end
407412

408-
409413
initialize_affects(cb::SymbolicDiscreteCallback) = cb.initialize
410414
function initialize_affects(cbs::Vector{SymbolicDiscreteCallback})
411415
mapreduce(initialize_affects, vcat, cbs, init = Equation[])
@@ -418,10 +422,13 @@ end
418422

419423
function namespace_callback(cb::SymbolicDiscreteCallback, s)::SymbolicDiscreteCallback
420424
function namespace_affects(af)
421-
return af isa AbstractVector ? namespace_affect.(af, Ref(s)) : namespace_affect(af, s)
425+
return af isa AbstractVector ? namespace_affect.(af, Ref(s)) :
426+
namespace_affect(af, s)
422427
end
423-
SymbolicDiscreteCallback(namespace_condition(condition(cb), s), namespace_affects(affects(cb)),
424-
reinitializealg=cb.reinitializealg, initialize=namespace_affects(initialize_affects(cb)), finalize=namespace_affects(finalize_affects(cb)))
428+
SymbolicDiscreteCallback(
429+
namespace_condition(condition(cb), s), namespace_affects(affects(cb)),
430+
reinitializealg = cb.reinitializealg, initialize = namespace_affects(initialize_affects(cb)),
431+
finalize = namespace_affects(finalize_affects(cb)))
425432
end
426433

427434
SymbolicDiscreteCallbacks(cb::Pair) = SymbolicDiscreteCallback[SymbolicDiscreteCallback(cb)]
@@ -698,7 +705,7 @@ function generate_single_rootfinding_callback(
698705
end
699706
end
700707
user_initfun = isnothing(affect_function.initialize) ? SciMLBase.INITIALIZE_DEFAULT :
701-
(c, u, t, i) -> affect_function.initialize(i)
708+
(c, u, t, i) -> affect_function.initialize(i)
702709
if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing &&
703710
(save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
704711
initfn = let save_idxs = save_idxs
@@ -715,7 +722,8 @@ function generate_single_rootfinding_callback(
715722
return ContinuousCallback(
716723
cond, affect_function.affect, affect_function.affect_neg, rootfind = cb.rootfind,
717724
initialize = initfn,
718-
finalize = isnothing(affect_function.finalize) ? SciMLBase.FINALIZE_DEFAULT : (c, u, t, i) -> affect_function.finalize(i),
725+
finalize = isnothing(affect_function.finalize) ? SciMLBase.FINALIZE_DEFAULT :
726+
(c, u, t, i) -> affect_function.finalize(i),
719727
initializealg = reinitialization_alg(cb))
720728
end
721729

@@ -742,8 +750,8 @@ function generate_vector_rootfinding_callback(
742750
affect_neg::Union{Function, Nothing},
743751
initialize::Union{Function, Nothing},
744752
finalize::Union{Function, Nothing}}[
745-
compile_affect_fn(cb, sys, dvs, ps, kwargs)
746-
for cb in cbs]
753+
compile_affect_fn(cb, sys, dvs, ps, kwargs)
754+
for cb in cbs]
747755

748756
cond = function (out, u, t, integ)
749757
rf_ip(out, u, parameter_values(integ), t)
@@ -789,31 +797,37 @@ function generate_vector_rootfinding_callback(
789797
end
790798
initialize = nothing
791799
if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
792-
initialize = handle_optional_setup_fn(map((cb, fn) -> begin
793-
if (save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
794-
let save_idxs = save_idxs
795-
if !isnothing(fn.initialize)
796-
(i) -> begin
797-
fn.initialize(i)
798-
for idx in save_idxs
799-
SciMLBase.save_discretes!(i, idx)
800-
end
801-
end
802-
else
803-
(i) -> begin
804-
for idx in save_idxs
805-
SciMLBase.save_discretes!(i, idx)
800+
initialize = handle_optional_setup_fn(
801+
map(
802+
(cb, fn) -> begin
803+
if (save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
804+
let save_idxs = save_idxs
805+
if !isnothing(fn.initialize)
806+
(i) -> begin
807+
fn.initialize(i)
808+
for idx in save_idxs
809+
SciMLBase.save_discretes!(i, idx)
810+
end
811+
end
812+
else
813+
(i) -> begin
814+
for idx in save_idxs
815+
SciMLBase.save_discretes!(i, idx)
816+
end
817+
end
806818
end
807819
end
820+
else
821+
fn.initialize
808822
end
809-
end
810-
else
811-
fn.initialize
812-
end
813-
end, cbs, affect_functions), SciMLBase.INITIALIZE_DEFAULT)
814-
823+
end,
824+
cbs,
825+
affect_functions),
826+
SciMLBase.INITIALIZE_DEFAULT)
827+
815828
else
816-
initialize = handle_optional_setup_fn(map(fn -> fn.initialize, affect_functions), SciMLBase.INITIALIZE_DEFAULT)
829+
initialize = handle_optional_setup_fn(
830+
map(fn -> fn.initialize, affect_functions), SciMLBase.INITIALIZE_DEFAULT)
817831
end
818832

819833
finalize = handle_optional_setup_fn(
@@ -833,10 +847,13 @@ function compile_affect_fn(cb, sys::AbstractODESystem, dvs, ps, kwargs)
833847
if eq_neg_aff === eq_aff
834848
affect_neg = affect
835849
else
836-
affect_neg = _compile_optional_affect(NULL_AFFECT, eq_neg_aff, cb, sys, dvs, ps; kwargs...)
850+
affect_neg = _compile_optional_affect(
851+
NULL_AFFECT, eq_neg_aff, cb, sys, dvs, ps; kwargs...)
837852
end
838-
initialize = _compile_optional_affect(NULL_AFFECT, initialize_affects(cb), cb, sys, dvs, ps; kwargs...)
839-
finalize = _compile_optional_affect(NULL_AFFECT, finalize_affects(cb), cb, sys, dvs, ps; kwargs...)
853+
initialize = _compile_optional_affect(
854+
NULL_AFFECT, initialize_affects(cb), cb, sys, dvs, ps; kwargs...)
855+
finalize = _compile_optional_affect(
856+
NULL_AFFECT, finalize_affects(cb), cb, sys, dvs, ps; kwargs...)
840857
(affect = affect, affect_neg = affect_neg, initialize = initialize, finalize = finalize)
841858
end
842859

@@ -928,7 +945,6 @@ function compile_affect(affect::FunctionalAffect, cb, sys, dvs, ps; kwargs...)
928945
compile_user_affect(affect, cb, sys, dvs, ps; kwargs...)
929946
end
930947

931-
932948
function _compile_optional_affect(default, aff, cb, sys, dvs, ps; kwargs...)
933949
if isnothing(aff) || aff == default
934950
return nothing
@@ -942,13 +958,15 @@ function generate_timed_callback(cb, sys, dvs, ps; postprocess_affect_expr! = no
942958
as = compile_affect(affects(cb), cb, sys, dvs, ps; expression = Val{false},
943959
postprocess_affect_expr!, kwargs...)
944960

945-
user_initfun = _compile_optional_affect(NULL_AFFECT, initialize_affects(cb), cb, sys, dvs, ps; kwargs...)
946-
user_finfun = _compile_optional_affect(NULL_AFFECT, finalize_affects(cb), cb, sys, dvs, ps; kwargs...)
961+
user_initfun = _compile_optional_affect(
962+
NULL_AFFECT, initialize_affects(cb), cb, sys, dvs, ps; kwargs...)
963+
user_finfun = _compile_optional_affect(
964+
NULL_AFFECT, finalize_affects(cb), cb, sys, dvs, ps; kwargs...)
947965
if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing &&
948966
(save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
949-
initfn = let
967+
initfn = let
950968
save_idxs = save_idxs
951-
initfun=user_initfun
969+
initfun = user_initfun
952970
function (cb, u, t, integrator)
953971
if !isnothing(initfun)
954972
initfun(integrator)
@@ -959,17 +977,21 @@ function generate_timed_callback(cb, sys, dvs, ps; postprocess_affect_expr! = no
959977
end
960978
end
961979
else
962-
initfn = isnothing(user_initfun) ? SciMLBase.INITIALIZE_DEFAULT : (_,_,_,i) -> user_initfun(i)
980+
initfn = isnothing(user_initfun) ? SciMLBase.INITIALIZE_DEFAULT :
981+
(_, _, _, i) -> user_initfun(i)
963982
end
964-
finfun = isnothing(user_finfun) ? SciMLBase.FINALIZE_DEFAULT : (_,_,_,i) -> user_finfun(i)
983+
finfun = isnothing(user_finfun) ? SciMLBase.FINALIZE_DEFAULT :
984+
(_, _, _, i) -> user_finfun(i)
965985
if cond isa AbstractVector
966986
# Preset Time
967987
return PresetTimeCallback(
968-
cond, as; initialize = initfn, finalize=finfun, initializealg = reinitialization_alg(cb))
988+
cond, as; initialize = initfn, finalize = finfun,
989+
initializealg = reinitialization_alg(cb))
969990
else
970991
# Periodic
971992
return PeriodicCallback(
972-
as, cond; initialize = initfn, finalize=finfun, initializealg = reinitialization_alg(cb))
993+
as, cond; initialize = initfn, finalize = finfun,
994+
initializealg = reinitialization_alg(cb))
973995
end
974996
end
975997

@@ -983,8 +1005,10 @@ function generate_discrete_callback(cb, sys, dvs, ps; postprocess_affect_expr! =
9831005
as = compile_affect(affects(cb), cb, sys, dvs, ps; expression = Val{false},
9841006
postprocess_affect_expr!, kwargs...)
9851007

986-
user_initfun = _compile_optional_affect(NULL_AFFECT, initialize_affects(cb), cb, sys, dvs, ps; kwargs...)
987-
user_finfun = _compile_optional_affect(NULL_AFFECT, finalize_affects(cb), cb, sys, dvs, ps; kwargs...)
1008+
user_initfun = _compile_optional_affect(
1009+
NULL_AFFECT, initialize_affects(cb), cb, sys, dvs, ps; kwargs...)
1010+
user_finfun = _compile_optional_affect(
1011+
NULL_AFFECT, finalize_affects(cb), cb, sys, dvs, ps; kwargs...)
9881012
if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing &&
9891013
(save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
9901014
initfn = let save_idxs = save_idxs, initfun = user_initfun
@@ -998,11 +1022,14 @@ function generate_discrete_callback(cb, sys, dvs, ps; postprocess_affect_expr! =
9981022
end
9991023
end
10001024
else
1001-
initfn = isnothing(user_initfun) ? SciMLBase.INITIALIZE_DEFAULT : (_,_,_,i) -> user_initfun(i)
1025+
initfn = isnothing(user_initfun) ? SciMLBase.INITIALIZE_DEFAULT :
1026+
(_, _, _, i) -> user_initfun(i)
10021027
end
1003-
finfun = isnothing(user_finfun) ? SciMLBase.FINALIZE_DEFAULT : (_,_,_,i) -> user_finfun(i)
1028+
finfun = isnothing(user_finfun) ? SciMLBase.FINALIZE_DEFAULT :
1029+
(_, _, _, i) -> user_finfun(i)
10041030
return DiscreteCallback(
1005-
c, as; initialize = initfn, finalize = finfun, initializealg = reinitialization_alg(cb))
1031+
c, as; initialize = initfn, finalize = finfun,
1032+
initializealg = reinitialization_alg(cb))
10061033
end
10071034
end
10081035

0 commit comments

Comments
 (0)