rref

rref.py

"""Implementation of RREF (Reduced Row Echelon Form)."""

import functools
import hashlib
import time
from abc import ABC, abstractmethod
from textwrap import dedent
from typing import Any, Iterator, Optional

import numpy as np
import numpy.typing as npt
import sympy as sp
from pydantic import BaseModel, Field
from tqdm import tqdm

np.random.seed(42)
np.set_printoptions(
    precision=2,
    suppress=True,
    edgeitems=16,
    linewidth=180,
)

EPSILON = 1e-6

LOG_ENABLED = False
DELAY = False
PROMPT = False

DELAY_TIME = 0.25
PROMPT_MESSAGE = "--- hit enter to continue ---"


## Python list utilities to push from the left or the right
def pushl(lst: list, x: Any) -> list:
    """Push an element to the left of a list."""
    lst.insert(0, x)
    return lst


def pushr(lst: list, x: Any) -> list:
    """Push an element to the right of a list."""
    lst.append(x)
    return lst


def array_hash(array: npt.NDArray[np.float64]) -> str:
    """Hash a numpy array."""
    return hashlib.sha256(array.tobytes()).hexdigest()


def log_matrix(
    matrix: npt.NDArray[np.float64], message: str, enabled: bool = LOG_ENABLED, prompt: bool = PROMPT
) -> None:
    """Log a matrix."""
    full_message = dedent(f"""
DELAY: {DELAY}
MESSAGE: {message}
MATRIX:
{matrix}
""")
    if enabled:
        print(full_message)
    if DELAY:
        time.sleep(DELAY_TIME)
    if prompt:
        input(PROMPT_MESSAGE)


class ElementaryMatrix(BaseModel, ABC, frozen=True):
    """Base class for elementary matrices."""

    n: int = Field(..., description="Size of the matrix (n x n)")

    class Config:  # noqa: D106
        arbitrary_types_allowed = True

    @abstractmethod
    def as_matrix(self) -> npt.NDArray[np.float64]:
        """Get the matrix representation of the elementary matrix."""
        raise NotImplementedError("Subclasses must implement as_matrix()")

    @abstractmethod
    def det(self) -> float:
        """Get the determinant of the elementary matrix."""
        raise NotImplementedError("Subclasses must implement det()")

    @abstractmethod
    def inverse(self) -> "ElementaryMatrix":
        """Get the inverse of the elementary matrix."""
        raise NotImplementedError("Subclasses must implement inverse()")


class RowMultiplication(ElementaryMatrix, frozen=True):
    """Elementary matrix that multiplies a row by a scalar."""

    i: int = Field(..., description="The index of the row to multiply")
    value: float = Field(..., description="The scalar to multiply the row by")

    def model_post_init(self, *args: object, **kwargs: object) -> None:
        """Post-initialization hook."""
        assert 0 <= self.i < self.n, "RowMultiplication: Index must be in the range [0, n)"
        assert self.value != 0, "RowMultiplication: Value must be non-zero"

    def as_matrix(self) -> npt.NDArray[np.float64]:
        """Get the matrix representation of the elementary matrix."""
        matrix = np.eye(self.n, dtype=np.float64)
        matrix[self.i, self.i] = self.value
        return matrix

    def det(self) -> float:
        """Get the determinant of the elementary row multiplication matrix."""
        return self.value

    def inverse(self) -> "RowMultiplication":
        """Get the inverse of the elementary row multiplication matrix."""
        return RowMultiplication(n=self.n, i=self.i, value=1 / self.value)


class RowSwitch(ElementaryMatrix, frozen=True):
    """Elementary matrix that swaps two rows."""

    i: int = Field(..., description="The index of the first row to swap")
    j: int = Field(..., description="The index of the second row to swap")

    def model_post_init(self, *args: object, **kwargs: object) -> None:
        """Post-initialization hook."""
        assert 0 <= self.i < self.n, "RowSwitch: Index must be in the range [0, n)"
        assert 0 <= self.j < self.n, "RowSwitch: Index must be in the range [0, n)"
        assert self.i != self.j, "RowSwitch: Indices must be different"

    def as_matrix(self) -> npt.NDArray[np.float64]:
        """Get the matrix representation of the elementary matrix."""
        matrix = np.eye(self.n, dtype=np.float64)
        matrix[[self.i, self.j]] = matrix[[self.j, self.i]]
        return matrix

    def det(self) -> float:
        """Get the determinant of the elementary row switch matrix."""
        return -1.0

    def inverse(self) -> "RowSwitch":
        """Get the inverse of the elementary row switch matrix."""
        return self


