add discrete_events field for ODEs

Author: isaacsas

src/structural_transformation/codegen.jl

@@ -1,6 +1,6 @@
 using LinearAlgebra
 
-using ModelingToolkit: isdifferenceeq, has_continuous_events, process_events
+using ModelingToolkit: isdifferenceeq, process_events
 
 const MAX_INLINE_NLSOLVE_SIZE = 8
 

src/systems/abstractsystem.jl

@@ -1039,6 +1039,7 @@ function Base.hash(sys::AbstractSystem, s::UInt)
     end
     s = foldr(hash, get_observed(sys), init = s)
     s = foldr(hash, get_continuous_events(sys), init = s)
+    s = foldr(hash, get_discrete_events(sys), init = s)
     s = hash(independent_variables(sys), s)
     return s
 end
@@ -1066,16 +1067,17 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem; name::Symbol = nam
     sts = union(get_states(basesys), get_states(sys))
     ps = union(get_ps(basesys), get_ps(sys))
     obs = union(get_observed(basesys), get_observed(sys))
-    evs = union(get_continuous_events(basesys), get_continuous_events(sys))
+    cevs = union(get_continuous_events(basesys), get_continuous_events(sys))
+    devs = union(get_discrete_events(basesys), get_discrete_events(sys))
     defs = merge(get_defaults(basesys), get_defaults(sys)) # prefer `sys`
     syss = union(get_systems(basesys), get_systems(sys))
 
     if length(ivs) == 0
         T(eqs, sts, ps, observed = obs, defaults = defs, name = name, systems = syss,
-          continuous_events = evs)
+          continuous_events = cevs, discrete_events = devs)
     elseif length(ivs) == 1
         T(eqs, ivs[1], sts, ps, observed = obs, defaults = defs, name = name,
-          systems = syss, continuous_events = evs)
+          systems = syss, continuous_events = cevs, discrete_events = devs)
     end
 end
 

src/systems/callbacks.jl

@@ -5,8 +5,7 @@ has_continuous_events(sys::AbstractSystem) = isdefined(sys, :continuous_events)
 
 has_discrete_events(sys::AbstractSystem) = isdefined(sys, :discrete_events)
 function get_discrete_events(sys::AbstractSystem)
-    has_discrete_events(sys) ||
-        error("Systems of type $(typeof(sys)) do not support discrete events.")
+    has_discrete_events(sys) || return SymbolicDiscreteCallback[]
     getfield(sys, :discrete_events)
 end
 
@@ -62,10 +61,10 @@ affect_equations(cb::SymbolicContinuousCallback) = cb.affect
 function affect_equations(cbs::Vector{SymbolicContinuousCallback})
     reduce(vcat, [affect_equations(cb) for cb in cbs])
 end
-namespace_equation(cb::SymbolicContinuousCallback, s)::SymbolicContinuousCallback = SymbolicContinuousCallback(namespace_equation.(equations(cb),
-                                                                                                                                   (s,)),
-                                                                                                               namespace_equation.(affect_equations(cb),
-                                                                                                                                   (s,)))
+function namespace_equation(cb::SymbolicContinuousCallback, s)::SymbolicContinuousCallback
+    SymbolicContinuousCallback(namespace_equation.(equations(cb), (s,)),
+                               namespace_equation.(affect_equations(cb), (s,)))
+end
 
 function continuous_events(sys::AbstractSystem)
     obs = get_continuous_events(sys)
@@ -124,6 +123,11 @@ function namespace_equation(cb::SymbolicDiscreteCallback, s)::SymbolicDiscreteCa
                              namespace_equation.(affect_equations(cb), Ref(s)))
 end
 
+SymbolicDiscreteCallbacks(cb::SymbolicDiscreteCallback) = [cb]
+SymbolicDiscreteCallbacks(cbs::Vector{<:SymbolicDiscreteCallback}) = cbs
+SymbolicDiscreteCallbacks(cbs::Vector) = SymbolicDiscreteCallback.(cbs)
+SymbolicDiscreteCallbacks(::Nothing) = SymbolicDiscreteCallback[]
+
 function discrete_events(sys::AbstractSystem)
     obs = get_discrete_events(sys)
     systems = get_systems(sys)

