more tests

Author: TorkelE

src/systems/callbacks.jl

@@ -36,8 +36,12 @@ end
 function infer_discrete_parameters(affects)
     discrete_parameters = Set()
     for affect in affects
-        infer_discrete_parameters!(discrete_parameters, affect.lhs)
-        infer_discrete_parameters!(discrete_parameters, affect.rhs)
+        if affect isa Equation
+            infer_discrete_parameters!(discrete_parameters, affect.lhs)
+            infer_discrete_parameters!(discrete_parameters, affect.rhs)
+        elseif affect isa NamedTuple
+            haskey(affect, :modified) && union!(discrete_parameters, affect.modified)
+        end
     end
     return collect(discrete_parameters)
 end
@@ -49,6 +53,12 @@ function infer_discrete_parameters!(discrete_parameters, expr)
     # Change this coming line to a Symbolic append type of thing.
     union!(discrete_parameters, filter(ModelingToolkit.isparameter, dynamic_symvars))
 end
+
+# When updating vector variables, the affect side can be a vector.
+function infer_discrete_parameters!(discrete_parameters, expr_vec::Vector)
+    foreach(expr -> infer_discrete_parameters!(discrete_parameters, expr), expr_vec)
+end
+
 # Functions for replacing a Pre-call with a `1.0` (removing its content from an expression).
 is_precall(expr) = iscall(expr) ? operation(expr) isa Pre : false
 precall_to_1(expr) = (is_precall(expr) ? 1.0 : expr)

test/symbolic_events.jl

@@ -1441,9 +1441,7 @@ end
     @test_nowarn ODEProblem(sys, [], (0.0, 1.0))
 end
 
-
 @testset "Automatic inference of `discrete_parameters`" begin
-
     # Basic case, checks for both types of events (in combination and isolation).
     let
         # Creates models with continuous, discrete, or both types of events
@@ -1519,4 +1517,52 @@ end
         @test sol.ps[k1] == sol.ps[k2] == sol.ps[k3] == sol.ps[k4] == sol.ps[k5] == 0.0
         @test sol.ps[d1] == sol.ps[d2] == sol.ps[d3] == sol.ps[d4] == sol.ps[d5] == 0.0
     end
+
+    # Checks that everything works for vector-valued parameters and variables.
+    let
+        # Creates the model using the macro. 
+        @mtkmodel VectorParams begin
+            @parameters begin
+                k(t)[1:2] = [1, 1]
+                kup = 2.0
+            end
+            @variables begin
+                X(t)[1:2] = [4.0, 4.0]
+            end
+            @equations begin
+                D(X[1]) ~ -k[1]*X[1] + k[2]*X[2]
+                D(X[2]) ~ k[1]*X[1] - k[2]*X[2]
+            end
+            @continuous_events begin
+                (k[2] ~ t) => [k[1] ~ Pre(k[1] + kup)]
+            end
+        end
+        @mtkcompile model = VectorParams()
+
+        # Simulates the model. Checks that the correct values are achieved.
+        prob = ODEProblem(model, [], (0.0, 100.0))
+        sol = solve(prob, Rosenbrock23())
+        @test sol.ps[model.kup]  == 2.0
+        @test sol.ps[model.k[1]] == 3.0
+        @test sol.ps[model.k[2]] == 1.0
+        @test sol[model.X[1]][end] ≈ 2.0 atol = 1e-8 rtol = 1e-8
+        @test sol[model.X[2]][end] ≈ 6.0 atol = 1e-8 rtol = 1e-8
+    end
+
+    # Checks for a functional affect.
+    let
+        # Creates model.
+        @variables X(t) = 5.0
+        @parameters p = 2.0 d(t) = 1.0
+        eqs = [D(X) ~ p - d * X]
+        affect!(mod, obs, ctx, integ) = return (; d = 2.0)
+        cevent = [t ~ 1.0] => (f = affect!, modified = (; d))
+        @mtkcompile sys = System(eqs, t; continuous_events = [cevent])
+
+        # Simualtes the model and checks that values is correct.
+        sol = solve(ODEProblem(sys, [], (0.0, 100.0)), Rosenbrock23())
+        @test sol[X][end] ≈ 1.0 atol = 1e-8 rtol = 1e-8
+        @test sol.ps[p] == 2.0
+        @test sol.ps[d] == [1.0, 2.0]
+    end
 end