class RowAddition(ElementaryMatrix, frozen=True):
    """Elementary matrix that adds a row j multiplied by a scalar m to row i."""

    i: int = Field(..., description="The index of the row to add")
    j: int = Field(..., description="The index of the column to add to")
    value: float = Field(..., description="The scalar to multiply the row to add by")

    def model_post_init(self, *args: object, **kwargs: object) -> None:
        """Post-initialization hook."""
        assert 0 <= self.i < self.n, "RowAddition: Index must be in the range [0, n)"
        assert 0 <= self.j < self.n, "RowAddition: Index must be in the range [0, n)"
        assert self.i != self.j, "RowAddition: Indices must be different"

    def as_matrix(self) -> npt.NDArray[np.float64]:
        """Get the matrix representation of the elementary matrix."""
        matrix = np.eye(self.n, dtype=np.float64)
        matrix[self.i, self.j] = self.value
        return matrix

    def det(self) -> float:
        """Get the determinant of the elementary row addition matrix."""
        return 1.0

    def inverse(self) -> "RowAddition":
        """Get the inverse of the elementary row addition matrix."""
        return RowAddition(n=self.n, i=self.i, j=self.j, value=-self.value)


LEFT_APPLY_CACHE: dict[tuple[ElementaryMatrix, str], npt.NDArray[np.float64]] = {}
RIGHT_APPLY_CACHE: dict[tuple[ElementaryMatrix, str], npt.NDArray[np.float64]] = {}


