Merge pull request #1384 from Keno/kf/state_selection

Partial State Selection

Author: YingboMa

docs/src/tutorials/tearing_parallelism.md

@@ -178,8 +178,8 @@ investigate what this means:
 
 ```julia
 using ModelingToolkit.BipartiteGraphs
-big_rc = initialize_system_structure(expand_connections(big_rc))
-inc_org = BipartiteGraphs.incidence_matrix(structure(big_rc).graph)
+ts = TearingState(expand_connections(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)
 ```
@@ -190,7 +190,7 @@ The figure on the left is the original incidence matrix of the algebraic equatio
 Notice that the original formulation of the model has dependencies between different
 equations, and so the full set of equations must be solved together. That exposes
 no parallelism. However, the Block Lower Triangular (BLT) transformation exposes
-independent blocks. This is then further impoved by the tearing process, which
+independent blocks. This is then further improved by the tearing process, which
 removes 90% of the equations and transforms the nonlinear equations into 50
 independent blocks *which can now all be solved in parallel*. The conclusion
 is that, your attempts to parallelize are neigh: performing parallelism after

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/bipartite_graph.jl

@@ -306,6 +306,25 @@ function Graphs.add_edge!(g::BipartiteGraph, edge::BipartiteEdge, md=NO_METADATA
     return true  # edge successfully added
 end
 
+Graphs.rem_edge!(g::BipartiteGraph, i::Integer, j::Integer) =
+    Graphs.rem_edge!(g, BipartiteEdge(i, j))
+function Graphs.rem_edge!(g::BipartiteGraph, edge::BipartiteEdge)
+    @unpack fadjlist, badjlist = g
+    s, d = src(edge), dst(edge)
+    (has_𝑠vertex(g, s) && has_𝑑vertex(g, d)) || error("edge ($edge) out of range.")
+    @inbounds list = fadjlist[s]
+    index = searchsortedfirst(list, d)
+    @inbounds (index <= length(list) && list[index] == d) || error("graph does not have edge $edge")
+    deleteat!(list, index)
+    g.ne -= 1
+    if badjlist isa AbstractVector
+        @inbounds list = badjlist[d]
+        index = searchsortedfirst(list, s)
+        deleteat!(list, index)
+    end
+    return true  # edge successfully deleted
+end
+
 function Graphs.add_vertex!(g::BipartiteGraph{T}, type::VertType) where T
     if type === DST
         if g.badjlist isa AbstractVector
@@ -322,10 +341,23 @@ function Graphs.add_vertex!(g::BipartiteGraph{T}, type::VertType) where T
 end
 
 function set_neighbors!(g::BipartiteGraph, i::Integer, new_neighbors::AbstractVector)
-    old_nneighbors = length(g.fadjlist[i])
+    old_neighbors = g.fadjlist[i]
+    old_nneighbors = length(old_neighbors)
     new_nneighbors = length(new_neighbors)
     g.fadjlist[i] = new_neighbors
     g.ne += new_nneighbors - old_nneighbors
+    if isa(g.badjlist, AbstractVector)
+        for n in old_neighbors
+            @inbounds list = g.badjlist[n]
+            index = searchsortedfirst(list, i)
+            deleteat!(list, index)
+        end
+        for n in new_neighbors
+            @inbounds list = g.badjlist[n]
+            index = searchsortedfirst(list, i)
+            insert!(list, index, i)
+        end
+    end
 end
 
 ###

src/structural_transformation/StructuralTransformations.jl

@@ -10,23 +10,24 @@ 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,
                        empty_substitutions, get_substitutions,
-                       get_structure, get_iv, independent_variables,
-                       has_structure, defaults, InvalidSystemException,
+                       get_tearing_state, get_iv, independent_variables,
+                       has_tearing_state, defaults, InvalidSystemException,
                        ExtraEquationsSystemException,
                        ExtraVariablesSystemException,
                        get_postprocess_fbody, vars!,
                        IncrementalCycleTracker, add_edge_checked!, topological_sort,
-                       invalidate_cache!, Substitutions
+                       invalidate_cache!, Substitutions, get_or_construct_tearing_state
 
 using ModelingToolkit.BipartiteGraphs
 import .BipartiteGraphs: invview
 using Graphs
 using ModelingToolkit.SystemStructures
+using ModelingToolkit.SystemStructures: algeqs
 
 using ModelingToolkit.DiffEqBase
 using ModelingToolkit.StaticArrays
@@ -38,10 +39,10 @@ using SparseArrays
 
 using NonlinearSolve
 
-export tearing, dae_index_lowering, check_consistency
-export tearing_assignments, tearing_substitution
+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, find_solvables!
+export tearing_assignments, tearing_substitution
 export torn_system_jacobian_sparsity
 export full_equations
 
@@ -50,6 +51,7 @@ include("pantelides.jl")
 include("bipartite_tearing/modia_tearing.jl")
 include("tearing.jl")
 include("symbolics_tearing.jl")
+include("partial_state_selection.jl")
 include("codegen.jl")
 
 end # module

src/structural_transformation/bipartite_tearing/modia_tearing.jl

@@ -13,6 +13,14 @@
 #
 ################################################
 
+function try_assign_eq!(ict::IncrementalCycleTracker, vj::Integer, eq::Integer)
+    G = ict.graph
+    add_edge_checked!(ict, Iterators.filter(!=(vj), 𝑠neighbors(G.graph, eq)), vj) do G
+        G.matching[vj] = eq
+        G.ne += length(𝑠neighbors(G.graph, eq)) - 1
+    end
+end
+
 """
     (eSolved, vSolved, eResidue, vTear) = tearEquations!(td, Gsolvable, es, vs)
 
@@ -33,10 +41,7 @@ function tearEquations!(ict::IncrementalCycleTracker, Gsolvable, es::Vector{Int}
     for eq in es  # iterate only over equations that are not in eSolvedFixed
         for vj in Gsolvable[eq]
             if G.matching[vj] === unassigned && (vj in vActive)
-                r = add_edge_checked!(ict, Iterators.filter(!=(vj), 𝑠neighbors(G.graph, eq)), vj) do G
-                    G.matching[vj] = eq
-                    G.ne += length(𝑠neighbors(G.graph, eq)) - 1
-                end
+                r = try_assign_eq!(ict, vj, eq)
                 r && break
             end
         end
@@ -45,6 +50,15 @@ function tearEquations!(ict::IncrementalCycleTracker, Gsolvable, es::Vector{Int}
     return ict
 end
 
+function tear_graph_block_modia!(var_eq_matching, graph, solvable_graph, eqs, vars)
+    ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph); dir=:in)
+    tearEquations!(ict, solvable_graph.fadjlist, eqs, vars)
+    for var in vars
+        var_eq_matching[var] = ict.graph.matching[var]
+    end
+    return nothing
+end
+
 """
     tear_graph_modia(sys) -> sys
 
@@ -58,13 +72,10 @@ function tear_graph_modia(graph::BipartiteGraph, solvable_graph::BipartiteGraph;
     for vars in var_sccs
         filtered_vars = filter(varfilter, vars)
         ieqs = Int[var_eq_matching[v] for v in filtered_vars if var_eq_matching[v] !== unassigned]
-
-        ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph); dir=:in)
-        tearEquations!(ict, solvable_graph.fadjlist, ieqs, filtered_vars)
-
         for var in vars
