Merge pull request #3760 from AayushSabharwal/as/linear-scc

feat: use `LinearProblem` for linear SCCs in `SCCNonlinearProblem`

Author: ChrisRackauckas

Project.toml

@@ -147,7 +147,7 @@ REPL = "1"
 RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
-SCCNonlinearSolve = "1.0.0"
+SCCNonlinearSolve = "1.4.0"
 SciMLBase = "2.108.0"
 SciMLPublic = "1.0.0"
 SciMLStructures = "1.7"
@@ -156,7 +156,7 @@ Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "0.1.0, 1, 2"
 SparseArrays = "1"
 SpecialFunctions = "1, 2"
-StaticArrays = "0.10, 0.11, 0.12, 1.0"
+StaticArrays = "1.9.14"
 StochasticDelayDiffEq = "1.10"
 StochasticDiffEq = "6.72.1"
 SymbolicIndexingInterface = "0.3.39"

src/problems/linearproblem.jl

@@ -1,3 +1,42 @@
+struct LinearFunction{iip, I} <: SciMLBase.AbstractSciMLFunction{iip}
+    interface::I
+    A::AbstractMatrix
+    b::AbstractVector
+end
+
+function LinearFunction{iip}(
+        sys::System; expression = Val{false}, check_compatibility = true,
+        sparse = false, eval_expression = false, eval_module = @__MODULE__,
+        checkbounds = false, cse = true, kwargs...) where {iip}
+    check_complete(sys, LinearProblem)
+    check_compatibility && check_compatible_system(LinearProblem, sys)
+
+    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
+    else
+        symbolic_interface = SciMLBase.SymbolicLinearInterface(
+            update_A, update_b, sys, observedfun, nothing)
+    end
+
+    return LinearFunction{iip, typeof(symbolic_interface)}(symbolic_interface, A, b)
+end
+
 function SciMLBase.LinearProblem(sys::System, op; kwargs...)
     SciMLBase.LinearProblem{true}(sys, op; kwargs...)
 end
@@ -9,14 +48,14 @@ 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}
+        eval_module = @__MODULE__, u0_constructor = identity, u0_eltype = nothing,
+        kwargs...) where {iip}
     check_complete(sys, LinearProblem)
     check_compatibility && check_compatible_system(LinearProblem, sys)
 
-    _, u0,
+    f, u0,
     p = process_SciMLProblem(
-        EmptySciMLFunction{iip}, sys, op; check_length, expression,
+        LinearFunction{iip}, sys, op; check_length, expression,
         build_initializeprob = false, symbolic_u0 = true, u0_constructor, u0_eltype,
         kwargs...)
 
@@ -33,45 +72,38 @@ function SciMLBase.LinearProblem{iip}(
     u0_eltype = something(u0_eltype, floatT)
 
     u0_constructor = get_p_constructor(u0_constructor, u0Type, u0_eltype)
+    symbolic_interface = f.interface
+    A,
+    b = get_A_b_from_LinearFunction(
+        sys, f, p; eval_expression, eval_module, expression, u0_constructor, sparse)
 
-    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)
+    kwargs = (; u0, process_kwargs(sys; kwargs...)..., f = symbolic_interface)
+    args = (; A, b, p)
 
+    return maybe_codegen_scimlproblem(expression, LinearProblem{iip}, args; kwargs...)
+end
+
+function get_A_b_from_LinearFunction(
+        sys::System, f::LinearFunction, p; eval_expression = false,
+        eval_module = @__MODULE__, expression = Val{false}, u0_constructor = identity,
+        u0_eltype = float, sparse = false)
+    @unpack A, b, interface = f
     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))
+        A = u0_constructor(u0_eltype.(get_A(p)))
+        b = u0_constructor(u0_eltype.(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))
+        A = u0_constructor(u0_eltype.(interface.update_A!(p)))
+        b = u0_constructor(u0_eltype.(interface.update_b!(p)))
+    end
+    if sparse
+        A = SparseArrays.sparse(A)
     end
 
-    kwargs = (; u0, process_kwargs(sys; kwargs...)..., f = symbolic_interface)
-    args = (; A, b, p)
-
-    return maybe_codegen_scimlproblem(expression, LinearProblem{iip}, args; kwargs...)
+    return A, b
 end
 
 # For remake

