From 69ab9322b0124900d1904110632ba70fddf65883 Mon Sep 17 00:00:00 2001
From: Christian Gutsche <christian.gutsche@mailbox.tu-dresden.de>
Date: Wed, 17 Jan 2024 10:47:14 +0100
Subject: [PATCH] adding event support and docs, test rebase

---
 docs/src/basics/MTKModel_Connector.md | 58 ++++++++++++++++++++++++++-
 src/systems/model_parsing.jl          | 13 ++++--
 test/model_parsing.jl                 | 30 ++++++++++++++
 3 files changed, 96 insertions(+), 5 deletions(-)

diff --git a/docs/src/basics/MTKModel_Connector.md b/docs/src/basics/MTKModel_Connector.md
index ac7f542152..5a9fec08a4 100644
--- a/docs/src/basics/MTKModel_Connector.md
+++ b/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:
 
@@ -177,11 +179,65 @@ getdefault(model_c3.model_a.k_array[2])
 
   - List all the equations here
 
+
 #### `@defaults` begin block
 
   - 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.
@@ -380,4 +436,4 @@ Using ternary operator or if-elseif-else statement, conditional initial guesses
         p = flag ? 1 : 2
     end
 end
-```
+```
\ No newline at end of file
diff --git a/src/systems/model_parsing.jl b/src/systems/model_parsing.jl
index 9da2de6695..1ec9c3692d 100644
--- a/src/systems/model_parsing.jl
+++ b/src/systems/model_parsing.jl
@@ -118,11 +118,15 @@ 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...)]))))
+    !(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...)]))))
+    !(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)
@@ -779,6 +783,7 @@ 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
diff --git a/test/model_parsing.jl b/test/model_parsing.jl
index dcacbc4d20..4f9f62ea25 100644
--- a/test/model_parsing.jl
+++ b/test/model_parsing.jl
@@ -426,6 +426,36 @@ end
     @test A.structure[:components] == [[:cc, :C]]
 end
 
+@testset "Event handeling 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, 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