feat: add LinearProblem codegen

Author: AayushSabharwal

src/ModelingToolkit.jl

@@ -188,6 +188,7 @@ include("problems/jumpproblem.jl")
 include("problems/initializationproblem.jl")
 include("problems/sccnonlinearproblem.jl")
 include("problems/bvproblem.jl")
+include("problems/linearproblem.jl")
 
 include("modelingtoolkitize/common.jl")
 include("modelingtoolkitize/odeproblem.jl")

src/problems/compatibility.jl

@@ -169,3 +169,12 @@ function check_no_equations(sys::System, T)
         """))
     end
 end
+
+function check_affine(sys::System, T)
+    if !isaffine(sys)
+        throw(SystemCompatibilityError("""
+        A non-affine system cannot be used to construct a `$T`. Consider a
+        `NonlinearProblem` instead.
+        """))
+    end
+end

src/problems/docs.jl

@@ -391,3 +391,32 @@ $PROBLEM_INTERNALS_HEADER
 
 $PROBLEM_INTERNAL_KWARGS
 """ SciMLBase.IntervalNonlinearProblem
+
+@doc """
+    SciMLBase.LinearProblem(sys::System, op; kwargs...)
+    SciMLBase.LinearProblem{iip}(sys::System, op; kwargs...)
+
+Build a `LinearProblem` given a system `sys` and operating point `op`. `iip` is a boolean
+indicating whether the problem should be in-place. The operating point should be an
+iterable collection of key-value pairs mapping variables/parameters in the system to the
+(initial) values they should take in `LinearProblem`. Any values not provided will
+fallback to the corresponding default (if present).
+
+Note that since `u0` is optional for `LinearProblem`, values of unknowns do not need to be
+specified in `op` to create a `LinearProblem`. In such a case, `prob.u0` will be `nothing`
+and attempting to symbolically index the problem with an unknown, observable, or expression
+depending on unknowns/observables will error.
+
+Updating the parameters automatically updates the `A` and `b` arrays.
+
+# Keyword arguments
+
+$PROBLEM_KWARGS
+$(prob_fun_common_kwargs(LinearProblem, false))
+
+All other keyword arguments are forwarded to the $func constructor.
+
+$PROBLEM_INTERNALS_HEADER
+
+$PROBLEM_INTERNAL_KWARGS
+""" SciMLBase.LinearProblem

src/problems/linearproblem.jl

@@ -0,0 +1,98 @@
+function SciMLBase.LinearProblem(sys::System, op; kwargs...)
+    SciMLBase.LinearProblem{true}(sys, op; kwargs...)
+end
+
+function SciMLBase.LinearProblem(sys::System, op::StaticArray; kwargs...)
+    SciMLBase.LinearProblem{false}(sys, op; kwargs...)
+end
+
+function SciMLBase.LinearProblem{iip}(
+        sys::System, op; check_length = true, expression = Val{false},
+        check_compatibility = true, sparse = false, eval_expression = false,
+        eval_module = @__MODULE__, checkbounds = false, cse = true,
+        u0_constructor = identity, u0_eltype = nothing, kwargs...) where {iip}
+    check_complete(sys, LinearProblem)
+    check_compatibility && check_compatible_system(LinearProblem, sys)
+
+    _, u0, p = process_SciMLProblem(
+        EmptySciMLFunction{iip}, sys, op; check_length, expression,
+        build_initializeprob = false, symbolic_u0 = true, u0_constructor, u0_eltype,
+        kwargs...)
+
+    if any(x -> symbolic_type(x) != NotSymbolic(), u0)
+        u0 = nothing
+    end
+
+    u0Type = typeof(op)
+    floatT = if u0 === nothing
+        calculate_float_type(op, u0Type)
+    else
+        eltype(u0)
+    end
+    u0_eltype = something(u0_eltype, floatT)
+
+    u0_constructor = get_p_constructor(u0_constructor, u0Type, u0_eltype)
+
+    A, b = calculate_A_b(sys; sparse)
+    update_A = generate_update_A(sys, A; expression, wrap_gfw = Val{true}, eval_expression,
+        eval_module, checkbounds, cse, kwargs...)
+    update_b = generate_update_b(sys, b; expression, wrap_gfw = Val{true}, eval_expression,
+        eval_module, checkbounds, cse, kwargs...)
+    observedfun = ObservedFunctionCache(
+        sys; steady_state = false, expression, eval_expression, eval_module, checkbounds,
+        cse)
+
+    if expression == Val{true}
+        symbolic_interface = quote
+            update_A = $update_A
+            update_b = $update_b
+            sys = $sys
+            observedfun = $observedfun
+            $(SciMLBase.SymbolicLinearInterface)(
+                update_A, update_b, sys, observedfun, nothing)
+        end
+        get_A = build_explicit_observed_function(
+            sys, A; param_only = true, eval_expression, eval_module)
+        if sparse
+            get_A = SparseArrays.sparse ∘ get_A
+        end
+        get_b = build_explicit_observed_function(
+            sys, b; param_only = true, eval_expression, eval_module)
+        A = u0_constructor(get_A(p))
+        b = u0_constructor(get_b(p))
+    else
+        symbolic_interface = SciMLBase.SymbolicLinearInterface(
+            update_A, update_b, sys, observedfun, nothing)
+        A = u0_constructor(update_A(p))
+        b = u0_constructor(update_b(p))
+    end
+
+    kwargs = (; u0, process_kwargs(sys; kwargs...)..., f = symbolic_interface)
+    args = (; A, b, p)
+
+    return maybe_codegen_scimlproblem(expression, LinearProblem{iip}, args; kwargs...)
+end
+
+# For remake
+function SciMLBase.get_new_A_b(
+        sys::AbstractSystem, f::SciMLBase.SymbolicLinearInterface, p, A, b; kw...)
+    if ArrayInterface.ismutable(A)
+        f.update_A!(A, p)
+        f.update_b!(b, p)
+    else
+        # The generated function has both IIP and OOP variants
+        A = StaticArraysCore.similar_type(A)(f.update_A!(p))
+        b = StaticArraysCore.similar_type(b)(f.update_b!(p))
+    end
+    return A, b
+end
+
+function check_compatible_system(T::Type{LinearProblem}, sys::System)
+    check_time_independent(sys, T)
+    check_affine(sys, T)
+    check_not_dde(sys)
+    check_no_cost(sys, T)
+    check_no_constraints(sys, T)
+    check_no_jumps(sys, T)
+    check_no_noise(sys, T)
+end

