feat: accept array equations on the NonlinearProblem path

Author: utkuyilmaz1903

docs/src/API/model_building.md

@@ -160,7 +160,11 @@ It is also possible (though not always advisable) to build numerical problems fr
 passing them through `mtkcompile`. To do this, the system must first be marked as "complete" via
 the `complete` function. This process is used to indicate that a system will not be modified
 further and allows ModelingToolkit to perform any necessary preprocessing to it. `mtkcompile`
-calls `complete` internally.
+calls `complete` internally. `NonlinearProblem` and `NonlinearLeastSquaresProblem`
+accept unscalarized array equations with scalar unknowns (`collect(u)`). Each
+array equation contributes one residual row per element and may be mixed with
+scalar equations. Use `mtkcompile` before requesting `jac = true` or
+`sparse = true`.
 
 ```@docs
 complete

docs/src/tutorials/nonlinear.md

@@ -41,3 +41,10 @@ Jacobian function:
 prob = NonlinearProblem(ns, vcat(guesses, ps), jac = true)
 sol = solve(prob, NewtonRaphson())
 ```
+
+!!! note "Array residuals"
+
+    `NonlinearProblem` and `NonlinearLeastSquaresProblem` accept unscalarized
+    array equations after [`complete`](@ref) when unknowns are scalar
+    (`collect(u)`). Use [`mtkcompile`](@ref) before requesting `jac = true` or
+    `sparse = true`. Time-dependent systems are converted to steady state.

lib/ModelingToolkitBase/src/problems/nonlinearproblem.jl

@@ -28,20 +28,31 @@ function SciMLBase.NonlinearFunction{iip, spec}(
     ) where {iip, spec, E}
     check_complete(sys, NonlinearFunction)
     opts.check_compatibility && check_compatible_system(NonlinearFunction, sys)
+    check_array_unknowns(unknowns(sys))
 
     (; u0, p, jac, sparse, analytic, simplify, initialization_data) = opts
     codegen_opts = opts.codegen
 
+    if (jac || sparse) && has_array_equations(equations(sys))
+        throw(
+            ArgumentError(
+                "Array residuals do not support `jac = true` or `sparse = true`. " *
+                    "Call `mtkcompile` first."
+            )
+        )
+    end
+
     f = generate_rhs(sys, codegen_opts)
 
     if spec === SciMLBase.FunctionWrapperSpecialize && iip
         if u0 === nothing || p === nothing
             error("u0, and p must be specified for FunctionWrapperSpecialize on NonlinearFunction.")
         end
+        resid = resid_prototype === nothing ? u0 : resid_prototype
         if E
-            f = :($(SciMLBase.wrapfun_iip)($f, ($u0, $u0, $p)))
+            f = :($(SciMLBase.wrapfun_iip)($f, ($resid, $u0, $p)))
         else
-            f = SciMLBase.wrapfun_iip(f, (u0, u0, p))
+            f = SciMLBase.wrapfun_iip(f, (resid, u0, p))
         end
     end
 
@@ -181,8 +192,8 @@ end
     _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        NonlinearFunction{_iip}, sys, op;
-        check_length, expression, kwargs...
+        NonlinearFunction{_iip, spec}, sys, op;
+        check_length, check_compatibility, expression, kwargs...
     )
 
     if lb === nothing && ub === nothing

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -3277,13 +3277,30 @@ function Base.eltype(::Type{<:TreeIterator{ModelingToolkitBase.AbstractSystem}})
     return ModelingToolkitBase.AbstractSystem
 end
 
-function check_array_equations_unknowns(eqs, dvs)
-    if any(eq -> eq isa Equation && Symbolics.isarraysymbolic(eq.lhs), eqs)
-        throw(ArgumentError("The system has array equations. Call `mtkcompile` to handle such equations or scalarize them manually."))
+function has_array_equations(eqs)
+    return any(eq -> eq isa Equation && SU.is_array_shape(SU.shape(eq.lhs)), eqs)
+end
+
+const ARRAY_EQUATIONS_ERROR = "The system has array equations. Call `mtkcompile` to handle such equations or scalarize them manually."
+const ARRAY_UNKNOWNS_ERROR = "The system has array unknowns. Call `mtkcompile` to handle this or scalarize them manually with `collect(u)`."
+
+function check_array_equations(eqs)
+    if has_array_equations(eqs)
+        throw(ArgumentError(ARRAY_EQUATIONS_ERROR))
     end