src/systems/diffeqs/odesystem.jl

@@ -102,6 +102,12 @@ struct ODESystem <: AbstractODESystem
     """
     continuous_events::Vector{SymbolicContinuousCallback}
     """
+    discrete_events: A `Vector{SymbolicDiscreteCallback}` that models events. Symbolic
+    analog to `SciMLBase.DiscreteCallback` that exectues an affect when a given condition is
+    true at the end of an integration step.
+    """
+    discrete_events::Vector{SymbolicDiscreteCallback}
+    """
     tearing_state: cache for intermediate tearing state
     """
     tearing_state::Any
@@ -112,19 +118,20 @@ struct ODESystem <: AbstractODESystem
 
     function ODESystem(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad,
                        jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults,
-                       torn_matching, connector_type, connections, preface, events,
-                       tearing_state = nothing, substitutions = nothing;
+                       torn_matching, connector_type, connections, preface, cevents,
+                       devents, tearing_state = nothing, substitutions = nothing;
                        checks::Bool = true)
         if checks
             check_variables(dvs, iv)
             check_parameters(ps, iv)
             check_equations(deqs, iv)
-            check_equations(equations(events), iv)
+            check_equations(equations(cevents), iv)
             all_dimensionless([dvs; ps; iv]) || check_units(deqs)
         end
         new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, torn_matching,
-            connector_type, connections, preface, events, tearing_state, substitutions)
+            connector_type, connections, preface, cevents, devents, tearing_state,
+            substitutions)
     end
 end
 
@@ -139,6 +146,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
                    connector_type = nothing,
                    preface = nothing,
                    continuous_events = nothing,
+                   discrete_events = nothing,
                    checks = true)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
@@ -172,9 +180,11 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         throw(ArgumentError("System names must be unique."))
     end
     cont_callbacks = SymbolicContinuousCallbacks(continuous_events)
+    disc_callbacks = SymbolicDiscreteCallbacks(discrete_events)
     ODESystem(deqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac,
               ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing,
-              connector_type, nothing, preface, cont_callbacks, checks = checks)
+              connector_type, nothing, preface, cont_callbacks, disc_callbacks,
+              checks = checks)
 end
 
 function ODESystem(eqs, iv = nothing; kwargs...)
@@ -244,6 +254,7 @@ function flatten(sys::ODESystem, noeqs = false)
                          parameters(sys),
                          observed = observed(sys),
                          continuous_events = continuous_events(sys),
+                         discrete_events = discrete_events(sys),
                          defaults = defaults(sys),
                          name = nameof(sys),
                          checks = false)

src/systems/nonlinear/nonlinearsystem.jl

@@ -83,9 +83,13 @@ function NonlinearSystem(eqs, states, ps;
                          systems = NonlinearSystem[],
                          connector_type = nothing,
                          continuous_events = nothing, # this argument is only required for ODESystems, but is added here for the constructor to accept it without error
+                         discrete_events = nothing,   # this argument is only required for ODESystems, but is added here for the constructor to accept it without error
                          checks = true)
     continuous_events === nothing || isempty(continuous_events) ||
         throw(ArgumentError("NonlinearSystem does not accept `continuous_events`, you provided $continuous_events"))
+    discrete_events === nothing || isempty(discrete_events) ||
+        throw(ArgumentError("NonlinearSystem does not accept `discrete_events`, you provided $discrete_events"))
+
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     # Move things over, but do not touch array expressions

test/runtests.jl

@@ -39,7 +39,7 @@ println("Last test requires gcc available in the path!")
 @testset "Serialization" begin include("serialization.jl") end
 @safetestset "print_tree" begin include("print_tree.jl") end
 @safetestset "error_handling" begin include("error_handling.jl") end
-@safetestset "Callbacks" begin include("root_equations.jl") end
+@safetestset "root_equations" begin include("root_equations.jl") end
 @safetestset "state_selection" begin include("state_selection.jl") end
 @safetestset "Modelingtoolkitize Test" begin include("modelingtoolkitize.jl") end
 @safetestset "ControlSystem Test" begin include("controlsystem.jl") end