feat: accept array equations on the implicit-DAE path

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -3286,17 +3286,36 @@ function check_array_equations_unknowns(eqs, dvs)
     end
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Number of scalar residual rows the equations stand for. An array equation contributes one
+row per element, so it cannot be counted as a single equation.
+"""
+count_equation_rows(eqs) = sum(equation_row_count, eqs; init = 0)
+
+equation_row_count(eq) = 1
+
+function equation_row_count(eq::Equation)
+    # A residual may be written with the array on either side, as `D(u[2:4]) ~ rhs` or
+    # `0 ~ rhs`.
+    sh = SU.shape(eq.lhs)
+    SU.is_array_shape(sh) || (sh = SU.shape(eq.rhs))
+    return SU.is_array_shape(sh) ? prod(length, sh) : 1
+end
+
 function check_eqs_u0(eqs, dvs, u0; check_length = true, kwargs...)
+    neqs = count_equation_rows(eqs)
     if u0 !== nothing
         if check_length
-            if !(length(eqs) == length(dvs) == length(u0))
-                throw(ArgumentError("Equations ($(length(eqs))), unknowns ($(length(dvs))), and initial conditions ($(length(u0))) are of different lengths."))
+            if !(neqs == length(dvs) == length(u0))
+                throw(ArgumentError("Equations ($(neqs)), unknowns ($(length(dvs))), and initial conditions ($(length(u0))) are of different lengths."))
             end
         elseif length(dvs) != length(u0)
             throw(ArgumentError("Unknowns ($(length(dvs))) and initial conditions ($(length(u0))) are of different lengths."))
         end
-    elseif check_length && (length(eqs) != length(dvs))
-        throw(ArgumentError("Equations ($(length(eqs))) and Unknowns ($(length(dvs))) are of different lengths."))
+    elseif check_length && (neqs != length(dvs))
+        throw(ArgumentError("Equations ($(neqs)) and Unknowns ($(length(dvs))) are of different lengths."))
     end
     return nothing
 end

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -37,6 +37,84 @@ $GENERATE_X_KWARGS
 
 All other keyword arguments are forwarded to [`build_function_wrapper`](@ref).
 """
