Banded matrix in Python
Matrix stored only within a fixed band around the diagonal.
- File
banded_matrix.py, 166 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/banded-matrix/. The authors are listed in prov/HE5FHC.jsonl.
"""A banded matrix in diagonal storage, with matvec, band LU and band LU with partial pivoting.
Entry a[i][j] with -q <= j - i <= p ... is kept in ab[q + i - j][j], so row d of `ab`
holds one diagonal: d = q is the main diagonal, d < q the superdiagonals, d > q the
subdiagonals. This is the LAPACK general band layout, written with 0-based indices.
Only +, -, * and / are used, so entries may be fractions.Fraction.
"""
class SingularMatrix(ValueError):
pass
class BandedMatrix:
"""n by n matrix with a_ij = 0 whenever i - j > p or j - i > q."""
def __init__(self, n, p, q, zero=0):
if n < 0 or p < 0 or q < 0:
raise ValueError("n, p and q must be nonnegative")
self.n, self.p, self.q = n, p, q
self.zero = zero
self.ab = [[zero] * n for _ in range(p + q + 1)]
@staticmethod
def from_dense(a, p, q):
n = len(a)
zero = 0 * a[0][0] if n else 0
m = BandedMatrix(n, p, q, zero)
for i in range(n):
for j in range(n):
if a[i][j] != 0:
m.set(i, j, a[i][j])
return m
def in_band(self, i, j):
return 0 <= i < self.n and 0 <= j < self.n and -self.q <= i - j <= self.p
def get(self, i, j):
if not (0 <= i < self.n and 0 <= j < self.n):
raise IndexError((i, j))
return self.ab[self.q + i - j][j] if self.in_band(i, j) else self.zero
def set(self, i, j, v):
if not self.in_band(i, j):
if v != 0:
raise ValueError(f"entry ({i}, {j}) lies outside the band")
return
self.ab[self.q + i - j][j] = v
def to_dense(self):
return [[self.get(i, j) for j in range(self.n)] for i in range(self.n)]
def copy(self, q=None):
"""Copy, optionally with a wider upper bandwidth (room for fill from pivoting)."""
m = BandedMatrix(self.n, self.p, self.q if q is None else q, self.zero)
for i in range(self.n):
for j in range(max(0, i - self.p), min(self.n, i + self.q + 1)):
m.set(i, j, self.get(i, j))
return m
def matvec(self, x):
"""y = A x in 2 n (p + q + 1) flops at most."""
n, p, q = self.n, self.p, self.q
y = []
for i in range(n): # @line M1
s = 0 * x[0]
for j in range(max(0, i - p), min(n, i + q + 1)): # @line M2
s += self.ab[q + i - j][j] * x[j] # @line M3
y.append(s)
return y
def lu(self):
"""Band LU without pivoting. Returns a BandLU; raises SingularMatrix on a zero pivot."""
w = self.copy()
n, p, q = w.n, w.p, w.q
for k in range(n): # @line L1
piv = w.get(k, k)
if piv == 0: # @line L2
raise SingularMatrix(f"zero pivot at step {k}")
for i in range(k + 1, min(k + p, n - 1) + 1): # @line L3
w.set(i, k, w.get(i, k) / piv) # @line L4
for j in range(k + 1, min(k + q, n - 1) + 1): # @line L5
ukj = w.get(k, j)
for i in range(k + 1, min(k + p, n - 1) + 1): # @line L6
w.set(i, j, w.get(i, j) - w.get(i, k) * ukj) # @line L7
return BandLU(w, None) # @line L8
def lu_pivoted(self):
"""Band LU with partial pivoting: P A = L U, with U of upper bandwidth p + q.
The interchanges are recorded in piv and applied step by step in solve, as in LINPACK;
the multipliers of step k stay in column k. Raises SingularMatrix if A is singular.
"""
n, p, q = self.n, self.p, self.q
w = self.copy(q=p + q)
piv = []
for k in range(n):
last = min(k + p, n - 1)
r = max(range(k, last + 1), key=lambda i: abs(w.get(i, k)))
if w.get(r, k) == 0:
raise SingularMatrix(f"no nonzero pivot in column {k}")
piv.append(r)
right = min(k + p + q, n - 1)
if r != k:
for j in range(k, right + 1):
a, b = w.get(k, j), w.get(r, j)
w.set(k, j, b)
w.set(r, j, a)
for i in range(k + 1, last + 1):
w.set(i, k, w.get(i, k) / w.get(k, k))
for j in range(k + 1, right + 1):
ukj = w.get(k, j)
if ukj != 0:
for i in range(k + 1, last + 1):
w.set(i, j, w.get(i, j) - w.get(i, k) * ukj)
return BandLU(w, piv)
def check_invariant(self):
assert len(self.ab) == self.p + self.q + 1
assert all(len(row) == self.n for row in self.ab)
# Slots of ab that do not correspond to matrix entries must hold zero.
for d in range(self.p + self.q + 1):
for j in range(self.n):
i = d - self.q + j
if not 0 <= i < self.n:
assert self.ab[d][j] == 0
return True
class BandLU:
"""Factors stored in band form: multipliers below the diagonal, U on and above it."""
def __init__(self, w, piv):
self.w, self.piv = w, piv
def solve(self, b):
w, piv = self.w, self.piv
n, p, q = w.n, w.p, w.q
y = list(b)
for k in range(n): # @line S1
if piv is not None and piv[k] != k:
y[k], y[piv[k]] = y[piv[k]], y[k]
for i in range(k + 1, min(k + p, n - 1) + 1): # @line S2
y[i] -= w.get(i, k) * y[k]
x = [None] * n
for i in range(n - 1, -1, -1): # @line S3
s = y[i]
for j in range(i + 1, min(i + q, n - 1) + 1):
s -= w.get(i, j) * x[j]
x[i] = s / w.get(i, i) # @line S4
return x
def lower(self):
"""Unit lower triangular L (no pivoting only), as a dense list of rows."""
if self.piv is not None:
raise ValueError("L is not formed explicitly for the pivoted factorization")
n = self.w.n
one = self.w.zero + 1
return [[one if i == j else (self.w.get(i, j) if j < i else self.w.zero) for j in range(n)] for i in range(n)]
def upper(self):
n = self.w.n
return [[self.w.get(i, j) if j >= i else self.w.zero for j in range(n)] for i in range(n)]
test_banded_matrix.py 188 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/banded-matrix/. The authors are listed in prov/HE5FHC.jsonl.
import random
from fractions import Fraction
import pytest
from banded_matrix import BandedMatrix, SingularMatrix
U = Fraction(1, 2**53)
def gamma(m):
return m * U / (1 - m * U)
def exact_solve(a, b):
n = len(b)
m = [[Fraction(v) for v in row] + [Fraction(bi)] for row, bi in zip(a, b)]
for k in range(n):
p = next(i for i in range(k, n) if m[i][k] != 0)
m[k], m[p] = m[p], m[k]
for i in range(k + 1, n):
f = m[i][k] / m[k][k]
for j in range(k, n + 1):
m[i][j] -= f * m[k][j]
x = [Fraction(0)] * n
for i in range(n - 1, -1, -1):
x[i] = (m[i][n] - sum(m[i][j] * x[j] for j in range(i + 1, n))) / m[i][i]
return x
def dense_lu_nopivot(a):
"""Doolittle LU in exact arithmetic: the oracle for the band factorization."""
n = len(a)
w = [[Fraction(v) for v in row] for row in a]
for k in range(n):
for i in range(k + 1, n):
w[i][k] /= w[k][k]
for j in range(k + 1, n):
w[i][j] -= w[i][k] * w[k][j]
return w
def matmul(x, y):
n = len(x)
return [[sum(x[i][k] * y[k][j] for k in range(n)) for j in range(n)] for i in range(n)]
def random_band(rng, n, p, q, dominant=False):
a = [[rng.uniform(-1, 1) for j in range(n)] for i in range(n)]
a = [[a[i][j] if (i - j <= p and j - i <= q) else 0.0 for j in range(n)] for i in range(n)]
if dominant:
for i in range(n):
a[i][i] = rng.choice([-1, 1]) * (sum(abs(a[i][j]) for j in range(n) if j != i) + rng.uniform(0.1, 1))
return a
def test_page_example():
a = [[2, -1, 0, 0], [-1, 2, -1, 0], [0, -1, 2, -1], [0, 0, -1, 2]]
a = [[Fraction(v) for v in row] for row in a]
m = BandedMatrix.from_dense(a, 1, 1)
assert m.ab == [[0, -1, -1, -1], [2, 2, 2, 2], [-1, -1, -1, 0]]
f = m.lu()
assert [f.w.get(i, i) for i in range(4)] == [2, Fraction(3, 2), Fraction(4, 3), Fraction(5, 4)]
assert [f.w.get(i + 1, i) for i in range(3)] == [Fraction(-1, 2), Fraction(-2, 3), Fraction(-3, 4)]
assert f.solve([Fraction(1)] * 4) == [2, 3, 3, 2]
b = BandedMatrix.from_dense([[Fraction(v) for v in r] for r in [[1, 1, 0], [4, 1, 1], [0, 4, 1]]], 1, 1)
g = b.lu_pivoted()
assert g.piv == [1, 2, 2]
assert g.upper() == [[4, 1, 1], [0, 4, 1], [0, 0, Fraction(-7, 16)]]
assert g.solve([Fraction(1), Fraction(2), Fraction(3)]) == [Fraction(2, 7), Fraction(5, 7), Fraction(1, 7)]
def test_access_and_invariant():
m = BandedMatrix(5, 1, 2)
m.set(3, 2, 7.0)
m.set(1, 3, 2.0)
assert m.get(3, 2) == 7.0 and m.get(1, 3) == 2.0 and m.get(4, 0) == 0
with pytest.raises(ValueError):
m.set(3, 0, 1.0)
m.set(3, 0, 0) # zero outside the band is fine
with pytest.raises(IndexError):
m.get(5, 0)
assert m.check_invariant()
def test_edges():
assert BandedMatrix.from_dense([[4.0]], 0, 0).lu().solve([2.0]) == [0.5]
d = BandedMatrix.from_dense([[2.0, 0, 0], [0, -4.0, 0], [0, 0, 0.5]], 0, 0)
assert d.lu().solve([1.0, 1.0, 1.0]) == [0.5, -0.25, 2.0]
lo = [[Fraction(v) for v in r] for r in [[2, 0, 0], [1, 3, 0], [0, -1, 5]]]
assert BandedMatrix.from_dense(lo, 1, 0).lu().solve([Fraction(1)] * 3) == exact_solve(lo, [1, 1, 1])
with pytest.raises(SingularMatrix):
BandedMatrix.from_dense([[0.0, 1.0], [1.0, 0.0]], 1, 1).lu()
assert BandedMatrix.from_dense([[0.0, 1.0], [1.0, 0.0]], 1, 1).lu_pivoted().solve([1.0, 2.0]) == [2.0, 1.0]
sing = [[1.0, 1.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 1.0]]
with pytest.raises(SingularMatrix):
BandedMatrix.from_dense(sing, 1, 1).lu_pivoted()
def test_matvec_against_dense():
rng = random.Random(1)
for _ in range(200):
n, p, q = rng.randint(1, 9), rng.randint(0, 3), rng.randint(0, 3)
a = random_band(rng, n, p, q)
x = [Fraction(rng.randint(-5, 5)) for _ in range(n)]
af = [[Fraction(v) for v in row] for row in a]
m = BandedMatrix.from_dense(af, p, q)
assert m.matvec(x) == [sum(af[i][j] * x[j] for j in range(n)) for i in range(n)]
assert m.to_dense() == af and m.check_invariant()
def test_exact_lu_equals_dense_elimination():
rng = random.Random(2)
for _ in range(150):
n, p, q = rng.randint(1, 8), rng.randint(0, 3), rng.randint(0, 3)
a = [[Fraction(v).limit_denominator(50) for v in row] for row in random_band(rng, n, p, q, dominant=True)]
f = BandedMatrix.from_dense(a, p, q).lu()
w = dense_lu_nopivot(a)
for i in range(n):
for j in range(n):
assert f.w.get(i, j) == w[i][j] # no fill outside the band: w is zero there
assert matmul(f.lower(), f.upper()) == a
def test_float_backward_error_bound():
"""|A - L U| <= gamma_{min(p,q)+1} |L| |U| for the computed factors (500 random matrices)."""
rng = random.Random(2027)
for _ in range(500):
n, p, q = rng.randint(1, 10), rng.randint(0, 3), rng.randint(0, 3)
a = random_band(rng, n, p, q, dominant=rng.random() < 0.7)
try:
f = BandedMatrix.from_dense(a, p, q).lu()
except SingularMatrix:
continue
lo = [[Fraction(v) for v in row] for row in f.lower()]
up = [[Fraction(v) for v in row] for row in f.upper()]
lu = matmul(lo, up)
alu = matmul([[abs(v) for v in r] for r in lo], [[abs(v) for v in r] for r in up])
g = gamma(min(p, q) + 1)
for i in range(n):
for j in range(n):
assert abs(Fraction(a[i][j]) - lu[i][j]) <= g * alu[i][j]
def test_pivoted_exact_and_fill_bound():
rng = random.Random(3)
for _ in range(150):
n, p, q = rng.randint(1, 8), rng.randint(0, 3), rng.randint(0, 3)
a = [[Fraction(v).limit_denominator(20) for v in row] for row in random_band(rng, n, p, q)]
b = [Fraction(rng.randint(-5, 5)) for _ in range(n)]
try:
x = exact_solve(a, b)
except StopIteration:
with pytest.raises(SingularMatrix):
BandedMatrix.from_dense(a, p, q).lu_pivoted()
continue
f = BandedMatrix.from_dense(a, p, q).lu_pivoted()
assert f.solve(b) == x
up = f.upper()
assert all(up[i][j] == 0 for i in range(n) for j in range(n) if j - i > p + q)
def test_pivoted_float_backward_error_and_tridiagonal_growth():
rng = random.Random(4)
for _ in range(300):
n, p, q = rng.randint(2, 12), rng.randint(1, 3), rng.randint(1, 3)
a = random_band(rng, n, p, q)
try:
f = BandedMatrix.from_dense(a, p, q).lu_pivoted()
except SingularMatrix:
continue
b = [rng.uniform(-1, 1) for _ in range(n)]
x = f.solve(b)
r = [Fraction(bi) - sum(Fraction(a[i][j]) * Fraction(x[j]) for j in range(n)) for i, bi in enumerate(b)]
an = max(sum(abs(Fraction(v)) for v in row) for row in a)
eta = max(abs(v) for v in r) / (an * max(abs(Fraction(v)) for v in x) + max(abs(Fraction(v)) for v in b))
assert eta <= 100 * (p + q + 1) * U
if p == q == 1:
growth = max(abs(v) for row in f.upper() for v in row) / max(abs(v) for row in a for v in row)
assert growth <= 2 * (1 + 1e-12)
def test_storage_is_linear_in_n():
m = BandedMatrix(10000, 2, 3)
assert sum(len(row) for row in m.ab) == 6 * 10000
What the next level needs
- the tests did not run in this build