Give partition names

Author: YingboMa

examples/electrical_components.jl

@@ -2,7 +2,8 @@ using Test
 using ModelingToolkit, OrdinaryDiffEq
 
 # Basic electric components
-const t = Sym{ModelingToolkit.Parameter{Real}}(:t)
+#const t = Sym{ModelingToolkit.Parameter{Real}}(:t)
+@parameters t
 function Pin(;name)
     @variables v(t) i(t)
     ODESystem(Equation[], t, [v, i], [], name=name, defaults=[v=>1.0, i=>1.0])

src/structural_transformation/codegen.jl

@@ -40,26 +40,27 @@ function torn_system_jacobian_sparsity(sys)
     # dependencies.
     avars2dvars = Dict{Int,Set{Int}}()
     c = 0
-    for (_, _, teqs, tvars) in partitions
+    for partition in partitions
+        @unpack e_residual, v_residual = partition
         # initialization
-        for tvar in tvars
+        for tvar in v_residual
             avars2dvars[tvar] = Set{Int}()
         end
-        for teq in teqs
+        for teq in e_residual
             c += 1
             for var in 𝑠neighbors(graph, teq)
                 # Skip the tearing variables in the current partition, because
                 # we are computing them from all the other states.
-                LightGraphs.insorted(var, tvars) && continue
+                LightGraphs.insorted(var, v_residual) && continue
                 deps = get(avars2dvars, var, nothing)
                 if deps === nothing # differential variable
                     @assert !isalgvar(s, var)
-                    for tvar in tvars
+                    for tvar in v_residual
                         push!(avars2dvars[tvar], var)
                     end
                 else # tearing variable from previous partitions
                     @assert isalgvar(s, var)
-                    for tvar in tvars
+                    for tvar in v_residual
                         union!(avars2dvars[tvar], avars2dvars[var])
                     end
                 end
@@ -97,22 +98,22 @@ function partitions_dag(s::SystemStructure)
     @unpack partitions, graph = s
 
     # `partvars[i]` contains all the states that appear in `partitions[i]`
-    partvars = map(partitions) do (_, _, reqs, tvars)
+    partvars = map(partitions) do partition
         ipartvars = Set{Int}()
-        for req in reqs
+        for req in partition.e_residual
             union!(ipartvars, 𝑠neighbors(graph, req))
         end
         ipartvars
     end
 
     I, J = Int[], Int[]
     n = length(partitions)
-    for i in 1:n
+    for (i, partition) in enumerate(partitions)
         for j in i+1:n
             # The only way for a later partition `j` to depend on an earlier
             # partition `i` is when `partvars[j]` contains one of tearing
             # variables of partition `i`.
-            if !isdisjoint(partvars[j], partitions[i][4])
+            if !isdisjoint(partvars[j], partition.v_residual)
                 # j depends on i
                 push!(I, i)
                 push!(J, j)
@@ -170,8 +171,8 @@ function get_torn_eqs_vars(sys)
     vars = s.fullvars
     eqs = equations(sys)
 
-    torn_eqs  = map(idxs-> eqs[idxs], map(x->x[3], partitions))
-    torn_vars = map(idxs->vars[idxs], map(x->x[4], partitions))
+    torn_eqs  = map(idxs-> eqs[idxs], map(x->x.e_residual, partitions))
+    torn_vars = map(idxs->vars[idxs], map(x->x.v_residual, partitions))
 
     gen_nlsolve.((sys,), torn_eqs, torn_vars)
 end