+
+"""
+Treat a derivative of an array-valued expression as a leaf, so that
+[`expand_array_derivatives`](@ref) collects `D(u[2:4])` itself rather than descending into
+it. Scalar variables are not atomic here, so nothing else is collected.
+"""
+struct ArrayDerivativeIsAtomic end
+
+function (::ArrayDerivativeIsAtomic)(ex::SymbolicT)
+    return isdifferential(ex) && SU.is_array_shape(SU.shape(ex))
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Rewrite derivatives of array-valued expressions, such as `D(u[2:4])`, into arrays of the
+corresponding scalar derivatives. Implicit-DAE codegen binds scalar `D(uᵢ)` terms to
+elements of the `du` argument, and a derivative of a slice matches none of them.
+
+Takes every residual at once so that the search and substitution caches are shared.
+"""
+function expand_array_derivatives(rhss::Vector{SymbolicT})
+    terms = Set{SymbolicT}()
+    for rhs in rhss
+        SU.search_variables!(terms, rhs; is_atomic = ArrayDerivativeIsAtomic())
+    end
+    isempty(terms) && return rhss
+
+    subs = Dict{SymbolicT, SymbolicT}()
+    for term in terms
+        op = operation(term)
+        arg = only(arguments(term))
+        sh = SU.shape(arg)
+        # Preserve the shape: a derivative of a 2D slice must expand to a 2D array of
+        # scalar derivatives, or it will not broadcast against the surrounding slices.
+        sz = ntuple(i -> length(sh[i]), length(sh))
+        els = [op(arg[idx]) for idx in SU.stable_eachindex(arg)]
+        subs[term] = SU.Const{VartypeT}(reshape(els, sz))
+    end
+
+    subber = SU.IRSubstituter{false}(IRStructure{VartypeT}(), subs)
+    return map(subber, rhss)
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Assemble residuals into a single array expression, where each residual writes to a
+contiguous region of the output. An array-valued residual stands for one output row per
+element, and writing it as a region keeps the array computation intact instead of
+scalarizing it into one expression per row.
+
+Returns `rhss` unchanged when every residual is scalar.
+"""
+function array_residual_maker(rhss::Vector{SymbolicT})
+    any(rhs -> SU.is_array_shape(SU.shape(rhs)), rhss) || return rhss
+
+    regions = Vector{Vector{UnitRange{Int}}}(undef, length(rhss))
+    values = similar(rhss)
+    offset = 0
+    for (i, rhs) in enumerate(rhss)
+        sh = SU.shape(rhs)
+        if SU.is_array_shape(sh)
+            n = prod(length, sh)
+            # Elements become consecutive output rows, so a rank > 1 residual is flattened
+            # rather than given a multidimensional region. `vec` here stays symbolic.
+            values[i] = length(sh) == 1 ? rhs : vec(rhs)
+        else
+            n = 1
+            # `ArrayMaker` regions only accept array-valued entries.
+            values[i] = SU.Const{VartypeT}([rhs])
+        end
+        regions[i] = [(offset + 1):(offset + n)]
+        offset += n
+    end
+    return SU.ArrayMaker{VartypeT}(regions, values)
+end
+
 function generate_rhs(
         sys::System, opts::GeneratedFunctionOptions;
         implicit_dae::Bool = false, scalar::Bool = false,
@@ -57,6 +135,7 @@ function generate_rhs(
     t = get_iv(sys)
     ddvs = nothing
     extra_assignments = Assignment[]
+    assemble_residuals = false
 
     # used for DAEProblem and ImplicitDiscreteProblem
     if implicit_dae
@@ -83,7 +162,11 @@ function generate_rhs(
         else
             D = Differential(t)
             ddvs = map(D, dvs)
-            rhss = [_iszero(eq.lhs) ? eq.rhs : eq.rhs - eq.lhs for eq in eqs]
+            rhss = SymbolicT[_iszero(eq.lhs) ? eq.rhs : eq.rhs - eq.lhs for eq in eqs]
+            # Rewrite array derivatives to the scalar ones bound to the `du` argument.
+            # Assembly into a single array happens below, after assertions.
+            rhss = expand_array_derivatives(rhss)
+            assemble_residuals = true
         end
     else
         if !override_discrete && !is_discrete_system(sys)
@@ -94,7 +177,18 @@ function generate_rhs(
     end
 
     if !isempty(assertions(sys)) && !isempty(rhss)
-        rhss[end] += unwrap(get_assertions_expr(sys))
+        assertion_expr = unwrap(get_assertions_expr(sys))
+        # An array-valued residual stands for several output rows, and `+` is not defined
+        # between a symbolic array and a scalar, so add the assertion to each of its rows.
+        rhss[end] = if SU.is_array_shape(SU.shape(rhss[end]))
+            unwrap(wrap(rhss[end]) .+ assertion_expr)
+        else
+            rhss[end] + assertion_expr
+        end
+    end
+
+    if assemble_residuals
+        rhss = array_residual_maker(rhss)
     end
 
     # TODO: add an optional check on the ordering of observed equations

lib/ModelingToolkitBase/src/systems/ir_info.jl

@@ -82,6 +82,10 @@ function get_ir_info(sys::System)
     elseif is_time_dependent(sys)
         for (i, eq) in enumerate(eqs)
             isdiffeq(eq) || continue
+            # A derivative of an array slice, as in `D(u[2:4]) ~ rhs`, has no scalar
+            # `toterm` name to key a substitution on. Implicit-DAE codegen expands such
+            # derivatives into their scalar elements before building the residual.
+            SU.is_array_shape(SU.shape(eq.lhs)) && continue
             ttk = default_toterm(eq.lhs)
             isequal(ttk, eq.rhs) && continue
 

lib/ModelingToolkitBase/src/systems/problem_utils.jl

@@ -2154,7 +2154,9 @@ function __process_SciMLProblem(
     iv = has_iv(sys) ? get_iv(sys) : nothing
     eqs = equations(sys)
 
-    check_array_equations_unknowns(eqs, dvs)
+    # Implicit-DAE codegen expands an array equation into one output row per element, so
+    # array equations are usable there. Every other problem type still needs `mtkcompile`.
+    implicit_dae || check_array_equations_unknowns(eqs, dvs)
 
     op = build_operating_point(sys, op; fast_path = true)
 

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 SU.is_array_shape(SU.shape(arg))
+                for idx in SU.stable_eachindex(arg)
+                    push!(diffvars, arg[idx])
+                end
+            else
+                push!(diffvars, arg)
+            end
         end
         empty!(vars)
     end

lib/ModelingToolkitBase/test/array_equation_dae.jl

@@ -0,0 +1,148 @@
+using ModelingToolkitBase, Test
+using ModelingToolkitBase: unwrap, complete, unknowns
+using Symbolics
+using SciMLBase
+using OrdinaryDiffEqBDF: DFBDF
+using DiffEqBase: BrownFullBasicInit
+
+# A system whose interior is written as one array equation over slices, as produced by a
+# finite-difference PDE discretization that does not scalarize.
+function heat_array_system(n)
+    @independent_variables t
+    @variables u(t)[1:n]
+    D = Differential(t)
+    dx = 1 / (n - 1)
+    # Residual (cardinalized) form, as a finite-difference discretization emits it: the
+    # derivative sits inside the expression rather than being the equation's whole LHS.
+    lap = (u[1:(n - 2)] .- 2 .* u[2:(n - 1)] .+ u[3:n]) ./ dx^2
+    interior = broadcast(-, D(u[2:(n - 1)]), lap) ~ zeros(n - 2)
+    eqs = [interior, u[1] ~ 0.0, u[n] ~ 0.0]
+    @named sys = System(eqs, t, collect(u), [])
+    return complete(sys), u, t, D
+end
+
+@testset "array equations reach DAEProblem" begin
+    n = 11
+    sys, u, t, D = heat_array_system(n)
+    xs = range(0.0, 1.0, length = n)
+    op = vcat(
+        [u[i] => sinpi(xs[i]) for i in 1:n],
+        [D(u[i]) => 0.0 for i in 1:n]
+    )
+
+    prob = DAEProblem(sys, op, (0.0, 0.1); build_initializeprob = false)
+
+    # one output row per element of the array equation, not one per equation
+    @test length(prob.u0) == n
+    @test prob.u0 isa Vector{Float64}
+
+    # the interior points are differential, the two boundary points algebraic
+    @test prob.differential_vars !== nothing
+    @test count(prob.differential_vars) == n - 2
+
+    # the residual evaluates: no `Differential` survives into the generated code
+    out = zeros(n)
+    du = zeros(n)
+    prob.f(out, du, prob.u0, prob.p, 0.0)
+    @test all(isfinite, out)
+    # with du = 0 the interior residual is minus the Laplacian, which is nonzero here
+    @test any(!iszero, out)
+end
+
+@testset "other problem types still require scalarized equations" begin
+    n = 11
+    sys, u, t, D = heat_array_system(n)
+    op = [u[i] => 0.0 for i in 1:n]
+    # ODEProblem cannot consume array equations; the guard must remain in place
+    @test_throws Exception ODEProblem(sys, op, (0.0, 0.1); build_initializeprob = false)
+end
+
+@testset "array-equation DAE solves to the analytic solution" begin
+    n = 21
+    sys, u, t, D = heat_array_system(n)
+    xs = range(0.0, 1.0, length = n)
+    op = vcat(
+        [u[i] => sinpi(xs[i]) for i in 1:n],
+        [D(u[i]) => 0.0 for i in 1:n]
+    )
+    tend = 0.1
+    prob = DAEProblem(sys, op, (0.0, tend); build_initializeprob = false)
+    # `du0` above is not consistent; the solver's own DAE initialization supplies it.
+    sol = solve(
+        prob, DFBDF(); initializealg = BrownFullBasicInit(),
+        reltol = 1.0e-8, abstol = 1.0e-8, saveat = [tend]
+    )
+    @test SciMLBase.successful_retcode(sol)
+    exact = [exp(-pi^2 * tend) * sinpi(xi) for xi in xs]
+    # second-order spatial discretization on 21 points
+    @test maximum(abs.(sol.u[end] .- exact)) < 5.0e-3
+end
+
+@testset "array equations over a 2D slice keep their shape" begin
+    # A derivative of a 2D slice must substitute a 2D array of scalar derivatives; a
+    # flattened one does not broadcast against the surrounding slices and codegen fails
+    # with a DimensionMismatch.
+    n = 6
+    @independent_variables t
+    @variables w(t)[1:n, 1:n]
+    D = Differential(t)
+    dx = 1 / (n - 1)
+    inner = 2:(n - 1)
+    lap = (
+        w[1:(n - 2), inner] .+ w[3:n, inner] .+ w[inner, 1:(n - 2)] .+
+            w[inner, 3:n] .- 4 .* w[inner, inner]
+    ) ./ dx^2
+    eqs = Equation[broadcast(-, D(w[inner, inner]), lap) ~ zeros(n - 2, n - 2)]
+    for i in 1:n
+        push!(eqs, w[i, 1] ~ 0.0)
+        push!(eqs, w[i, n] ~ 0.0)
+    end
+    for j in inner
+        push!(eqs, w[1, j] ~ 0.0)
+        push!(eqs, w[n, j] ~ 0.0)
+    end
+    @named sys2d = System(eqs, t, vec(collect(w)), [])
+    sys2d = complete(sys2d)
+
+    op = vcat(
+        [w[i, j] => 0.25 for i in 1:n, j in 1:n] |> vec,
+        [D(w[i, j]) => 0.0 for i in 1:n, j in 1:n] |> vec
+    )
+    prob = DAEProblem(sys2d, op, (0.0, 0.01); build_initializeprob = false)
+    @test length(prob.u0) == n * n
+    out = zeros(n * n)
+    prob.f(out, zeros(n * n), prob.u0, prob.p, 0.0)
+    @test all(isfinite, out)
+end
+
+@testset "array equations written as `D(slice) ~ rhs`" begin
+    # The residual form above puts the derivative inside the expression. The equivalent
+    # `D(u[2:n-1]) ~ rhs` form has no scalar `toterm` name for its LHS, which the
+    # derivative-substitution machinery has to skip rather than trip over.
+    n = 11
+    @independent_variables t
+    @variables u(t)[1:n]
+    D = Differential(t)
+    dx = 1 / (n - 1)
+    lap = (u[1:(n - 2)] .- 2 .* u[2:(n - 1)] .+ u[3:n]) ./ dx^2
+    eqs = [D(u[2:(n - 1)]) ~ lap, u[1] ~ 0.0, u[n] ~ 0.0]
+    @named sys = System(eqs, t, collect(u), [])
+    sys = complete(sys)
+
+    xs = range(0.0, 1.0, length = n)
+    op = vcat([u[i] => sinpi(xs[i]) for i in 1:n], [D(u[i]) => 0.0 for i in 1:n])
+    prob = DAEProblem(sys, op, (0.0, 0.1); build_initializeprob = false)
+    @test length(prob.u0) == n
+
+    # the residual matches the analytic derivative of the initial condition
+    out = zeros(n)
+    du = zeros(n)
+    du[2:(n - 1)] .= [-pi^2 * sinpi(x) for x in xs[2:(n - 1)]]
+    prob.f(out, du, prob.u0, prob.p, 0.0)
+    @test maximum(abs, out) < 1.0e-1
+
+    sol = solve(prob, DFBDF(); initializealg = BrownFullBasicInit(), reltol = 1.0e-8,
+        abstol = 1.0e-8, saveat = [0.1])
+    @test SciMLBase.successful_retcode(sol)
+    @test maximum(abs, sol.u[end] .- [exp(-pi^2 * 0.1) * sinpi(x) for x in xs]) < 1.0e-2
+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,32 @@ 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]
+
+    # a slice of rank 2 records every element, not just the first column
+    @variables z(t)[1:3, 1:2]
+    twod = collect_differential_variables(Dt(z[1:2, 1:2]) ~ z[1:2, 1:2])
+    @test twod == Set(Any[unwrap(z[i, j]) for i in 1:2, j in 1:2])
+end
+
 @testset "parse_variable with scalarized arrays" begin
     @variables scalarized_x(t)[1:2]
     @parameters scalarized_p[1:2]