An algorithm to generate LaTeX Gaussian-eliminations

gausselimination.py

from enum import Enum
from fractions import Fraction

import numpy as np


def override_str(self):
    if self._denominator == 1:
        return str(self._numerator)
    else:
        return '-\\frac{%s}{%s}' % (
        abs(self._numerator), abs(self._denominator)) if self._numerator * self._denominator < 0 \
            else '\\frac{%s}{%s}' % (self._numerator, self._denominator)


# is there a better way to override methods of a class?
Fraction.__str__ = override_str


class OpTypes(Enum):
    SWAP = 0
    ADD = 1
    MULT = 2


operation_example = {
    'type': OpTypes.SWAP,
    'row1': 1,
    'row2': 2
}


def gaussian_elimination(matrix):
    # matrix = matrix.astype(np.float)
    m, n = matrix.shape
    p_row = 0
    p_col = 0
    while p_row < m and p_col < n:
        current_col = matrix[:, p_col]
        max_el = np.argmax(current_col[p_row:]) + p_row
        # if it's 0 we do nothing
        if current_col[max_el] == 0:
            p_col += 1
            continue

        if max_el != p_row:
            # we do a swap, !!!
            matrix[[max_el, p_row]] = matrix[[p_row, max_el]]

        s_row = matrix[p_row]

        for i in range(m):
            if i == p_row:
                continue
            f = matrix[i, p_col] / matrix[p_row, p_col]
            # we add the row !!!
            matrix[i] -= f * s_row

        # we set the pivot to 1 !!!
        matrix[p_row] = s_row / matrix[p_row, p_col]
        p_col += 1
        p_row += 1

    return matrix


class GaussianEliminationGenerator:
    def __init__(self, matrix, delim=0, prefix='(A|\mathbf{b})'):
        self.matrix = matrix
        self.delim = delim
        self.prefix = prefix
        self.m, self.n = matrix.shape
        self.p_row = 0
        self.p_col = 0
        self.latex_code = []

    @staticmethod
    def generate_latex_matrix(matrix, delim=0):
        m, n = matrix.shape
        code = f"\\begin{{pmatrix}}[{'c'*delim+'|'+'c'*(n-delim)}]" if delim else f"\\begin{{pmatrix}}[{'c'*n}]"
        code += ' \\\\ '.join([' & '.join([str(i) for i in row]) for row in matrix])
        code += f" \\end{{pmatrix}}"
        return code

    @staticmethod
    def generate_latex_swap_row(row1, row2):
        return f'R_{{{row1}}} \\leftrightarrow R_{{{row2}}}'

    @staticmethod
    def generate_latex_add_row(multiple, row):
        if multiple == 0:
            return ''
        if multiple < 0:
            return f'+({multiple})R_{{{row}}}'
        if multiple == 1:
            return f'+R_{{{row}}}'
        if multiple > 0:
            return f'+{multiple}R_{{{row}}}'

    @staticmethod
    def generate_latex_multiply(multiple):
        return f'\cdot {multiple}' if multiple > 0 else f'\cdot ({multiple})'

    @staticmethod
    def generate_latex_row_operations(m):
        backslash = '\\'
        return f'\\begin{{matrix*}}[l] {(backslash + "scriptstyle %s " + backslash * 2) * m} \\end{{matrix*}}'

    @staticmethod
    def pick_best_pivot(column):
        absed = abs(column)
        if 1 in absed:
            return np.where(absed == 1)[0][0]
        nonzeroed = absed[np.nonzero(absed)]
        if nonzeroed.any():
            return np.where(absed == nonzeroed)[0][0]
        return 0

    def _iterate(self):
        operations = []

        current_col = self.matrix[:, self.p_col]
        sliced = current_col[self.p_row:]
        # a human will always pick the row which has a 1 as the pivot
        best_el = GaussianEliminationGenerator.pick_best_pivot(sliced) + self.p_row
        # if it's 0 we do nothing
        if current_col[best_el] == 0:
            self.p_col += 1
            if self.p_col < self.n:
                return self._iterate()
            else:
                return operations

        if best_el != self.p_row:
            # we do a swap, !!!
            self.matrix[[best_el, self.p_row]] = self.matrix[[self.p_row, best_el]]
            operations.append({
                'type': OpTypes.SWAP,
                'args': {
                    'row1': best_el,
                    'row2': self.p_row
                }
            })
            return operations

        s_row = self.matrix[self.p_row]

        for i in range(self.m):
            if i == self.p_row:
                continue
            f = self.matrix[i, self.p_col] / self.matrix[self.p_row, self.p_col]
            # we add the row !!!
            self.matrix[i] -= f * s_row
            operations.append({
                'type': OpTypes.ADD,
                'args': {
                    'row1': self.p_row,
                    'row2': i,
                    'multiple': -f
                }
            })

        # we set the pivot to 1 !!!
        multiple = 1 / self.matrix[self.p_row, self.p_col]
        if round(multiple, 2) != 1:
            self.matrix[self.p_row] = s_row * multiple
            operations.append({
                'type': OpTypes.MULT,
                'args': {
                    'row1': self.p_row,
                    'multiple': multiple
                }
            })
        self.p_col += 1
        self.p_row += 1

        return operations

    def generate(self):
        begin = GaussianEliminationGenerator.generate_latex_matrix(self.matrix, self.delim)
        self.latex_code.append(f'{self.prefix}=&' + begin)
        while self.p_row < self.m and self.p_col < self.n:
            current = GaussianEliminationGenerator.generate_latex_matrix(self.matrix, self.delim)

            operations = self._iterate()
            row_ops_latex = GaussianEliminationGenerator.generate_latex_row_operations(self.m)
            subs = [''] * self.m
            for op in operations:
                if op['type'] == OpTypes.SWAP:
                    row1 = op['args']['row1'] + 1
                    row2 = op['args']['row2'] + 1
                    row_swap_latex = GaussianEliminationGenerator.generate_latex_swap_row(row1, row2)
                    subs[op['args']['row1']] = row_swap_latex
                if op['type'] == OpTypes.ADD:
                    row = op['args']['row1'] + 1
                    pos = op['args']['row2']
                    multiple = op['args']['multiple']
                    row_add_latex = GaussianEliminationGenerator.generate_latex_add_row(multiple, row)
                    subs[pos] = row_add_latex
                if op['type'] == OpTypes.MULT:
                    pos = op['args']['row1']
                    multiple = op['args']['multiple']
                    row_multiply_latex = GaussianEliminationGenerator.generate_latex_multiply(multiple)
                    subs[pos] = row_multiply_latex

            current += row_ops_latex % tuple(subs)
            self.latex_code.append('&' + current)
        final = GaussianEliminationGenerator.generate_latex_matrix(self.matrix, self.delim)
        self.latex_code.append(f'{self.prefix}\'=&' + final)
        return '\\\\\n'.join(self.latex_code)


# Example matrix we wish to reduce
a = np.array([[Fraction(1, 9), Fraction(8, 9), Fraction(4, 9), 5],
              [Fraction(4, 9), -Fraction(4, 9), Fraction(7, 9), 1],
              [-Fraction(8, 9), Fraction(1, 9), -Fraction(4, 9), 2]
])

a = a + Fraction()
print(GaussianEliminationGenerator(a, delim=3, prefix=r'(A\mid \mathbf{0})').generate())

# todo
#  add (A|b)= prefix
#  make the algorithm calculate more humanly