feat: accept array equations on the implicit-DAE path

[ChrisRackauckas-Claude]

Note

Draft — please ignore until reviewed by @ChrisRackauckas.
Stacked on #4982 (collect_operator_variables), which supplies the correct
differential_vars this relies on.

Why

A finite-difference PDE discretization can emit its interior as one array equation over
slices instead of one scalar equation per grid point, which makes the symbolic system
size independent of grid resolution (SciML/MethodOfLines.jl#428). Measured on the 1D heat
equation, generating code directly from such a system: 1 614 expression nodes at every
grid size, versus 20 646 → 117 446 growing for the scalarized form, with identical
numerics.

mtkcompile scalarizes array equations, so that structure is erased before codegen.
ODEProblem fundamentally needs it to: it requires D(x) = f(x), and isolating the
derivative out of a residual is structural simplification.

The implicit-DAE path does not have that requirement — a residual D(u) - f ~ 0 is exactly
what DAEProblem consumes.

What changed

Four changes, all inert outside the implicit-DAE path with array equations:

check_array_equations_unknowns	skipped when implicit_dae; every other problem type keeps the guard
check_eqs_u0	skipped only when a DAE system actually contains array equations, where an equation count is not comparable with an unknown count — scalarized DAEs keep the check
expand_array_derivatives	rewrites D(u[2:4]) into the scalar derivatives bound to the du argument, which a derivative of a slice otherwise matches none of
flatten_array_residuals	expands an array-valued residual entry into one output row per element, indexing the shared expression rather than scalarizing it

Both codegen helpers return their input unchanged when nothing is array-valued, so existing
implicit-DAE systems generate identical code.

Evidence

Before this, DAEProblem on such a system failed in sequence: the array-equation guard,
then Equations (3), unknowns (21), then a literal Differential surviving into the
generated code (MethodError: no method matching (::Differential)(::Vector{Float64})).

After, on the heat equation with 3 equations and 21 unknowns:

• 21 output rows produced from 3 equations (81 from 3 at the larger size)
• differential_vars = 19/21 and 79/81 — interior differential, boundaries algebraic
• residual evaluates finitely
• end-to-end solve with DFBDF + BrownFullBasicInit: max|u - exact| = 7.57e-4
  against exp(-π²t)sin(πx), the expected second-order spatial error on 21 points
• ODEProblem on the same system still throws

Test Summary: | Pass  Total    Time
arrdae        |   11     11  4m16.8s

Includes a 2D case: expand_array_derivatives originally flattened the scalarized
derivative with vec, which produced a length-81 vector against 9×9 surrounding slices
and failed with DimensionMismatch. Every array equation of two or more dimensions was
affected, and the 1D tests could not catch it because vec is a no-op there. Fixed, with
a regression test.

Notes for review

• DAEProblem still needs build_initializeprob=false here: MTK's initialization system
  builds a time-independent system and rejects the Differential inside the array
  equation. The solver's own BrownFullBasicInit covers that case, but it is only valid
  when no initialization equation constrains a derivative or an algebraic variable, so
  automating the choice needs care. Not addressed in this PR.
• Array equations only survive System construction in residual form. Written
  naturally as D(u[2:n-1]) ~ rhs, diff2term fails with
  Can only build StableIndex{Int} from indexed symbolic — it cannot name a derivative
  of a slice. Worth knowing for anyone generating such systems.
• I have not measured whether CSE hoists the shared array expression across the expanded
  rows. If it does not, the stencil is recomputed per row, which would be worse than
  scalarizing — so the performance claim above is about the symbolic representation,
  not yet about generated-code quality.
• Tests are lib/ModelingToolkitBase/test/array_equation_dae.jl; I ran that file and
  variable_utils.jl, not the full suite.

🤖 Generated with Claude Code

https://claude.ai/code/session_01S1dxUEPGysx27SSSpxtiQ9

[AayushSabharwal]

This entire function can be made more efficient. Instead of making half a dozen comments, I'll compile it into this one. _array_derivative_terms! is just SU.search_variables! with is_atomic = x -> isdifferential(x) && SU.is_array_shape(SU.shape(x)) to get all array derivative applications. This removes the entire function and makes it more efficient. The buffer should be a Set{SymbolicT}(). subs should be concretely typed. Instead of the weird collect(scalarize(wrap( this should use SU.stable_eachindex to iterate over the symbolic array and manually build an array_literal. wrap is dynamic dispatch and doesn't infer, as is every function after that point. Substitution should use IRSubstituter. There's no need to wrap the argument to substitute and then unwrap it.

If this function also accepted rhss in its entirety instead of one element at a time, the caching in search_variables! and IRSubstituter is also shared.

[ChrisRackauckas-Claude]

Done — rewritten along all of those lines. _array_derivative_terms! is gone, replaced by search_variables! with an ArrayDerivativeIsAtomic callable (isdifferential(ex) && SU.is_array_shape(SU.shape(ex))); since scalar variables are not atomic under that predicate, nothing else is collected and no post-filter is needed. Buffer is Set{SymbolicT}, subs is Dict{SymbolicT, SymbolicT}, the elements come from SU.stable_eachindex instead of collect(scalarize(wrap(...))), substitution goes through IRSubstituter{false}, and the wrap/unwrap round trip is gone.

It now takes rhss in its entirety, so the search_variables! and IRSubstituter caches are shared across residuals.

One detail worth flagging: the elements have to be reshaped back to the argument shape before building the literal. A derivative of a 2D slice must expand to a 2D array or it stops broadcasting against the surrounding slices — that was a real bug in the previous revision of this PR. I verified stable_eachindex iterates column-major and matches scalarize exactly, including for rank 2.

[AayushSabharwal]

This entire function is redundant. generate_rhs should build an ArrayMaker expression where each equation writes to the appropriate contiguous region of the output. That will codegen much better than this approach. For example, a 1D PDE with a body of length n would generate something equivalent to

@makearray out[1:(n+2)] begin
  out[1:1] => left_boundary
  out[2:(n+1)] => body
  out[(n+2):(n+2)] => right_boundary
end

Except generated programmatically with the constructor instead of with a macro.

[ChrisRackauckas-Claude]

Done — flatten_array_residuals is gone, replaced by array_residual_maker, which builds a single ArrayMaker where each residual writes to a contiguous region, exactly as in your sketch but constructed programmatically.

Two constraints I hit that shape the implementation:

• ArrayMaker rejects scalar values (AssertionError: sT <: AbstractArray), so a scalar residual is lifted to a length-1 array literal via SU.Const{VartypeT}([rhs]) and given a length-1 region.
• The output is a flat residual vector, so a rank > 1 residual cannot be given a multidimensional region. There is no reshape for symbolic arrays, and vec(wrap(x)) scalarizes — but vec on the unwrapped BasicSymbolic stays symbolic (vec((w(t))[2:3, 4:5]), shape 1:4), so that is what is used.

It does codegen much better, as you said. Generated DAEFunction expression size for a 1D heat system, before vs after:

n	flatten_array_residuals	ArrayMaker
20	19,141	11,948
200	140,095	53,035
2000	1,393,803	480,776

Still linear in n — the remaining growth is the unknowns/argument handling, not the residual body — but the constant is ~2.9x smaller at n=2000.

[AayushSabharwal]

This check can still work if it is refactored slightly to not assume each equation is a scalar.

[ChrisRackauckas-Claude]

Done — the check runs on this path now rather than being skipped. check_eqs_u0 counts residual rows instead of equations: a new count_equation_rows/equation_row_count sums prod(length, shape) for an array equation and 1 otherwise, taking the shape from whichever side carries it (D(u[2:4]) ~ rhs or 0 ~ rhs).

Since array equations previously could not reach check_eqs_u0 at all, this is a strict generalization for every other caller. GROUP=InterfaceI passes: 1595 passed, 0 failed, 5 pre-existing broken.

[AayushSabharwal]

SU.is_array_shape(SU.shape(arg)) is a more efficient check.

[ChrisRackauckas-Claude]

Done, using SU.is_array_shape(SU.shape(arg)).

[ChrisRackauckas-Claude]

Pushed as d6640a0, rebased onto d9c3814 so the accepted suggestion is preserved (that commit applied the stable_eachindex half of the collect_operator_variables comment; the rebase also applies the is_array_shape half).

Replies are on each thread. Summary:

• _array_derivative_terms! removed in favour of search_variables! with an array-derivative is_atomic; concretely typed Set/Dict, stable_eachindex, IRSubstituter, no wrap/unwrap round trip, and it takes all of rhss so the caches are shared.
• flatten_array_residuals removed in favour of an ArrayMaker with one contiguous region per residual. Generated expression size drops ~2.9x at n=2000.
• check_eqs_u0 now counts residual rows, so the check runs on the implicit-DAE path instead of being skipped.
• collect_operator_variables uses is_array_shape/stable_eachindex.

One thing I found while doing this, which is not from the review. The natural D(u[2:n-1]) ~ rhs form crashed with ArgumentError: Can only build StableIndex{Int} from indexed symbolic, thrown from diff2term via ir_info.jl: a derivative of a slice has no scalar toterm name to key a substitution on. The existing tests only covered the residual form (broadcast(-, D(...), lap) ~ 0), which is what MethodOfLines emits, so it was not caught. ir_info now skips array-valued LHSs — the implicit-DAE path has already expanded those derivatives to their scalar elements by then — and there is a new testset covering the form end to end. Happy to split that into its own PR if you would rather keep this one to the review changes.

Verification, all local:

array equations reach DAEProblem                        |    6      6
other problem types still require scalarized equations  |    1      1
array-equation DAE solves to the analytic solution      |    2      2
array equations over a 2D slice keep their shape        |    2      2
array equations written as `D(slice) ~ rhs`             |    4      4

GROUP=InterfaceI: InterfaceI | 1595  5(broken)  1600  56m00.2s
     Testing ModelingToolkitBase tests passed

runic --check is clean. I have not run the full suite beyond InterfaceI.

[ChrisRackauckas-Claude]

Correction to the codegen numbers above. I measured length(string(DAEFunction(...))), but that string embeds the System, whose show prints all n unknowns. So roughly 4n of what I counted was the system's printed repr, not generated code. Both columns were inflated the same way, so the comparison was directionally right, but it understated the improvement and it made me wrongly say the new residual body is "still linear in n". It is not.

Counting generated code only (everything before the embedded sys = Model sys:), and splitting out the du argument reads:

n	old total	old non-du	new total	new non-du
20	237	201	134	98
200	1,677	1,281	494	98
2000	16,077	12,081	4,094	98

The ArrayMaker residual body is O(1) — exactly 98 lines at every size, with the broadcast count fixed at 10. The old flatten_array_residuals body was O(n) (12,081 lines at n=2000). So the improvement is ~123x on the residual body, not the ~2.9x I quoted.

What is still O(n), and why. The remaining growth is entirely 2(n-2) element reads of the du argument (local var"##cse#k" = __mtk_arg_1[i]). The state vector does not have this problem: u is reconstructed once as reshape(view(___mtkunknowns___, 1:n), (n,)), so u[2:n-1] is a genuine O(1) slice.

The asymmetry comes from array_variable_assignments being filtered by required_arrvars, which only contains array variables appearing unindexed in the expression. expand_array_derivatives rewrites D(u[2:n-1]) into an array literal of scalar D(u[i]) terms, so D(u) never appears as an array, no view is emitted for the du buffer, and each element is read individually.

I tried the obvious fix — substituting D(u)[2:n-1] instead of the literal, so that D(u) lands in required_arrvars. That does not work: indexing an array derivative normalizes straight back, D(u)[2:4] becomes D(u[2:4]), so the slice-of-array-derivative form does not survive as an expression. Making the du side O(1) would need codegen to bind the du buffer as an array and slice the bound local, i.e. a change in build_function_wrapper rather than in this expansion. Happy to look at that separately if you think it is worth it — it would make the whole generated function O(1) in the grid size.

[AayushSabharwal]

Just typeassert that shape(eq.lhs)::SU.ShapeVecT and return prod(length, sh; init = 1)

[AayushSabharwal]

This shouldn't be a callable singleton struct instead of a normal function?

[AayushSabharwal]

Using IRStructureSearchBuffer here is much faster

[AayushSabharwal]

Suggested change

	sz = ntuple(i -> length(sh[i]), length(sh))
	sz = size(arg)

[AayushSabharwal]

A better method is to build arrargs = Symbolics.SArgsT() and sizehint! it. Then, push sz and each of op(arg[idx]). This allows building Symbolics.STerm(SU.array_literal, arrargs; type = symtype(arg), shape = sh). This avoids the dynamic dispatch reshape.

[AayushSabharwal]

Why not just reuse get_irstructure(sys)? That way the expression is also prepopulated for codegen to use.

[AayushSabharwal]

Shouldn't need the map iirc. Could also use map! and avoid the allocation.

[AayushSabharwal]

Suggested change

	sh = SU.shape(rhs)
	sh = SU.shape(rhs)::SU.ShapeVecT

For now, since this code assumes it anyway. We'll have to lift this restriction to allow resizing the PDE without redoing codegen.

[AayushSabharwal]

A future PR should probably make assertions codegen to isoutofdomain?