-    return if any(x -> Symbolics.isarraysymbolic(x), dvs)
-        throw(ArgumentError("The system has array unknowns. Call `mtkcompile` to handle this or scalarize them manually."))
+    return nothing
+end
+
+function check_array_unknowns(dvs)
+    if any(Symbolics.isarraysymbolic, dvs)
+        throw(ArgumentError(ARRAY_UNKNOWNS_ERROR))
     end
+    return nothing
+end
+
+function check_array_equations_unknowns(eqs, dvs)
+    check_array_equations(eqs)
+    return check_array_unknowns(dvs)
 end
 
 """

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -15,15 +15,6 @@ const EXPERIMENTAL_WARNING = """
     This API is experimental and may change in a future non-breaking release.
 """
 
-"""
-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.
-"""
-function array_derivative_is_atomic(ex::SymbolicT)
-    return isdifferential(ex) && SU.is_array_shape(SU.shape(ex))
-end
-
 """
     $(TYPEDSIGNATURES)
 
@@ -43,25 +34,7 @@ function expand_array_derivatives!(rhss::Vector{SymbolicT}, ir::IRStructure{Vart
     end
     isempty(terms) && return rhss
 
-    subs = Dict{SymbolicT, SymbolicT}()
-    for term in terms
-        op = operation(term)
-        arg = only(arguments(term))
-        sh = SU.shape(arg)::SU.ShapeVecT
-        # 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.
-        arrargs = Symbolics.SArgsT()
-        sizehint!(arrargs, prod(length, sh; init = 1) + 1)
-        push!(arrargs, SU.Const{VartypeT}(size(arg)))
-        for idx in SU.stable_eachindex(arg)
-            push!(arrargs, op(arg[idx]))
-        end
-        subs[term] = Symbolics.STerm(
-            SU.array_literal, arrargs; type = symtype(arg), shape = sh
-        )
-    end
-
-    subber = SU.IRSubstituter{false}(ir, subs)
+    subber = SU.IRSubstituter{false}(ir, array_derivative_expansion_map(terms))
     map!(subber, rhss, rhss)
     return rhss
 end
@@ -100,6 +73,88 @@ function array_residual_maker(rhss::Vector{SymbolicT})
     return SU.ArrayMaker{VartypeT}(regions, values)
 end
 
+"""
+    residual_eltype(x)
+
+Return the numeric element type contributed by `x`, or `Union{}`.
+"""
+function residual_eltype(x)
+    x = SciMLBase.unwrap_parameters(x)
+    if x isa Number
+        return typeof(x)
+    elseif x isa AbstractArray
+        T = eltype(x)
+        return T === Any ? Union{} : T
+    elseif SciMLStructures.isscimlstructure(x)
+        tun = first(SciMLStructures.canonicalize(SciMLStructures.Tunable(), x))
+        return residual_eltype(tun)
+    else
+        return Union{}
+    end
+end
+
+"""
+    similar_for_residual(prototype, extras...)
+
+Return an allocator using `prototype` and the promoted runtime element type.
+"""
+function similar_for_residual(prototype, extras...)
+    T = residual_eltype(prototype)
+    T === Union{} && (T = Float64)
+    for x in extras
+        T = promote_type(T, residual_eltype(x))
+    end
+    return sz -> similar(prototype, T, sz)
+end
+
+function residual_allocator_arg(arg)
+    name = arg isa DestructuredArgs ? arg.name : arg
+    name isa SymbolicT && return name
+    name isa Symbol || throw(
+        ArgumentError(
+            "Cannot form a residual allocator argument from $(typeof(arg))."
+        )
+    )
+    return SSym(name; type = Any, shape = SU.ShapeVecT())
+end
+
+function residual_allocator_term(args)
+    return STerm(
+        similar_for_residual, SArgsT((map(residual_allocator_arg, args)...,));
+        type = SU.FnType{Tuple, Any, Any},
+        shape = SU.ShapeVecT(),
+    )
+end
+
+function inject_similar_for_residual(body, alloc_term)
+    if body isa Let
+        return Let(
+            body.pairs, inject_similar_for_residual(body.body, alloc_term),
+            body.let_block
+        )
+    elseif body isa SymbolicT && Code.supports_with_allocator(body)
+        # The public helper would wrap the symbolic allocator in `Const`.
+        return STerm(
+            Code.with_allocator,
+            SArgsT((alloc_term, body));
+            type = SU.symtype(body),
+            shape = SU.shape(body),
+        )
+    else
+        return body
+    end
+end
+
+function wrap_oop_similar_for_residual(fn)
+    fn isa Func || return fn
+    isempty(fn.args) && return fn
+    alloc_term = residual_allocator_term(fn.args)
+    return Func(
+        fn.args, fn.kwargs, inject_similar_for_residual(fn.body, alloc_term),
+        fn.pre
+    )
+end
+
 """
     $(TYPEDSIGNATURES)
 
@@ -121,6 +176,10 @@ $GENERATE_X_KWARGS
   by default, which leaves the standard codegen path byte-identical.
 
 All other keyword arguments are forwarded to [`build_function_wrapper`](@ref).
+
+Time-independent systems use `_iszero(lhs) ? rhs : rhs - lhs` and
+`array_residual_maker`. No `du`. Out-of-place ArrayMaker residuals promote
+their element type from the function arguments.
 """
 function generate_rhs(
         sys::System, opts::GeneratedFunctionOptions;
@@ -175,6 +234,9 @@ function generate_rhs(
             expand_array_derivatives!(rhss, get_irstructure(sys))
             assemble_residuals = true
         end
+    elseif !is_time_dependent(sys)
+        rhss = SymbolicT[_iszero(eq.lhs) ? eq.rhs : eq.rhs - eq.lhs for eq in eqs]
+        assemble_residuals = true
     else
         if !override_discrete && !is_discrete_system(sys)
             check_operator_variables(eqs, Differential)
@@ -223,10 +285,20 @@ function generate_rhs(
         (; p_end = (t === nothing ? length(args) : length(args) - 1) - length(extra_args))
 
     u_arg = scalar ? -1 : (implicit_dae ? 2 : 1)
+    codegen_opts = opts.codegen
+    if !implicit_dae && assemble_residuals && rhss isa SymbolicT &&
+            Code.supports_with_allocator(rhss)
+        # `ArrayMaker` otherwise allocates a `Float64` buffer out of place.
+        oop_wrap, iip_wrap = codegen_opts.wrap_code
+        codegen_opts = setproperties(
+            codegen_opts,
+            (; wrap_code = (wrap_oop_similar_for_residual ∘ oop_wrap, iip_wrap))
+        )
+    end
     res = build_function_wrapper(
         sys, rhss, collect(Any, args), BuildFunctionWrapperOptions(;
             p_start, extra_assignments, u_arg, n_param_buffers, p_end_kw...,
-            codegen_function_options = opts.codegen
+            codegen_function_options = codegen_opts
         )
     )
     nargs = length(args) - length(p) + 1

lib/ModelingToolkitBase/src/systems/problem_utils.jl

@@ -2198,9 +2198,13 @@ function __process_SciMLProblem(
     iv = has_iv(sys) ? get_iv(sys) : nothing
     eqs = equations(sys)
 
-    # 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)
+    if !implicit_dae
+        if constructor <: NonlinearFunction
+            check_array_unknowns(dvs)
+        else
+            check_array_equations_unknowns(eqs, dvs)
+        end
+    end
 
     op = build_operating_point(sys, op; fast_path = true)
 
@@ -2318,17 +2322,20 @@ function __process_SciMLProblem(
         du0 = nothing
     end
 
-    if constructor <: NonlinearFunction && length(dvs) != length(eqs)
-        kwargs = merge(
-            kwargs,
-            (;
-                resid_prototype = u0_constructor(
-                    calculate_resid_prototype(
-                        length(eqs), u0, p
-                    )
-                ),
+    if constructor <: NonlinearFunction
+        nrows = count_equation_rows(eqs)
+        if length(dvs) != nrows
+            kwargs = merge(
+                kwargs,
+                (;
+                    resid_prototype = u0_constructor(
+                        calculate_resid_prototype(
+                            nrows, u0, p
+                        )
+                    ),
+                )
             )
-        )
+        end
     end
 
     f = constructor(

lib/ModelingToolkitBase/src/systems/system.jl

@@ -1345,11 +1345,11 @@ end
 
 Given a time-dependent system `sys` of ODEs, convert it to a time-independent system of
 nonlinear equations that solve for the steady-state of the unknowns. This is done by
-replacing every derivative `D(x)` of an unknown `x` with zero. Note that this process
-does not retain noise equations, brownian terms, jumps or costs associated with `sys`.
-All other information such as initial conditions, bindings, guesses, observed and
-initialization equations are retained. The independent variable of `sys` becomes a
-parameter of the returned system.
+replacing every derivative `D(x)` of an unknown `x` with zero. Array derivatives are
+expanded first. Note that this process does not retain noise equations, brownian terms,
+jumps or costs associated with `sys`. All other information such as initial conditions,
+bindings, guesses, observed and initialization equations are retained. The independent
+variable of `sys` becomes a parameter of the returned system.
 
 If `sys` is hierarchical (it contains subsystems) this transformation will be applied
 recursively to all subsystems. The output system will be marked as `complete` if and only
@@ -1362,7 +1362,7 @@ function NonlinearSystem(sys::System)
     if !is_time_dependent(sys)
         throw(ArgumentError("`NonlinearSystem` constructor expects a time-dependent `System`"))
     end
-    eqs = equations(sys)
+    eqs = expand_array_derivatives(equations(sys))
     obs = observed(sys)
     D = Differential(get_iv(sys))
     subrules = Dict([D(x) => 0.0 for x in unknowns(sys)])

lib/ModelingToolkitBase/src/utils.jl

@@ -688,6 +688,44 @@ isoperator(::Type{op}) where {op <: SU.Operator} = Base.Fix2(isoperator, op)
 isdifferential(expr) = isoperator(expr, Differential)
 isdiffeq(eq) = isdifferential(eq.lhs) || isoperator(eq.lhs, Shift)
 
+function array_derivative_is_atomic(ex::SymbolicT)
+    return isdifferential(ex) && SU.is_array_shape(SU.shape(ex))
+end
+
+function array_derivative_expansion(term::SymbolicT)
+    op = operation(term)
+    arg = only(arguments(term))
+    sh = SU.shape(arg)::SU.ShapeVecT
+    # Preserve rank for broadcasts against surrounding slices.
+    arrargs = Symbolics.SArgsT()
+    sizehint!(arrargs, prod(length, sh; init = 1) + 1)
+    push!(arrargs, SU.Const{VartypeT}(size(arg)))
+    for idx in SU.stable_eachindex(arg)
+        push!(arrargs, op(arg[idx]))
+    end
+    return Symbolics.STerm(
+        SU.array_literal, arrargs; type = symtype(arg), shape = sh
+    )
+end
+
+function array_derivative_expansion_map(terms)
+    subs = Dict{SymbolicT, SymbolicT}()
+    for term in terms
+        subs[term] = array_derivative_expansion(term)
+    end
+    return subs
+end
+
+function expand_array_derivatives(eqs::Vector{Equation})
+    terms = Set{SymbolicT}()
+    for eq in eqs
+        SU.search_variables!(terms, eq; is_atomic = array_derivative_is_atomic)
+    end
+    isempty(terms) && return eqs
+    subs = array_derivative_expansion_map(terms)
+    return map(eq -> substitute(eq, subs), eqs)
+end
+
 isvariable(x::Num)::Bool = isvariable(value(x))
 function isvariable(x)
     x isa SymbolicT || return false