Post State-Selection cleanup

Author: Keno

docs/src/tutorials/tearing_parallelism.md

@@ -211,8 +211,8 @@ investigate what this means:
 
 ```julia
 using ModelingToolkit.BipartiteGraphs
-big_rc = initialize_system_structure(big_rc)
-inc_org = BipartiteGraphs.incidence_matrix(structure(big_rc).graph)
+ts = TearingState(big_rc)
+inc_org = BipartiteGraphs.incidence_matrix(ts.graph)
 blt_org = StructuralTransformations.sorted_incidence_matrix(big_rc, only_algeqs=true, only_algvars=true)
 blt_reduced = StructuralTransformations.sorted_incidence_matrix(sys, only_algeqs=true, only_algvars=true)
 ```

src/ModelingToolkit.jl

@@ -193,7 +193,7 @@ export calculate_factorized_W, generate_factorized_W
 export calculate_hessian, generate_hessian
 export calculate_massmatrix, generate_diffusion_function
 export stochastic_integral_transform
-export initialize_system_structure
+export TearingState, StateSelectionState
 export generate_difference_cb
 
 export BipartiteGraph, equation_dependencies, variable_dependencies

src/structural_transformation/StructuralTransformations.jl

@@ -10,12 +10,12 @@ using SymbolicUtils.Rewriters
 using SymbolicUtils: similarterm, istree
 
 using ModelingToolkit
-using ModelingToolkit: ODESystem, AbstractSystem,var_from_nested_derivative, Differential,
+using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Differential,
                        states, equations, vars, Symbolic, diff2term, value,
                        operation, arguments, Sym, Term, simplify, solve_for,
                        isdiffeq, isdifferential, isinput,
-                       get_structure, get_iv, independent_variables,
-                       get_structure, defaults, InvalidSystemException,
+                       get_iv, independent_variables,
+                       defaults, InvalidSystemException,
                        ExtraEquationsSystemException,
                        ExtraVariablesSystemException,
                        get_postprocess_fbody, vars!,
@@ -25,6 +25,7 @@ using ModelingToolkit.BipartiteGraphs
 import .BipartiteGraphs: invview
 using Graphs
 using ModelingToolkit.SystemStructures
+using ModelingToolkit.SystemStructures: algeqs
 
 using ModelingToolkit.DiffEqBase
 using ModelingToolkit.StaticArrays
@@ -38,7 +39,7 @@ using NonlinearSolve
 
 export tearing, partial_state_selection, dae_index_lowering, check_consistency
 export build_torn_function, build_observed_function, ODAEProblem
-export sorted_incidence_matrix
+export sorted_incidence_matrix, pantelides!, tearing_reassemble
 
 include("utils.jl")
 include("pantelides.jl")

src/structural_transformation/codegen.jl

@@ -2,9 +2,9 @@ using LinearAlgebra
 
 const MAX_INLINE_NLSOLVE_SIZE = 8
 
-function torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs)
-    s = structure(sys)
-    @unpack fullvars, graph = s
+function torn_system_jacobian_sparsity(state, var_eq_matching, var_sccs)
+    fullvars = state.fullvars
+    graph = state.structure.graph
 
     # The sparsity pattern of `nlsolve(f, u, p)` w.r.t `p` is difficult to
     # determine in general. Consider the "simplest" case, a linear system. We
@@ -59,12 +59,12 @@ function torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs)
                 Graphs.insorted(var, v_residual) && continue
                 deps = get(avars2dvars, var, nothing)
                 if deps === nothing # differential variable
-                    @assert !isalgvar(s, var)
+                    @assert !isalgvar(state.structure, var)
                     for tvar in v_residual
                         push!(avars2dvars[tvar], var)
                     end
                 else # tearing variable from previous partitions
-                    @assert isalgvar(s, var)
+                    @assert isalgvar(state.structure, var)
                     for tvar in v_residual
                         union!(avars2dvars[tvar], avars2dvars[var])
                     end
@@ -73,18 +73,19 @@ function torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs)
         end
     end
 
-    dvrange = diffvars_range(s)
+    dvrange = diffvars_range(state.structure)
     dvar2idx = Dict(v=>i for (i, v) in enumerate(dvrange))
     I = Int[]; J = Int[]
     eqidx = 0
