sparse_elim.py

from math import gcd
from fractions import Fraction
from collections import defaultdict
from sparse_row import SparseRow
from stop_watch import StopWatch

def lcm(a, *args):
    for b in args:
        a *= b // gcd(a,b)
    return a

class EquationIndex:
    """
    EquationIndex is a helper identifier, any new equation obtains
    a new variable EquationIndex with coefficient 1.
    This allows tracing the origin of a derived equation, currently used
    for determining the common denominator
    = how much did we have to divide to obtain a derived equation
    """
    def __init__(self, equation):
        self.equation = equation

"""
ElimMatrix keeps a set of non-redundant linear equations (SparseRows).
User can add a new linear equation using function
  add(equation)
or check whether a certain equation can be derived from the equations added so far
using the function
  check(equation)
It also automatically detects if two variables are forced to be zero,
or proportional to each other, and communicates this fact back in the form of pairs
which are forced to be equal, returned by the "add" function.
Besides that, ElimMatrix also offer method
  get_inverse(x)
which finds a variable y (if possible) such that y == -x.

Variables which are non-zero, and not substituted by a proportion of other variable
are kept in a "matrix", rows of which are equations derived from the input equations.
Every row in the matrix must have a "pivot" variable such that there is no other
equation in the matrix containing this variable.
"""
class ElimMatrix:
    def __init__(self):
        self.rows = dict() # pivot -> row
        self.cols = defaultdict(set) # varname -> pivot set
        self.value_to_var = dict() # value_key -> pivot variable proportional to that
        self.zeroes = dict() # varname -> equation(size=1)
        self.proportional_to = dict() # varname -> (varname, equation(size=2) )
        self.root_to_proportions = dict() # varname -> size, dict(proportion -> varnames)

    def query(self, query_r):
        query_r = SparseRow(query_r)
        self._eliminate(query_r)
        if not all(isinstance(key, EquationIndex) for key in query_r.keys()):
            return 0
        return self._least_denom(query_r)

    def add(self, added):
        #assert(self.check_consistency())
        #print("elim.add({})".format(added))
        new_r = SparseRow(added)
        self._elim_by_proportions(new_r)
        eq_pure = [
            (x,coef) for (x,coef) in new_r.items()
            if not isinstance(x, EquationIndex)
        ]
        if len(eq_pure) <= 2 and any(x not in self.cols for x,coef in eq_pure):
            if not eq_pure: return False, () # Already known equation
            eqi = EquationIndex(added)
            new_r[eqi] = Fraction(1)
            to_glue_out = []
            if len(eq_pure) == 1:
                (x,coef), = eq_pure
                new_r *= -1/coef
                self._add_zero(x,new_r, to_glue_out)
            else:
                (x,cx),(y,cy) = eq_pure
                if x in self.cols or (y not in self.cols and
                                      self._get_root_to_prop(x)[0] > self._get_root_to_prop(y)[0]):
                    x,cx,y,cy = y,cy,x,cx
                new_r *= -1/cx
                self._add_proportion(x,y, new_r, to_glue_out)
            return True, to_glue_out

        self._elim_by_matrix(new_r)

        pivot_candidates = list(filter(
            lambda x: not isinstance(x, EquationIndex), new_r.keys(),
        ))
        if not pivot_candidates: return False, () # Already known equation

        # label equation
        eqi = EquationIndex(added)
        new_r[eqi] = Fraction(1)

        # select pivot
        if len(pivot_candidates) == 2:
            def cand_cost(cand):
                return 3*len(self.cols[cand]) + self._get_root_to_prop(cand)[0]
        else:
            def cand_cost(cand):
                return len(self.cols[cand])
        pivot = min(pivot_candidates, key = cand_cost)
        new_r *= -1 / new_r[pivot]

        to_glue_out = []
        cols_to_update = [
            (self.cols[x], x)
            for x in pivot_candidates
            if x != pivot
        ]
        #_, cols_to_update_proj2 = zip(*cols_to_update)
        #print("C", cols_to_update_proj2)

        # update matrix, compute glued
        main_col = set(self.cols[pivot])
        for ri in main_col:
            row = self.rows.get(ri)
            coef = row[pivot]
            self._deactivate_row(ri, row)

            row.iadd_coef(coef, new_r) # the essential command

            if self._activate_row(ri, row, to_glue_out):
                # kept in matrix
                for col, ci in cols_to_update:
                    if ci in row: col.add(ri)
                    else: col.discard(ri)
            else:
                # removed from matrix
                for col, ci in cols_to_update:
                    col.discard(ri)

        # add new row
        self.rows[pivot] = new_r
        if self._activate_row(pivot, new_r, to_glue_out):
            # add new row to columns
            self.cols[pivot] = { pivot }
            for col, ci in cols_to_update: col.add(pivot)

        return True, to_glue_out

    def get_inverse(self, x):
        r, eq_rx = self.proportional_to.get(x, (x, None))
        if eq_rx is None:
            ratio = Fraction(-1)
            denom_x = 1
        else:
            ratio = -eq_rx[r]
            denom_x = self._least_denom(eq_rx)            
        r_dict = self._get_root_to_prop(r)[1]
        y_list = r_dict.get(ratio, None)
        if y_list is None: return None, 0

        y = y_list[0]

        # get denominator
        _,eq_ry = self.proportional_to.get(y, (x, None))
        if eq_ry is None: denom_y = 1
        else: denom_y = self._least_denom(eq_ry)

        return y, lcm(denom_x, denom_y)

    # debug functions
    def print_self(self):
        keys = set(self.cols.keys())
        pivots = set(self.rows.keys())
        print("Pivots:", pivots)
        print("Variables:", keys)
        print("------------")
        keys = sorted(pivots)+sorted(keys-pivots)
        print(' '.join("{:^7}".format(k) for k in keys))
        for _,r in sorted(self.rows.items()):
            print(' '.join("{:>3}/{:<3}".format(
                r[k].numerator, r[k].denominator
            ) if r[k] != 0 else 7*' ' for k in keys))
        print(".....")
        for x,(y,eq) in self.proportional_to.items():
            print("  ({}) -> {}*({})".format(x, eq[y], y))
        for x in self.zeroes.keys():
            print("  ({}) -> 0".format(x))
        print("------------")

    def check_consistency(self): # debug function, verification of internal consistency
        ok = True
        for ri, r in self.rows.items():
            if r[ri] != -1:
                print("pivot {} has wrong coefficient in its row: {}".format(ri, r[ri]))
                ok = False
            for ci in r.keys():
                if isinstance(ci, EquationIndex): continue
                if ri not in self.cols[ci]:
                    print("[{}, {}] not in column".format(ri, ci))
                    ok = False
        for ci, col in self.cols.items():
            for ri in col:
                if ri not in self.rows or ci not in self.rows[ri]:
                    print("[{}, {}] not in row".format(ri, ci))
                    ok = False
        for x,eq in self.zeroes.items():
            if eq[x] != -1:
                print("zero {} has wrong coefficient: {}".format(ri, r[ri]))
                ok = False
            if any(y != x and not isinstance(y, EquationIndex) for y in eq.keys()):
                print("other variables in zero equation {}: {}".format(x, eq))
                ok = False
        for x,(y,eq) in self.proportional_to.items():
            if eq[x] != -1 or eq[y] == 0:
                print("proportion {} -> {} has wrong coefficients: {}".format(x,y, r))
                ok = False
            if any(z != x and z != y and not isinstance(z, EquationIndex) for z in eq.keys()):
                print("other variables in proportion equation {} -> {}: {}".format(x, y, eq))
                ok = False

        return ok

    # helper functions
    def _eliminate(self, row): # in place elimination of a single row
        self._elim_by_proportions(row)
        self._elim_by_matrix(row)

    def _elim_by_proportions(self, row):
        updates = []
        for var, coef in row.items():
            _,update = self.proportional_to.get(var, (None, None))
            if update is None: update = self.zeroes.get(var)
            if update is not None: updates.append((coef, update))
        for coef, update in updates:
            row += coef*update
    def _elim_by_matrix(self, row):
        updates = []
        for var, coef in row.items():
            update = self.rows.get(var)
            if update is not None: updates.append((coef, update))
        for coef, update in updates:
            row += coef*update        
            
    def _least_denom(self, row): # common denominator of the equation indices
        res = 1
        for key, coef in row.items():
            if not isinstance(key, EquationIndex): continue
            denom = coef.denominator
            res *= denom // gcd(res, denom)
        return res

    def _add_zero(self, x, eq, to_glue_out):
        if x in self.cols: del self.cols[x]
        if self.zeroes:
            y,eq2 = next(iter(self.zeroes.items()))
            denom = lcm(self._least_denom(eq),self._least_denom(eq2))
            to_glue_out.append((x,y, denom))
        self.zeroes[x] = eq
        if x in self.root_to_proportions:
            _, d = self.root_to_proportions[x]
            for q,l in d.items():
                for y in l:
                    if y != x:
                        eq2 = self.proportional_to[y][1] + q*eq
                        self.zeroes[y] = eq2
                        del self.proportional_to[y]
                        if q != 1:
                            denom = lcm(self._least_denom(eq),self._least_denom(eq2))
                            to_glue_out.append((x,y, denom))
            del self.root_to_proportions[x]

    def _get_root_to_prop(self, x):
        res = self.root_to_proportions.get(x, None)
        if res is None:
            res = 1, {Fraction(1) : [x]}
            self.root_to_proportions[x] = res
        return res

    def _add_proportion(self, x, y, eq_yx, to_glue_out):
        if x in self.cols: del self.cols[x]
        if 22 in eq_yx and 13 in eq_yx and 26 in eq_yx:
            raise Exception()
        # get data
        x_size, x_dict = self._get_root_to_prop(x)
        y_size, y_dict = self._get_root_to_prop(y)

        ratio_yx = eq_yx[y]
        denom_yx = self._least_denom(eq_yx)
        for ratio_xz, z_list in x_dict.items():
            ratio_yz = ratio_xz*ratio_yx

            # update proportional_to
            for z in z_list:
                if z == x: eq_yz = eq_yx
                else:
                    # x,eq2 ==
                    _,eq_xz = self.proportional_to[z]
                    # eq_yz = eq_xz + ratio_xz * eq_yx
                    eq_xz.iadd_coef(ratio_xz, eq_yx)
                    eq_yz = eq_xz
                self.proportional_to[z] = y, eq_yz

            # update y_dict
            zz_list = y_dict.setdefault(ratio_yz, z_list)
            if zz_list is not z_list:
                z = z_list[0]
                zz = zz_list[0]
                zz_list.extend(z_list)

                # get the denominator of the equation stating z == zz
                if z == x: denom_xz = 1
                else:
                    _, eq_xz = self.proportional_to[z]
                    denom_xz = self._least_denom(eq_xz)
                if zz == y: denom_yzz = 1
                else:
                    _, eq_yzz = self.proportional_to[zz]
                    denom_yzz = self._least_denom(eq_yzz)
                denom = lcm(denom_xz, denom_yzz, ratio_yz.denominator * denom_yx)

                # update to_glue
                to_glue_out.append((z, zz, denom))

        self.root_to_proportions[y] = x_size+y_size, y_dict
        del self.root_to_proportions[x]

    def _row_valkey(self, pivot, row):
        relevant_items = tuple(
            (x,coef) for (x,coef) in row.items()
            if not isinstance(x, EquationIndex) and x != pivot
        )
        if not relevant_items: return ()
        _,coef0 = min(relevant_items)
        return frozenset((x, coef/coef0) for (x,coef) in relevant_items)

    # remove from self.rows and self.cols
    def _remove_row(self, pivot, row):
        del self.rows[pivot]
        for x,coef in row.items():
            if not isinstance(x, EquationIndex) and x in self.cols:
                self.cols[x].discard(pivot)

    # updates only self.value_to_var, not self.rows nor self.cols
    def _deactivate_row(self, pivot, row):
        valkey = self._row_valkey(pivot, row)
        del self.value_to_var[valkey]

    def _activate_row(self, x, row, to_glue_out):
        valkey = self._row_valkey(x, row)
        if len(valkey) <= 1:
            if len(valkey) == 0:
                self._add_zero(x, row, to_glue_out)
            else:
                (y,_), = valkey
                self._add_proportion(x, y, row, to_glue_out)
            self._remove_row(x, row)
            return False
        else:
            y = self.value_to_var.setdefault(valkey, x)
            if y is x: return True

            # x and y are proportional

            # swap if more efficient
            cost_x, _ = self._get_root_to_prop(x)
            cost_y, _ = self._get_root_to_prop(y)
            if cost_x > cost_y:
                x,y = y,x
                eq_x = self.rows[x]
                eq_y = row
                self.value_to_var[valkey] = y
                preserve_x = True
            else:
                eq_x = row
                eq_y = self.rows[y]
                preserve_x = False

            # get equation stating x = coef*y
            z,_ = next(iter(valkey))
            eq = eq_x + eq_y * (-eq_x[z] / eq_y[z])

            # add as a proportion, remove from matrix
            self._add_proportion(x, y, eq, to_glue_out)
            self._remove_row(x, eq_x)

            return preserve_x

