fix: add reinitalizealg back

Author: vyudu

ext/MTKFMIExt.jl

@@ -93,7 +93,7 @@ with the name `namespace__variable`.
 - `name`: The name of the system.
 """
 function MTK.FMIComponent(::Val{Ver}; fmu = nothing, tolerance = 1e-6,
-        communication_step_size = nothing, type, name) where {Ver}
+        communication_step_size = nothing, type, name, reinitializealg = nothing) where {Ver}
     if Ver != 2 && Ver != 3
         throw(ArgumentError("FMI Version must be `2` or `3`"))
     end
@@ -238,7 +238,7 @@ function MTK.FMIComponent(::Val{Ver}; fmu = nothing, tolerance = 1e-6,
         finalize_affect = MTK.FunctionalAffect(fmiFinalize!, [], [wrapper], [])
         step_affect = MTK.FunctionalAffect(Returns(nothing), [], [], [])
         instance_management_callback = MTK.SymbolicDiscreteCallback(
-            (t != t - 1), step_affect; finalize = finalize_affect)
+            (t != t - 1), step_affect; finalize = finalize_affect, reinitializealg)
 
         push!(params, wrapper)
         append!(observed, der_observed)
@@ -279,7 +279,7 @@ function MTK.FMIComponent(::Val{Ver}; fmu = nothing, tolerance = 1e-6,
             fmiCSStep!; observed = cb_observed, modified = cb_modified, ctx = _functor)
         instance_management_callback = MTK.SymbolicDiscreteCallback(
             communication_step_size, step_affect; initialize = initialize_affect,
-            finalize = finalize_affect)
+            finalize = finalize_affect, reinitializealg)
 
         # guarded in case there are no outputs/states and the variable is `[]`.
         symbolic_type(__mtk_internal_o) == NotSymbolic() || push!(params, __mtk_internal_o)

src/systems/callbacks.jl

@@ -200,7 +200,9 @@ Affects (i.e. `affect` and `affect_neg`) can be specified as either:
     + `ctx` is a user-defined context object passed to `f!` when invoked. This value is aliased for each problem.
 * A [`ImperativeAffect`](@ref); refer to its documentation for details.
 
-DAEs will automatically be reinitialized.
+`reinitializealg` is used to set how the system will be reinitialized after the callback. 
+- Symbolic affects have reinitialization built in. In this case the algorithm will default to SciMLBase.NoInit(), and should **not** be provided.
+- Functional and imperative affects will default to SciMLBase.CheckInit(), which will error if the system is not properly reinitialized after the callback. If your system is a DAE, pass in an algorithm like SciMLBase.BrownBasicFullInit() to properly re-initialize.
 
 Initial and final affects can also be specified identically to positive and negative edge affects. Initialization affects
 will run as soon as the solver starts, while finalization affects will be executed after termination.
@@ -212,6 +214,7 @@ struct SymbolicContinuousCallback <: AbstractCallback
     initialize::Union{Affect, Nothing}
     finalize::Union{Affect, Nothing}
     rootfind::Union{Nothing, SciMLBase.RootfindOpt}
+    reinitializealg::SciMLBase.DAEInitializationAlgorithm
 
     function SymbolicContinuousCallback(
             conditions::Union{Equation, Vector{Equation}},
@@ -221,13 +224,21 @@ struct SymbolicContinuousCallback <: AbstractCallback
             initialize = nothing,
             finalize = nothing,
             rootfind = SciMLBase.LeftRootFind,
+            reinitializealg = nothing,
             iv = nothing,
             algeeqs = Equation[])
         conditions = (conditions isa AbstractVector) ? conditions : [conditions]
+
+        if isnothing(reinitializealg)
+            any(a -> (a isa FunctionalAffect || a isa ImperativeAffect), [affect, affect_neg, initialize, finalize]) ? 
+                reinitializealg = SciMLBase.CheckInit() : 
+                reinitializealg = SciMLBase.NoInit()
+        end
+
         new(conditions, make_affect(affect; iv, algeeqs, discrete_parameters),
             make_affect(affect_neg; iv, algeeqs, discrete_parameters),
             make_affect(initialize; iv, algeeqs, discrete_parameters), make_affect(
-                finalize; iv, algeeqs, discrete_parameters), rootfind)
+                finalize; iv, algeeqs, discrete_parameters), rootfind, reinitializealg)
     end # Default affect to nothing
 end
 
@@ -424,16 +435,22 @@ struct SymbolicDiscreteCallback <: AbstractCallback
     affect::Union{Affect, Nothing}
     initialize::Union{Affect, Nothing}
     finalize::Union{Affect, Nothing}
+    reinitializealg::SciMLBase.DAEInitializationAlgorithm
 
     function SymbolicDiscreteCallback(
             condition, affect = nothing;
             initialize = nothing, finalize = nothing, iv = nothing,
-            algeeqs = Equation[], discrete_parameters = Any[])
+            algeeqs = Equation[], discrete_parameters = Any[], reinitializealg = nothing)
         c = is_timed_condition(condition) ? condition : value(scalarize(condition))
 
+        if isnothing(reinitializealg)
+            any(a -> (a isa FunctionalAffect || a isa ImperativeAffect), [affect, affect_neg, initialize, finalize]) ? 
+                reinitializealg = SciMLBase.CheckInit() : 
+                reinitializealg = SciMLBase.NoInit()
+        end
         new(c, make_affect(affect; iv, algeeqs, discrete_parameters),
             make_affect(initialize; iv, algeeqs, discrete_parameters),
-            make_affect(finalize; iv, algeeqs, discrete_parameters))
+            make_affect(finalize; iv, algeeqs, discrete_parameters), reinitializealg)
     end # Default affect to nothing
 end
 
@@ -525,7 +542,8 @@ function Base.hash(cb::AbstractCallback, s::UInt)
     !is_discrete(cb) && (s = hash(affect_negs(cb), s))
     s = hash(initialize_affects(cb), s)
     s = hash(finalize_affects(cb), s)
-    !is_discrete(cb) ? hash(cb.rootfind, s) : s
+    !is_discrete(cb) && (s = hash(cb.rootfind, s))
+    hash(cb.reinitializealg, s)
 end
 
 ###########################
@@ -562,7 +580,7 @@ end
 function Base.:(==)(e1::AbstractCallback, e2::AbstractCallback)
     (is_discrete(e1) === is_discrete(e2)) || return false
     (isequal(e1.conditions, e2.conditions) && isequal(e1.affect, e2.affect) &&
-     isequal(e1.initialize, e2.initialize) && isequal(e1.finalize, e2.finalize)) ||
+     isequal(e1.initialize, e2.initialize) && isequal(e1.finalize, e2.finalize)) && isequal(e1.reinitializealg, e2.reinitializealg) ||
         return false
     is_discrete(e1) ||
         (isequal(e1.affect_neg, e2.affect_neg) && isequal(e1.rootfind, e2.rootfind))
@@ -664,15 +682,15 @@ function generate_continuous_callbacks(sys::AbstractSystem, dvs = unknowns(sys),
         ps = parameters(sys; initial_parameters = true); kwargs...)
     cbs = continuous_events(sys)
     isempty(cbs) && return nothing
-    cb_classes = Dict{SciMLBase.RootfindOpt, Vector{SymbolicContinuousCallback}}()
+    cb_classes = Dict{Tuple{SciMLBase.RootfindOpt, SciMLBase.DAEReinitializationAlg}, Vector{SymbolicContinuousCallback}}()
 
     # Sort the callbacks by their rootfinding method
     for cb in cbs
-        _cbs = get!(() -> SymbolicContinuousCallback[], cb_classes, cb.rootfind)
+        _cbs = get!(() -> SymbolicContinuousCallback[], cb_classes, (cb.rootfind, cb.reinitializealg))
         push!(_cbs, cb)
     end
-    cb_classes = sort!(OrderedDict(cb_classes))
-    compiled_callbacks = [generate_callback(cb, sys; kwargs...) for (rf, cb) in cb_classes]
+    sort!(OrderedDict(cb_classes), by = cb -> cb.rootfind)
+    compiled_callbacks = [generate_callback(cb, sys; kwargs...) for ((rf, reinit), cb) in cb_classes]
     if length(compiled_callbacks) == 1
         return only(compiled_callbacks)
     else
@@ -741,7 +759,7 @@ function generate_callback(cbs::Vector{SymbolicContinuousCallback}, sys; kwargs.
 
     return VectorContinuousCallback(
         trigger, affect, affect_neg, length(eqs); initialize, finalize,
-        rootfind = cbs[1].rootfind, initializealg = SciMLBase.NoInit())
+        rootfind = cbs[1].rootfind, initializealg = cbs[1].reinitializealg)
 end
 
 function generate_callback(cb, sys; kwargs...)
@@ -768,16 +786,16 @@ function generate_callback(cb, sys; kwargs...)
     if is_discrete(cb)
         if is_timed && conditions(cb) isa AbstractVector
             return PresetTimeCallback(trigger, affect; initialize,
-                finalize, initializealg = SciMLBase.NoInit())
+                finalize, initializealg = cb.reinitializealg)
         elseif is_timed
-            return PeriodicCallback(affect, trigger; initialize, finalize, initializealg = SciMLBase.NoInit())
+            return PeriodicCallback(affect, trigger; initialize, finalize, initializealg = cb.reinitializealg)
         else
             return DiscreteCallback(trigger, affect; initialize,
-                finalize, initializealg = SciMLBase.NoInit())
+                finalize, initializealg = cb.reinitializealg)
         end
     else
         return ContinuousCallback(trigger, affect, affect_neg; initialize, finalize,
-            rootfind = cb.rootfind, initializealg = SciMLBase.NoInit())
+            rootfind = cb.rootfind, initializealg = cb.reinitializealg)
     end
 end
 

src/systems/model_parsing.jl

@@ -64,8 +64,6 @@ function _model_macro(mod, name, expr, isconnector)
     push!(exprs.args, :(systems = ODESystem[]))
     push!(exprs.args, :(equations = Union{Equation, Vector{Equation}}[]))
     push!(exprs.args, :(defaults = Dict{Num, Union{Number, Symbol, Function}}()))
-    push!(exprs.args, :(disc_events = []))
-    push!(exprs.args, :(cont_events = []))
 
     Base.remove_linenums!(expr)
     for arg in expr.args
@@ -107,8 +105,6 @@ function _model_macro(mod, name, expr, isconnector)
     push!(exprs.args, :(push!(parameters, $(ps...))))
     push!(exprs.args, :(push!(systems, $(comps...))))
     push!(exprs.args, :(push!(variables, $(vs...))))
-    push!(exprs.args, :(push!(disc_events, $(d_evts...))))
-    push!(exprs.args, :(push!(cont_events, $(c_evts...))))
 
     gui_metadata = isassigned(icon) > 0 ? GUIMetadata(GlobalRef(mod, name), icon[]) :
                    GUIMetadata(GlobalRef(mod, name))
@@ -120,7 +116,7 @@ function _model_macro(mod, name, expr, isconnector)
 
     sys = :($ODESystem($(flatten_equations)(equations), $iv, variables, parameters;
         name, description = $description, systems, gui_metadata = $gui_metadata,
-        defaults, continuous_events = cont_events, discrete_events = disc_events))
+        defaults))
 
     if length(ext) == 0
         push!(exprs.args, :(var"#___sys___" = $sys))
@@ -131,6 +127,18 @@ function _model_macro(mod, name, expr, isconnector)
     isconnector && push!(exprs.args,
         :($Setfield.@set!(var"#___sys___".connector_type=$connector_type(var"#___sys___"))))
 
+    !isempty(c_evts) && push!(exprs.args,
+        :($Setfield.@set!(var"#___sys___".continuous_events=$SymbolicContinuousCallback.([
+            $(c_evts...)
+        ]))))
+
+    @show d_evts
+    !isempty(d_evts) && push!(exprs.args,
+        :($Setfield.@set!(var"#___sys___".discrete_events=$SymbolicDiscreteCallback.([
+            $(d_evts...)
+        ]))))
+
+
     f = if length(where_types) == 0
         :($(Symbol(:__, name, :__))(; name, $(kwargs...)) = $exprs)
     else

test/fmi/fmi.jl

@@ -157,7 +157,7 @@ end
     @testset "v2, CS" begin
         fmu = loadFMU(joinpath(FMU_DIR, "SimpleAdder.fmu"); type = :CS)
         @named adder = MTK.FMIComponent(
-            Val(2); fmu, type = :CS, communication_step_size = 1e-6)
+                                        Val(2); fmu, type = :CS, communication_step_size = 1e-6, reinitializealg = BrownFullBasicInit())
         @test MTK.isinput(adder.a)
         @test MTK.isinput(adder.b)
         @test MTK.isoutput(adder.out)
@@ -209,7 +209,7 @@ end
     @testset "v3, CS" begin
         fmu = loadFMU(joinpath(FMU_DIR, "StateSpace.fmu"); type = :CS)
         @named sspace = MTK.FMIComponent(
-            Val(3); fmu, communication_step_size = 1e-6, type = :CS)
+                                         Val(3); fmu, communication_step_size = 1e-6, type = :CS, reinitializealg = BrownFullBasicInit())
         @test MTK.isinput(sspace.u)
         @test MTK.isoutput(sspace.y)
         @test !MTK.isinput(sspace.x) && !MTK.isoutput(sspace.x)

test/symbolic_events.jl

@@ -1204,7 +1204,7 @@ end
     @mtkmodel DECAY begin
         @parameters begin
             unrelated[1:2] = zeros(2)
-            k = 0.0
+            k(t) = 0.0
         end
         @variables begin
             x(t) = 10.0
@@ -1213,7 +1213,7 @@ end
             D(x) ~ -k * x
         end
         @discrete_events begin
-            (t == 1.0) => [k ~ 1.0]
+            (t == 1.0) => [k ~ 1.0], discrete_parameters = [k]
         end
     end
     @mtkbuild decay = DECAY()
@@ -1338,7 +1338,4 @@ end
     sol = solve(prob, FBDF())
     @test prob.ps[g] == sol.ps[g]
 end
-# TODO: test:
-# - Functional affects reinitialize correctly
 # - explicit equation of t in a functional affect
-# - reinitialization after affects