diff --git a/lib/ModelingToolkitTearing/src/tearingstate.jl b/lib/ModelingToolkitTearing/src/tearingstate.jl index 3b8afda..062e90a 100644 --- a/lib/ModelingToolkitTearing/src/tearingstate.jl +++ b/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 diff --git a/lib/ModelingToolkitTearing/test/runtests.jl b/lib/ModelingToolkitTearing/test/runtests.jl index bf11637..32ab664 100644 --- a/lib/ModelingToolkitTearing/test/runtests.jl +++ b/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