if __name__ == "__main__":

    elim = ElimMatrix()

    #elim.add(SparseRow({'A': Fraction(3, 2), 'B': Fraction(-1, 1)}))
    #elim.add(SparseRow({'C': Fraction(3, 2), 'D': Fraction(-1, 1)}))
    #elim.add(SparseRow({'A': Fraction(1, 1), 'C': Fraction(1, 1)}))
    #print(elim.get_inverse('D'))
    #print(elim.proportional_to)

    elim.add({0: Fraction(1, 1)})
    elim.add({0: Fraction(1, 1)})
    elim.add({3: Fraction(2, 1), 4: Fraction(-2, 1)})
    elim.add({6: Fraction(-1, 1), 4: Fraction(1, 1)})
    elim.add({6: Fraction(2, 1), 8: Fraction(-2, 1)})
    elim.add({3: Fraction(-1, 1), 8: Fraction(1, 1)})
    elim.add({3: Fraction(-1, 1), 10: Fraction(1, 1)})
    elim.add({13: Fraction(2, 1), 3: Fraction(-1, 1)})
    elim.add({13: Fraction(2, 1), 17: Fraction(-1, 1)})
    elim.add({3: Fraction(-1, 1), 17: Fraction(1, 1)})
    elim.add({12: Fraction(-1, 1), 18: Fraction(1, 1)})
    elim.add({13: Fraction(2, 1), 21: Fraction(-1, 1)})
    elim.add({3: Fraction(-1, 1), 21: Fraction(1, 1)})
    elim.add({12: Fraction(-1, 1), 22: Fraction(1, 1)})
    elim.add({25: Fraction(-1, 1), 26: Fraction(1, 1)})
    elim.add({27: Fraction(1, 1), 25: Fraction(-1, 1), 28: Fraction(-1, 1)})
    elim.add({28: Fraction(1, 1), 29: Fraction(-1, 1)})
    elim.add({28: Fraction(-1, 1), 29: Fraction(1, 1)})
    elim.add({30: Fraction(2, 1), 3: Fraction(-1, 1)})
    elim.add({30: Fraction(2, 1), 33: Fraction(-1, 1)})
    elim.add({3: Fraction(-1, 1), 33: Fraction(1, 1)})
    elim.add({12: Fraction(-1, 1), 34: Fraction(1, 1)})
    elim.add({12: Fraction(-1, 1), 25: Fraction(1, 1)})
    elim.add({3: Fraction(2, 1), 35: Fraction(-2, 1)})
    elim.add({6: Fraction(-1, 1), 35: Fraction(1, 1)})
    elim.add({6: Fraction(-1, 1), 37: Fraction(1, 1)})
    elim.add({43: Fraction(2, 1), 44: Fraction(-1, 1)})
    elim.add({45: Fraction(2, 1), 3: Fraction(-1, 1)})
    elim.add({30: Fraction(-1, 1), 45: Fraction(1, 1)})
    elim.add({43: Fraction(1, 1), 45: Fraction(1, 1), 46: Fraction(-1, 1)})
    elim.add({48: Fraction(-1, 1), 46: Fraction(1, 1)})
    elim.add({50: Fraction(2, 1), 51: Fraction(-1, 1)})
    elim.add({52: Fraction(2, 1), 3: Fraction(-1, 1)})
    elim.add({13: Fraction(-1, 1), 52: Fraction(1, 1)})
    elim.add({50: Fraction(1, 1), 52: Fraction(1, 1), 53: Fraction(-1, 1)})
    elim.add({55: Fraction(-1, 1), 53: Fraction(1, 1)})
    elim.add({44: Fraction(2, 1), 58: Fraction(-2, 1)})
    elim.add({60: Fraction(-1, 1), 58: Fraction(1, 1)})
    elim.add({51: Fraction(2, 1), 62: Fraction(-2, 1)})
    elim.add({64: Fraction(-1, 1), 62: Fraction(1, 1)})
    elim.add({66: Fraction(-1, 1), 68: Fraction(1, 1)})
    elim.add({70: Fraction(-1, 1), 72: Fraction(1, 1)})
    elim.add({6: Fraction(-1, 1), 80: Fraction(1, 1)})
    elim.add({6: Fraction(-1, 1), 83: Fraction(1, 1)})
    elim.add({44: Fraction(-1, 1), 48: Fraction(1, 1), 87: Fraction(-1, 1)})
    elim.add({48: Fraction(-1, 1), 3: Fraction(1, 1), 88: Fraction(-1, 1)})
    elim.add({87: Fraction(-1, 1), 88: Fraction(1, 1)})
    elim.add({70: Fraction(-1, 1), 89: Fraction(1, 1)})
    elim.add({51: Fraction(-1, 1), 55: Fraction(1, 1), 92: Fraction(-1, 1)})
    elim.add({55: Fraction(-1, 1), 3: Fraction(1, 1), 93: Fraction(-1, 1)})
    elim.add({92: Fraction(-1, 1), 93: Fraction(1, 1)})
    elim.add({66: Fraction(-1, 1), 94: Fraction(1, 1)})
    elim.add({100: Fraction(-1, 1), 101: Fraction(1, 1), 102: Fraction(-1, 1)})
    elim.add({103: Fraction(2, 1), 64: Fraction(-1, 1)})
    elim.add({105: Fraction(2, 1), 107: Fraction(-1, 1)})
    elim.add({103: Fraction(-1, 1), 105: Fraction(1, 1), 108: Fraction(-1, 1)})
    elim.add({102: Fraction(-1, 1), 108: Fraction(1, 1)})
    elim.add({110: Fraction(-1, 1), 111: Fraction(1, 1), 112: Fraction(-1, 1)})
    elim.add({113: Fraction(2, 1), 115: Fraction(-1, 1)})
    elim.add({116: Fraction(2, 1), 60: Fraction(-1, 1)})
    elim.add({113: Fraction(-1, 1), 116: Fraction(1, 1), 118: Fraction(-1, 1)})
    elim.add({112: Fraction(-1, 1), 118: Fraction(1, 1)})
    elim.add({70: Fraction(-1, 1), 119: Fraction(1, 1)})
    elim.add({66: Fraction(-1, 1), 120: Fraction(1, 1)})
    elim.add({115: Fraction(-1, 1), 77: Fraction(1, 1), 122: Fraction(-1, 1)})
    elim.add({77: Fraction(-1, 1), 60: Fraction(1, 1), 123: Fraction(-1, 1)})
    elim.add({122: Fraction(-1, 1), 123: Fraction(1, 1)})
    elim.add({0: Fraction(1, 1)})
    elim.add({64: Fraction(-1, 1), 77: Fraction(1, 1), 125: Fraction(-1, 1)})
    elim.add({77: Fraction(-1, 1), 107: Fraction(1, 1), 126: Fraction(-1, 1)})
    print('A')
    elim.check_consistency()
    print('B')
    elim.add({125: Fraction(-1, 1), 126: Fraction(1, 1)})
    print('C')
    elim.check_consistency()