+    aeqs = algeqs(state.structure)
     for ieq in 𝑠vertices(graph)
-        isalgeq(s, ieq) && continue
+        ieq in aeqs && continue
         eqidx += 1
         for ivar in 𝑠neighbors(graph, ieq)
-            if isdiffvar(s, ivar)
+            if isdiffvar(state.structure, ivar)
                 push!(I, eqidx)
                 push!(J, dvar2idx[ivar])
-            elseif isalgvar(s, ivar)
+            elseif isalgvar(state.structure, ivar)
                 for dvar in avars2dvars[ivar]
                     push!(I, eqidx)
                     push!(J, dvar2idx[dvar])
@@ -170,18 +171,18 @@ function build_torn_function(
         isdiffeq(eq) && push!(rhss, eq.rhs)
     end
 
-    s = structure(sys)
-    @unpack fullvars = s
-    var_eq_matching, var_sccs = algebraic_variables_scc(sys)
+    state = TearingState(sys)
+    fullvars = state.fullvars
+    var_eq_matching, var_sccs = algebraic_variables_scc(state)
 
-    states = map(i->s.fullvars[i], diffvars_range(s))
+    states = map(i->fullvars[i], diffvars_range(state.structure))
     mass_matrix_diag = ones(length(states))
     torn_expr = []
     defs = defaults(sys)
 
     needs_extending = false
     for scc in var_sccs
-        torn_vars = [s.fullvars[var] for var in scc if var_eq_matching[var] !== unassigned]
+        torn_vars = [fullvars[var] for var in scc if var_eq_matching[var] !== unassigned]
         torn_eqs = [eqs[var_eq_matching[var]] for var in scc if var_eq_matching[var] !== unassigned]
         isempty(torn_eqs) && continue
         if length(torn_eqs) <= max_inlining_size
@@ -224,18 +225,18 @@ function build_torn_function(
     if expression
         expr, states
     else
-        observedfun = let sys = sys, dict = Dict()
+        observedfun = let state = state, dict = Dict()
             function generated_observed(obsvar, u, p, t)
                 obs = get!(dict, value(obsvar)) do
-                    build_observed_function(sys, obsvar, var_eq_matching, var_sccs, checkbounds=checkbounds)
+                    build_observed_function(state, obsvar, var_eq_matching, var_sccs, checkbounds=checkbounds)
                 end
                 obs(u, p, t)
             end
         end
 
         ODEFunction{true}(
                           @RuntimeGeneratedFunction(expr),
-                          sparsity = torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs),
+                          sparsity = torn_system_jacobian_sparsity(state, var_eq_matching, var_sccs),
                           syms = syms,
                           observed = observedfun,
                           mass_matrix = mass_matrix,
@@ -261,7 +262,7 @@ function find_solve_sequence(sccs, vars)
 end
 
 function build_observed_function(
-        sys, ts, var_eq_matching, var_sccs;
+        state, ts, var_eq_matching, var_sccs;
         expression=false,
         output_type=Array,
         checkbounds=true
@@ -273,12 +274,14 @@ function build_observed_function(
     ts = Symbolics.scalarize.(value.(ts))
 
     vars = Set()
+    sys = state.sys
     foreach(Base.Fix1(vars!, vars), ts)
     ivs = independent_variables(sys)
     dep_vars = collect(setdiff(vars, ivs))
 
-    s = structure(sys)
-    @unpack fullvars, graph = s
+    fullvars = state.fullvars
+    s = state.structure
+    graph = s.graph
     diffvars = map(i->fullvars[i], diffvars_range(s))
     algvars = map(i->fullvars[i], algvars_range(s))
 

src/structural_transformation/pantelides.jl

@@ -2,9 +2,10 @@
 ### Reassemble: structural information -> system
 ###
 
-function pantelides_reassemble(sys::ODESystem, eq_to_diff, assign)
-    s = structure(sys)
-    @unpack fullvars, var_to_diff = s
+function pantelides_reassemble(state::TearingState, var_eq_matching)
+    fullvars = state.fullvars
+    @unpack var_to_diff, eq_to_diff = state.structure
+    sys = state.sys
     # Step 1: write derivative equations
     in_eqs = equations(sys)
     out_eqs = Vector{Any}(undef, nv(eq_to_diff))
@@ -58,53 +59,29 @@ function pantelides_reassemble(sys::ODESystem, eq_to_diff, assign)
     end
 
     final_vars = unique(filter(x->!(operation(x) isa Differential), fullvars))
-    final_eqs = map(identity, filter(x->value(x.lhs) !== nothing, out_eqs[sort(filter(x->x !== unassigned, assign))]))
+    final_eqs = map(identity, filter(x->value(x.lhs) !== nothing, out_eqs[sort(filter(x->x !== unassigned, var_eq_matching))]))
 
     @set! sys.eqs = final_eqs
     @set! sys.states = final_vars
-    @set! sys.structure = nothing
     return sys
 end
 
 """
-    pantelides!(sys::ODESystem; kwargs...)
+    pantelides!(state::TransformationState; kwargs...)
 
 Perform Pantelides algorithm.
 """
-function pantelides!(sys::ODESystem; maxiters = 8000)
-    find_solvables!(sys)
-    s = structure(sys)
-    # D(j) = assoc[j]
-    @unpack graph, var_to_diff = s
-    # N.B.: var_derivative! and eq_derivative! are defined in symbolics_tearing.jl
-    return (sys, pantelides!(PantelidesSetup(sys, graph, var_to_diff))...)
-end
-
-struct PantelidesSetup{T}
-    system::T
-    graph::BipartiteGraph
-    var_to_diff::DiffGraph
-    eq_to_diff::DiffGraph
-    var_eq_matching::Matching
-end
-
-function PantelidesSetup(sys::T, graph, var_to_diff) where {T}
-    neqs = nsrcs(graph)
-    nvars = nv(var_to_diff)
-    var_eq_matching = Matching(nvars)
-    eq_to_diff = DiffGraph(neqs)
-    PantelidesSetup{T}(sys, graph, var_to_diff, eq_to_diff, var_eq_matching)
-end
-
-function pantelides!(p::PantelidesSetup; maxiters = 8000)
-    @unpack graph, var_to_diff, eq_to_diff, var_eq_matching = p
+function pantelides!(state::TransformationState; maxiters = 8000)
+    @unpack graph, var_to_diff, eq_to_diff = state.structure
     neqs = nsrcs(graph)
     nvars = nv(var_to_diff)
     vcolor = falses(nvars)
     ecolor = falses(neqs)
+    var_eq_matching = Matching(nvars)
     neqs′ = neqs
     for k in 1:neqs′
         eq′ = k
+        isempty(𝑠neighbors(graph, eq′)) && continue
         pathfound = false
         # In practice, `maxiters=8000` should never be reached, otherwise, the
         # index would be on the order of thousands.
@@ -128,7 +105,7 @@ function pantelides!(p::PantelidesSetup; maxiters = 8000)
 
                 add_edge!(var_to_diff, var, add_vertex!(var_to_diff))
                 push!(var_eq_matching, unassigned)
-                var_derivative!(p, eq)
+                var_derivative!(state, var)
             end
 
             for eq in eachindex(ecolor); ecolor[eq] || continue
@@ -138,7 +115,7 @@ function pantelides!(p::PantelidesSetup; maxiters = 8000)
                 # the new equation is created by differentiating `eq`
                 eq_diff = add_vertex!(eq_to_diff)
                 add_edge!(eq_to_diff, eq, eq_diff)
-                eq_derivative!(p, eq)
+                eq_derivative!(state, eq)
             end
 
             for var in eachindex(vcolor); vcolor[var] || continue
@@ -150,7 +127,7 @@ function pantelides!(p::PantelidesSetup; maxiters = 8000)
         end # for _ in 1:maxiters
         pathfound || error("maxiters=$maxiters reached! File a bug report if your system has a reasonable index (<100), and you are using the default `maxiters`. Try to increase the maxiters by `pantelides(sys::ODESystem; maxiters=1_000_000)` if your system has an incredibly high index and it is truly extremely large.")
     end # for k in 1:neqs′
-    return var_eq_matching, eq_to_diff
+    return var_eq_matching
 end
 
 """
@@ -161,8 +138,8 @@ DAE. `kwargs` are forwarded to [`pantelides!`](@ref). End users are encouraged t
 instead, which calls this function internally.
 """
 function dae_index_lowering(sys::ODESystem; kwargs...)
-    s = get_structure(sys)
-    (s isa SystemStructure) || (sys = initialize_system_structure(sys))
-    sys, var_eq_matching, eq_to_diff = pantelides!(sys; kwargs...)
-    return pantelides_reassemble(sys, eq_to_diff, var_eq_matching)
+    state = TearingState(sys)
+    find_solvables!(state)
+    var_eq_matching = pantelides!(state; kwargs...)
+    return pantelides_reassemble(state, var_eq_matching)
 end

src/structural_transformation/partial_state_selection.jl

@@ -1,12 +1,8 @@
-function partial_state_selection_graph!(sys::ODESystem)
-    s = get_structure(sys)
-    (s isa SystemStructure) || (sys = initialize_system_structure(sys))
-    s = structure(sys)
-    find_solvables!(sys; allow_symbolic=true)
-    @set! s.graph = complete(s.graph)
-    @set! sys.structure = s
-    var_eq_matching, eq_to_diff = pantelides!(PantelidesSetup(sys, s.graph, s.var_to_diff))
-    (sys, partial_state_selection_graph!(s.graph, s.solvable_graph, s.var_to_diff, var_eq_matching, eq_to_diff)...)
+function partial_state_selection_graph!(state::TransformationState)
+    find_solvables!(state; allow_symbolic=true)
+    var_eq_matching = complete(pantelides!(state))
+    complete!(state.structure)
+    partial_state_selection_graph!(state.structure, var_eq_matching)
 end
 
 function ascend_dg(xs, dg, level)
@@ -31,7 +27,9 @@ function ascend_dg_all(xs, dg, level, maxlevel)
     return r
 end
 
-function pss_graph_modia!(graph, solvable_graph, var_eq_matching, var_to_diff, eq_to_diff, varlevel, inv_varlevel, inv_eqlevel)
+function pss_graph_modia!(structure::SystemStructure, var_eq_matching, varlevel, inv_varlevel, inv_eqlevel)
+    @unpack eq_to_diff, var_to_diff, graph, solvable_graph = structure
+
     # var_eq_matching is a maximal matching on the top-differentiated variables.
     # Find Strongly connected components. Note that after pantelides, we expect
     # a balanced system, so a maximal matching should be possible.
@@ -44,7 +42,8 @@ function pss_graph_modia!(graph, solvable_graph, var_eq_matching, var_to_diff, e
         end
 
         # Now proceed level by level from lowest to highest and tear the graph.
-        eqs = [var_eq_matching[var] for var in vars]
+        eqs = [var_eq_matching[var] for var in vars if var_eq_matching[var] !== unassigned]
+        isempty(eqs) && continue
         maxlevel = level = maximum(map(x->inv_eqlevel[x], eqs))
         old_level_vars = ()
         ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph, complete(Matching(ndsts(graph)))); dir=:in)
@@ -97,7 +96,8 @@ function pss_graph_modia!(graph, solvable_graph, var_eq_matching, var_to_diff, e
 end
 
 struct SelectedState; end
-function partial_state_selection_graph!(graph, solvable_graph, var_to_diff, var_eq_matching, eq_to_diff)
+function partial_state_selection_graph!(structure::SystemStructure, var_eq_matching)
+    @unpack eq_to_diff, var_to_diff, graph, solvable_graph = structure
     eq_to_diff = complete(eq_to_diff)
 
     inv_eqlevel = map(1:nsrcs(graph)) do eq
@@ -134,9 +134,8 @@ function partial_state_selection_graph!(graph, solvable_graph, var_to_diff, var_
         end
     end
 
-    var_eq_matching = pss_graph_modia!(graph, solvable_graph,
-        complete(var_eq_matching), var_to_diff, eq_to_diff, varlevel, inv_varlevel,
-        inv_eqlevel)
+    var_eq_matching = pss_graph_modia!(structure,
+        complete(var_eq_matching), varlevel, inv_varlevel, inv_eqlevel)
 
-    var_eq_matching, eq_to_diff
+    var_eq_matching
 end