Merge pull request #2427 from cgutsche/dev-add-event-support-mtkmodel

Add event support mtkmodel

Author: ChrisRackauckas

docs/src/basics/MTKModel_Connector.md

@@ -32,6 +32,8 @@ equations.
   - `@parameters`: for specifying the symbolic parameters
   - `@structural_parameters`: for specifying non-symbolic parameters
   - `@variables`: for specifying the unknowns
+  - `@continuous_events`: for specifying a list of continuous events
+  - `@discrete_events`: for specifying a list of discrete events
 
 Let's explore these in more detail with the following example:
 
@@ -182,6 +184,59 @@ getdefault(model_c3.model_a.k_array[2])
   - Default values can be passed as pairs.
   - This is equivalent to passing `defaults` argument to `ODESystem`.
 
+#### `@continuous_events` begin block
+
+  - Defining continuous events as described [here](https://docs.sciml.ai/ModelingToolkit/stable/basics/Events/#Continuous-Events).
+  - If this block is not defined in the model, no continuous events will be added.
+
+```@example mtkmodel-example
+using ModelingToolkit
+
+@mtkmodel M begin
+    @parameters begin
+        k
+    end
+    @variables begin
+        x(t)
+        y(t)
+    end
+    @equations begin
+        x ~ k * D(x)
+        D(y) ~ -k
+    end
+    @continuous_events begin
+        [x ~ 1.5] => [x ~ 5, y ~ 5]
+        [t ~ 4] => [x ~ 10]
+    end
+end
+```
+
+#### `@discrete_events` begin block
+
+  - Defining discrete events as described [here](https://docs.sciml.ai/ModelingToolkit/stable/basics/Events/#Discrete-events-support).
+  - If this block is not defined in the model, no discrete events will be added.
+
+```@example mtkmodel-example
+using ModelingToolkit
+
+@mtkmodel M begin
+    @parameters begin
+        k
+    end
+    @variables begin
+        x(t)
+        y(t)
+    end
+    @equations begin
+        x ~ k * D(x)
+        D(y) ~ -k
+    end
+    @discrete_events begin
+        (t == 1.5) => [x ~ x + 5, y ~ 5]
+    end
+end
+```
+
 #### A begin block
 
   - Any other Julia operations can be included with dedicated begin blocks.

src/systems/model_parsing.jl

@@ -48,6 +48,8 @@ function _model_macro(mod, name, expr, isconnector)
     eqs = Expr[]
     icon = Ref{Union{String, URI}}()
     ps, sps, vs, = [], [], []
+    c_evts = []
+    d_evts = []
     kwargs = Set()
     where_types = Expr[]
 
@@ -61,7 +63,7 @@ function _model_macro(mod, name, expr, isconnector)
     for arg in expr.args
         if arg.head == :macrocall
             parse_model!(exprs.args, comps, ext, eqs, icon, vs, ps,
-                sps, dict, mod, arg, kwargs, where_types)
+                sps, c_evts, d_evts, dict, mod, arg, kwargs, where_types)
         elseif arg.head == :block
             push!(exprs.args, arg)
         elseif arg.head == :if
@@ -116,6 +118,16 @@ function _model_macro(mod, name, expr, isconnector)
     isconnector && push!(exprs.args,
         :($Setfield.@set!(var"#___sys___".connector_type=$connector_type(var"#___sys___"))))
 
+    !(c_evts == []) && push!(exprs.args,
+        :($Setfield.@set!(var"#___sys___".continuous_events=$SymbolicContinuousCallback.([
+            $(c_evts...)
+        ]))))
+
+    !(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
@@ -124,6 +136,7 @@ function _model_macro(mod, name, expr, isconnector)
             :($(Symbol(:__, name, :__))(; name, $(kwargs...))), where_types...)
         :($f_with_where = $exprs)
     end
+
     :($name = $Model($f, $dict, $isconnector))
 end
 
@@ -341,7 +354,7 @@ function get_var(mod::Module, b)
     end
 end
 
-function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps,
+function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps, c_evts, d_evts,
         dict, mod, arg, kwargs, where_types)
     mname = arg.args[1]
     body = arg.args[end]
@@ -359,6 +372,10 @@ function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps,
         parse_equations!(exprs, eqs, dict, body)
     elseif mname == Symbol("@constants")
         parse_constants!(exprs, dict, body, mod)
+    elseif mname == Symbol("@continuous_events")
+        parse_continuous_events!(c_evts, dict, body)
+    elseif mname == Symbol("@discrete_events")
+        parse_discrete_events!(d_evts, dict, body)
     elseif mname == Symbol("@icon")
         isassigned(icon) && error("This model has more than one icon.")
         parse_icon!(body, dict, icon, mod)
@@ -753,6 +770,24 @@ function parse_equations!(exprs, eqs, dict, body)
     end
 end
 
+function parse_continuous_events!(c_evts, dict, body)
+    dict[:continuous_events] = []
+    Base.remove_linenums!(body)
+    for arg in body.args
+        push!(c_evts, arg)
+        push!(dict[:continuous_events], readable_code.(c_evts)...)
+    end
+end
+
+function parse_discrete_events!(d_evts, dict, body)
+    dict[:discrete_events] = []
+    Base.remove_linenums!(body)
+    for arg in body.args
+        push!(d_evts, arg)
+        push!(dict[:discrete_events], readable_code.(d_evts)...)
+    end
+end
+
 function parse_icon!(body::String, dict, icon, mod)
     icon_dir = get(ENV, "MTK_ICONS_DIR", joinpath(DEPOT_PATH[1], "mtk_icons"))
     dict[:icon] = icon[] = if isfile(body)

test/model_parsing.jl

@@ -426,6 +426,36 @@ end
     @test A.structure[:components] == [[:cc, :C]]
 end
 
+@testset "Event handling in MTKModel" begin
+    @mtkmodel M begin
+        @variables begin
+            x(t)
+            y(t)
+            z(t)
+        end
+        @equations begin
+            x ~ -D(x)
+            D(y) ~ 0
+            D(z) ~ 0
+        end
+        @continuous_events begin
+            [x ~ 1.5] => [x ~ 5, y ~ 1]
+        end
+        @discrete_events begin
+            (t == 1.5) => [x ~ x + 5, z ~ 2]
+        end
+    end
+
+    @mtkbuild model = M()
+    u0 = [model.x => 10, model.y => 0, model.z => 0]
+
+    prob = ODEProblem(model, u0, (0, 5.0))
+    sol = solve(prob, Tsit5(), tstops = [1.5])
+
+    @test isequal(sol[model.y][end], 1.0)
+    @test isequal(sol[model.z][end], 2.0)
+end
+
 # Ensure that modules consisting MTKModels with component arrays and icons of
 # `Expr` type and `unit` metadata can be precompiled.
 module PrecompilationTest