# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Linear algebra primitives with multi-backend support.
This module provides backend-agnostic linear algebra operations
for NumPy, PyTorch, and JAX, with robust fallbacks.
Decompositions
--------------
- QR decomposition (with column pivoting option)
- Cholesky decomposition (with safe fallback)
- SVD (singular value decomposition)
- Eigendecomposition for symmetric matrices
Solvers
-------
- solve_triangular: Efficient triangular system solver
- solve_cholesky: Solve via Cholesky factorization
- solve_qr: Solve via QR decomposition
- lstsq: Least squares solution
Matrix Operations
-----------------
- safe_inverse: Matrix inverse with regularization
- woodbury_inverse: Efficient update of inverse
- quadratic_form: x'Ax computation
- log_determinant: Log of matrix determinant
Examples
--------
>>> from aurora.core.linalg import safe_cholesky, solve_cholesky
>>>
>>> # Robust Cholesky with fallback
>>> L = safe_cholesky(A)
>>> x = solve_cholesky(L, b)
>>>
>>> # QR for least squares
>>> Q, R = qr_decomposition(X)
>>> beta = solve_qr(Q, R, y)
"""
from __future__ import annotations
from typing import Literal, cast
import numpy as np
# =============================================================================
# Decompositions
# =============================================================================
[docs]
def qr_decomposition(
A: np.ndarray,
*,
mode: Literal["reduced", "complete", "r", "raw"] = "reduced",
pivoting: bool = False,
) -> tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, np.ndarray]:
"""QR decomposition with optional column pivoting.
Decomposes A = QR where Q is orthogonal and R is upper triangular.
With pivoting, AP = QR where P is a permutation matrix.
Parameters
----------
A : ndarray of shape (m, n)
Matrix to decompose.
mode : {'reduced', 'complete', 'r', 'raw'}, default='reduced'
- 'reduced': Return Q (m, k) and R (k, n), where k = min(m, n)
- 'complete': Return Q (m, m) and R (m, n)
- 'r': Return only R
- 'raw': Return (H, tau) for advanced use
pivoting : bool, default=False
If True, use column pivoting and return permutation indices.
Returns
-------
Q : ndarray
Orthogonal matrix (unless mode='r' or 'raw').
R : ndarray
Upper triangular matrix.
P : ndarray (only if pivoting=True)
Permutation indices such that A[:, P] = Q @ R.
Examples
--------
>>> A = np.array([[1, 2], [3, 4], [5, 6]])
>>> Q, R = qr_decomposition(A)
>>> np.allclose(A, Q @ R)
True
>>> # With pivoting for rank-deficient matrices
>>> Q, R, P = qr_decomposition(A, pivoting=True)
>>> np.allclose(A[:, P], Q @ R)
True
"""
if pivoting:
from scipy.linalg import qr as scipy_qr
Q, R, P = scipy_qr(A, mode="economic", pivoting=True)
return Q, R, P
else:
return cast(
"tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, np.ndarray]",
np.linalg.qr(A, mode=mode),
)
[docs]
def safe_cholesky(
A: np.ndarray,
*,
lower: bool = True,
max_tries: int = 5,
jitter: float = 1e-6,
) -> np.ndarray:
"""Cholesky decomposition with robustness to near-singular matrices.
Attempts standard Cholesky, and if it fails, adds progressively
larger diagonal jitter until successful.
Parameters
----------
A : ndarray of shape (n, n)
Symmetric positive semi-definite matrix.
lower : bool, default=True
If True, return lower triangular L. If False, return upper.
max_tries : int, default=5
Maximum number of attempts with increasing jitter.
jitter : float, default=1e-6
Initial diagonal jitter to add on failure.
Returns
-------
L : ndarray of shape (n, n)
Cholesky factor such that A ≈ L @ L.T.
Raises
------
np.linalg.LinAlgError
If decomposition fails after max_tries attempts.
Examples
--------
>>> A = np.array([[4, 2], [2, 1.001]]) # Nearly singular
>>> L = safe_cholesky(A)
>>> np.allclose(A, L @ L.T, atol=1e-5)
True
"""
A = np.asarray(A)
for i in range(max_tries):
try:
if i > 0:
# Add jitter to diagonal
A = A + (jitter * (10 ** (i - 1))) * np.eye(A.shape[0])
L = np.linalg.cholesky(A)
if not lower:
L = L.T
return L
except np.linalg.LinAlgError:
if i == max_tries - 1:
raise np.linalg.LinAlgError(
f"Cholesky decomposition failed after {max_tries} attempts "
f"with jitter up to {jitter * (10 ** (i - 1)):.2e}"
) from None
# Should not reach here
raise np.linalg.LinAlgError("Cholesky decomposition failed")
[docs]
def svd(
A: np.ndarray,
*,
full_matrices: bool = False,
compute_uv: bool = True,
) -> tuple[np.ndarray, np.ndarray, np.ndarray] | np.ndarray:
"""Singular Value Decomposition.
Decomposes A = U @ S @ V.T where U and V are orthogonal and S is diagonal.
Parameters
----------
A : ndarray of shape (m, n)
Matrix to decompose.
full_matrices : bool, default=False
If True, return full U (m, m) and Vh (n, n).
If False, return reduced U (m, k) and Vh (k, n).
compute_uv : bool, default=True
If True, return U, S, Vh. If False, return only S.
Returns
-------
U : ndarray
Left singular vectors (if compute_uv=True).
S : ndarray
Singular values (sorted in descending order).
Vh : ndarray
Right singular vectors transposed (if compute_uv=True).
Examples
--------
>>> A = np.random.randn(5, 3)
>>> U, S, Vh = svd(A)
>>> np.allclose(A, U @ np.diag(S) @ Vh)
True
"""
if compute_uv:
U, S, Vh = np.linalg.svd(A, full_matrices=full_matrices, compute_uv=True)
return U, S, Vh
return np.asarray(np.linalg.svd(A, full_matrices=full_matrices, compute_uv=False))
[docs]
def eigh(
A: np.ndarray,
*,
subset_by_index: tuple[int, int] | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Eigendecomposition for symmetric/Hermitian matrices.
Parameters
----------
A : ndarray of shape (n, n)
Symmetric matrix.
subset_by_index : tuple of (lo, hi), optional
Only compute eigenvalues/vectors in index range [lo, hi].
Returns
-------
eigenvalues : ndarray of shape (n,) or (hi-lo+1,)
Eigenvalues in ascending order.
eigenvectors : ndarray of shape (n, n) or (n, hi-lo+1)
Eigenvectors as columns.
Examples
--------
>>> A = np.array([[2, 1], [1, 2]])
>>> vals, vecs = eigh(A)
>>> np.allclose(A @ vecs, vecs @ np.diag(vals))
True
"""
if subset_by_index is not None:
from scipy.linalg import eigh as scipy_eigh
w, v = scipy_eigh(A, subset_by_index=subset_by_index)
return np.asarray(w), np.asarray(v)
else:
w, v = np.linalg.eigh(A)
return np.asarray(w), np.asarray(v)
# =============================================================================
# Solvers
# =============================================================================
[docs]
def solve_triangular(
A: np.ndarray,
b: np.ndarray,
*,
lower: bool = True,
trans: bool = False,
) -> np.ndarray:
"""Solve a triangular system Ax = b.
Parameters
----------
A : ndarray of shape (n, n)
Triangular matrix.
b : ndarray of shape (n,) or (n, k)
Right-hand side.
lower : bool, default=True
True if A is lower triangular, False if upper.
trans : bool, default=False
If True, solve A.T @ x = b instead.
Returns
-------
x : ndarray
Solution to the system.
"""
from scipy.linalg import solve_triangular as scipy_solve_tri
trans_arg = "T" if trans else "N"
return np.asarray(scipy_solve_tri(A, b, lower=lower, trans=trans_arg))
[docs]
def solve_cholesky(
L: np.ndarray,
b: np.ndarray,
*,
lower: bool = True,
) -> np.ndarray:
"""Solve a system using pre-computed Cholesky factor.
Solves A x = b where A = L @ L.T (lower=True) or A = U.T @ U (lower=False).
Parameters
----------
L : ndarray of shape (n, n)
Cholesky factor.
b : ndarray of shape (n,) or (n, k)
Right-hand side.
lower : bool, default=True
True if L is lower triangular.
Returns
-------
x : ndarray
Solution to the system.
Examples
--------
>>> A = np.array([[4, 2], [2, 5]])
>>> L = safe_cholesky(A)
>>> b = np.array([1, 2])
>>> x = solve_cholesky(L, b)
>>> np.allclose(A @ x, b)
True
"""
# Solve L @ y = b
y = solve_triangular(L, b, lower=lower, trans=False)
# Solve L.T @ x = y
x = solve_triangular(L, y, lower=lower, trans=True)
return x
[docs]
def solve_qr(
Q: np.ndarray,
R: np.ndarray,
b: np.ndarray,
) -> np.ndarray:
"""Solve a least squares problem using QR factors.
Solves min ||Ax - b||² where A = QR.
Parameters
----------
Q : ndarray of shape (m, k)
Orthogonal factor.
R : ndarray of shape (k, n)
Upper triangular factor.
b : ndarray of shape (m,)
Right-hand side.
Returns
-------
x : ndarray of shape (n,)
Least squares solution.
"""
# x = R^{-1} @ Q.T @ b
Qtb = Q.T @ b
x = solve_triangular(R, Qtb, lower=False)
return x
[docs]
def lstsq(
A: np.ndarray,
b: np.ndarray,
*,
rcond: float | None = None,
) -> tuple[np.ndarray, np.ndarray, int, np.ndarray]:
"""Least squares solution to Ax = b.
Parameters
----------
A : ndarray of shape (m, n)
Design matrix.
b : ndarray of shape (m,) or (m, k)
Right-hand side.
rcond : float, optional
Cutoff ratio for small singular values.
Returns
-------
x : ndarray
Least squares solution.
residuals : ndarray
Sum of squared residuals (if m > n).
rank : int
Effective rank of A.
s : ndarray
Singular values of A.
"""
return np.linalg.lstsq(A, b, rcond=rcond)
# =============================================================================
# Matrix Operations
# =============================================================================
[docs]
def safe_inverse(
A: np.ndarray,
*,
rcond: float = 1e-15,
) -> np.ndarray:
"""Compute matrix inverse with regularization for near-singular matrices.
Uses SVD-based pseudoinverse with condition number thresholding.
Parameters
----------
A : ndarray of shape (n, n)
Matrix to invert.
rcond : float, default=1e-15
Cutoff for small singular values. Singular values smaller than
rcond * largest_singular_value are set to zero.
Returns
-------
A_inv : ndarray of shape (n, n)
(Pseudo)inverse of A.
Examples
--------
>>> A = np.array([[1, 2], [2, 4.001]]) # Nearly singular
>>> A_inv = safe_inverse(A)
"""
return np.linalg.pinv(A, rcond=rcond)
[docs]
def woodbury_inverse(
A_inv: np.ndarray,
U: np.ndarray,
C: np.ndarray,
V: np.ndarray,
) -> np.ndarray:
"""Woodbury matrix identity for efficient inverse updates.
Computes (A + UCV)^{-1} given A^{-1} efficiently.
The Woodbury identity states:
(A + UCV)^{-1} = A^{-1} - A^{-1}U(C^{-1} + VA^{-1}U)^{-1}VA^{-1}
Parameters
----------
A_inv : ndarray of shape (n, n)
Inverse of A.
U : ndarray of shape (n, k)
Update matrix.
C : ndarray of shape (k, k)
Core matrix.
V : ndarray of shape (k, n)
Update matrix.
Returns
-------
result : ndarray of shape (n, n)
(A + UCV)^{-1}
Examples
--------
>>> A = np.eye(3)
>>> U = np.array([[1], [0], [0]])
>>> C = np.array([[1]])
>>> V = np.array([[1, 0, 0]])
>>> new_inv = woodbury_inverse(np.eye(3), U, C, V)
"""
# A^{-1}U
AiU = A_inv @ U
# VA^{-1}
VAi = V @ A_inv
# C^{-1} + VA^{-1}U
inner = np.linalg.inv(C) + VAi @ U
# (C^{-1} + VA^{-1}U)^{-1}
inner_inv = np.linalg.inv(inner)
# A^{-1} - A^{-1}U(...)VA^{-1}
return np.asarray(A_inv - AiU @ inner_inv @ VAi)
[docs]
def log_determinant(
A: np.ndarray,
*,
method: Literal["cholesky", "svd", "auto"] = "auto",
) -> float:
"""Compute log of matrix determinant in a numerically stable way.
Parameters
----------
A : ndarray of shape (n, n)
Positive definite matrix.
method : {'cholesky', 'svd', 'auto'}, default='auto'
Method to use. 'auto' tries Cholesky first, falls back to SVD.
Returns
-------
logdet : float
Log of determinant of A.
Examples
--------
>>> A = np.array([[2, 1], [1, 2]])
>>> logdet = log_determinant(A)
>>> np.isclose(logdet, np.log(np.linalg.det(A)))
True
"""
if method == "cholesky" or method == "auto":
try:
L = np.linalg.cholesky(A)
return float(2 * np.sum(np.log(np.diag(L))))
except np.linalg.LinAlgError:
if method == "cholesky":
raise
# Fall through to SVD
# SVD method
_, s, _ = np.linalg.svd(A)
return float(np.sum(np.log(s)))
[docs]
def matrix_rank(
A: np.ndarray,
*,
tol: float | None = None,
) -> int:
"""Compute matrix rank.
Parameters
----------
A : ndarray
Input matrix.
tol : float, optional
Threshold below which singular values are considered zero.
Returns
-------
rank : int
Numerical rank of the matrix.
"""
return int(np.linalg.matrix_rank(A, tol=tol))
[docs]
def condition_number(
A: np.ndarray,
*,
p: float | Literal["fro", "nuc"] | None = 2,
) -> float:
"""Compute matrix condition number.
Parameters
----------
A : ndarray
Input matrix.
p : {1, 2, inf, 'fro'}, default=2
Order of the norm.
Returns
-------
cond : float
Condition number of A.
"""
return float(np.linalg.cond(A, p=p))
[docs]
def weighted_condition_number(
X: np.ndarray,
weights: np.ndarray | None = None,
) -> float:
"""Compute condition number of weighted design matrix X^T W X.
Parameters
----------
X : ndarray, shape (n, p)
Design matrix.
weights : ndarray, shape (n,), optional
Observation weights. If None, computes condition of X^T X.
Returns
-------
kappa : float
Condition number κ = σ_max / σ_min of X^T W X.
"""
if weights is not None:
W_sqrt = np.sqrt(weights)
XtWX = (W_sqrt[:, None] * X).T @ (W_sqrt[:, None] * X)
else:
XtWX = X.T @ X
try:
s = np.linalg.svd(XtWX, compute_uv=False)
if s[-1] == 0:
return np.inf
return float(s[0] / s[-1])
except np.linalg.LinAlgError:
return np.inf
__all__ = [
# Decompositions
"qr_decomposition",
"safe_cholesky",
"svd",
"eigh",
# Solvers
"solve_triangular",
"solve_cholesky",
"solve_qr",
"lstsq",
# Matrix operations
"safe_inverse",
"woodbury_inverse",
"quadratic_form",
"log_determinant",
"matrix_rank",
"condition_number",
"weighted_condition_number",
]