fix: make vars search through arrays of symbolics

SciML · Oct 11, 2024 · b83a251 · b83a251
1 parent 816fde7
commit b83a251
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/src/utils.jl b/src/utils.jl
@@ -385,6 +385,12 @@ function vars!(vars, O; op = Differential)
     if isvariable(O)
         return push!(vars, O)
     end
+    if symbolic_type(O) == NotSymbolic() && O isa AbstractArray
+        for arg in O
+            vars!(vars, arg; op)
+        end
+        return vars
+    end
     !iscall(O) && return vars
 
     operation(O) isa op && return push!(vars, O)

diff --git a/test/variable_utils.jl b/test/variable_utils.jl
@@ -1,5 +1,5 @@
 using ModelingToolkit, Test
-using ModelingToolkit: value
+using ModelingToolkit: value, vars
 using SymbolicUtils: <ₑ
 @parameters α β δ
 expr = (((1 / β - 1) + δ) / α)^(1 / (α - 1))
@@ -33,3 +33,11 @@ aov = ModelingToolkit.collect_applied_operators(eq, Differential)
 
 ts = collect_ivs([eq])
 @test ts == Set([t])
+
+@testset "vars searching through array of symbolics" begin
+    fn(x, y) = sum(x) + y
+    @register_symbolic fn(x::AbstractArray, y)
+    @variables x y z
+    res = vars(fn([x, y], z))
+    @test length(res) == 3
+end