fix: fix unscalarized array passed to discrete_parameters

Author: AayushSabharwal

src/systems/callbacks.jl

@@ -904,7 +904,12 @@ function compile_equational_affect(
             obseqs, Dict([p => unPre(p) for p in parameters(affsys)]))
         rhss = map(x -> x.rhs, update_eqs)
         lhss = map(x -> x.lhs, update_eqs)
-        is_p = [lhs in Set(ps_to_update) for lhs in lhss]
+        update_ps_set = Set(ps_to_update)
+        is_p = map(lhss) do lhs
+            lhs in update_ps_set ||
+                iscall(lhs) && operation(lhs) === getindex &&
+                arguments(lhs)[1] in update_ps_set
+        end
         is_u = [lhs in Set(dvs_to_update) for lhs in lhss]
         dvs = unknowns(sys)
         ps = parameters(sys)

test/symbolic_events.jl

@@ -1495,3 +1495,43 @@ end
     @test v_sol[v] ≈ [vini, -vini, vini]
     @test M_sol[M] ≈ [Mini, -Mini, Mini]
 end
+
+@testset "Issue#3990: Scalarized array passed to `discrete_parameters` of symbolic affect" begin
+    N = 2
+    @parameters v(t)[1:N]
+    @parameters M(t)[1:N, 1:N]
+
+    @variables x(t)
+
+    Mini = rand(N, N) ./ (N^2)
+    vini = vec(sum(Mini, dims = 1))
+
+    v_eq = [D(x) ~ x * Symbolics.scalarize(sum(v))]
+    M_eq = [D(x) ~ x * Symbolics.scalarize(sum(M))]
+
+    v_event = ModelingToolkit.SymbolicDiscreteCallback(
+        1.0,
+        [v ~ -Pre(v)],
+        discrete_parameters = collect(v)
+    )
+
+    M_event = ModelingToolkit.SymbolicDiscreteCallback(
+        1.0,
+        [M ~ -Pre(M)],
+        discrete_parameters = vec(collect(M))
+    )
+
+    @mtkcompile v_sys = System(v_eq, t; discrete_events = v_event)
+    @mtkcompile M_sys = System(M_eq, t; discrete_events = M_event)
+
+    u0p0_map = Dict(x => 1.0, M => Mini, s => sMini, v => vini)
+
+    v_prob = ODEProblem(v_sys, u0p0_map, (0.0, 2.5))
+    M_prob = ODEProblem(M_sys, u0p0_map, (0.0, 2.5))
+
+    v_sol = solve(v_prob, Tsit5())
+    M_sol = solve(M_prob, Tsit5())
+
+    @test v_sol[v] ≈ [vini, -vini, vini]
+    @test M_sol[M] ≈ [Mini, -Mini, Mini]
+end