src/problems/sccnonlinearproblem.jl

@@ -10,7 +10,7 @@ end
 
 function CacheWriter(sys::AbstractSystem, buffer_types::Vector{TypeT},
         exprs::Dict{TypeT, Vector{Any}}, solsyms, obseqs::Vector{Equation};
-        eval_expression = false, eval_module = @__MODULE__, cse = true)
+        eval_expression = false, eval_module = @__MODULE__, cse = true, sparse = false)
     ps = parameters(sys; initial_parameters = true)
     rps = reorder_parameters(sys, ps)
     obs_assigns = [eq.lhs ← eq.rhs for eq in obseqs]
@@ -39,9 +39,22 @@ end
 struct SCCNonlinearFunction{iip} end
 
 function SCCNonlinearFunction{iip}(
-        sys::System, _eqs, _dvs, _obs, cachesyms; eval_expression = false,
+        sys::System, _eqs, _dvs, _obs, cachesyms, op; eval_expression = false,
         eval_module = @__MODULE__, cse = true, kwargs...) where {iip}
     ps = parameters(sys; initial_parameters = true)
+    subsys = System(
+        _eqs, _dvs, ps; observed = _obs, name = nameof(sys), defaults = defaults(sys))
+    @set! subsys.parameter_dependencies = parameter_dependencies(sys)
+    if get_index_cache(sys) !== nothing
+        @set! subsys.index_cache = subset_unknowns_observed(
+            get_index_cache(sys), sys, _dvs, getproperty.(_obs, (:lhs,)))
+        @set! subsys.complete = true
+    end
+    # generate linear problem instead
+    if isaffine(subsys)
+        return LinearFunction{iip}(
+            subsys; eval_expression, eval_module, cse, cachesyms, kwargs...)
+    end
     rps = reorder_parameters(sys, ps)
 
     obs_assignments = [eq.lhs ← eq.rhs for eq in _obs]
@@ -54,14 +67,6 @@ function SCCNonlinearFunction{iip}(
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
     f = GeneratedFunctionWrapper{(2, 2, is_split(sys))}(f_oop, f_iip)
 
-    subsys = System(_eqs, _dvs, ps; observed = _obs,
-        parameter_dependencies = parameter_dependencies(sys), name = nameof(sys))
-    if get_index_cache(sys) !== nothing
-        @set! subsys.index_cache = subset_unknowns_observed(
-            get_index_cache(sys), sys, _dvs, getproperty.(_obs, (:lhs,)))
-        @set! subsys.complete = true
-    end
-
     return NonlinearFunction{iip}(f; sys = subsys)
 end
 
@@ -70,7 +75,7 @@ function SciMLBase.SCCNonlinearProblem(sys::System, args...; kwargs...)
 end
 
 function SciMLBase.SCCNonlinearProblem{iip}(sys::System, op; eval_expression = false,
-        eval_module = @__MODULE__, cse = true, kwargs...) where {iip}
+        eval_module = @__MODULE__, cse = true, u0_constructor = identity, kwargs...) where {iip}
     if !iscomplete(sys) || get_tearing_state(sys) === nothing
         error("A simplified `System` is required. Call `mtkcompile` on the system before creating an `SCCNonlinearProblem`.")
     end
@@ -113,7 +118,7 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::System, op; eval_expression = f
 
     _, u0,
     p = process_SciMLProblem(
-        EmptySciMLFunction{iip}, sys, op; eval_expression, eval_module, kwargs...)
+        EmptySciMLFunction{iip}, sys, op; eval_expression, eval_module, symbolic_u0 = true, kwargs...)
 
     explicitfuns = []
     nlfuns = []
@@ -224,28 +229,57 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::System, op; eval_expression = f
             get(cachevars, T, [])
         end)
         f = SCCNonlinearFunction{iip}(
-            sys, _eqs, _dvs, _obs, cachebufsyms; eval_expression, eval_module, cse, kwargs...)
+            sys, _eqs, _dvs, _obs, cachebufsyms, op;
+            eval_expression, eval_module, cse, kwargs...)
         push!(nlfuns, f)
     end
 