-            var_eq_matching[var] = ict.graph.matching[var]
+            var_eq_matching[var] = unassigned
         end
+        tear_graph_block_modia!(var_eq_matching, graph, solvable_graph, ieqs, filtered_vars)
     end
 
     return var_eq_matching

src/structural_transformation/codegen.jl

@@ -4,9 +4,9 @@ using ModelingToolkit: isdifferenceeq, has_continuous_events, generate_rootfindi
 
 const MAX_INLINE_NLSOLVE_SIZE = 8
 
-function torn_system_with_nlsolve_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs)
-    s = structure(sys)
-    @unpack fullvars, graph = s
+function torn_system_with_nlsolve_jacobian_sparsity(state, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs)
+    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
@@ -73,6 +73,7 @@ function torn_system_with_nlsolve_jacobian_sparsity(sys, var_eq_matching, var_sc
     nlsolve_vars_set = BitSet(nlsolve_vars)
 
     I = Int[]; J = Int[]
+    s = state.structure
     for ieq in 𝑠vertices(graph)
         nieq = get(eqs2idx, ieq, 0)
         nieq == 0 && continue
@@ -244,15 +245,15 @@ function build_torn_function(
         push!(rhss, eq.rhs)
     end
 
-    s = structure(sys)
-    @unpack fullvars = s
-    var_eq_matching, var_sccs = algebraic_variables_scc(sys)
+    state = get_or_construct_tearing_state(sys)
+    fullvars = state.fullvars
+    var_eq_matching, var_sccs = algebraic_variables_scc(state)
     condensed_graph = MatchedCondensationGraph(
-        DiCMOBiGraph{true}(complete(s.graph), complete(var_eq_matching)), var_sccs)
+        DiCMOBiGraph{true}(complete(state.structure.graph), complete(var_eq_matching)), var_sccs)
     toporder = topological_sort_by_dfs(condensed_graph)
     var_sccs = var_sccs[toporder]
 
-    states_idxs = collect(diffvars_range(s))
+    states_idxs = collect(diffvars_range(state.structure))
     mass_matrix_diag = ones(length(states_idxs))
 
     assignments, deps, sol_states = tearing_assignments(sys)
@@ -276,7 +277,7 @@ function build_torn_function(
         torn_eqs_idxs = [var_eq_matching[var] for var in torn_vars_idxs]
         isempty(torn_eqs_idxs) && continue
         if length(torn_eqs_idxs) <= max_inlining_size
-            nlsolve_expr = gen_nlsolve!(is_not_prepended_assignment, eqs[torn_eqs_idxs], s.fullvars[torn_vars_idxs], defs, assignments, (deps, invdeps), var2assignment, checkbounds=checkbounds)
+            nlsolve_expr = gen_nlsolve!(is_not_prepended_assignment, eqs[torn_eqs_idxs], fullvars[torn_vars_idxs], defs, assignments, (deps, invdeps), var2assignment, checkbounds=checkbounds)
             append!(torn_expr, nlsolve_expr)
             push!(nlsolve_scc_idxs, i)
         else
@@ -297,7 +298,7 @@ function build_torn_function(
         rhss
     )
 
-    states = s.fullvars[states_idxs]
+    states = fullvars[states_idxs]
     syms = map(Symbol, states_idxs)
 
     pre = get_postprocess_fbody(sys)
@@ -322,10 +323,10 @@ function build_torn_function(
     if expression
         expr, states
     else
-        observedfun = let sys=sys, dict=Dict(), assignments=assignments, deps=(deps, invdeps), sol_states=sol_states, var2assignment=var2assignment
+        observedfun = let state = state, dict=Dict(), assignments=assignments, deps=(deps, invdeps), sol_states=sol_states, var2assignment=var2assignment
             function generated_observed(obsvar, u, p, t)
                 obs = get!(dict, value(obsvar)) do
-                    build_observed_function(sys, obsvar, var_eq_matching, var_sccs,
+                    build_observed_function(state, obsvar, var_eq_matching, var_sccs,
                                             assignments, deps, sol_states, var2assignment,
                                             checkbounds=checkbounds,
                                            )
@@ -336,7 +337,7 @@ function build_torn_function(
 
         ODEFunction{true}(
                           @RuntimeGeneratedFunction(expr),
-                          sparsity = jacobian_sparsity ? torn_system_with_nlsolve_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs) : nothing,
+                          sparsity = jacobian_sparsity ? torn_system_with_nlsolve_jacobian_sparsity(state, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs) : nothing,
                           syms = syms,
                           observed = observedfun,
                           mass_matrix = mass_matrix,
@@ -362,7 +363,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,
         assignments,
         deps,
         sol_states,
@@ -379,12 +380,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))
 
@@ -416,8 +419,13 @@ function build_observed_function(
     if !isempty(subset)
         eqs = equations(sys)
 
-        torn_eqs  = map(i->map(v->eqs[var_eq_matching[v]], var_sccs[i]), subset)
-        torn_vars = map(i->map(v->fullvars[v], var_sccs[i]), subset)
+        nested_torn_vars_idxs = []
+        for iscc in subset
+            torn_vars_idxs = Int[var for var in var_sccs[iscc] if var_eq_matching[var] !== unassigned]
+            isempty(torn_vars_idxs) || push!(nested_torn_vars_idxs, torn_vars_idxs)
+        end
+        torn_eqs = [[eqs[var_eq_matching[i]] for i in idxs] for idxs in nested_torn_vars_idxs]
+        torn_vars = [fullvars[idxs] for idxs in nested_torn_vars_idxs]
         u0map = defaults(sys)
         assignments = copy(assignments)
         solves = map(zip(torn_eqs, torn_vars)) do (eqs, vars)