Merge pull request #4772 from SciML/as/index-with-var

fix: handle variable discover of `x[i]` where `i` is a variable

Author: AayushSabharwal

lib/ModelingToolkitBase/Project.toml

@@ -188,7 +188,7 @@ StochasticDiffEq = "6.82.0, 7"
 Sundials = "6"
 SymbolicIndexingInterface = "0.3.39"
 SymbolicUtils = "4.35.3"
-Symbolics = "7.31"
+Symbolics = "7.32.1"
 TaskLocalValues = "0.1.3"
 Test = "1"
 Tracker = "0.2.30"

lib/ModelingToolkitBase/src/utils.jl

@@ -954,8 +954,15 @@ function collect_var!(unknowns::OrderedSet{SymbolicT}, parameters::OrderedSet{Sy
         )
     end
     arr, isarr = split_indexed_var(var)
-    if isarr && SU.is_array_shape(SU.shape(var))
-        # `var` is indexed, and it is an array, so it must be a slice. Replace it with `arr`.
+    if isarr && (
+            # `var` is indexed, and it is an array, so it must be a slice. Replace it with `arr`.
+            SU.is_array_shape(SU.shape(var)) ||
+                # `var` is of the form `x[i]` where `i` is also a variable/expression
+                any(!SU.isconst, Iterators.drop(arguments(var), 1))
+        )
+        for arg in Iterators.drop(arguments(var), 1)
+            collect_vars!(unknowns, parameters, arg, iv)
+        end
         var = arr
     end
     check_scope_depth(getmetadata(arr, SymScope, LocalScope())::AllScopes, depth) || return nothing

lib/ModelingToolkitBase/test/system_building.jl

@@ -72,3 +72,9 @@ ModelingToolkitBase.validate_operator(::MyOp, args...; kws...) = nothing
     @named sys = System([D(x) ~ x + SU.term(MyOp(), [0.0]; type = Real, shape = SU.ShapeVecT())], t)
     @test issetequal(unknowns(sys), [x])
 end
+
+@testset "Variable discovery `x[i]` where `i` is a variable" begin
+    @variables x(t)[1:2] i(t)::Int
+    @named sys = System([x[i] ~ 0], t)
+    @test issetequal(unknowns(sys), [x, i])
+end