src/systems/codegen.jl

@@ -1130,3 +1130,90 @@ function build_explicit_observed_function(sys, ts;
         return f
     end
 end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Return matrix `A` and vector `b` such that the system `sys` can be represented as
+`A * x = b` where `x` is `unknowns(sys)`. Errors if the system is not affine.
+
+# Keyword arguments
+
+- `sparse`: return a sparse `A`.
+"""
+function calculate_A_b(sys::System; sparse = false)
+    rhss = [eq.rhs for eq in full_equations(sys)]
+    dvs = unknowns(sys)
+
+    A = Matrix{Any}(undef, length(rhss), length(dvs))
+    b = Vector{Any}(undef, length(rhss))
+    for (i, rhs) in enumerate(rhss)
+        # mtkcompile makes this `0 ~ rhs` which typically ends up giving
+        # unknowns negative coefficients. If given the equations `A * x ~ b`
+        # it will simplify to `0 ~ b - A * x`. Thus this negation usually leads
+        # to more comprehensible user API.
+        resid = -rhs
+        for (j, var) in enumerate(dvs)
+            p, q, islinear = Symbolics.linear_expansion(resid, var)
+            if !islinear
+                throw(ArgumentError("System is not linear. Equation $((0 ~ rhs)) is not linear in unknown $var."))
+            end
+            A[i, j] = p
+            resid = q
+        end
+        # negate beucause `resid` is the residual on the LHS
+        b[i] = -resid
+    end
+
+    @assert all(Base.Fix1(isassigned, A), eachindex(A))
+    @assert all(Base.Fix1(isassigned, A), eachindex(b))
+
+    if sparse
+        A = SparseArrays.sparse(A)
+    end
+    return A, b
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Given a system `sys` and the `A` from [`calculate_A_b`](@ref) generate the function that
+updates `A` given the parameter object.
+
+# Keyword arguments
+
+$GENERATE_X_KWARGS
+
+All other keyword arguments are forwarded to [`build_function_wrapper`](@ref).
+"""
+function generate_update_A(sys::System, A::AbstractMatrix; expression = Val{true},
+        wrap_gfw = Val{false}, eval_expression = false, eval_module = @__MODULE__, kwargs...)
+    ps = reorder_parameters(sys)
+
+    res = build_function_wrapper(sys, A, ps...; p_start = 1, expression = Val{true},
+        similarto = typeof(A), kwargs...)
+    return maybe_compile_function(expression, wrap_gfw, (1, 1, is_split(sys)), res;
+        eval_expression, eval_module)
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Given a system `sys` and the `b` from [`calculate_A_b`](@ref) generate the function that
+updates `b` given the parameter object.
+
+# Keyword arguments
+
+$GENERATE_X_KWARGS
+
+All other keyword arguments are forwarded to [`build_function_wrapper`](@ref).
+"""
+function generate_update_b(sys::System, b::AbstractVector; expression = Val{true},
+        wrap_gfw = Val{false}, eval_expression = false, eval_module = @__MODULE__, kwargs...)
+    ps = reorder_parameters(sys)
+
+    res = build_function_wrapper(sys, b, ps...; p_start = 1, expression = Val{true},
+        similarto = typeof(b), kwargs...)
+    return maybe_compile_function(expression, wrap_gfw, (1, 1, is_split(sys)), res;
+        eval_expression, eval_module)
+end