fix: properly eliminate identically zero variables in alias_elimination!

Author: AayushSabharwal

src/systems/alias_elimination.jl

@@ -409,6 +409,107 @@ function find_perfect_aliases!(
     return aliases
 end
 
+"""
+    $TYPEDSIGNATURES
+
+If a row in `mm` only contains a single variable, that variable is identically zero.
+Eliminate such any such zero variables which can be identified from `mm`. Return:
+- `zero_vars`: A `Set{Int}` of variables identified to be identically zero.
+- `zero_eqs`: A `Set{Int}` of equations identifying such variables to be zero.
+- `eqs_to_substitute`: A `Set{Int}` of equations in `state` which contain variables
+  in `zero_vars`. These equations should be substituted to remove the zero variables.
+
+Also update `mm` in-place to remove such zero variables/equations.
+"""
+function eliminate_zero_variables!(state::TearingState, mm::CLIL.SparseMatrixCLIL)
+    # Inverse mapping of `mm.nzrows`
+    nzrow_to_idx = Dict{Int, Int}()
+    sizehint!(nzrow_to_idx, length(mm.nzrows))
+    for (i, eqidx) in enumerate(mm.nzrows)
+        nzrow_to_idx[eqidx] = i
+    end
+    # Variables we can identify are zero via `mm`
+    zero_vars = Set{Int}()
+    # Equations in `mm` of the form `var ~ 0`
+    zero_eqs = Set{Int}()
+    # Equations that `zero_vars` are incident on and need to be substituted
+    eqs_to_substitute = Set{Int}()
+    # Queue of indices in `mm.nzrows` containing rows to check for being
+    # `var ~ 0`
+    queue = Queue{Int}()
+    # Initially check all rows with just one element
+    for i in eachindex(mm.nzrows)
+        isone(length(mm.row_cols[i])) && push!(queue, i)
+    end
+
+    while !isempty(queue)
+        row_i = popfirst!(queue)
+        eqidx = mm.nzrows[row_i]
+        # Skip rows we already processed in case they show up twice
+        eqidx in zero_eqs && continue
+
+        # If a row only contains one non-zero element, that element is also zero.
+        # We filter `rcol` and `rval` to remove any zero elements. This could be
+        # done in the `for nbor in nbors` loop below. However, it's possible we process
+        # two zero variables both present in the same row before processing that row.
+        # This approach avoids the row being processed by each of the two zero variables
+        # individually in their `nbors` loop. We delay removing all zero elements to
+        # the latest possible time to aggregate updates.
+        rcol = mm.row_cols[row_i]
+        rval = mm.row_vals[row_i]
+        ptr = firstindex(rcol)
+        nnz = 0
+        for i in eachindex(rcol)
+            rcol[i] in zero_vars && continue
+            rcol[ptr] = rcol[i]
+            rval[ptr] = rval[i]
+            ptr = nextind(rcol, ptr)
+            nnz += 1
+        end
+        resize!(rcol, nnz)
+        resize!(rval, nnz)
+        # `mm` does have rows with empty `rcol` on entry to this function, but they
+        # should never be reachable by this search.
+        @assert !iszero(nnz) "`mm` should not be singular at this point!"
+        isone(nnz) || continue
+
+        zvar = rcol[1]
+        push!(zero_vars, zvar)
+        push!(zero_eqs, eqidx)
+
+        nbors = 𝑑neighbors(state.structure.graph, zvar)
+        for nbor in nbors
+            nzrow_idx = get(nzrow_to_idx, nbor, 0)
+            # If `nbor` is not present in `mm`, it must be substituted later
+            if iszero(nzrow_idx)
+                push!(eqs_to_substitute, nbor)
+                continue
+            end
+            # If it is present, process it
+            push!(queue, nzrow_idx)
+        end
+    end
+
+    # Remove zero equations from `mm`
+    if !isempty(zero_eqs)
+        ptr = firstindex(mm.nzrows)
+        for i in eachindex(mm.nzrows)
+            eqidx = mm.nzrows[i]
+            eqidx in zero_eqs && continue
+            mm.nzrows[ptr] = mm.nzrows[i]
+            mm.row_cols[ptr] = mm.row_cols[i]
+            mm.row_vals[ptr] = mm.row_vals[i]
+            ptr = nextind(mm.nzrows, ptr)
+        end
+        nnz = length(mm.nzrows) - length(zero_eqs)
+        resize!(mm.nzrows, nnz)
+        resize!(mm.row_cols, nnz)
+        resize!(mm.row_vals, nnz)
+    end
+
+    return zero_vars, zero_eqs, eqs_to_substitute
+end
+
 function alias_elimination!(
         state::TearingState; fully_determined = true,
         print_underconstrained_variables = false, kwargs...
@@ -438,6 +539,10 @@ function alias_elimination!(
     fullvars_to_idx = Dict{SymbolicT, Int}(Iterators.map(reverse, enumerate(fullvars)))
     original_eqs = state.original_eqs
 
+    zero_vars, zero_eqs, eqs_to_substitute = eliminate_zero_variables!(state, mm)
+    append!(vars_to_rm, zero_vars)
+    append!(eqs_to_rm, zero_eqs)
+
     for (ieq, eq) in enumerate(mm.nzrows)
         rcol = mm.row_cols[ieq]
         rval = mm.row_vals[ieq]
@@ -475,6 +580,18 @@ function alias_elimination!(
         end
     end
 
+    subrules = Dict{SymbolicT, SymbolicT}()
+    sizehint!(subrules, length(zero_vars))
+    for v in zero_vars
+        subrules[fullvars[v]] = Symbolics.COMMON_ZERO
+        push!(state.additional_observed, fullvars[v] ~ Symbolics.COMMON_ZERO)
+    end
+    subber = SU.IRSubstituter{true}(get_irstructure(sys), subrules)
+    for ieq in eqs_to_substitute
+        eqs[ieq] = subber(eqs[ieq])
+        original_eqs[ieq] = subber(original_eqs[ieq])
+    end
+
     @set! sys.eqs = eqs
     state.sys = sys