+    u0_eltype = Union{}
+    for x in u0
+        symbolic_type(x) == NotSymbolic() || continue
+        u0_eltype = typeof(x)
+        break
+    end
+    if u0_eltype == Union{}
+        u0_eltype = Float64
+    end
+    u0_eltype = float(u0_eltype)
+
     if !isempty(cachetypes)
         templates = map(cachetypes, cachesizes) do T, n
             # Real refers to `eltype(u0)`
             if T == Real
-                T = eltype(u0)
+                T = u0_eltype
             elseif T <: Array && eltype(T) == Real
-                T = Array{eltype(u0), ndims(T)}
+                T = Array{u0_eltype, ndims(T)}
             end
             BufferTemplate(T, n)
         end
         p = rebuild_with_caches(p, templates...)
     end
 
+    # yes, `get_p_constructor` since this is only used for `LinearProblem` and
+    # will retain the shape of `A`
+    u0_constructor = get_p_constructor(u0_constructor, typeof(u0), u0_eltype)
     subprobs = []
-    for (f, vscc) in zip(nlfuns, var_sccs)
+    for (i, (f, vscc)) in enumerate(zip(nlfuns, var_sccs))
         _u0 = SymbolicUtils.Code.create_array(
             typeof(u0), eltype(u0), Val(1), Val(length(vscc)), u0[vscc]...)
-        prob = NonlinearProblem(f, _u0, p)
+        symbolic_idxs = findall(x -> symbolic_type(x) != NotSymbolic(), _u0)
+        explicitfuns[i](p, subprobs)
+        if f isa LinearFunction
+            _u0 = isempty(symbolic_idxs) ? _u0 : zeros(u0_eltype, length(_u0))
+            _u0 = u0_eltype.(_u0)
+            symbolic_interface = f.interface
+            A,
+            b = get_A_b_from_LinearFunction(
+                sys, f, p; eval_expression, eval_module, u0_constructor, u0_eltype)
+            prob = LinearProblem{iip}(A, b, p; f = symbolic_interface, u0 = _u0)
+        else
+            isempty(symbolic_idxs) || throw(MissingGuessError(dvs[vscc], _u0))
+            _u0 = u0_eltype.(_u0)
+            prob = NonlinearProblem(f, _u0, p)
+        end
         push!(subprobs, prob)
     end
 
@@ -255,5 +289,5 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::System, op; eval_expression = f
     @set! sys.eqs = new_eqs
     @set! sys.index_cache = subset_unknowns_observed(
         get_index_cache(sys), sys, new_dvs, getproperty.(obs, (:lhs,)))
-    return SCCNonlinearProblem(subprobs, explicitfuns, p, true; sys)
+    return SCCNonlinearProblem(Tuple(subprobs), Tuple(explicitfuns), p, true; sys)
 end

src/systems/codegen.jl

@@ -1189,10 +1189,11 @@ $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...)
+        wrap_gfw = Val{false}, eval_expression = false, eval_module = @__MODULE__, cachesyms = (), kwargs...)
     ps = reorder_parameters(sys)
 
-    res = build_function_wrapper(sys, A, ps...; p_start = 1, expression = Val{true},
+    res = build_function_wrapper(
+        sys, A, ps..., cachesyms...; 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)
@@ -1211,10 +1212,11 @@ $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...)
+        wrap_gfw = Val{false}, eval_expression = false, eval_module = @__MODULE__, cachesyms = (), kwargs...)
     ps = reorder_parameters(sys)
 
