Restore rule-based flattening of array dependent variables (#101)

#98 rewrote `chain_flatten_array_variables` to avoid a QA ignore for
`Chain` and `Prewalk`. The rewrite keyed its replacement map on the array
operation alone, so every `u[i]` overwrote the previous entry and all
indices collapsed onto the last one, and it built replacements with
`Symbolics.variable(...; T = symtype(name))`, which yields an array-valued
variable rather than a callable. The resulting `Arr` reached `get_ops`, so
discretizing any `PDESystem` written against `@variables (u(..))[1:N]`
failed with `MethodError: operation(::Symbolics.Arr{Num, 1})`.

The ignore was avoidable without touching this code: `Chain` and `Prewalk`
are public in `SymbolicUtils.Rewriters`, which owns them, and were only
non-public via the parent `SymbolicUtils`. Import them from the owner and
restore the previous implementation, which shipped in 0.1.30.

Fixes #100



Claude-Session: https://claude.ai/code/session_0117aJvytDT3y2jMi5Ng6Q5m

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: 3 people

src/PDEBase.jl

@@ -14,7 +14,9 @@ using SciMLBase: AbstractDiscretization, AbstractDiscretizationMetadata, Nonline
     ODEFunction, ODEProblem
 import SciMLPublic: @public
 using Symbolics: @variables, Differential, Equation, Num, unwrap
-using SymbolicUtils: arguments, getmetadata, iscall, operation, substitute, unwrap_const
+using SymbolicUtils: @rule, arguments, getmetadata, iscall, operation, substitute,
+    unwrap_const
+using SymbolicUtils.Rewriters: Chain, Prewalk
 using SymbolicIndexingInterface: is_time_dependent
 using TermInterface: maketerm, metadata
 

src/make_pdesys_compatible.jl

@@ -26,36 +26,19 @@ function _replace_ops(term, op_map)
     return term
 end
 
-function _replace_indexed_array_variables(term, replacements)
-    term = safe_unwrap(term)
-    iscall(term) || return term
-
-    op = operation(term)
-    args = arguments(term)
-    if op === getindex
-        array = first(args)
-        if iscall(array)
-            replacement = get(replacements, operation(array), nothing)
-            replacement === nothing || return replacement(arguments(array)...)
-        end
-    end
-
-    new_args = _replace_indexed_array_variables.(args, Ref(replacements))
-    return all(isequal.(args, new_args)) ? term : maketerm(typeof(term), op, new_args, metadata(term))
-end
-
 function chain_flatten_array_variables(dvs)
-    replacements = Dict()
+    rs = []
     for dv in dvs
         dv = safe_unwrap(dv)
         if isequal(operation(dv), getindex)
             name = operation(arguments(dv)[1])
             idxs = arguments(dv)[2:end]
             fullname = Symbol(string(name) * "_" * string(idxs))
-            replacements[name] = Symbolics.variable(fullname; T = SymbolicUtils.symtype(name))
+            newop = (@variables $fullname(..))[1]
+            push!(rs, @rule getindex($(name)(~~a), idxs...) => newop(~a...))
         end
     end
-    return isempty(replacements) ? identity : term -> _replace_indexed_array_variables(term, replacements)
+    return isempty(rs) ? identity : Prewalk(Chain(rs))
 end
 
 function apply_lhs_rhs(f, eqs)

test/array_variables.jl

@@ -0,0 +1,29 @@
+using PDEBase
+using ModelingToolkit
+using Symbolics
+using SymbolicUtils
+using Test
+
+@testset "Flattening array dependent variables" begin
+    N = 3
+    @parameters x
+    @variables (u(..))[1:N]
+
+    dvs = collect([u(x)[i] for i in 1:N])
+    ch = PDEBase.chain_flatten_array_variables(dvs)
+    flattened = map(dv -> ch(PDEBase.safe_unwrap(dv)), dvs)
+
+    # Each `u[i]` needs its own replacement; keying only on `u` collapses them all.
+    @test length(unique(string.(flattened))) == N
+
+    for f in flattened
+        # The replacement stays a call in the independent variables so downstream
+        # `operation`/`arguments` keep working. An `Arr` here means the index was
+        # never resolved away.
+        @test !(f isa Symbolics.Arr)
+        @test iscall(f)
+        @test isequal(only(arguments(f)), PDEBase.safe_unwrap(x))
+    end
+
+    @test PDEBase.chain_flatten_array_variables([]) === identity
+end