reframr/linalg.py

import math
import site
import sys
from pathlib import Path

_VENDOR_ROOT = Path(__file__).resolve().parent.parent / ".vendor"
for _vendor_path in (_VENDOR_ROOT / "python", _VENDOR_ROOT / "sitepkgs"):
    if _vendor_path.exists():
        vendor_text = str(_vendor_path)
        if vendor_text not in sys.path:
            sys.path.insert(0, vendor_text)

try:
    import numpy as np
except ModuleNotFoundError:
    user_site = site.getusersitepackages()
    if user_site and user_site not in sys.path:
        sys.path.append(user_site)
    try:
        import numpy as np
    except ModuleNotFoundError:
        np = None

if np is not None and not hasattr(np, "asarray"):
    np = None

Matrix = list[list[float]]
Vector = list[float]
SUMPROD = getattr(math, "sumprod", None)


def zeros(rows: int, cols: int) -> Matrix:
    return [[0.0 for _ in range(cols)] for _ in range(rows)]


def zeros_vector(size: int) -> Vector:
    return [0.0 for _ in range(size)]


def identity(size: int) -> Matrix:
    matrix = zeros(size, size)
    for index in range(size):
        matrix[index][index] = 1.0
    return matrix


def copy_matrix(matrix: Matrix) -> Matrix:
    return [row[:] for row in matrix]


def transpose(matrix: Matrix) -> Matrix:
    if not matrix:
        return []
    if np is not None:
        return np.asarray(matrix, dtype=np.float64).T.tolist()
    return [list(column) for column in zip(*matrix)]


def matvec(matrix: Matrix, vector: Vector) -> Vector:
    if np is not None:
        return (np.asarray(matrix, dtype=np.float64) @ np.asarray(vector, dtype=np.float64)).tolist()
    if SUMPROD is not None:
        return [SUMPROD(row, vector) for row in matrix]
    return [sum(value * vector[idx] for idx, value in enumerate(row)) for row in matrix]


def matmul(left: Matrix, right: Matrix) -> Matrix:
    if not left or not right:
        return []
    if np is not None:
        return (np.asarray(left, dtype=np.float64) @ np.asarray(right, dtype=np.float64)).tolist()
    right_t = transpose(right)
    if SUMPROD is not None:
        return [[SUMPROD(row, column) for column in right_t] for row in left]
    return [
        [sum(a * b for a, b in zip(row, column)) for column in right_t]
        for row in left
    ]


def add_matrices(left: Matrix, right: Matrix) -> Matrix:
    return [
        [left[row][col] + right[row][col] for col in range(len(left[row]))]
        for row in range(len(left))
    ]


def subtract_matrices(left: Matrix, right: Matrix) -> Matrix:
    return [
        [left[row][col] - right[row][col] for col in range(len(left[row]))]
        for row in range(len(left))
    ]


def scale_matrix(matrix: Matrix, scalar: float) -> Matrix:
    return [[scalar * value for value in row] for row in matrix]


def dot(left: Vector, right: Vector) -> float:
    if np is not None:
        return float(np.dot(np.asarray(left, dtype=np.float64), np.asarray(right, dtype=np.float64)))
    if SUMPROD is not None:
        return SUMPROD(left, right)
    return sum(a * b for a, b in zip(left, right))


def norm(vector: Vector) -> float:
    return math.sqrt(dot(vector, vector))


def outer(left: Vector, right: Vector) -> Matrix:
    if np is not None:
        return np.outer(np.asarray(left, dtype=np.float64), np.asarray(right, dtype=np.float64)).tolist()
    return [[a * b for b in right] for a in left]


def mean(values: Vector) -> float:
    return sum(values) / len(values) if values else 0.0


def trace(matrix: Matrix) -> float:
    return sum(matrix[index][index] for index in range(min(len(matrix), len(matrix[0]))))


def covariance_matrix(samples: list[Vector]) -> Matrix:
    if not samples:
        return []
    if np is not None:
        sample_array = np.asarray(samples, dtype=np.float64)
        centered = sample_array - sample_array.mean(axis=0, keepdims=True)
        denominator = max(len(samples) - 1, 1)
        return ((centered.T @ centered) / denominator).tolist()

    feature_count = len(samples[0])
    sample_count = len(samples)
    means = [
        sum(sample[feature] for sample in samples) / sample_count
        for feature in range(feature_count)
    ]
    covariance = zeros(feature_count, feature_count)
    for sample in samples:
        centered = [sample[index] - means[index] for index in range(feature_count)]
        for row in range(feature_count):
            for col in range(feature_count):
                covariance[row][col] += centered[row] * centered[col]

    denominator = max(sample_count - 1, 1)
    return scale_matrix(covariance, 1.0 / denominator)


