fix: record array elements in collect_operator_variables

lib/ModelingToolkitBase/src/utils.jl

@@ -714,7 +714,17 @@ function collect_operator_variables(eqs::Vector{Equation}, ::Type{op}) where {op
         SU.search_variables!(vars, eq; is_atomic = OperatorIsAtomic{op}())
         for v in vars
             isoperator(v, op) || continue
-            push!(diffvars, arguments(v)[1])
+            arg = arguments(v)[1]
+            # An operator applied to an array variable or slice, such as `D(u[2:4])`,
+            # names the array rather than its elements. Callers test membership of the
+            # scalar unknowns, so record the elements.
+            if symtype(arg) <: AbstractArray
+                for el in vec(collect(Symbolics.scalarize(wrap(arg))))
+                    push!(diffvars, unwrap(el))
+                end
+            else
+                push!(diffvars, arg)
+            end
         end
         empty!(vars)
     end

lib/ModelingToolkitBase/test/variable_utils.jl

@@ -1,5 +1,5 @@
 using ModelingToolkitBase, Test
-using ModelingToolkitBase: value, parse_variable
+using ModelingToolkitBase: value, parse_variable, unwrap
 using SymbolicUtils: <ₑ
 import SymbolicUtils as SU
 
@@ -44,6 +44,27 @@ aov = ModelingToolkitBase.collect_applied_operators(eq, Differential)
 ts = collect_ivs([eq])
 @test ts == Set([t])
 
+@testset "collect_differential_variables with array variables" begin
+    # A derivative of an array variable or of a slice names the array, not its elements.
+    # Callers such as `DAEProblem`'s `differential_vars` test membership of the scalar
+    # unknowns, so the elements have to be recorded.
+    @variables w(t)[1:4]
+    Dt = Differential(t)
+
+    whole = collect_differential_variables(Dt(w) ~ w)
+    @test whole == Set(Any[unwrap(el) for el in collect(w)])
+
+    sliced = collect_differential_variables(Dt(w[2:3]) ~ w[1:2])
+    @test sliced == Set(Any[unwrap(w[2]), unwrap(w[3])])
+
+    # scalar derivatives are unaffected
+    @test collect_differential_variables(Dt(w[1]) ~ w[2]) == Set(Any[unwrap(w[1])])
+
+    # and the elements are exactly what a `differential_vars` style membership test needs
+    sts = [unwrap(el) for el in collect(w)]
+    @test map(Base.Fix2(in, sliced), sts) == [false, true, true, false]
+end
+
 @testset "parse_variable with scalarized arrays" begin
     @variables scalarized_x(t)[1:2]
     @parameters scalarized_p[1:2]