LU decomposition in Python
Factors a matrix into a unit lower triangular and an upper triangular factor, reusable across many right-hand sides.
- File
lu_decomposition.py, 107 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/lu-decomposition/. The authors are listed in prov/EH2SQ6.jsonl.
"""Exact LU decomposition with row swaps, P A = L U, over Q or over Z/pZ."""
from fractions import Fraction
from typing import Optional, Union
Number = Union[int, Fraction]
Matrix = list[list[Number]]
def _norm(x: Number, p: Optional[int]) -> Number:
return Fraction(x) if p is None else int(x) % p
def _div(x: Number, y: Number, p: Optional[int]) -> Number:
if p is None:
return Fraction(x) / Fraction(y)
return int(x) * pow(int(y), -1, p) % p
def lu_decomposition(a: Matrix, p: Optional[int] = None) -> tuple[list[int], Matrix, Matrix]:
"""Return (perm, L, U) with P A = L U.
a is an m by n matrix (a list of m rows of length n) over Q (p is None,
entries are int or Fraction) or over Z/pZ for a prime p. Row r of P A is
row perm[r] of A. L is m by m unit lower triangular, U is m by n in row
echelon form, and its number of nonzero rows is the rank of A. The first
nonzero entry of a column is the pivot, so the result is deterministic.
Cost: O(m n min(m, n)) field operations.
"""
m = len(a)
n = len(a[0]) if m else 0
u = [[_norm(x, p) for x in row] for row in a] # @line L1
one, zero = _norm(1, p), _norm(0, p)
low = [[one if i == j else zero for j in range(m)] for i in range(m)]
perm = list(range(m))
r = 0
for c in range(n): # @line L2
if r >= m: # @line L3
break
piv = next((i for i in range(r, m) if u[i][c] != 0), None) # @line L4
if piv is None:
continue
if piv != r: # @line L5
u[r], u[piv] = u[piv], u[r]
perm[r], perm[piv] = perm[piv], perm[r]
for k in range(r):
low[r][k], low[piv][k] = low[piv][k], low[r][k]
for i in range(r + 1, m): # @line L6
if u[i][c] == 0:
continue
f = _div(u[i][c], u[r][c], p) # @line L7
low[i][r] = f
row_i, row_r = u[i], u[r]
row_i[c] = zero # @line L8
for j in range(c + 1, n):
if p is None:
row_i[j] = row_i[j] - f * row_r[j]
else:
row_i[j] = (row_i[j] - f * row_r[j]) % p
r += 1 # @line L9
return perm, low, u # @line L10
def lu_solve(perm: list[int], low: Matrix, u: Matrix, b: list[Number], p: Optional[int] = None) -> list[Number]:
"""Solve A x = b from P A = L U, for a square nonsingular A.
Forward substitution gives y with L y = P b, and back substitution gives
x with U x = y. Raises ZeroDivisionError or ValueError when U has a zero
diagonal entry, that is when A is singular. Cost: O(n^2) field operations.
"""
n = len(perm)
y: list[Number] = []
for i in range(n): # @line L11
s = _norm(b[perm[i]], p)
for k in range(i):
s -= low[i][k] * y[k]
y.append(_norm(s, p))
x: list[Number] = [_norm(0, p)] * n
for i in range(n - 1, -1, -1): # @line L12
s = y[i]
for k in range(i + 1, n):
s -= u[i][k] * x[k]
x[i] = _div(s, u[i][i], p)
return x # @line L13
def lu_determinant(a: Matrix, p: Optional[int] = None) -> Number:
"""Return det A for a square matrix, as sign(P) times the product of the diagonal of U."""
perm, _, u = lu_decomposition(a, p)
sign = 1
seen = [False] * len(perm)
for i in range(len(perm)):
if not seen[i]:
j, length = i, 0
while not seen[j]:
seen[j] = True
j = perm[j]
length += 1
if length % 2 == 0:
sign = -sign
d = _norm(sign, p)
for i in range(len(u)):
d = _norm(d * u[i][i], p)
return d
test_lu_decomposition.py 130 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/lu-decomposition/. The authors are listed in prov/EH2SQ6.jsonl.
import itertools
import random
from fractions import Fraction
import pytest
from lu_decomposition import lu_decomposition, lu_determinant, lu_solve
def leibniz(a, p=None):
n = len(a)
total = 0
for s in itertools.permutations(range(n)):
inv = sum(1 for i in range(n) for j in range(i + 1, n) if s[i] > s[j])
t = -1 if inv % 2 else 1
for i in range(n):
t *= a[i][s[i]]
total += t
return total if p is None else total % p
def brute_rank(a, p=None):
m = len(a)
n = len(a[0]) if m else 0
for k in range(min(m, n), 0, -1):
for rows in itertools.combinations(range(m), k):
for cols in itertools.combinations(range(n), k):
if leibniz([[a[i][j] for j in cols] for i in rows], p) != 0:
return k
return 0
def matmul(x, y, p=None):
inner = len(y)
cols = len(y[0]) if y else 0
out = [[sum((x[i][k] * y[k][j] for k in range(inner)), Fraction(0)) for j in range(cols)] for i in range(len(x))]
return out if p is None else [[int(v) % p for v in row] for row in out]
def check(a, p=None):
m = len(a)
n = len(a[0]) if m else 0
perm, low, u = lu_decomposition(a, p)
assert sorted(perm) == list(range(m))
for i in range(m):
assert low[i][i] == 1
assert all(low[i][j] == 0 for j in range(i + 1, m))
pa = [[a[perm[i]][j] if p is None else a[perm[i]][j] % p for j in range(n)] for i in range(m)]
assert matmul(low, u, p) == pa
lead = [next((j for j in range(n) if u[i][j] != 0), None) for i in range(m)]
rank = sum(1 for x in lead if x is not None)
assert all(x is None for x in lead[rank:])
assert all(lead[i] < lead[i + 1] for i in range(rank - 1))
return perm, low, u, rank
def test_worked_example():
a = [[2, 1, 1, 0], [4, 2, 3, 1], [-2, 0, 1, 3], [6, 3, 5, 4]]
perm, low, u, rank = check(a)
assert perm == [0, 2, 1, 3]
assert low == [[1, 0, 0, 0], [-1, 1, 0, 0], [2, 0, 1, 0], [3, 0, 2, 1]]
assert u == [[2, 1, 1, 0], [0, 1, 2, 3], [0, 0, 1, 1], [0, 0, 0, 2]]
assert lu_solve(perm, low, u, [3, 8, 0, 13]) == [1, -1, 2, 0]
assert lu_determinant(a) == -4 == leibniz(a)
def test_edges():
assert lu_decomposition([]) == ([], [], [])
assert lu_decomposition([[]]) == ([0], [[1]], [[]])
assert lu_decomposition([[5]]) == ([0], [[1]], [[5]])
assert lu_determinant([]) == 1
assert lu_solve(*lu_decomposition([[3]]), [6]) == [2]
assert check([[0, 0], [0, 0]])[3] == 0
perm, low, u, rank = check([[0, 1], [1, 1]])
assert perm == [1, 0] and lu_determinant([[0, 1], [1, 1]]) == -1
assert check([[0, 0, 1], [0, 2, 3], [0, 4, 5]])[3] == 2
assert check([[1, 2, 3, 4]])[3] == 1
assert check([[0], [0], [7]])[3] == 1
with pytest.raises(ZeroDivisionError):
lu_solve(*lu_decomposition([[1, 2], [2, 4]]), [1, 2])
def test_large_entries():
rng = random.Random(7)
for _ in range(20):
n = rng.randrange(1, 6)
a = [[rng.randrange(-2**80, 2**80) for _ in range(n)] for _ in range(n)]
perm, low, u, rank = check(a)
assert lu_determinant(a) == leibniz(a)
if rank == n:
x = [rng.randrange(-2**40, 2**40) for _ in range(n)]
b = [sum(r[j] * x[j] for j in range(n)) for r in a]
assert lu_solve(perm, low, u, b) == x
def test_random_over_q():
rng = random.Random(2024)
for _ in range(600):
m, n = rng.randrange(0, 5), rng.randrange(0, 5)
a = [[rng.choice([0, 0, 1, -1, 2, rng.randrange(-9, 10)]) for _ in range(n)] for _ in range(m)]
if m and rng.random() < 0.3:
a[rng.randrange(m)] = [0] * n
perm, low, u, rank = check(a)
if m and n:
assert rank == brute_rank(a)
if m == n:
assert lu_determinant(a) == leibniz(a)
if rank == n and n:
b = [rng.randrange(-9, 10) for _ in range(n)]
x = lu_solve(perm, low, u, b)
assert [sum(a[i][j] * x[j] for j in range(n)) for i in range(n)] == b
@pytest.mark.parametrize("p", [2, 3, 7])
def test_random_mod_p(p):
rng = random.Random(p)
for _ in range(300):
m, n = rng.randrange(1, 5), rng.randrange(1, 5)
a = [[rng.randrange(p) for _ in range(n)] for _ in range(m)]
perm, low, u, rank = check(a, p)
assert rank == brute_rank(a, p)
if m == n:
assert lu_determinant(a, p) == leibniz(a, p)
if rank == n:
b = [rng.randrange(p) for _ in range(n)]
x = lu_solve(perm, low, u, b, p)
assert [sum(a[i][j] * x[j] for j in range(n)) % p for i in range(n)] == b
What the next level needs
- the tests did not run in this build