class ElementaryMatrixOperations(BaseModel, frozen=False):
    """Operations on elementary matrices.

    Left elementary matrices operate on the rows of a matrix and are applied from right to left.

    Right elementary matrices operate on the columns of a matrix and are applied from left to right.
    """

    left_matrices: list[ElementaryMatrix] = Field(default_factory=list, description="The list of left matrices")
    right_matrices: list[ElementaryMatrix] = Field(default_factory=list, description="The list of right matrices")

    def add_operations(self, other: "ElementaryMatrixOperations") -> "ElementaryMatrixOperations":
        """Add two elementary matrix operations together.

        The provided 'other' operations will be applied after the current operations.
        """
        return ElementaryMatrixOperations(
            left_matrices=other.left_matrices + self.left_matrices,
            right_matrices=self.right_matrices + other.right_matrices,
        )

    def left_matrix(self) -> npt.NDArray[np.float64]:
        """Apply the left operations to the identity matrix."""
        return functools.reduce(
            lambda x, y: x @ y,
            iter(m.as_matrix() for m in self.left_matrices),
        )

    def append_left_matrix(self, elementary_matrix: ElementaryMatrix) -> "ElementaryMatrixOperations":
        """Append a left matrix to the identity matrix."""
        return ElementaryMatrixOperations(
            left_matrices=pushl(self.left_matrices, elementary_matrix),
            right_matrices=self.right_matrices,
        )

    def right_matrix(self) -> npt.NDArray[np.float64]:
        """Apply the right operations to the identity matrix."""
        return functools.reduce(
            lambda x, y: x @ y,
            iter(m.as_matrix() for m in self.right_matrices),
        )

    def append_right_matrix(self, elementary_matrix: ElementaryMatrix) -> "ElementaryMatrixOperations":
        """Append a right matrix to the identity matrix."""
        return ElementaryMatrixOperations(
            left_matrices=self.left_matrices,
            right_matrices=pushr(self.right_matrices, elementary_matrix),
        )

    def apply(self, matrix: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        """Apply the operations to a matrix."""
        input_type = matrix.dtype

        def left_inner_apply(
            a: npt.NDArray[np.float64], left_elementary_matrix: ElementaryMatrix
        ) -> npt.NDArray[np.float64]:
            key = (left_elementary_matrix, array_hash(a))
            if key in LEFT_APPLY_CACHE:
                return LEFT_APPLY_CACHE[key]
            result = left_elementary_matrix.as_matrix() @ a
            LEFT_APPLY_CACHE[key] = result
            return result

        def right_inner_apply(
            a: npt.NDArray[np.float64], right_elementary_matrix: ElementaryMatrix
        ) -> npt.NDArray[np.float64]:
            key = (right_elementary_matrix, array_hash(a))
            if key in RIGHT_APPLY_CACHE:
                return RIGHT_APPLY_CACHE[key]
            result = right_elementary_matrix.as_matrix() @ a
            RIGHT_APPLY_CACHE[key] = result
            return result

        base = matrix
        for left_matrix in reversed(self.left_matrices):
            base = left_inner_apply(base, left_matrix)
        for right_matrix in self.right_matrices:
            base = right_inner_apply(base, right_matrix)

        base = base.astype(input_type)
        return base


def first_nonident_rowcol(matrix: npt.NDArray[np.float64]) -> Optional[int]:
    """Find the first non-identity row or column in a matrix."""
    assert matrix.ndim == 2, "Matrix must be 2D"
    m, n = matrix.shape
    diag_size = min(m, n)

    for i in range(diag_size):
        diag_i_is_one = matrix[i, i] == 1
        rest_row_is_zero = np.all(matrix[i + 1 :, i] == 0)
        rest_col_is_zero = np.all(matrix[i, i + 1 :] == 0)
        if not diag_i_is_one or not rest_col_is_zero or not rest_row_is_zero:
            return i
    return None


def matrix_first_nonzero_from_coord(
    matrix: npt.NDArray[np.float64], coord: tuple[int, int]
) -> Optional[tuple[int, int]]:
    """Find the first non-zero element in a matrix from a given coordinate.

    This is a breadth-first search (BFS) algorithm that explores the matrix in a grid-like manner.
    Only looking in the (>x, >y) direction.
    """
    assert matrix.ndim == 2, "Matrix must be 2D"
    m, n = matrix.shape
    assert 0 <= coord[0] < m, "Index must be in the range [0, m)"
    assert 0 <= coord[1] < n, "Index must be in the range [0, n)"

    stack = [coord]
    visited = set()

    while stack:
        curr_i, curr_j = stack.pop()
        if (curr_i, curr_j) in visited:
            continue

        visited.add((curr_i, curr_j))

        if curr_i >= coord[0] and curr_j >= coord[1] and not np.isclose(matrix[curr_i, curr_j], 0.0, atol=EPSILON):
            return (curr_i, curr_j)

        ## Only look in the (>x, >y) directions
        for ni, nj in [(curr_i + 1, curr_j), (curr_i, curr_j + 1)]:
            if ni < m and nj < n and (ni, nj) not in visited:
                stack.append((ni, nj))

    return None


def make_ident_rowcol(
    matrix: npt.NDArray[np.float64], i: int, operations: Optional[ElementaryMatrixOperations] = None
) -> Optional[ElementaryMatrixOperations]:
    """Make the i-th row and column of a matrix identity."""
    assert matrix.ndim == 2, "Matrix must be 2D"
    m, n = matrix.shape
    diag_size = min(m, n)
    assert 0 <= i < diag_size, "Index must be in the range [0, diag_size)"

    if operations is None:
        operations = ElementaryMatrixOperations()

    ## We need to ensure that (i,i) is non-zero before proceeding.
    ## If it is zero, we need to find the nearest non-zero element w.r.t. (i, i)
    ## and switch the row/column to make it non-zero.
    ## exit if we can't do it (meaning the rest of thematrix is 0)
    latest_matrix = operations.apply(matrix)
    if np.isclose(latest_matrix[i, i], 0.0, atol=EPSILON):
        coord = matrix_first_nonzero_from_coord(latest_matrix, (i, i))

        ## If we can't find a non-zero element past (i, i) then the rest of the matrix is 0
        ## and we're done, we've done all we can
        if coord is None:
            return operations

        ## switch row/col if coord is non-zero w.r.t. (i, i)
        nonzero_row_idx, nonzero_col_idx = coord
        if nonzero_row_idx != i:
            operations = operations.append_left_matrix(RowSwitch(n=m, i=i, j=nonzero_row_idx))
            log_matrix(operations.apply(matrix), f"Switching rows: {i} <-> {nonzero_row_idx}")
        if nonzero_col_idx != i:
            operations = operations.append_right_matrix(RowSwitch(n=n, i=i, j=nonzero_col_idx))
            log_matrix(operations.apply(matrix), f"Switching columns: {i} <-> {nonzero_col_idx}")

    ## (i,i) is non zero, if it's non one we scale it down to 1 by multiplying by inverse
    latest_matrix = operations.apply(matrix)
    assert not np.isclose(
        latest_matrix[i, i], 0.0
    ), f"Diagonal element {i} should be non-zero, got {round(latest_matrix[i, i], 2)}"
    if not np.isclose(latest_matrix[i, i], 1.0, atol=EPSILON):
        operations = operations.append_right_matrix(RowMultiplication(n=n, i=i, value=1 / latest_matrix[i, i]))
        log_matrix(operations.apply(matrix), f"Scaling diagonal element: {i} -> 1")

    ## Now we apply RowAddition matrices to make the rest of the 'row' zero (fixed column)
    latest_matrix = operations.apply(matrix)
    for j in range(i + 1, diag_size):
        if not np.isclose(latest_matrix[i, j], 0.0, atol=EPSILON):
            operations = operations.append_right_matrix(RowAddition(n=n, i=i, j=j, value=-latest_matrix[i, j]))
            log_matrix(operations.apply(matrix), f"Shifting row coord: {(i, j)} = {round(latest_matrix[i, j], 2)} -> 0")

    ## Same for the columns, make the rest of the 'column' zero (fixed row)
    latest_matrix = operations.apply(matrix)
    for j in range(i + 1, diag_size):
        if not np.isclose(latest_matrix[j, i], 0.0, atol=EPSILON):
            operations = operations.append_left_matrix(RowAddition(n=m, i=j, j=i, value=-latest_matrix[j, i]))
            log_matrix(
                operations.apply(matrix), f"Shifting column coord: {(j, i)} = {round(latest_matrix[j, i], 2)} -> 0"
            )

    return operations


def make_ident_lower_triangular_ops(
    matrix: npt.NDArray[np.float64], operations: Optional[ElementaryMatrixOperations] = None
) -> ElementaryMatrixOperations:
    """Make the matrix lower triangular via row operations."""
    assert matrix.ndim == 2, "Matrix must be 2D"
    m, n = matrix.shape
    diag_size = min(m, n)

    if operations is None:
        operations = ElementaryMatrixOperations()

    latest_matrix = operations.apply(matrix)
    for col_idx in range(diag_size):
        for row_idx in range(col_idx + 1, m):
            if not np.isclose(latest_matrix[row_idx, col_idx], 0.0, atol=EPSILON):
                operations = operations.append_left_matrix(
                    RowAddition(
                        n=m,
                        i=row_idx,
                        j=col_idx,
                        value=-(latest_matrix[row_idx, col_idx] / latest_matrix[col_idx, col_idx]),
                    )
                )
                log_matrix(
                    operations.apply(matrix),
                    f"Shifting column coord: {(row_idx, col_idx)} = {round(latest_matrix[row_idx, col_idx], 2)} -> 0",
                )
                latest_matrix = operations.apply(matrix)

    return operations


def make_ident_diagonal_ops(
    matrix: npt.NDArray[np.float64], operations: Optional[ElementaryMatrixOperations] = None
) -> ElementaryMatrixOperations:
    """Make the diagonal elements of a matrix identity."""
    assert matrix.ndim == 2, "Matrix must be 2D"
    m, n = matrix.shape
    diag_size = min(m, n)

    if operations is None:
        operations = ElementaryMatrixOperations()

    latest_matrix = operations.apply(matrix)
    for i in range(diag_size):
        if not np.isclose(latest_matrix[i, i], 0.0, atol=EPSILON):
            operations = operations.append_left_matrix(RowMultiplication(n=m, i=i, value=1 / latest_matrix[i, i]))
            log_matrix(
                operations.apply(matrix),
                f"Scaling diagonal element: {i}, {latest_matrix[i, i]} -> 1",
            )

    return operations


def make_ident_upper_triangular_ops(
    matrix: npt.NDArray[np.float64], operations: Optional[ElementaryMatrixOperations] = None
) -> ElementaryMatrixOperations:
    """Make the matrix upper triangular via row operations."""
    assert matrix.ndim == 2, "Matrix must be 2D"
    m, n = matrix.shape
    diag_size = min(m, n)

    if operations is None:
        operations = ElementaryMatrixOperations()

    latest_matrix = operations.apply(matrix)
    for col_idx in range(diag_size):
        for row_idx in range(col_idx - 1, -1, -1):  # decrement to zero
            if not np.isclose(latest_matrix[row_idx, col_idx], 0.0, atol=EPSILON):
                operations = operations.append_left_matrix(
                    RowAddition(n=m, i=row_idx, j=col_idx, value=-latest_matrix[row_idx, col_idx])
                )
                log_matrix(
                    operations.apply(matrix),
                    f"Shifting column coord: {(row_idx, col_idx)} = {round(latest_matrix[row_idx, col_idx], 2)} -> 0",
                )
                latest_matrix = operations.apply(matrix)

    return operations


def staircase_form_ops(
    matrix: npt.NDArray[np.float64], operations: Optional[ElementaryMatrixOperations] = None
) -> ElementaryMatrixOperations:
    """Compute the staircase form of a matrix."""
    assert matrix.ndim == 2, "Matrix must be 2D"
    m, n = matrix.shape

    if operations is None:
        operations = ElementaryMatrixOperations()

    def count_leading_zeros(row: npt.NDArray[np.float64]) -> int:
        """Count the number of leading zeros in a row."""
        idxs = np.nonzero(np.abs(row) > EPSILON)[0]
        return int(idxs[0]) if idxs.size > 0 else len(row)

    ## Sort row indices based on leading zeros
    latest_matrix = operations.apply(matrix)
    sorted_row_indices = list(range(m))
    sorted_row_indices.sort(key=lambda i: count_leading_zeros(latest_matrix[i, :]))

    ## Convert that sorted order into row swaps
    current_positions = list(range(m))
    for sorted_pos, row_idx in enumerate(sorted_row_indices):
        ## Where is the row_idx currently in our matrix?
        current_pos_of_desired_row = current_positions.index(row_idx)
        if current_pos_of_desired_row != sorted_pos:
            operations = operations.append_left_matrix(RowSwitch(n=m, i=current_pos_of_desired_row, j=sorted_pos))
            log_matrix(
                operations.apply(matrix),
                f"Switching rows: {current_pos_of_desired_row} <-> {sorted_pos}",
            )

            ## Update local tracking of current positions
            current_positions[current_pos_of_desired_row], current_positions[sorted_pos] = (
                current_positions[sorted_pos],
                current_positions[current_pos_of_desired_row],
            )

    return operations


def smith_normal_form_ops(matrix: npt.NDArray[np.float64]) -> ElementaryMatrixOperations:
    """Compute the Smith normal form of a matrix."""
    assert matrix.ndim == 2, "Matrix must be 2D"

    n = matrix.shape[0]
    operations = ElementaryMatrixOperations()
    for row_idx in range(n):
        next_operations = make_ident_rowcol(matrix, row_idx, operations)
        if next_operations is None:
            return operations
        operations = next_operations
    return operations


def rref_ops(matrix: npt.NDArray[np.float64]) -> ElementaryMatrixOperations:
    """Compute the reduced row echelon form of a matrix."""
    assert matrix.ndim == 2, "Matrix must be 2D"
    operations = ElementaryMatrixOperations()
    operations = make_ident_lower_triangular_ops(matrix, operations)
    operations = make_ident_diagonal_ops(matrix, operations)
    operations = make_ident_upper_triangular_ops(matrix, operations)
    operations = staircase_form_ops(matrix, operations)
    return operations


def det(matrix: npt.NDArray[np.float64]) -> float:
    """Compute the determinant of a matrix."""
    assert matrix.ndim == 2, "Matrix must be 2D"
    assert matrix.shape[0] == matrix.shape[1], "Matrix must be square"

    operations = rref_ops(matrix)
    assert len(operations.right_matrices) == 0, "Right matrices should be empty"
    dets: Iterator[float] = iter(m.inverse().det() for m in operations.left_matrices)
    return functools.reduce(
        lambda x, y: x * y,
        dets,
    )


def test_fuzz_rref() -> None:
    """Test the rref function with a random matrix."""
    print("Fuzz testing square matrices")
    tests = 300
    max_size = 10

    def test_single_matrix() -> bool:
        n = np.random.randint(1, max_size)
        base_matrix = np.random.random((n, n)) * 20 - 10
        operations = rref_ops(base_matrix)
        out_mat = np.round(operations.left_matrix() @ base_matrix / EPSILON) * EPSILON
        target_mat = np.array(sp.Matrix(base_matrix).rref()[0]).astype(np.float64)
        # assert np.allclose(out_mat, target_mat, atol=EPSILON), "RREF does not match numpy.linalg.rref"
        return np.allclose(out_mat, target_mat, atol=EPSILON)

    results = list(tqdm((test_single_matrix() for _ in range(tests)), total=tests, desc="Testing RREF Square Matrices"))
    print(f"Passed {sum(results)}/{tests} tests")

    tests = 300
    max_rows = 10
    max_cols = 10

    def test_single_matrix() -> bool:
        n = np.random.randint(1, max_rows)
        m = np.random.randint(1, max_cols)
        base_matrix = np.random.random((n, m)) * 20 - 10
        operations = rref_ops(base_matrix)
        out_mat = np.round(operations.left_matrix() @ base_matrix / EPSILON) * EPSILON
        target_mat = np.array(sp.Matrix(base_matrix).rref()[0]).astype(np.float64)
        return np.allclose(out_mat, target_mat, atol=EPSILON)

    results = list(
        tqdm((test_single_matrix() for _ in range(tests)), total=tests, desc="Testing RREF Rectangular Matrices")
    )
    print(f"Passed {sum(results)}/{tests} tests")

    tests = 300
    max_rows = 10
    max_cols = 10
    max_row_dependencies = 5
    max_col_dependencies = 5

    def test_single_matrix() -> bool:
        n = np.random.randint(2, max_rows)
        m = np.random.randint(2, max_cols)
        base_matrix = np.random.random((n, m)) * 20 - 10
        row_dependencies = np.random.randint(0, max_row_dependencies)
        col_dependencies = np.random.randint(0, max_col_dependencies)
        for _ in range(row_dependencies):
            scalar = np.random.random()
            rand_row_a = np.random.randint(0, n - 1)
            rand_row_b = np.random.randint(0, n - 1)
            base_matrix[rand_row_a, :] = base_matrix[rand_row_a, :] + scalar * base_matrix[rand_row_b, :]
        for _ in range(col_dependencies):
            scalar = np.random.random()
            rand_col_a = np.random.randint(0, m - 1)
            rand_col_b = np.random.randint(0, m - 1)
            base_matrix[:, rand_col_a] = base_matrix[:, rand_col_a] + scalar * base_matrix[:, rand_col_b]
        operations = rref_ops(base_matrix)
        out_mat = np.round(operations.left_matrix() @ base_matrix / EPSILON) * EPSILON
        target_mat = np.array(sp.Matrix(base_matrix).rref()[0]).astype(np.float64)
        return np.allclose(out_mat, target_mat, atol=EPSILON)

    results = list(
        tqdm(
            (test_single_matrix() for _ in range(tests)),
            total=tests,
            desc="Testing RREF Rectangular Matrices With Linear Dependencies",
        )
    )
    print(f"Passed {sum(results)}/{tests} tests")