Merge pull request #1843 from SciML/myb/complete_ag

Compute `complete_ag` that contains irreducible alias information

Author: YingboMa

src/bipartite_graph.jl

@@ -63,6 +63,13 @@ end
 function Base.setindex!(m::Matching{U}, v::Union{Integer, U}, i::Integer) where {U}
     if m.inv_match !== nothing
         oldv = m.match[i]
+        # TODO: maybe default Matching to always have an `inv_match`?
+
+        # To maintain the invariant that `m.inv_match[m.match[i]] == i`, we need
+        # to unassign the matching at `m.inv_match[v]` if it exists.
+        if v isa Int && (iv = m.inv_match[v]) isa Int
+            m.match[iv] = unassigned
+        end
         if isa(oldv, Int)
             @assert m.inv_match[oldv] == i
             m.inv_match[oldv] = unassigned

src/structural_transformation/pantelides.jl

@@ -83,20 +83,27 @@ function pantelides!(state::TransformationState, ag::Union{AliasGraph, Nothing}
     ecolor = falses(neqs)
     var_eq_matching = Matching(nvars)
     neqs′ = neqs
-    nnonemptyeqs = count(eq -> !isempty(𝑠neighbors(graph, eq)), 1:neqs′)
+    nnonemptyeqs = count(eq -> !isempty(𝑠neighbors(graph, eq)) && eq_to_diff[eq] === nothing,
+                         1:neqs′)
 
     # Allow matching for the highest differentiated variable that
     # currently appears in an equation (or implicit equation in a side ag)
     varwhitelist = falses(nvars)
     for var in 1:nvars
-        if var_to_diff[var] === nothing
+        if var_to_diff[var] === nothing && !varwhitelist[var]
             while isempty(𝑑neighbors(graph, var)) && (ag === nothing || !haskey(ag, var))
                 var′ = invview(var_to_diff)[var]
                 var′ === nothing && break
                 var = var′
             end
-            if !isempty(𝑑neighbors(graph, var)) || (ag !== nothing && haskey(ag, var))
-                varwhitelist[var] = true
+            if !isempty(𝑑neighbors(graph, var))
+                if ag !== nothing && haskey(ag, var)
+                    # TODO: remove lower diff vars from whitelist
+                    c, a = ag[var]
+                    iszero(c) || (varwhitelist[a] = true)
+                else
+                    varwhitelist[var] = true
+                end
             end
         end
     end
@@ -107,6 +114,7 @@ function pantelides!(state::TransformationState, ag::Union{AliasGraph, Nothing}
 
     for k in 1:neqs′
         eq′ = k
+        eq_to_diff[eq′] === nothing || continue
         isempty(𝑠neighbors(graph, eq′)) && continue
         pathfound = false
         # In practice, `maxiters=8000` should never be reached, otherwise, the

src/structural_transformation/partial_state_selection.jl

@@ -98,7 +98,10 @@ function pss_graph_modia!(structure::SystemStructure, var_eq_matching, varlevel,
         end
     end
     for var in 1:ndsts(graph)
-        if varlevel[var] !== 0 && var_eq_matching[var] === unassigned
+        dv = var_to_diff[var]
+        # If `var` is not algebraic (not differentiated nor a dummy derivative),
+        # then it's a SelectedState
+        if !(dv === nothing || (varlevel[dv] !== 0 && var_eq_matching[dv] === unassigned))
             var_eq_matching[var] = SelectedState()
         end
     end
@@ -154,11 +157,13 @@ function partial_state_selection_graph!(structure::SystemStructure, var_eq_match
     var_eq_matching
 end
 
-function dummy_derivative_graph!(state::TransformationState, jac = nothing; kwargs...)
+function dummy_derivative_graph!(state::TransformationState, jac = nothing,
+                                 (ag, diff_va) = (nothing, nothing);
+                                 kwargs...)
     state.structure.solvable_graph === nothing && find_solvables!(state; kwargs...)
-    var_eq_matching = complete(pantelides!(state))
+    var_eq_matching = complete(pantelides!(state, ag))
     complete!(state.structure)
-    dummy_derivative_graph!(state.structure, var_eq_matching, jac)
+    dummy_derivative_graph!(state.structure, var_eq_matching, jac, (ag, diff_va))
 end
 
 function compute_diff_level(diff_to_x)
@@ -178,7 +183,8 @@ function compute_diff_level(diff_to_x)
     return xlevel, maxlevel
 end
 
-function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, jac)
+function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, jac,
+                                 (ag, diff_va))
     @unpack eq_to_diff, var_to_diff, graph = structure
     diff_to_eq = invview(eq_to_diff)
     diff_to_var = invview(var_to_diff)
@@ -242,6 +248,18 @@ function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, ja
             vars = [diff_to_var[var] for var in vars if diff_to_var[var] !== nothing]
         end
     end
+    if diff_va !== nothing
+        n_dummys = length(dummy_derivatives)
+        needed = count(x -> x isa Int, diff_to_eq) - n_dummys
+        n = 0
+        for v in diff_va
+            c, a = ag[v]
+            n += 1
+            push!(dummy_derivatives, iszero(c) ? v : a)
+            needed == n && break
+            continue
+        end
+    end
 
     dummy_derivatives_set = BitSet(dummy_derivatives)
     # We can eliminate variables that are not a selected state (differential
@@ -251,6 +269,9 @@ function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, ja
         dummy_derivatives_set = dummy_derivatives_set
 
         v -> begin
+            if ag !== nothing
+                haskey(ag, v) && return false
+            end
             dv = var_to_diff[v]
             dv === nothing || dv in dummy_derivatives_set
         end
@@ -266,7 +287,8 @@ function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, ja
                                        Union{Unassigned, SelectedState};
                                        varfilter = can_eliminate)
     for v in eachindex(var_eq_matching)
-        can_eliminate(v) && continue
+        dv = var_to_diff[v]
+        (dv === nothing || dv in dummy_derivatives_set) && continue
         var_eq_matching[v] = SelectedState()
     end
 

src/systems/abstractsystem.jl

@@ -405,7 +405,7 @@ function namespace_assignment(eq::Assignment, sys)
     Assignment(_lhs, _rhs)
 end
 
-function namespace_expr(O, sys, n = nameof(sys)) where {T}
+function namespace_expr(O, sys, n = nameof(sys))
     ivs = independent_variables(sys)
     O = unwrap(O)
     if any(isequal(O), ivs)

src/systems/alias_elimination.jl

@@ -6,7 +6,10 @@ const KEEP = typemin(Int)
 
 function alias_eliminate_graph!(state::TransformationState)
     mm = linear_subsys_adjmat(state)
-    size(mm, 1) == 0 && return AliasGraph(ndsts(state.structure.graph)), mm, BitSet() # No linear subsystems
+    if size(mm, 1) == 0
+        ag = AliasGraph(ndsts(state.structure.graph))
+        return ag, mm, ag, mm, BitSet() # No linear subsystems
+    end
 
     @unpack graph, var_to_diff = state.structure
 
@@ -39,7 +42,7 @@ alias_elimination(sys) = alias_elimination!(TearingState(sys; quick_cancel = tru
 function alias_elimination!(state::TearingState)
     sys = state.sys
     complete!(state.structure)
-    ag, mm, updated_diff_vars = alias_eliminate_graph!(state)
+    ag, mm, complete_ag, complete_mm, updated_diff_vars = alias_eliminate_graph!(state)
     isempty(ag) && return sys
 
     fullvars = state.fullvars
@@ -543,6 +546,7 @@ function alias_eliminate_graph!(graph, var_to_diff, mm_orig::SparseMatrixCLIL)
     #
     nvars = ndsts(graph)
     ag = AliasGraph(nvars)
+    complete_ag = AliasGraph(nvars)
     mm, echelon_mm = simple_aliases!(ag, graph, var_to_diff, mm_orig)
 
     # Step 3: Handle differentiated variables
@@ -702,6 +706,7 @@ function alias_eliminate_graph!(graph, var_to_diff, mm_orig::SparseMatrixCLIL)
             while (iv = diff_to_var[v]) in zero_vars_set
                 v = iv
             end
+            complete_ag[v] = 0
             if diff_to_var[v] === nothing # `v` is reducible
                 dag[v] = 0
             end
@@ -729,6 +734,13 @@ function alias_eliminate_graph!(graph, var_to_diff, mm_orig::SparseMatrixCLIL)
     # Step 4: Merge dag and ag
     removed_aliases = BitSet()
     merged_ag = AliasGraph(nvars)
+    for (v, (c, a)) in dag
+        complete_ag[v] = c => a
+    end
+    for (v, (c, a)) in ag
+        (processed[v] || (!iszero(a) && processed[a])) && continue
+        complete_ag[v] = c => a
+    end
     for (v, (c, a)) in dag
         # D(x) ~ D(y) cannot be removed if x and y are not aliases
         if v != a && !iszero(a)
@@ -789,7 +801,8 @@ function alias_eliminate_graph!(graph, var_to_diff, mm_orig::SparseMatrixCLIL)
         update_graph_neighbors!(graph, ag)
     end
 
-    return ag, mm, updated_diff_vars
+    complete_mm = reduce!(copy(echelon_mm), mm_orig, complete_ag, size(echelon_mm, 1))
+    return ag, mm, complete_ag, complete_mm, updated_diff_vars
 end
 
 function update_graph_neighbors!(graph, ag)

test/structural_transformation/index_reduction.jl

@@ -119,9 +119,7 @@ sol = solve(prob_auto, Rodas5());
 #plot(sol, vars=(D(x), y))
 
 let pss_pendulum2 = partial_state_selection(pendulum2)
-    # This currently selects `T` rather than `x` at top level. Needs tearing priorities to fix.
-    @test length(equations(pss_pendulum2)) == 4
-    @test length(equations(ModelingToolkit.ode_order_lowering(pss_pendulum2))) == 4
+    @test length(equations(pss_pendulum2)) <= 6
 end
 
 eqs = [D(x) ~ w,