Merge pull request #142 from JuliaComputing/as/index-with-var

fix: handle `x[i]` where `i` is a variable in `TearingState` constructor

Author: AayushSabharwal

lib/ModelingToolkitTearing/src/tearingstate.jl

@@ -249,12 +249,26 @@ function TearingState(sys::System, source_info::Union{Nothing, MTKBase.EquationS
 
             if !in(v, dvs)
                 arr, isidx = MTKBase.split_indexed_var(v)
-                # Handles cases list `D(x[1:2])`
-                if isidx && SU.is_array_shape(SU.shape(v)) && SU.shape(v) isa SU.ShapeVecT
-                    for idx in SU.stable_eachindex(v)
-                        push!(auxiliary_vars, v[idx])
+                if isidx
+                    # Handles cases list `D(x[1:2])`.
+                    if SU.is_array_shape(SU.shape(v)) && SU.shape(v) isa SU.ShapeVecT
+                        for idx in SU.stable_eachindex(v)
+                            push!(auxiliary_vars, v[idx])
+                        end
+                        continue
+                    end
+                    # Handles cases like `x[i]` where `i` is a variable.
+                    if any(!SU.isconst, Iterators.drop(arguments(v), 1)) && SU.shape(arr) isa SU.ShapeVecT
+                        for idx in SU.stable_eachindex(arr)
+                            push!(auxiliary_vars, arr[idx])
+                        end
+                        for arg in Iterators.drop(arguments(v), 1)
+                            SU.search_variables!(
+                                auxiliary_vars, arg; is_atomic = MTKBase.OperatorIsAtomic{SU.Operator}()
+                            )
+                        end
+                        continue
                     end
-                    continue
                 end
                 isvalid = @match v begin
                     BSImpl.Term(; f, args) => f isa Shift || isempty(args) || f isa SU.Operator && is_transparent_operator(f)::Bool

lib/ModelingToolkitTearing/test/runtests.jl

@@ -493,3 +493,23 @@ end
         @test ts_deferred.eqs_source == [[:a], [:a], [:a], [:b]]
     end
 end
+
+@testset "`TearingState` ctor for `x[i]` where `i` is a variable" begin
+    N = 3
+    dt = 0.1
+    clock1 = Clock(dt)
+    k = ShiftIndex(clock1)
+
+    @variables u(t) y(t) i(t)::Int
+    @variables x(t)[1:N] = zeros(N)
+
+    eqs = Equation[
+        u ~ sin(t);
+        [x[i + 1](k) ~ u(k - i) for i in 0:(N - 1)];
+        y(k) ~ x(k)[Symbolics.unwrap(i)];
+        i(k) ~ ifelse(i(k - 1) + 1 > N, 1, i(k - 1) + 1)
+    ]
+
+    @named cl = System(eqs, t)
+    @test_nowarn TearingState(cl)
+end