@@ -243,9 +244,9 @@ given a set of `vars`, find the groups of equations we need to solve for
 to obtain the solution to `vars`
 """
 function find_solve_sequence(partitions, vars)
-    subset = filter(x -> !isdisjoint(x[4], vars), partitions)
+    subset = filter(x -> !isdisjoint(x.v_residual, vars), partitions)
     isempty(subset) && return []
-    vars′ = mapreduce(x->x[4], union, subset)
+    vars′ = mapreduce(x->x.v_residual, union, subset)
     if vars′ == vars
         return subset
     else
@@ -289,8 +290,8 @@ function build_observed_function(
     if !isempty(subset)
         eqs = equations(sys)
 
-        torn_eqs  = map(idxs-> eqs[idxs[3]], subset)
-        torn_vars = map(idxs->fullvars[idxs[4]], subset)
+        torn_eqs  = map(idxs-> eqs[idxs.e_residual], subset)
+        torn_vars = map(idxs->fullvars[idxs.v_residual], subset)
 
         solves = gen_nlsolve.((sys,), torn_eqs, torn_vars)
     else

src/structural_transformation/tearing.jl

@@ -8,15 +8,13 @@ function tear_graph(sys)
     s = structure(sys)
     @unpack graph, solvable_graph, assign, inv_assign, scc = s
 
-    partitions = map(scc) do c
+    @set! sys.structure.partitions = map(scc) do c
         ieqs = filter(eq->isalgeq(s, eq), c)
         vars = inv_assign[ieqs]
 
         td = TraverseDAG(graph.fadjlist, length(assign))
-        e_solved, v_solved, e_residue, v_tear = tearEquations!(td, solvable_graph.fadjlist, ieqs, vars)
+        SystemPartition(tearEquations!(td, solvable_graph.fadjlist, ieqs, vars)...)
     end
-
-    @set! sys.structure.partitions = partitions
     return sys
 end
 
@@ -37,7 +35,8 @@ function tearing_reassemble(sys; simplify=false)
     active_eqs  = trues(ns)
     active_vars = trues(nd)
     rvar2req = Vector{Int}(undef, nd)
-    for (ith_scc, (e_solved, v_solved, e_residue, v_tear)) in enumerate(partitions)
+    for (ith_scc, partition) in enumerate(partitions)
+        @unpack e_solved, v_solved, e_residual, v_residual = partition
         for ii in eachindex(e_solved)
             ieq = e_solved[ii]; ns -= 1
             iv = v_solved[ii]; nd -= 1
@@ -161,18 +160,19 @@ function tearing_reassemble(sys; simplify=false)
     ### update partitions
     newpartitions = similar(partitions, 0)
     emptyintvec = Int[]
-    for ii in eachindex(partitions)
-        _, _, og_e_residue, og_v_tear = partitions[ii]
-        isempty(og_v_tear) && continue
-        e_residue = similar(og_e_residue)
-        v_tear = similar(og_v_tear)
-        for ii in eachindex(og_e_residue)
-            e_residue[ii] = eq_reidx[og_e_residue[ii]]
-            v_tear[ii] = var_reidx[og_v_tear[ii]]
+    for (ii, partition) in enumerate(partitions)
+        @unpack e_residual, v_residual = partition
+        isempty(v_residual) && continue
+        new_e_residual = similar(e_residual)
+        new_v_residual = similar(v_residual)
+        for ii in eachindex(e_residual)
+            new_e_residual[ii] = eq_reidx[ e_residual[ii]]
+            new_v_residual[ii] = var_reidx[v_residual[ii]]
         end
         # `emptyintvec` is aliased to save memory
         # We need them for type stability
-        push!(newpartitions, (emptyintvec, emptyintvec, e_residue, v_tear))
+        newpart = SystemPartition(emptyintvec, emptyintvec, new_e_residual, new_v_residual)
+        push!(newpartitions, newpart)
     end
 
     obseqs = solvars .~ rhss

src/structural_transformation/utils.jl

@@ -262,21 +262,21 @@ function reordered_matrix(sys, partitions=structure(sys).partitions)
     I, J = Int[], Int[]
     ii = 0
     M = Int[]
-    for (e_solved, v_solved, e_residue, v_tear) in partitions
-        append!(M, v_solved)
-        append!(M, v_tear)
+    for partition in partitions
+        append!(M, partition.v_solved)
+        append!(M, partition.v_residual)
     end
     M = inverse_mapping(vcat(M, setdiff(1:nvars, M)))
-    for (e_solved, v_solved, e_residue, v_tear) in partitions
-        for es in e_solved
+    for partition in partitions
+        for es in partition.e_solved
             isdiffeq(eqs[es]) && continue
             ii += 1
             js = [M[x] for x in 𝑠neighbors(graph, es) if isalgvar(s, x)]
             append!(I, fill(ii, length(js)))
             append!(J, js)
         end
 
-        for er in e_residue
+        for er in partition.e_residual
             isdiffeq(eqs[er]) && continue
             ii += 1
             js = [M[x] for x in 𝑠neighbors(graph, er) if isalgvar(s, x)]

src/systems/abstractsystem.jl

@@ -244,7 +244,7 @@ function Base.getproperty(sys::AbstractSystem, name::Symbol; namespace=true)
         end
     end
 
-    throw(error("Variable $name does not exist"))
+    throw(ArgumentError("Variable $name does not exist"))
 end
 
 function Base.setproperty!(sys::AbstractSystem, prop::Symbol, val)
@@ -525,7 +525,7 @@ function structural_simplify(sys::AbstractSystem)
     sys = initialize_system_structure(alias_elimination(sys))
     check_consistency(structure(sys))
     if sys isa ODESystem
-        sys = sort_states(dae_index_lowering(sys))
+        sys = dae_index_lowering(sys)
     end
     sys = tearing(sys)
     fullstates = [map(eq->eq.lhs, observed(sys)); states(sys)]

src/systems/systemstructure.jl

@@ -37,12 +37,26 @@ for v in 𝑣vertices(graph); active_𝑣vertices[v] || continue
 end
 =#
 
-export SystemStructure, initialize_system_structure, find_linear_equations
+export SystemStructure, SystemPartition
+export initialize_system_structure, find_linear_equations
 export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq
 export dervars_range, diffvars_range, algvars_range
 
 @enum VariableType::Int8 DIFFERENTIAL_VARIABLE ALGEBRAIC_VARIABLE DERIVATIVE_VARIABLE
 
+Base.@kwdef struct SystemPartition
+    e_solved::Vector{Int}
+    v_solved::Vector{Int}
+    e_residual::Vector{Int}
+    v_residual::Vector{Int}
+end
+
+function Base.:(==)(s1::SystemPartition, s2::SystemPartition)
+    tup1 = (s1.e_solved, s1.v_solved, s1.e_residual, s1.v_residual)
+    tup2 = (s2.e_solved, s2.v_solved, s2.e_residual, s2.v_residual)
+    tup1 == tup2
+end
+
 Base.@kwdef struct SystemStructure
     fullvars::Vector
     vartype::Vector{VariableType}
@@ -55,7 +69,7 @@ Base.@kwdef struct SystemStructure
     assign::Vector{Int}
     inv_assign::Vector{Int}
     scc::Vector{Vector{Int}}
-    partitions::Vector{NTuple{4, Vector{Int}}}
+    partitions::Vector{SystemPartition}
 end
 
 isdervar(s::SystemStructure, var::Integer) = s.vartype[var] === DERIVATIVE_VARIABLE
@@ -168,7 +182,7 @@ function initialize_system_structure(sys)
         assign = Int[],
         inv_assign = Int[],
         scc = Vector{Int}[],
-        partitions = NTuple{4, Vector{Int}}[],
+        partitions = SystemPartition[],
     )
     return sys
 end

test/structural_transformation/tearing.jl

@@ -53,7 +53,7 @@ tornsys = tearing(sys)
 sss = structure(tornsys)
 @unpack graph, solvable_graph, assign, partitions = sss
 @test graph.fadjlist == [[1]]
-@test partitions == [([], [], [1], [1])]
+@test partitions == [StructuralTransformations.SystemPartition([], [], [1], [1])]
 
 # Before:
 #      u1  u2  u3  u4  u5