fix: fix collect_var

Author: vyudu

src/systems/callbacks.jl

@@ -301,7 +301,9 @@ function make_affect(affect::Vector{Equation}; discrete_parameters::AbstractVect
     # get accessed parameters p from Pre(p) in the callback parameters
     accessed_params = filter(isparameter, map(unPre, collect(pre_params)))
     union!(accessed_params, sys_params)
-    # add unknowns to the map
+
+    # add scalarized unknowns to the map.
+    _dvs = reduce(vcat, map(scalarize, _dvs), init = Any[])
     for u in _dvs
         aff_map[u] = u
     end
@@ -616,7 +618,8 @@ function compile_condition(
     end
 
     if !is_discrete(cbs)
-        condit = [cond.lhs - cond.rhs for cond in condit]
+        condit = reduce(vcat, flatten_equations(condit))
+        condit = condit isa AbstractVector ? [c.lhs - c.rhs for c in condit] : [condit.lhs - condit.rhs]
     end
 
     fs = build_function_wrapper(sys,

src/utils.jl

@@ -595,6 +595,7 @@ end
 
 function collect_var!(unknowns, parameters, var, iv; depth = 0)
     isequal(var, iv) && return nothing
+    var = unwrap(var)
     check_scope_depth(getmetadata(var, SymScope, LocalScope()), depth) || return nothing
     if iscalledparameter(var)
         callable = getcalledparameter(var)

test/odesystem.jl

@@ -1032,24 +1032,26 @@ prob = ODEProblem(sys, [x => 1.0], (0.0, 10.0))
 @test_nowarn solve(prob, Tsit5())
 
 # Issue#2383
-@variables x(t)[1:3]
-@parameters p[1:3, 1:3]
-eqs = [
-    D(x) ~ p * x
-]
-@mtkbuild sys = ODESystem(eqs, t; continuous_events = [[norm(x) ~ 3.0] => [x ~ ones(3)]])
-# array affect equations used to not work
-prob1 = @test_nowarn ODEProblem(sys, [x => ones(3)], (0.0, 10.0), [p => ones(3, 3)])
-sol1 = @test_nowarn solve(prob1, Tsit5())
-
-# array condition equations also used to not work
-@mtkbuild sys = ODESystem(
-    eqs, t; continuous_events = [[x ~ sqrt(3) * ones(3)] => [x ~ ones(3)]])
-# array affect equations used to not work
-prob2 = @test_nowarn ODEProblem(sys, [x => ones(3)], (0.0, 10.0), [p => ones(3, 3)])
-sol2 = @test_nowarn solve(prob2, Tsit5())
-
-@test sol1 ≈ sol2
+@testset "Arrays in affect/condition equations" begin
+    @variables x(t)[1:3]
+    @parameters p[1:3, 1:3]
+    eqs = [
+        D(x) ~ p * x
+    ]
+    @mtkbuild sys = ODESystem(eqs, t; continuous_events = [[norm(x) ~ 3.0] => [x ~ ones(3)]])
+    # array affect equations used to not work
+    prob1 = @test_nowarn ODEProblem(sys, [x => ones(3)], (0.0, 10.0), [p => ones(3, 3)])
+    sol1 = @test_nowarn solve(prob1, Tsit5())
+    
+    # array condition equations also used to not work
+    @mtkbuild sys = ODESystem(
+        eqs, t; continuous_events = [[x ~ sqrt(3) * ones(3)] => [x ~ ones(3)]])
+    # array affect equations used to not work
+    prob2 = @test_nowarn ODEProblem(sys, [x => ones(3)], (0.0, 10.0), [p => ones(3, 3)])
+    sol2 = @test_nowarn solve(prob2, Tsit5())
+    
+    @test sol1.u ≈ sol2.u[2:end]
+end
 
 # Requires fix in symbolics for `linear_expansion(p * x, D(y))`
 @test_skip begin
@@ -1196,10 +1198,12 @@ end
 end
 
 # Namespacing of array variables
-@variables x(t)[1:2]
-@named sys = ODESystem(Equation[], t)
-@test getname(unknowns(sys, x)) == :sys₊x
-@test size(unknowns(sys, x)) == size(x)
+@testset "Namespacing of array variables" begin
+    @variables x(t)[1:2]
+    @named sys = ODESystem(Equation[], t)
+    @test getname(unknowns(sys, x)) == :sys₊x
+    @test size(unknowns(sys, x)) == size(x)
+end
 
 # Issue#2667 and Issue#2953
 @testset "ForwardDiff through ODEProblem constructor" begin
@@ -1537,8 +1541,7 @@ end
 @testset "Observed variables dependent on discrete parameters" begin
     @variables x(t) obs(t)
     @parameters c(t)
-    @mtkbuild sys = ODESystem(
-        [D(x) ~ c * cos(x), obs ~ c], t, [x], [c]; discrete_events = [1.0 => [c ~ c + 1]])
+    @mtkbuild sys = ODESystem([D(x) ~ c * cos(x), obs ~ c], t, [x], [c]; discrete_events = [SymbolicDiscreteCallback(1.0 => [c ~ Pre(c) + 1], discrete_parameters = [c])])
     prob = ODEProblem(sys, [x => 0.0], (0.0, 2pi), [c => 1.0])
     sol = solve(prob, Tsit5())
     @test sol[obs] ≈ 1:7
@@ -1598,15 +1601,15 @@ end
 # Test `isequal`
 @testset "`isequal`" begin
     @variables X(t)
-    @parameters p d
+    @parameters p d(t)
     eq = D(X) ~ p - d * X
 
     osys1 = complete(ODESystem([eq], t; name = :osys))
     osys2 = complete(ODESystem([eq], t; name = :osys))
     @test osys1 == osys2 # true
 
-    continuous_events = [[X ~ 1.0] => [X ~ X + 5.0]]
-    discrete_events = [5.0 => [d ~ d / 2.0]]
+    continuous_events = [[X ~ 1.0] => [X ~ Pre(X) + 5.0]]
+    discrete_events = [SymbolicDiscreteCallback(5.0 => [d ~ d / 2.0], discrete_parameters = [d])]
 
     osys1 = complete(ODESystem([eq], t; name = :osys, continuous_events))
     osys2 = complete(ODESystem([eq], t; name = :osys))