fix: fix initialization and finalization affects

Author: vyudu

src/structural_transformation/bipartite_tearing/modia_tearing.jl

@@ -96,7 +96,7 @@ function tear_graph_modia(structure::SystemStructure, isder::F = nothing,
     ieqs = Int[]
     filtered_vars = BitSet()
     free_eqs = free_equations(graph, var_sccs, var_eq_matching, varfilter)
-    is_overdetemined = !isempty(free_eqs)
+    is_overdetermined = !isempty(free_eqs)
     for vars in var_sccs
         for var in vars
             if varfilter(var)
@@ -112,7 +112,7 @@ function tear_graph_modia(structure::SystemStructure, isder::F = nothing,
             filtered_vars, isder)
         # If the systems is overdetemined, we cannot assume the free equations
         # will not form algebraic loops with equations in the sccs.
-        if !is_overdetemined
+        if !is_overdetermined
             vargraph.ne = 0
             for var in vars
                 vargraph.matching[var] = unassigned
@@ -121,7 +121,7 @@ function tear_graph_modia(structure::SystemStructure, isder::F = nothing,
         empty!(ieqs)
         empty!(filtered_vars)
     end
-    if is_overdetemined
+    if is_overdetermined
         free_vars = findall(x -> !(x isa Int), var_eq_matching)
         tear_graph_block_modia!(var_eq_matching, ict, solvable_graph, free_eqs,
             BitSet(free_vars), isder)

src/structural_transformation/utils.jl

@@ -218,6 +218,7 @@ function find_eq_solvables!(state::TearingState, ieq, to_rm = Int[], coeffs = no
     all_int_vars = true
     coeffs === nothing || empty!(coeffs)
     empty!(to_rm)
+
     for j in 𝑠neighbors(graph, ieq)
         var = fullvars[j]
         isirreducible(var) && (all_int_vars = false; continue)

src/systems/callbacks.jl

@@ -217,7 +217,7 @@ struct SymbolicContinuousCallback <: AbstractCallback
     end # Default affect to nothing
 end
 
-SymbolicContinuousCallback(p::Pair, args...; kwargs...) = SymbolicContinuousCallback(p[1], p[2])
+SymbolicContinuousCallback(p::Pair, args...; kwargs...) = SymbolicContinuousCallback(p[1], p[2], args...; kwargs...)
 SymbolicContinuousCallback(cb::SymbolicContinuousCallback, args...; kwargs...) = cb
 
 make_affect(affect::Nothing; kwargs...) = nothing
@@ -395,7 +395,7 @@ Arguments:
 - algeeqs: Algebraic equations of the system that must be satisfied after the callback occurs.
 """
 struct SymbolicDiscreteCallback <: AbstractCallback
-    conditions::Union{Real, Vector{<:Real}, Vector{Equation}}
+    conditions::Any
     affect::Union{Affect, Nothing}
     initialize::Union{Affect, Nothing}
     finalize::Union{Affect, Nothing}
@@ -410,7 +410,7 @@ struct SymbolicDiscreteCallback <: AbstractCallback
     end # Default affect to nothing
 end
 
-SymbolicDiscreteCallback(p::Pair, args...; kwargs...) = SymbolicDiscreteCallback(p[1], p[2])
+SymbolicDiscreteCallback(p::Pair, args...; kwargs...) = SymbolicDiscreteCallback(p[1], p[2], args...; kwargs...)
 SymbolicDiscreteCallback(cb::SymbolicDiscreteCallback, args...; kwargs...) = cb
 
 """
@@ -630,7 +630,7 @@ end
 """
 Compile user-defined functional affect.
 """
-function compile_functional_affect(affect::FunctionalAffect, cb, sys; kwargs...)
+function compile_functional_affect(affect::FunctionalAffect, sys; kwargs...)
     dvs = unknowns(sys)
     ps = parameters(sys)
     dvs_ind = Dict(reverse(en) for en in enumerate(dvs))
@@ -639,11 +639,9 @@ function compile_functional_affect(affect::FunctionalAffect, cb, sys; kwargs...)
     if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
         p_inds = [(pind = parameter_index(sys, sym)) === nothing ? sym : pind
                   for sym in parameters(affect)]
-        save_idxs = get(ic.callback_to_clocks, cb, Int[])
     else
         ps_ind = Dict(reverse(en) for en in enumerate(ps))
         p_inds = map(sym -> get(ps_ind, sym, sym), parameters(affect))
-        save_idxs = Int[]
     end
     # HACK: filter out eliminated symbols. Not clear this is the right thing to do
     # (MTK should keep these symbols)
@@ -652,13 +650,9 @@ function compile_functional_affect(affect::FunctionalAffect, cb, sys; kwargs...)
     p = filter(x -> !isnothing(x[2]), collect(zip(parameters_syms(affect), p_inds))) |>
         NamedTuple
 
-    let u = u, p = p, user_affect = func(affect), ctx = context(affect),
-        save_idxs = save_idxs
-        function (integ)
+    let u = u, p = p, user_affect = func(affect), ctx = context(affect)
+        (integ) -> begin
             user_affect(integ, u, p, ctx)
-            for idx in save_idxs
-                SciMLBase.save_discretes!(integ, idx)
-            end
         end
     end
 end
@@ -670,6 +664,8 @@ function generate_continuous_callbacks(sys::AbstractSystem, dvs = unknowns(sys),
     cbs = continuous_events(sys)
     isempty(cbs) && return nothing
     cb_classes = Dict{SciMLBase.RootfindOpt, Vector{SymbolicContinuousCallback}}()
+
+    # Sort the callbacks by their rootfinding method
     for cb in cbs
         _cbs = get!(() -> SymbolicContinuousCallback[], cb_classes, cb.rootfind)
         push!(_cbs, cb)
@@ -709,12 +705,12 @@ function generate_callback(cbs::Vector{SymbolicContinuousCallback}, sys; kwargs.
     inits = []
     finals = []
     for cb in cbs
-        affect = compile_affect(cb.affect, cb, sys, default = (args...) -> ())
+        affect = compile_affect(cb.affect, cb, sys, default = nothing)
         push!(affects, affect)
-        affect_neg = (cb.affect_neg === cb.affect) ? affect : compile_affect(cb.affect_neg, cb, sys, default = (args...) -> ())
+        affect_neg = (cb.affect_neg == cb.affect) ? affect : compile_affect(cb.affect_neg, cb, sys, default = nothing)
         push!(affect_negs, affect_neg)
-        push!(inits, compile_affect(cb.initialize, cb, sys, default = nothing))
-        push!(finals, compile_affect(cb.finalize, cb, sys, default = nothing))
+        push!(inits, compile_affect(cb.initialize, cb, sys; default = nothing, is_init = true))
+        push!(finals, compile_affect(cb.finalize, cb, sys; default = nothing))
     end
 
     # Since there may be different number of conditions and affects,
@@ -746,10 +742,16 @@ function generate_callback(cb, sys; kwargs...)
 
     trigger = is_timed ? conditions(cb) : compile_condition(cb, sys, dvs, ps; kwargs...)
     affect = compile_affect(cb.affect, cb, sys, default = (args...) -> ())
-    affect_neg = hasfield(typeof(cb), :affect_neg) ?
-                 compile_affect(cb.affect_neg, cb, sys, default = affect) : nothing
-    initialize = compile_affect(cb.initialize, cb, sys, default = SciMLBase.INITIALIZE_DEFAULT)
-    finalize = compile_affect(cb.finalize, cb, sys, default = SciMLBase.FINALIZE_DEFAULT)
+    affect_neg = if is_discrete(cb)
+        nothing
+    else
+        (cb.affect == cb.affect_neg) ? affect : compile_affect(cb.affect_neg, cb, sys, default = nothing)
+    end
+    init = compile_affect(cb.initialize, cb, sys, default = SciMLBase.INITIALIZE_DEFAULT, is_init = true)
+    final = compile_affect(cb.finalize, cb, sys, default = SciMLBase.FINALIZE_DEFAULT)
+
+    initialize = isnothing(cb.initialize) ? init : ((c, u, t, i) -> init(i))
+    finalize = isnothing(cb.finalize) ? final : ((c, u, t, i) -> final(i))
 
     if is_discrete(cb)
         if is_timed && conditions(cb) isa AbstractVector
@@ -784,32 +786,81 @@ Notes
   - `kwargs` are passed through to `Symbolics.build_function`.
 """
 function compile_affect(
-        aff::Union{Nothing, Affect}, cb::AbstractCallback, sys::AbstractSystem; default = nothing, kwargs...)
-    isnothing(aff) && return default
-
+        aff::Union{Nothing, Affect}, cb::AbstractCallback, sys::AbstractSystem; default = nothing, is_init = false, kwargs...)
     save_idxs = if !(has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing)
         Int[]
     else
         get(ic.callback_to_clocks, cb, Int[])
     end
 
-    if aff isa AffectSystem
-        affsys = system(aff)
-        aff_map = aff_to_sys(aff)
-        sys_map = Dict([v => k for (k, v) in aff_map])
-        reinit = has_alg_eqs(sys)
-        ps_to_modify = discretes(aff)
-        dvs_to_modify = setdiff(unknowns(aff), getfield.(observed(sys), :lhs))
+    f = if isnothing(aff)
+        default 
+    elseif aff isa AffectSystem
+        compile_equational_affect(aff, sys)
+    elseif aff isa FunctionalAffect || aff isa ImperativeAffect
+        compile_functional_affect(aff, sys; kwargs...)
+    end
+    wrap_save_discretes(f, save_idxs; is_init)
+end
+
+# Init can be: user defined function, nothing, or INITIALIZE_DEFAULT
+function wrap_save_discretes(f, save_idxs; is_init = false)
+    if isempty(save_idxs) || f === SciMLBase.FINALIZE_DEFAULT || (isnothing(f) && !is_init)
+        return f
+    elseif f === SciMLBase.INITIALIZE_DEFAULT
+        let save_idxs = save_idxs
+            (c, u, t, i) -> begin
+                f(c, u, t, i)
+                for idx in save_idxs
+                    SciMLBase.save_discretes!(i, idx)
+                end
+            end
+        end
+    else
+        let save_idxs = save_idxs
+            (i) -> begin
+                isnothing(f) || f(i)
+                for idx in save_idxs
+                    SciMLBase.save_discretes!(i, idx)
+                end
+            end
+        end
+    end
+end
+
+"""
+Initialize and Finalize for VectorContinuousCallback.
+"""
+function compile_vector_optional_affect(funs, default)
+    all(isnothing, funs) && return default
+    return let funs = funs
+        function (cb, u, t, integ)
+            for func in funs
+                isnothing(func) ? continue : func(integ)
+            end
+        end
+    end
+end
+
+function compile_equational_affect(aff::AffectSystem, sys; kwargs...)
+    affsys = system(aff)
+    aff_map = aff_to_sys(aff)
+    sys_map = Dict([v => k for (k, v) in aff_map])
+    ps_to_modify = discretes(aff)
+    dvs_to_modify = setdiff(unknowns(aff), getfield.(observed(sys), :lhs))
+    #TODO: Add an optimization for systems without algebraic equations 
 
-        function affect!(integrator)
+    return let dvs_to_modify = dvs_to_modify, aff_map = aff_map, sys_map = sys_map, affsys = affsys, ps_to_modify = ps_to_modify
+
+        @inline function affect!(integrator)
             pmap = Pair[]
             for pre_p in parameters(affsys)
                 p = only(arguments(unwrap(pre_p)))
                 pval = isparameter(p) ? integrator.ps[p] : integrator[p]
                 push!(pmap, pre_p => pval)
             end
             guesses = Pair[u => integrator[aff_map[u]] for u in unknowns(affsys)]
-            affprob = ImplicitDiscreteProblem(affsys, Pair[], (0, 1), pmap; guesses, build_initializeprob = reinit)
+            affprob = ImplicitDiscreteProblem(affsys, Pair[], (0, 1), pmap; guesses, build_initializeprob = false)
 
             affsol = init(affprob, SimpleIDSolve())
             for u in dvs_to_modify
@@ -818,28 +869,9 @@ function compile_affect(
             for p in ps_to_modify
                 integrator.ps[p] = affsol[sys_map[p]]
             end
-            for idx in save_idxs
-                SciMLBase.save_discretes!(integrator, idx)
-            end
 
             sys isa JumpSystem && reset_aggregated_jumps!(integrator)
         end
-    elseif aff isa FunctionalAffect || aff isa ImperativeAffect
-        compile_functional_affect(aff, cb, sys; kwargs...)
-    end
-end
-
-"""
-Initialize and Finalize for VectorContinuousCallback.
-"""
-function compile_vector_optional_affect(funs, default)
-    all(isnothing, funs) && return default
-    return let funs = funs
-        function (cb, u, t, integ)
-            for func in funs
-                isnothing(func) ? continue : func(integ)
-            end
-        end
     end
 end
 

src/systems/discrete_system/implicit_discrete_system.jl

@@ -264,7 +264,7 @@ function flatten(sys::ImplicitDiscreteSystem, noeqs = false)
 end
 
 function generate_function(
-        sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
+        sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, cachesyms::Tuple = (), kwargs...)
     iv = get_iv(sys)
     # Algebraic equations get shifted forward 1, to match with differential equations
     exprs = map(equations(sys)) do eq
@@ -280,8 +280,9 @@ function generate_function(
 
     u_next = map(Shift(iv, 1), dvs)
     u = dvs
+    p = (reorder_parameters(sys, unwrap.(ps))..., cachesyms...)
     build_function_wrapper(
-        sys, exprs, u_next, u, ps..., iv; p_start = 3, extra_assignments, kwargs...)
+        sys, exprs, u_next, u, p..., iv; p_start = 3, extra_assignments, kwargs...)
 end
 
 function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)

src/systems/imperative_affect.jl

@@ -109,10 +109,6 @@ function namespace_affect(affect::ImperativeAffect, s)
         affect.skip_checks)
 end
 
-function compile_affect(affect::ImperativeAffect, cb, sys, dvs, ps; kwargs...)
-    compile_functional_affect(affect, cb, sys, dvs, ps; kwargs...)
-end
-
 function invalid_variables(sys, expr)
     filter(x -> !any(isequal(x), all_symbols(sys)), reduce(vcat, vars(expr); init = []))
 end
@@ -155,7 +151,7 @@ function check_assignable(sys, sym)
     end
 end
 
-function compile_functional_affect(affect::ImperativeAffect, cb, sys; kwargs...)
+function compile_functional_affect(affect::ImperativeAffect, sys; kwargs...)
     #=
     Implementation sketch:
         generate observed function (oop), should save to a component array under obs_syms
@@ -235,14 +231,8 @@ function compile_functional_affect(affect::ImperativeAffect, cb, sys; kwargs...)
 
     upd_funs = NamedTuple{mod_names}((setu.((sys,), first.(mod_pairs))...,))
 
-    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
-        save_idxs = get(ic.callback_to_clocks, cb, Int[])
-    else
-        save_idxs = Int[]
-    end
-
     let user_affect = func(affect), ctx = context(affect)
-        function (integ)
+        @inline function (integ)
             # update the to-be-mutated values; this ensures that if you do a no-op then nothing happens
             modvals = mod_og_val_fun(integ.u, integ.p, integ.t)
             upd_component_array = NamedTuple{mod_names}(modvals)
@@ -256,10 +246,6 @@ function compile_functional_affect(affect::ImperativeAffect, cb, sys; kwargs...)
 
             # write the new values back to the integrator
             _generated_writeback(integ, upd_funs, upd_vals)
-
-            for idx in save_idxs
-                SciMLBase.save_discretes!(integ, idx)
-            end
         end
     end
 end

src/systems/systemstructure.jl

@@ -688,6 +688,7 @@ function _structural_simplify!(state::TearingState, io; simplify = false,
         check_consistency = true, fully_determined = true, warn_initialize_determined = false,
         dummy_derivative = true,
         kwargs...)
+
     if fully_determined isa Bool
         check_consistency &= fully_determined
     else