feat: initial implementation of SCCNonlinearProblem codegen

Author: AayushSabharwal

Project.toml

@@ -123,7 +123,7 @@ REPL = "1"
 RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
-SciMLBase = "2.57.1"
+SciMLBase = "2.61"
 SciMLStructures = "1.0"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"

src/structural_transformation/bipartite_tearing/modia_tearing.jl

@@ -62,6 +62,15 @@ function tear_graph_block_modia!(var_eq_matching, ict, solvable_graph, eqs, vars
     return nothing
 end
 
+function build_var_eq_matching(structure::SystemStructure, ::Type{U} = Unassigned;
+        varfilter::F2 = v -> true, eqfilter::F3 = eq -> true) where {U, F2, F3}
+    @unpack graph, solvable_graph = structure
+    var_eq_matching = maximal_matching(graph, eqfilter, varfilter, U)
+    matching_len = max(length(var_eq_matching),
+        maximum(x -> x isa Int ? x : 0, var_eq_matching, init = 0))
+    return complete(var_eq_matching, matching_len), matching_len
+end
+
 function tear_graph_modia(structure::SystemStructure, isder::F = nothing,
         ::Type{U} = Unassigned;
         varfilter::F2 = v -> true,
@@ -78,10 +87,7 @@ function tear_graph_modia(structure::SystemStructure, isder::F = nothing,
     # find them here [TODO: It would be good to have an explicit example of this.]
 
     @unpack graph, solvable_graph = structure
-    var_eq_matching = maximal_matching(graph, eqfilter, varfilter, U)
-    matching_len = max(length(var_eq_matching),
-        maximum(x -> x isa Int ? x : 0, var_eq_matching, init = 0))
-    var_eq_matching = complete(var_eq_matching, matching_len)
+    var_eq_matching, matching_len = build_var_eq_matching(structure, U; varfilter, eqfilter)
     full_var_eq_matching = copy(var_eq_matching)
     var_sccs = find_var_sccs(graph, var_eq_matching)
     vargraph = DiCMOBiGraph{true}(graph, 0, Matching(matching_len))

src/systems/abstractsystem.jl

@@ -162,11 +162,12 @@ object.
 """
 function generate_custom_function(sys::AbstractSystem, exprs, dvs = unknowns(sys),
         ps = parameters(sys); wrap_code = nothing, postprocess_fbody = nothing, states = nothing,
-        expression = Val{true}, eval_expression = false, eval_module = @__MODULE__, kwargs...)
+        expression = Val{true}, eval_expression = false, eval_module = @__MODULE__,
+        cachesyms::Tuple = (), kwargs...)
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system.")
     end
-    p = reorder_parameters(sys, unwrap.(ps))
+    p = (reorder_parameters(sys, unwrap.(ps))..., cachesyms...)
     isscalar = !(exprs isa AbstractArray)
     if wrap_code === nothing
         wrap_code = isscalar ? identity : (identity, identity)
@@ -187,7 +188,7 @@ function generate_custom_function(sys::AbstractSystem, exprs, dvs = unknowns(sys
             postprocess_fbody,
             states,
             wrap_code = wrap_code .∘ wrap_mtkparameters(sys, isscalar) .∘
-                        wrap_array_vars(sys, exprs; dvs) .∘
+                        wrap_array_vars(sys, exprs; dvs, cachesyms) .∘
                         wrap_parameter_dependencies(sys, isscalar),
             expression = Val{true}
         )
@@ -199,7 +200,7 @@ function generate_custom_function(sys::AbstractSystem, exprs, dvs = unknowns(sys
             postprocess_fbody,
             states,
             wrap_code = wrap_code .∘ wrap_mtkparameters(sys, isscalar) .∘
-                        wrap_array_vars(sys, exprs; dvs) .∘
+                        wrap_array_vars(sys, exprs; dvs, cachesyms) .∘
                         wrap_parameter_dependencies(sys, isscalar),
             expression = Val{true}
         )
@@ -231,116 +232,51 @@ end
 
 function wrap_array_vars(
         sys::AbstractSystem, exprs; dvs = unknowns(sys), ps = parameters(sys),
-        inputs = nothing, history = false)
+        inputs = nothing, history = false, cachesyms::Tuple = ())
     isscalar = !(exprs isa AbstractArray)
-    array_vars = Dict{Any, AbstractArray{Int}}()
-    if dvs !== nothing
-        for (j, x) in enumerate(dvs)
-            if iscall(x) && operation(x) == getindex
-                arg = arguments(x)[1]
-                inds = get!(() -> Int[], array_vars, arg)
-                push!(inds, j)
-            end
-        end
-        for (k, inds) in array_vars
-            if inds == (inds′ = inds[1]:inds[end])
-                array_vars[k] = inds′
-            end
-        end
+    var_to_arridxs = Dict()
 
-        uind = 1
-    else
+    if dvs === nothing
         uind = 0
-    end
-    # values are (indexes, index of buffer, size of parameter)
-    array_parameters = Dict{Any, Tuple{AbstractArray{Int}, Int, Tuple{Vararg{Int}}}}()
-    # If for some reason different elements of an array parameter are in different buffers
-    other_array_parameters = Dict{Any, Any}()
-
-    hasinputs = inputs !== nothing
-    input_vars = Dict{Any, AbstractArray{Int}}()
-    if hasinputs
-        for (j, x) in enumerate(inputs)
-            if iscall(x) && operation(x) == getindex
-                arg = arguments(x)[1]
-                inds = get!(() -> Int[], input_vars, arg)
-                push!(inds, j)
-            end
-        end
-        for (k, inds) in input_vars
-            if inds == (inds′ = inds[1]:inds[end])
-                input_vars[k] = inds′
-            end
-        end
-    end
-    if has_index_cache(sys)
-        ic = get_index_cache(sys)
     else
-        ic = nothing
-    end
-    if ps isa Tuple && eltype(ps) <: AbstractArray
-        ps = Iterators.flatten(ps)
-    end
-    for p in ps
-        p = unwrap(p)
-        if iscall(p) && operation(p) == getindex
-            p = arguments(p)[1]
-        end
-        symtype(p) <: AbstractArray && Symbolics.shape(p) != Symbolics.Unknown() || continue
-        scal = collect(p)
-        # all scalarized variables are in `ps`
-        any(isequal(p), ps) || all(x -> any(isequal(x), ps), scal) || continue
-        (haskey(array_parameters, p) || haskey(other_array_parameters, p)) && continue
-
-        idx = parameter_index(sys, p)
-        idx isa Int && continue
-        if idx isa ParameterIndex
-            if idx.portion != SciMLStructures.Tunable()
+        uind = 1
+        for (i, x) in enumerate(dvs)
+            iscall(x) && operation(x) == getindex || continue
+            arg = arguments(x)[1]
+            inds = get!(() -> [], var_to_arridxs, arg)
+            push!(inds, (uind, i))
+        end
+    end
+    p_start = uind + 1 + (inputs !== nothing) + history
+    input_ind = inputs === nothing ? -1 : (p_start - 1)
+    rps = (reorder_parameters(sys, ps)..., cachesyms...)
+    for sym in reduce(vcat, rps; init = [])
+        iscall(sym) && operation(sym) == getindex || continue
+        arg = arguments(sym)[1]
+        if inputs !== nothing
+            idx = findfirst(isequal(sym), inputs)
+            if idx !== nothing
+                inds = get!(() -> [], var_to_arridxs, arg)
+                push!(inds, (input_ind, idx))
                 continue
             end
-            array_parameters[p] = (vec(idx.idx), 1, size(idx.idx))
-        else
-            # idx === nothing
-            idxs = map(Base.Fix1(parameter_index, sys), scal)
-            if first(idxs) isa ParameterIndex
-                buffer_idxs = map(Base.Fix1(iterated_buffer_index, ic), idxs)
-                if allequal(buffer_idxs)
-                    buffer_idx = first(buffer_idxs)
-                    if first(idxs).portion == SciMLStructures.Tunable()
-                        idxs = map(x -> x.idx, idxs)
-                    else
-                        idxs = map(x -> x.idx[end], idxs)
-                    end
-                else
-                    other_array_parameters[p] = scal
-                    continue
-                end
-            else
-                buffer_idx = 1
-            end
-
-            sz = size(idxs)
-            if vec(idxs) == idxs[begin]:idxs[end]
-                idxs = idxs[begin]:idxs[end]
-            elseif vec(idxs) == idxs[begin]:-1:idxs[end]
-                idxs = idxs[begin]:-1:idxs[end]
-            end
-            idxs = vec(idxs)
-            array_parameters[p] = (idxs, buffer_idx, sz)
         end
+        bufferidx = findfirst(buf -> any(isequal(sym), buf), rps)
+        idxinbuffer = findfirst(isequal(sym), rps[bufferidx])
+        inds = get!(() -> [], var_to_arridxs, arg)
+        push!(inds, (p_start + bufferidx - 1, idxinbuffer))
     end
 
-    inputind = if history
-        uind + 2
-    else
-        uind + 1
-    end
-    params_offset = if history && hasinputs
-        uind + 2
-    elseif history || hasinputs
-        uind + 1
-    else
-        uind
+    viewsyms = Dict()
+    splitsyms = Dict()
+    for (arrsym, idxs) in var_to_arridxs
+        length(idxs) == length(arrsym) || continue
+        # allequal(first, idxs) is a 1.11 feature
+        if allequal(Iterators.map(first, idxs))
+            viewsyms[arrsym] = (first(first(idxs)), reshape(last.(idxs), size(arrsym)))
+        else
+            splitsyms[arrsym] = reshape(idxs, size(arrsym))
+        end
     end
     if isscalar
         function (expr)
@@ -349,15 +285,11 @@ function wrap_array_vars(
                 [],
                 Let(
                     vcat(
-                        [k ← :(view($(expr.args[uind].name), $v)) for (k, v) in array_vars],
-                        [k ← :(view($(expr.args[inputind].name), $v))
-                         for (k, v) in input_vars],
-                        [k ← :(reshape(
-                             view($(expr.args[params_offset + buffer_idx].name), $idxs),
-                             $sz))
-                         for (k, (idxs, buffer_idx, sz)) in array_parameters],
-                        [k ← Code.MakeArray(v, symtype(k))
-                         for (k, v) in other_array_parameters]
+                        [sym ← :(view($(expr.args[i].name), $idxs))
+                         for (sym, (i, idxs)) in viewsyms],
+                        [sym ←
+                         MakeArray([expr.args[bufi].elems[vali] for (bufi, vali) in idxs],
+                             expr.args[idxs[1][1]]) for (sym, idxs) in splitsyms]
                     ),
                     expr.body,
                     false
@@ -371,15 +303,11 @@ function wrap_array_vars(
                 [],
                 Let(
                     vcat(
-                        [k ← :(view($(expr.args[uind].name), $v)) for (k, v) in array_vars],
-                        [k ← :(view($(expr.args[inputind].name), $v))
-                         for (k, v) in input_vars],
-                        [k ← :(reshape(
-                             view($(expr.args[params_offset + buffer_idx].name), $idxs),
-                             $sz))
-                         for (k, (idxs, buffer_idx, sz)) in array_parameters],
-                        [k ← Code.MakeArray(v, symtype(k))
-                         for (k, v) in other_array_parameters]
+                        [sym ← :(view($(expr.args[i].name), $idxs))
+                         for (sym, (i, idxs)) in viewsyms],
+                        [sym ←
+                         MakeArray([expr.args[bufi].elems[vali] for (bufi, vali) in idxs],
+                             expr.args[idxs[1][1]]) for (sym, idxs) in splitsyms]
                     ),
                     expr.body,
                     false
@@ -392,17 +320,11 @@ function wrap_array_vars(
                 [],
                 Let(
                     vcat(
-                        [k ← :(view($(expr.args[uind + 1].name), $v))
-                         for (k, v) in array_vars],
-                        [k ← :(view($(expr.args[inputind + 1].name), $v))
-                         for (k, v) in input_vars],
-                        [k ← :(reshape(
-                             view($(expr.args[params_offset + buffer_idx + 1].name),
-                                 $idxs),
-                             $sz))
-                         for (k, (idxs, buffer_idx, sz)) in array_parameters],
-                        [k ← Code.MakeArray(v, symtype(k))
-                         for (k, v) in other_array_parameters]
+                        [sym ← :(view($(expr.args[i + 1].name), $idxs))
+                         for (sym, (i, idxs)) in viewsyms],
+                        [sym ← MakeArray(
+                             [expr.args[bufi + 1].elems[vali] for (bufi, vali) in idxs],
+                             expr.args[idxs[1][1] + 1]) for (sym, idxs) in splitsyms]
                     ),
                     expr.body,
                     false