MTKTearing: keep array differential equations intact through structural simplification

[ChrisRackauckas-Claude]

Companion to SciML/ModelingToolkit.jl#5102. Bumps ModelingToolkitTearing to 1.21.0.

Array differential equations such as the MethodOfLines discretization D(u[2:(n-1)]) ~ lap(u) are currently always flattened by scalarize_tearing_state_eqs!, so nothing downstream of mtkcompile can ever see them as arrays again. This PR teaches the structural simplification pipeline to keep those equations as first-class units when that is valid while still running matching, Pantelides, dummy derivatives, tearing and alias elimination on them. It is not a switch that turns any algorithm off.

Design

Incidence stays scalar; the array structure is carried alongside

TearingState still scalarizes every equation into rows of the bipartite graph and fullvars still holds scalar variables. This is deliberate: for D(u[2:(n-1)]) ~ lap(u) each element depends on a different subset of u, and per-element incidence is exactly what index reduction and tearing need to make correct decisions (which u[k] are states, whether a boundary equation u[1] ~ 0 makes D(u[1]) a dummy derivative, and so on). Collapsing an array equation into one graph node would either lose that information or require every algorithm to reason about slices.

Instead, when the constructor scalarizes an eligible equation it records an ArrayEquationGroup:

mutable struct ArrayEquationGroup
    eq::Equation                 # canonical form `D(x[slice]) ~ rhs`
    lhs_vars::Vector{SymbolicT}  # D(x[k]) for each element, in row order
    dirty::Bool                  # rows can no longer be re-assembled
end

and two per-row vectors on TearingState: row_group[i] (index of the group of equation row i, 0 for scalar rows) and row_elem[i] (linear element index within the group). Rows appended later (eq_derivative!) are always scalar.

Eligibility (canonicalize_array_equation): one side is D(x) or D(x[slice]) for an array unknown x with a constant-index slice, the other side has no derivatives and matching shape. The residual forms emitted by MethodOfLines, D(x[slice]) .- f ~ 0 and D(x[slice]) .+ f ~ 0, are canonicalized to D(x[slice]) ~ ±f. Algebraic array equations and equations with derivatives on both sides are scalarized as before.

Every algorithm runs; any rewrite the array equation cannot represent marks the group dirty

The algorithms are unchanged. Wherever a pass modifies a row so that it is no longer "element k of D(x[slice]) ~ rhs solved for D(x[k])", the group is marked dirty (dirty_array_group!):

Pass	Where	Why it dirties
Pantelides / index reduction	eq_derivative!	an element was differentiated: the block is part of a higher-index structure
Consistency / redundancy removal	rm_eqs_vars!	an element row was deleted
Dummy derivatives	substitute_derivatives_algevars!	a dummy derivative was substituted into a row
Tearing / reassembly	codegen_equation! (algebraic and solved branches), inline linear SCC path	a row is solved for something other than its own D(x[k]), or a previously solved derivative was substituted into its RHS
Clock inference	system_subset, shift_discrete_system, substitute_sample_time	rows split across partitions / discrete systems
Brownian handling (MTK side)	__mtkcompile	the drift equation no longer matches the array equation

Substituting an eliminated alias or zero variable into a row does not dirty the group: the eliminated variable stays available as an observed equation and the array equation remains correct.

Integer-linear elimination cannot pivot on array-equation rows

linear_subsys_adjmat! leaves the rows of intact groups out of mm (is_intact_array_group_row). Otherwise alias elimination could pivot on D(x[1]) ~ -x[1] + y to eliminate y from y ~ sum(x), which would leave D(x[1]) matched to a rewritten equation. Their solvability is still recorded in solvable_graph, so tearing still sees D(x[k]) as solvable from its row. Rows of dirty groups are ordinary scalar rows and take part in the elimination. This is why the MTK side needs no changes to alias_elimination!.

Emission

At the end of reassembly blt_reorder_generated_equations! places the rows of each intact group contiguously and in element order (tagged with the earliest SCC among them; they are all differential equations of selected states so algebraic BLT order is unaffected). With preserve_array_equations = true on DefaultReassembleAlgorithm (also accepted as a mtkcompile keyword, which already forwards to the reassemble algorithm), collapse_array_equations! replaces each intact run by the single array equation. Unknowns stay scalar (x[k]) in the same order as the rows they replace, so mass matrices and the ODE/DAE code generators lay the equation out over the right slots. row_to_equation_indices maps graph rows to emitted equations for consumers such as map_variables_to_equations. Dirty groups are emitted scalarized exactly as today.

The default is preserve_array_equations = false, so the output of mtkcompile is unchanged by this PR. The opt-in exists only because explicit ODEProblem codegen for array equations is landing separately (SciML/ModelingToolkit.jl#5101); DAEProblem already handles them. Once that is in, the default can flip.

What still scalarizes, and why

• Any group whose element is differentiated by Pantelides (high-index systems, e.g. the Cartesian pendulum with array q, v): the differentiated row and its dummy-derivative substitutions are not representable as the original array equation.
• Rows solved for a variable other than their own D(x[k]), e.g. when tearing/state selection picks a different state and the row becomes algebraic in it.
• Array equations with derivatives on both sides, algebraic array equations, non-constant slices, and equations split across clock partitions or in discrete systems.
• Equations whose elements are removed as redundant.

In every case the fallback is exactly the current behaviour.

Tests

lib/ModelingToolkitTearing/test/runtests.jl, testset "Array equation groups":

• TearingState tracks eligible equations as groups with correct row_group/row_elem.
• MethodOfLines residual forms are canonicalized.
• Ineligible equations are not grouped.
• eq_derivative! dirties the group; the appended row is scalar.
• Intact group rows are excluded from mm but remain solvable for their derivatives; dirty groups take part; alias_elimination! leaves the array equation intact.
• preserve_array_equations (keyword and algorithm option) emits one array equation over scalar unknowns; row_to_equation_indices is correct; default output unchanged.
• A high-index system dirties its groups and is emitted scalarized.

Full MTKTearing test suite passes locally both against the companion MTK branch and against registered ModelingToolkit v11.42.0 / ModelingToolkitBase v1.69.0 (this PR does not depend on the MTK PR).

How to verify O(1) through mtkcompile

See SciML/ModelingToolkit.jl#5102 (test/structural_transformation/array_equations.jl): a heat equation in MethodOfLines residual form compiled with mtkcompile(heat; preserve_array_equations = true) yields one array equation plus observed boundary conditions, DAEProblem solves it to the same result as the scalarized system, and the Expr size of generate_rhs(...; implicit_dae = true) is identical for n = 24, 48, 96.

Made with Cursor

[ChrisRackauckas-Claude]

Local check: this branch’s ModelingToolkitTearing/test/runtests.jl passes against registered ModelingToolkit v11.42.0 + ModelingToolkitBase v1.69.0 (dev’d StateSelection/MTKTearing only).