-    res = build_function_wrapper(sys, b, ps...; p_start = 1, expression = Val{true},
+    res = build_function_wrapper(
+        sys, b, ps..., cachesyms...; 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)

test/initial_values.jl

@@ -202,7 +202,7 @@ end
     @variables a(t) b(t) c(t) d(t) e(t)
     eqs = [D(a) ~ b, D(b) ~ c, D(c) ~ d, D(d) ~ e, D(e) ~ 1]
     @mtkcompile sys = System(eqs, t)
-    @test_throws ["a(t)", "c(t)"] ODEProblem(
+    @test_throws ["d(t)", "c(t)"] ODEProblem(
         sys, [e => 2, a => b, b => a + 1, c => d, d => c + 1], (0, 1))
 end
 

test/initializationsystem.jl

@@ -1670,15 +1670,15 @@ end
            x[1] ~ 0.01exp(-1)
            x[2] ~ 0.01cos(t)]
 
-    @mtkbuild sys = ODESystem(eqs, t)
+    @mtkcompile sys = System(eqs, t)
     prob = ODEProblem(sys, [], (0.0, 1.0))
     sol = solve(prob, Tsit5())
     @test SciMLBase.successful_retcode(sol)
 end
 
 @testset "Defaults removed with ` => nothing` aren't retained" begin
     @variables x(t)[1:2]
-    @mtkbuild sys = System([D(x[1]) ~ -x[1], x[1] + x[2] ~ 3], t; defaults = [x[1] => 1])
+    @mtkcompile sys = System([D(x[1]) ~ -x[1], x[1] + x[2] ~ 3], t; defaults = [x[1] => 1])
     prob = ODEProblem(sys, [x[1] => nothing, x[2] => 1], (0.0, 1.0))
     @test SciMLBase.initialization_status(prob) == SciMLBase.FULLY_DETERMINED
 end
@@ -1696,7 +1696,7 @@ end
             D(x) ~ r * x
         end
     end
-    @mtkbuild sys = Foo(p = "a")
+    @mtkcompile sys = Foo(p = "a")
     prob = ODEProblem(sys, [], (0.0, 1.0))
     @test prob.p.nonnumeric[1] isa Vector{AbstractString}
     integ = init(prob)

test/scc_nonlinear_problem.jl

@@ -27,16 +27,21 @@ using ModelingToolkit: t_nounits as t, D_nounits as D
     @test_throws ["not compatible"] SCCNonlinearProblem(_model, [])
     model = mtkcompile(model)
     prob = NonlinearProblem(model, [u => zeros(8)])
-    sccprob = SCCNonlinearProblem(model, [u => zeros(8)])
+    sccprob = SCCNonlinearProblem(model, collect(u[1:5]) .=> zeros(5))
     sol1 = solve(prob, NewtonRaphson())
     sol2 = solve(sccprob, NewtonRaphson())
     @test SciMLBase.successful_retcode(sol1)
     @test SciMLBase.successful_retcode(sol2)
     @test sol1[u] ≈ sol2[u]
 
-    sccprob = SCCNonlinearProblem{false}(model, SA[u => zeros(8)])
+    sccprob = SCCNonlinearProblem{false}(model, SA[(collect(u[1:5]) .=> zeros(5))...])
     for prob in sccprob.probs
-        @test prob.u0 isa SVector
+        if prob isa LinearProblem
+            @test prob.A isa SMatrix
+            @test prob.b isa SVector
+        else
+            @test prob.u0 isa SVector
+        end
         @test !SciMLBase.isinplace(prob)
     end
 
@@ -91,8 +96,9 @@ end
     @mtkcompile sys = System(eqs, [u], [p1, p2])
     sccprob = SCCNonlinearProblem(sys, [u => u0, p1 => p[1], p2 => p[2][]])
     sccsol = solve(sccprob, SimpleNewtonRaphson(); abstol = 1e-9)
+    sccresid = prob.f(sccsol[u], (u0, p))
     @test SciMLBase.successful_retcode(sccsol)
-    @test norm(sccsol.resid) < norm(sol.resid)
+    @test norm(sccresid) < norm(sol.resid)
 
     # Test BLT sorted
     @test istril(StructuralTransformations.sorted_incidence_matrix(sys), 1)
@@ -173,9 +179,11 @@ end
                               0 ~ func(x[1], x[2]) * exp(x[3]) - x[4]^3 - 5
                               0 ~ func(x[1], x[2]) * exp(x[4]) - x[3]^3 - 4])
     sccprob = SCCNonlinearProblem(sys, [])
+    # since explicitfuns are called during problem construction
+    @test val[] == 1
     sccsol = solve(sccprob, NewtonRaphson())
     @test SciMLBase.successful_retcode(sccsol)
-    @test val[] == 1
+    @test val[] == 2
 end
 
 import ModelingToolkitStandardLibrary.Blocks as B