def solve_linear_system(matrix: Matrix, vector: Vector) -> Vector:
    if np is not None:
        return np.linalg.solve(
            np.asarray(matrix, dtype=np.float64),
            np.asarray(vector, dtype=np.float64),
        ).tolist()
    size = len(matrix)
    augmented = [matrix[row][:] + [vector[row]] for row in range(size)]

    for pivot_index in range(size):
        pivot_row = max(
            range(pivot_index, size),
            key=lambda row_index: abs(augmented[row_index][pivot_index]),
        )
        augmented[pivot_index], augmented[pivot_row] = augmented[pivot_row], augmented[pivot_index]

        pivot_value = augmented[pivot_index][pivot_index]
        if abs(pivot_value) < 1e-12:
            raise ValueError("Singular matrix encountered while solving linear system.")

        inverse_pivot = 1.0 / pivot_value
        augmented[pivot_index] = [value * inverse_pivot for value in augmented[pivot_index]]

        for row_index in range(size):
            if row_index == pivot_index:
                continue
            factor = augmented[row_index][pivot_index]
            augmented[row_index] = [
                augmented[row_index][col] - factor * augmented[pivot_index][col]
                for col in range(size + 1)
            ]

    return [augmented[row][-1] for row in range(size)]


def invert_matrix(matrix: Matrix) -> Matrix:
    if np is not None:
        return np.linalg.inv(np.asarray(matrix, dtype=np.float64)).tolist()
    size = len(matrix)
    inverse_columns = []
    for basis_index in range(size):
        basis_vector = [0.0 for _ in range(size)]
        basis_vector[basis_index] = 1.0
        inverse_columns.append(solve_linear_system(matrix, basis_vector))
    return transpose(inverse_columns)


def dominant_eigenpair_symmetric(
    matrix: Matrix,
    max_iterations: int = 64,
    tolerance: float = 1e-10,
) -> tuple[float, Vector]:
    size = len(matrix)
    if size == 0:
        return 0.0, []
    if np is not None:
        values, vectors = np.linalg.eigh(np.asarray(matrix, dtype=np.float64))
        index = int(np.argmax(values))
        eigenvalue = float(values[index])
        if eigenvalue <= tolerance:
            return 0.0, zeros_vector(size)
        return eigenvalue, vectors[:, index].astype(float).tolist()

    vector = [1.0 / math.sqrt(size) for _ in range(size)]
    for _ in range(max_iterations):
        next_vector = matvec(matrix, vector)
        next_norm = norm(next_vector)
        if next_norm < tolerance:
            return 0.0, zeros_vector(size)

        next_vector = [value / next_norm for value in next_vector]
        delta = max(abs(a - b) for a, b in zip(vector, next_vector))
        vector = next_vector
        if delta < tolerance:
            break

    eigenvalue = dot(vector, matvec(matrix, vector))
    return eigenvalue, vector


def top_k_eigenpairs_symmetric(matrix: Matrix, k: int) -> list[tuple[float, Vector]]:
    if np is not None and matrix:
        values, vectors = np.linalg.eigh(np.asarray(matrix, dtype=np.float64))
        ranked = sorted(
            (
                (float(values[index]), vectors[:, index].astype(float).tolist())
                for index in range(len(values))
                if float(values[index]) > 1e-9
            ),
            key=lambda item: item[0],
            reverse=True,
        )
        return ranked[: min(k, len(ranked))]
    working = copy_matrix(matrix)
    eigenpairs: list[tuple[float, Vector]] = []
    for _ in range(min(k, len(working))):
        eigenvalue, eigenvector = dominant_eigenpair_symmetric(working)
        if eigenvalue <= 1e-9 or not eigenvector:
            break
        eigenpairs.append((eigenvalue, eigenvector))
        deflation = scale_matrix(outer(eigenvector, eigenvector), eigenvalue)
        working = subtract_matrices(working, deflation)
    return eigenpairs


def softmax(logits: Vector) -> Vector:
    if not logits:
        return []
    if np is not None:
        values = np.asarray(logits, dtype=np.float64)
        shifted = np.exp(values - values.max())
        total = float(shifted.sum())
        if total == 0.0:
            return [1.0 / len(logits) for _ in logits]
        return (shifted / total).tolist()
    max_logit = max(logits)
    shifted = [math.exp(logit - max_logit) for logit in logits]
    total = sum(shifted)
    if total == 0.0:
        return [1.0 / len(logits) for _ in logits]
    return [value / total for value in shifted]