Preserve array differential equations with `mtkcompile(sys; scalarize_arrays = false)` by ChrisRackauckas-Claude · Pull Request #5100 · SciML/ModelingToolkit.jl

ChrisRackauckas-Claude · 2026-09-09T16:11:21Z

Summary

mtkcompile scalarizes every array equation before structural simplification, so for a finite-difference discretization such as D(u[2:(n - 1)]) ~ lap(u) the compiled system, the generated code and the compile time all grow with the number of grid points. MethodOfLines currently gets O(1) scaling only by skipping mtkcompile entirely and building a DAEProblem from the completed system.

This PR adds mtkcompile(sys; scalarize_arrays = false), which compiles the system without tearing/index reduction and keeps first-order array differential equations intact, so that DAEProblem (and, with #5101, ODEProblem) code generation is independent of the array length. It is the first of two ModelingToolkit PRs; the ODE code generation half is stacked on top of this one.

What changes

mtkcompile(sys; scalarize_arrays = false) (ModelingToolkitBase). Dispatches to the tearing-free compiler (__mtkcompile_no_tearing, the former __mtkcompile of ModelingToolkitBase; ModelingToolkit's __mtkcompile override returns early to it). That compiler now:
- keeps D(x) ~ rhs where x is an array of unknowns or a slice of one, ordering the unknowns so the elements of each differentiated slice form a contiguous block matching the rows of the equation;
- rewrites residual-form array equations D(x) .- f ~ 0 (what MethodOfLines emits) to D(x) ~ f;
- scalarizes array algebraic equations, so that trivially defined elements (boundary conditions) become observed as before;
- counts an array equation as one row per element in the balance check;
- records the choice under the ScalarizeArraysCtx metadata key (arrays_scalarized(sys)), which the initialization system inherits (mirrors HomotopyCtx).
- Systems with brownians/jumps/poissonians/noise throw an ArgumentError with scalarize_arrays = false.
Code generation keeps the array form:
- expand_array_derivatives! binds a derivative of a slice to a view of du when its scalar derivatives are consecutive in du, instead of an array_literal of them (fixes Preserve array derivatives in implicit DAE residual code generation #5097).
- Arrays whose elements are partly unknowns and partly observed are reconstructed from views of the argument buffers plus the observed elements (partial_array_reconstruction), instead of through an observed equation that lists every element. Guarded by !arrays_scalarized(sys), so the default path is unchanged.
- default_toterm handles derivatives of slices: D(u[2:4]) becomes uˍt(t)[2:4].
- full_equations scalarizes preserved equations so jacobians/mass matrices/sparsity keep one row per scalar equation.
Initialization:
- write_possibly_indexed_array! broadcasts a scalar written to an array (or slice) key, so the scalar default derivative guess for an array uˍt(t) no longer produces vec(::Bool) (fixes DAEProblem from a system with array-form differential equations fails with MethodError: no method matching vec(::Bool) #5093 — the reproducer from the issue is now a test).
- The initialization system registers the rows of a preserved derivative equation so that it is not added a second time in scalarized form, and get_initialization_problem_type counts equation rows and copes with a nothing tearing state (no SCC problem without tearing).
Tests: test/array_equation_dae.jl (previously not in runtests.jl) is wired in, plus new test/array_equations_preserved.jl.

Example

using ModelingToolkit, OrdinaryDiffEqBDF
using ModelingToolkit: t_nounits as t, D_nounits as D

n = 96
@variables u(t)[1:n]
dx = 1 / (n - 1)
lap = (u[1:(n - 2)] .- 2 .* u[2:(n - 1)] .+ u[3:n]) ./ dx^2
@named heat = System([D(u[2:(n - 1)]) ~ lap, u[1] ~ 0, u[n] ~ 0], t, [u], [])

sys = mtkcompile(heat; scalarize_arrays = false)
equations(sys)   # 1 equation: D(u[2:95]) ~ ...
unknowns(sys)    # u[2], ..., u[95]
observed(sys)    # u[1] ~ 0, u[96] ~ 0

prob = DAEProblem(sys, [u => sinpi.(range(0, 1, length = n))], (0.0, 0.1))
sol = solve(prob, DFBDF())

How to verify the O(1) scaling

test/array_equations_preserved.jl checks, for n = 24, 48, 96, that the compiled system has exactly one equation and that the generated in-place DAEProblem function (generate_rhs(sys, GeneratedFunctionOptions(; expression = Val{true}); implicit_dae = true)) has the same number of Expr nodes and no array_literal. The same file checks the solution against the scalarized compilation and the analytic solution of the heat equation, with and without an initialization problem, and a 2D slice.

Limitations

No tearing or index reduction: differential equations must be explicit (D(x) ~ f or D(x) .- f ~ 0), and algebraic equations that do not trivially define an unknown are kept as algebraic equations. Higher-order array equations are scalarized and order-reduced as before.
ODEProblem from such a system still requires the follow-up Build ODEProblem from array differential equations with complete only #5101; here it keeps throwing the existing "array equations" error.

Tests run locally

ModelingToolkitBase groups Initialization, SymbolicIndexingInterface, InterfaceI, InterfaceII: pass, except the pre-existing JumpSystem Test error (FieldError: type NamedTuple has no field abstol from affect_tolerance in callbacks.jl on an SSAIntegrator), which also fails on master with JumpProcesses 9.32.2 and is untouched here.
ModelingToolkit group InterfaceI: pass (1507 passed, 3 broken, as on master).
MethodOfLines test/Discretization/problem_construction.jl against this branch (including the new ODE path from Keep array equations through the compiled ODEProblem path MethodOfLines.jl#688): pass.

CI notes

The integration tests (Catalyst, MethodOfLines, ...) and the benchmark job fail at resolution with "Unsatisfiable requirements ... ModelingToolkitBase restricted to versions 1.70.0" because the ModelingToolkit-side changes (__mtkcompile early return, count_equation_rows) need the ModelingToolkitBase in this PR. They will resolve once ModelingToolkitBase 1.70.0 is registered.
downgrade-mtkbase (InterfaceII) and sublibrary-ci InterfaceII fail only on the pre-existing JumpSystem Test error above (the same job fails on Sundials 6.7.1 + MCI 0.3 source for x86 InterfaceI #5098).
sublibrary-ci QA fails only on JET: the QA environment now resolves JET 0.12.1, which reports ~245 "local variable may be undefined" findings in Moshi @match expansions across ModelingToolkitBase (files this PR does not touch); Recognize copied symbolic missing bindings #5073 fails the same way with 234 findings. The ModelingToolkit QA job (JET 0.11.6) passes.
InterfaceI (x86) fails to precompile Sundials/QuasiMonteCarlo on 32-bit, as on master (Sundials 6.7.1 + MCI 0.3 source for x86 InterfaceI #5098).
Documentation build fails on the pre-existing unbound_inputs @ref in interactive_simulation.md, fixed by Qualify unbound_inputs @ref in the interactive simulation tutorial #5099; the new @docs block here renders.

Related: #5093, #5097, SciML/MethodOfLines.jl (array-form DAEProblem path).

Made with Cursor

…_arrays = false)` Structural simplification scalarizes every array equation before matching, so for a finite-difference discretization such as `D(u[2:(n - 1)]) ~ lap(u)` the compiled system, its generated code and the compile time all grow with the number of grid points. The array form is only preserved today by skipping `mtkcompile` altogether and building a `DAEProblem` from the completed system. Add the `scalarize_arrays` keyword to `mtkcompile`. With `scalarize_arrays = false` the system is compiled by the tearing-free compiler in ModelingToolkitBase (`__mtkcompile_no_tearing`, which ModelingToolkit's `__mtkcompile` dispatches to), which now keeps first-order array differential equations intact: the differentiated array or slice becomes a contiguous block of the unknowns, residual-form equations `D(x) .- f ~ 0` are rewritten to `D(x) ~ f`, and array algebraic equations are scalarized so that trivially defined elements (boundary conditions) become observed. The compiled system records the choice under `ScalarizeArraysCtx` (`arrays_scalarized`), which the initialization system inherits. Code generation then keeps the array form all the way through: - `expand_array_derivatives!` binds a derivative of a slice to a `view` of `du` when its scalar derivatives are consecutive (SciML#5097), instead of an `array_literal` of them. - Arrays whose elements are partly unknowns and partly observed are reconstructed from views of the argument buffers plus the observed elements, rather than through an observed equation listing every element. - `default_toterm` handles derivatives of slices (`D(u[2:4])` becomes `uˍt(t)[2:4]`). - `write_possibly_indexed_array!` broadcasts a scalar written to an array key (SciML#5093), so the default derivative guess no longer fails for array derivative variables. - `full_equations` scalarizes preserved equations so that jacobians, mass matrices and sparsity patterns keep one row per scalar equation. Also wire the existing `array_equation_dae.jl` test into the test suite and add tests that the compiled heat equation has one equation and generated `DAEProblem` code of the same size for n = 24, 48, 96, solves correctly with and without an initialization problem, and matches the scalarized compilation. Co-authored-by: Cursor <cursoragent@cursor.com>

…, SciML#5097) Co-authored-by: Cursor <cursoragent@cursor.com>

…rownFullBasicInit `BrownFullBasicInit()` fails on the oldest supported SciMLBase/DiffEqBase combination (the downgrade CI job) because of an unrelated parameter despecialization bug, and the tests do not depend on it: the discrete Laplacian of the initial condition is the consistent initial derivative, so pass it explicitly and solve with `NoInit()`. Also assert the residual vanishes for that `du0`. Co-authored-by: Cursor <cursoragent@cursor.com>

The QA suite requires every `@public` name to be rendered in the docs. Co-authored-by: Cursor <cursoragent@cursor.com>

… QA list Co-authored-by: Cursor <cursoragent@cursor.com>

…eserving path `__mtkcompile_no_tearing` and `count_equation_rows` are internal helpers reached from ModelingToolkit, like `__mtkcompile` before them. Co-authored-by: Cursor <cursoragent@cursor.com>

ChrisRackauckas-Claude · 2026-09-09T20:24:48Z

Closing this as an intermediate step that is not the direction we want.

scalarize_arrays = false gets O(1) code by switching tearing and index reduction off for the whole system, and then bolts a separate compiler path (__mtkcompile_no_tearing, ScalarizeArraysCtx, arrays_scalarized) onto every consumer. That is turning algorithms off rather than making them array-aware, and it would have to be kept in sync with the real mtkcompile forever. The goal is for tearing and index reduction themselves to keep array equations intact, which is a different piece of work and will come as its own PR.

What is useful here has been split out:

Build ODEProblem from array differential equations with complete only #5101 has been rewritten from master (no longer stacked on this PR) and makes ODEProblem(complete(sys), ...) work for array differential equations directly, mirroring what DAEProblem already does, with generated code independent of the array length. No new keyword, no alternative compiler.
The default_toterm, write_possibly_indexed_array! (DAEProblem from a system with array-form differential equations fails with MethodError: no method matching vec(::Bool) #5093) and du-view (Preserve array derivatives in implicit DAE residual code generation #5097) fixes are only needed once array equations survive mtkcompile, so they belong with that work.

AayushSabharwal · 2026-09-09T19:55:58Z

+Returns `nothing` if the array cannot be reconstructed: some element is neither in a buffer
+nor has a fallback, or the array is not indexed from one.
+"""
+function partial_array_reconstruction(


I don't think this should be necessary. array_variable_assignments in codegen_utils.jl should handle this. If that generates suboptimal code (such as a large array_literal) we should fix it and make it use e.g. ArrayMaker.

Oh I see that's exactly what this is doing.

This function should be checked using Cthulhu/@code_typed to make sure the closures don't cause excessive boxing from that one Julia bug. It would be nice to just use explicit let captures anyway.

AayushSabharwal · 2026-09-09T20:05:32Z

+            push!(element_exprs, name)
+            name
+        end
+        push!(writes, Expr(:(=), Expr(:ref, result, Tuple(carts[i])...), value))


I guess the splat here is fine but it would be nice to just not have to build carts since it will never infer.

AayushSabharwal · 2026-09-09T20:06:17Z

+        push!(used_buffers, name)
+        src = Expr(:call, view, name, pos:(pos + len - 1))
+        if N > 1
+            src = Expr(:call, reshape, src, map(length, region)...)


Suggested change

src = Expr(:call, reshape, src, map(length, region)...)

src = Expr(:call, reshape, src)

append!(src.args, Iterators.map(length, region))

AayushSabharwal · 2026-09-09T20:06:31Z

+        if N > 1
+            src = Expr(:call, reshape, src, map(length, region)...)
+        end
+        dest = Expr(:call, view, result, region...)


Suggested change

dest = Expr(:call, view, result, region...)

dest = Expr(:call, view, result)

append!(dest.args, region)

AayushSabharwal · 2026-09-09T20:10:03Z

+        if !isempty(present)
+            lo = Tuple(carts[first(present)])
+            hi = lo
+            for i in present


I have a gut feeling that this loop would really benefit from not using CartesianIndices, or at least turning lo and hi into Vector{Int}/Memory{Int} so that the broadcasts infer. We can even make it operate in-place.

AayushSabharwal · 2026-09-09T20:12:11Z

+            end
+            region = ntuple(d -> lo[d]:hi[d], N)
+            lin = LinearIndices(sz)
+            block = [lin[c] for c in CartesianIndices(region)]


block's ndims won't infer. Does this need to be an Array{T, N} or can it just be a Vector{T}? If the latter, it's better to do block = Vector{T}(undef, some_length) and write to it in a loop.

CartesianIndices(region) may not be 1-indexed, but lin is. Won't this cause problems? I guess what this is trying to do is block = [prod(Tuple(c)) - prod(lo) for c in CartesianIndices(region)]?

AayushSabharwal · 2026-09-09T20:13:14Z

+    alloc = if isempty(used_buffers)
+        Expr(:call, Expr(:curly, Array, eltype_expr, N), :undef, sz...)
+    else
+        Expr(:call, similar, first(used_buffers), eltype_expr, sz...)
+    end


Suggested change

alloc = if isempty(used_buffers)

Expr(:call, Expr(:curly, Array, eltype_expr, N), :undef, sz...)

else

Expr(:call, similar, first(used_buffers), eltype_expr, sz...)

end

alloc = if isempty(used_buffers)

Expr(:call, Expr(:curly, Array, eltype_expr, N), :undef)

else

Expr(:call, similar, first(used_buffers), eltype_expr)

end

append!(alloc.args, Iterators.map(length, sh))

No splat, no need for sz.

AayushSabharwal · 2026-09-09T20:24:16Z

+    push!(body.args, result)
+    push!(assignments, Assignment(arrvar, body))
+    return assignments
+end


This function would be much easier to understand with some more comments.

AayushSabharwal · 2026-09-09T20:25:00Z

+"""
+function partial_array_reconstruction(
+        arrvar::SymbolicT, idxs::Vector{Tuple{Int, Int}}, argument_name, buffer_offset::Int,
+        element_fallback


Might want to @nospecialize the argument_name and element_fallback

AayushSabharwal · 2026-09-09T20:25:16Z

+slice); any other slice is read one element at a time. Slices with an element in no buffer
+are skipped.
+"""
+function array_slice_assignments(


The same general comments from above apply here.

ChrisRackauckas-Claude · 2026-09-09T22:48:58Z

Superseded checkpoint pointers (this PR remains closed as the wrong intermediate):

Build ODEProblem from array differential equations with complete only #5101 — complete → ODEProblem
CHECKPOINT: preserve array DEs through mtkcompile (ArrayEquationGroup / not skip-tearing) #5102 + MTKTearing: keep array differential equations intact through structural simplification JuliaComputing/StateSelection.jl#158 — array groups through tearing/IR (draft)
https://github.com/ChrisRackauckas-Claude/ModelingToolkit.jl/blob/array-aware-tearing/docs/src/internals/ARRAY_EQUATIONS_CHECKPOINT.md

ChrisRackauckas and others added 2 commits September 9, 2026 12:06

Test the DAE initialization and du view for array equations (SciML#5093…

3a39a0b

…, SciML#5097) Co-authored-by: Cursor <cursoragent@cursor.com>

This was referenced Sep 9, 2026

Build ODEProblem from array differential equations with complete only #5101

Closed

Keep array equations through the compiled ODEProblem path SciML/MethodOfLines.jl#688

Closed

ChrisRackauckas and others added 4 commits September 9, 2026 12:55

Document arrays_scalarized and ScalarizeArraysCtx in the API reference

e072087

The QA suite requires every `@public` name to be rendered in the docs. Co-authored-by: Cursor <cursoragent@cursor.com>

Approve the ScalarizeArraysCtx and arrays_scalarized reexports in the…

4232f34

… QA list Co-authored-by: Cursor <cursoragent@cursor.com>

Allow the qualified ModelingToolkitBase accesses used by the array-pr…

562ebe4

…eserving path `__mtkcompile_no_tearing` and `count_equation_rows` are internal helpers reached from ModelingToolkit, like `__mtkcompile` before them. Co-authored-by: Cursor <cursoragent@cursor.com>

ChrisRackauckas-Claude closed this Sep 9, 2026

AayushSabharwal reviewed Sep 9, 2026

View reviewed changes

This was referenced Sep 9, 2026

MTKTearing: keep array differential equations intact through structural simplification JuliaComputing/StateSelection.jl#158

Draft

CHECKPOINT: preserve array DEs through mtkcompile (ArrayEquationGroup / not skip-tearing) #5102

Draft

	src = Expr(:call, reshape, src, map(length, region)...)
	src = Expr(:call, reshape, src)
	append!(src.args, Iterators.map(length, region))

	dest = Expr(:call, view, result, region...)
	dest = Expr(:call, view, result)
	append!(dest.args, region)

Uh oh!

Conversation

ChrisRackauckas-Claude commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changes

Example

How to verify the O(1) scaling

Limitations

Tests run locally

CI notes

Uh oh!

ChrisRackauckas-Claude commented Sep 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ChrisRackauckas-Claude commented Sep 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ChrisRackauckas-Claude commented Sep 9, 2026 •

edited

Loading