Add back rank2 (highest diffed variables)

Co-authored-by: Keno Fischer <[email protected]>

Author: YingboMa

src/structural_transformation/StructuralTransformations.jl

@@ -54,6 +54,7 @@ export tearing_assignments, tearing_substitution
 export torn_system_jacobian_sparsity
 export full_equations
 export but_ordered_incidence, lowest_order_variable_mask, highest_order_variable_mask
+export computed_highest_diff_variables
 
 include("utils.jl")
 include("pantelides.jl")

src/structural_transformation/pantelides.jl

@@ -70,7 +70,7 @@ function pantelides_reassemble(state::TearingState, var_eq_matching)
 end
 
 """
-    computed_highest_diff_variables(var_to_diff)
+    computed_highest_diff_variables(structure)
 
 Computes which variables are the "highest-differentiated" for purposes of
 pantelides. Ordinarily this is relatively straightforward. However, in our

src/systems/alias_elimination.jl

@@ -264,7 +264,8 @@ function find_linear_variables(graph, linear_equations, var_to_diff, irreducible
     return linear_variables
 end
 
-function aag_bareiss!(graph, var_to_diff, mm_orig::SparseMatrixCLIL{T, Ti}) where {T, Ti}
+function aag_bareiss!(structure, mm_orig::SparseMatrixCLIL{T, Ti}) where {T, Ti}
+    @unpack graph, var_to_diff = structure
     mm = copy(mm_orig)
     linear_equations_set = BitSet(mm_orig.nzrows)
 
@@ -279,6 +280,7 @@ function aag_bareiss!(graph, var_to_diff, mm_orig::SparseMatrixCLIL{T, Ti}) wher
         v -> var_to_diff[v] === nothing === invview(var_to_diff)[v]
     end
     is_linear_variables = is_algebraic.(1:length(var_to_diff))
+    is_highest_diff = computed_highest_diff_variables(structure)
     for i in 𝑠vertices(graph)
         # only consider linear algebraic equations
         (i in linear_equations_set && all(is_algebraic, 𝑠neighbors(graph, i))) &&
@@ -291,25 +293,31 @@ function aag_bareiss!(graph, var_to_diff, mm_orig::SparseMatrixCLIL{T, Ti}) wher
 
     local bar
     try
-        bar = do_bareiss!(mm, mm_orig, is_linear_variables)
+        bar = do_bareiss!(mm, mm_orig, is_linear_variables, is_highest_diff)
     catch e
         e isa OverflowError || rethrow(e)
         mm = convert(SparseMatrixCLIL{BigInt, Ti}, mm_orig)
-        bar = do_bareiss!(mm, mm_orig, is_linear_variables)
+        bar = do_bareiss!(mm, mm_orig, is_linear_variables, is_highest_diff)
     end
 
     return mm, solvable_variables, bar
 end
 
-function do_bareiss!(M, Mold, is_linear_variables)
+function do_bareiss!(M, Mold, is_linear_variables, is_highest_diff)
     rank1r = Ref{Union{Nothing, Int}}(nothing)
+    rank2r = Ref{Union{Nothing, Int}}(nothing)
     find_pivot = let rank1r = rank1r
         (M, k) -> begin
             if rank1r[] === nothing
                 r = find_masked_pivot(is_linear_variables, M, k)
                 r !== nothing && return r
                 rank1r[] = k - 1
             end
+            if rank2r[] === nothing
+                r = find_masked_pivot(is_highest_diff, M, k)
+                r !== nothing && return r
+                rank2r[] = k - 1
+            end
             # TODO: It would be better to sort the variables by
             # derivative order here to enable more elimination
             # opportunities.
@@ -334,15 +342,19 @@ function do_bareiss!(M, Mold, is_linear_variables)
     bareiss_ops = ((M, i, j) -> nothing, myswaprows!,
         bareiss_update_virtual_colswap_mtk!, bareiss_zero!)
 
-    rank2, = bareiss!(M, bareiss_ops; find_pivot = find_and_record_pivot)
+    rank3, = bareiss!(M, bareiss_ops; find_pivot = find_and_record_pivot)
+    rank2 = something(rank2r[], rank3)
     rank1 = something(rank1r[], rank2)
-    (rank1, rank2, pivots)
+    (rank1, rank2, rank3, pivots)
 end
 
-function simple_aliases!(ils, graph, solvable_graph, eq_to_diff, var_to_diff)
-    ils, solvable_variables, (rank1, rank2, pivots) = aag_bareiss!(graph,
-        var_to_diff,
-        ils)
+function alias_eliminate_graph!(state::TransformationState, ils::SparseMatrixCLIL)
+    @unpack structure = state
+    @unpack graph, solvable_graph, var_to_diff, eq_to_diff = state.structure
+    # Step 1: Perform Bareiss factorization on the adjacency matrix of the linear
+    #         subsystem of the system we're interested in.
+    #
+    ils, solvable_variables, (rank1, rank2, rank3, pivots) = aag_bareiss!(structure, ils)
 
     ## Step 2: Simplify the system using the Bareiss factorization
     rk1vars = BitSet(@view pivots[1:rank1])
@@ -362,14 +374,6 @@ function simple_aliases!(ils, graph, solvable_graph, eq_to_diff, var_to_diff)
     return ils
 end
 
-function alias_eliminate_graph!(state::TransformationState, ils::SparseMatrixCLIL)
-    @unpack graph, solvable_graph, var_to_diff, eq_to_diff = state.structure
-    # Step 1: Perform Bareiss factorization on the adjacency matrix of the linear
-    #         subsystem of the system we're interested in.
-    #
-    return simple_aliases!(ils, graph, solvable_graph, eq_to_diff, var_to_diff)
-end
-
 function exactdiv(a::Integer, b)
     d, r = divrem(a, b)
     @assert r == 0