Chebyshev iteration in Python
Accelerated stationary iteration using Chebyshev polynomials and eigenvalue bounds.
- File
chebyshev_iteration.py, 97 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/chebyshev-iteration/. The authors are listed in prov/3W6N6C.jsonl.
"""Chebyshev iteration for A x = b when the spectrum of A lies in [lmin, lmax], 0 < lmin.
Three-term form of Saad, Algorithm 12.1. No inner products are needed for
the iteration itself; the residual norm is computed only for the stopping
test, every `check` steps. Works with floats or fractions.Fraction.
Standard library only.
"""
class CSR:
"""Compressed sparse row matrix: row i holds data[ptr[i]:ptr[i+1]]."""
def __init__(self, n, ptr, ind, data):
self.n, self.ptr, self.ind, self.data = n, ptr, ind, data
@classmethod
def from_dense(cls, a):
ptr, ind, data = [0], [], []
for row in a:
for j, v in enumerate(row):
if v != 0:
ind.append(j)
data.append(v)
ptr.append(len(ind))
return cls(len(a), ptr, ind, data)
def matvec(self, x):
out = []
for i in range(self.n):
acc = 0 * x[0]
for p in range(self.ptr[i], self.ptr[i + 1]):
acc += self.data[p] * x[self.ind[p]]
out.append(acc)
return out
def _mv(a, x):
if isinstance(a, CSR):
return a.matvec(x)
return [sum((aij * xj for aij, xj in zip(row, x)), 0 * x[0]) for row in a]
class ChebyshevResult:
def __init__(self, x, converged, iterations, status):
self.x = x
self.converged = converged
self.iterations = iterations
self.status = status # "converged", "maxit" or "diverged"
def chebyshev(a, b, lmin, lmax, x0=None, tol=1e-10, maxit=1000, check=1, trace=None):
"""Iterate until ||r||_2 <= tol ||b||_2 (squared norms are compared).
trace, if a list, receives x_k for k = 0, 1, ... (for exact tests).
Raises ValueError unless 0 < lmin <= lmax.
"""
if not (0 < lmin <= lmax):
raise ValueError("need 0 < lmin <= lmax; Chebyshev iteration needs a definite spectrum")
n = len(b)
zero = 0 * b[0]
x = [zero] * n if x0 is None else list(x0)
theta = (lmax + lmin) / 2 # @line L1
delta = (lmax - lmin) / 2
r = [bi - t for bi, t in zip(b, _mv(a, x))] # @line L2
d = [ri / theta for ri in r]
if delta != 0:
sigma = theta / delta
rho = 1 / sigma
bb = sum((t * t for t in b), zero)
target = tol * tol * bb
rr0 = sum((t * t for t in r), zero)
if trace is not None:
trace.append(list(x))
for k in range(maxit + 1): # @line L3
if k % check == 0 or k == maxit:
rr = sum((t * t for t in r), zero)
if rr <= target: # @line L4
return ChebyshevResult(x, True, k, "converged")
if rr > 1e12 * (rr0 + bb) and rr0 + bb > 0:
return ChebyshevResult(x, False, k, "diverged")
if k == maxit:
break
ad = _mv(a, d) # @line L5
x = [xi + di for xi, di in zip(x, d)] # @line L6
r = [ri - t for ri, t in zip(r, ad)] # @line L7
if trace is not None:
trace.append(list(x))
if delta == 0:
d = [ri / theta for ri in r]
continue
rho_new = 1 / (2 * sigma - rho) # @line L8
d = [rho_new * rho * di + (2 * rho_new / delta) * ri for di, ri in zip(d, r)] # @line L9
rho = rho_new
return ChebyshevResult(x, False, maxit, "maxit")
test_chebyshev_iteration.py 145 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/chebyshev-iteration/. The authors are listed in prov/3W6N6C.jsonl.
import math
import random
from fractions import Fraction as F
import pytest
from chebyshev_iteration import CSR, chebyshev
def cheb(k, t):
"""T_k(t) by the three-term recurrence (exact for Fractions)."""
a, b = 1, t
if k == 0:
return a
for _ in range(k - 1):
a, b = b, 2 * t * b - a
return b
def mv(a, x):
return [sum(p * q for p, q in zip(row, x)) for row in a]
def test_page_example_exact():
a = [[F(2), F(-1)], [F(-1), F(2)]]
b = [F(1), F(1)]
tr = []
res = chebyshev(a, b, F(1), F(3), tol=0, maxit=4, trace=tr)
errs = [1 - xk[0] for xk in tr]
assert errs == [1, F(1, 2), F(1, 7), F(1, 26), F(1, 97)]
assert all(xk[0] == xk[1] for xk in tr)
assert not res.converged and res.status == "maxit"
def test_one_by_one_and_equal_bounds():
res = chebyshev([[F(4)]], [F(2)], F(4), F(4), tol=0)
assert res.converged and res.iterations == 1 and res.x == [F(1, 2)]
# A = c I with exact bounds: one step
res = chebyshev([[F(3), 0], [0, F(3)]], [F(1), F(2)], F(3), F(3), tol=0)
assert res.converged and res.iterations == 1
def test_random_diagonal_error_is_chebyshev_polynomial():
rng = random.Random(31337)
for trial in range(500):
n = rng.randint(1, 4)
lam = [F(rng.randint(1, 40), rng.randint(1, 4)) for _ in range(n)]
lo = min(lam) - F(rng.randint(0, 2), 8)
if lo <= 0:
lo = min(lam)
hi = max(lam) + F(rng.randint(0, 3), 4)
if lo == hi:
hi = lo + 1
a = [[lam[i] if i == j else F(0) for j in range(n)] for i in range(n)]
xt = [F(rng.randint(-5, 5)) for _ in range(n)]
b = mv(a, xt)
k = rng.randint(1, 6)
tr = []
chebyshev(a, b, lo, hi, tol=0, maxit=k, trace=tr)
theta, delta = (hi + lo) / 2, (hi - lo) / 2
for j, xj in enumerate(tr):
tk = cheb(j, theta / delta)
for i in range(n):
expected = cheb(j, (theta - lam[i]) / delta) / tk * xt[i]
assert xt[i] - xj[i] == expected
# the error bound 1 / T_k(sigma) in the 2-norm
e2 = sum((xt[i] - xj[i]) ** 2 for i in range(n))
e02 = sum(t * t for t in xt)
assert e2 * tk * tk <= e02
def laplacian(n):
ptr, ind, data = [0], [], []
for i in range(n):
for j, v in ((i - 1, -1.0), (i, 2.0), (i + 1, -1.0)):
if 0 <= j < n:
ind.append(j)
data.append(v)
ptr.append(len(ind))
return CSR(n, ptr, ind, data)
def test_iteration_count_matches_theory():
n = 60
a = laplacian(n)
lmin = 2 - 2 * math.cos(math.pi / (n + 1))
lmax = 2 - 2 * math.cos(n * math.pi / (n + 1))
rng = random.Random(4)
xt = [rng.uniform(-1, 1) for _ in range(n)]
b = a.matvec(xt)
eps = 1e-8
res = chebyshev(a, b, lmin, lmax, tol=eps, maxit=5000)
kappa = lmax / lmin
q = (math.sqrt(kappa) - 1) / (math.sqrt(kappa) + 1)
# residual bound ||r_k|| <= 2 q^k ||r_0|| (symmetric A, exact bounds)
kbound = math.ceil(math.log(2 * math.sqrt(sum(t * t for t in b)) / (eps * math.sqrt(sum(t * t for t in b)))) / math.log(1 / q))
assert res.converged
assert res.iterations <= kbound
assert res.iterations >= kbound // 3 # the bound is not wildly pessimistic here
assert res.iterations <= math.ceil(math.sqrt(kappa) / 2 * math.log(2 / eps))
def test_underestimated_lmax_diverges():
n = 40
a = laplacian(n)
b = [1.0] * n
res = chebyshev(a, b, 0.01, 2.0, tol=1e-10, maxit=3000)
assert not res.converged and res.status == "diverged"
def test_underestimated_lmin_slower_but_ok():
n = 40
a = laplacian(n)
b = [1.0] * n
lmin = 2 - 2 * math.cos(math.pi / (n + 1))
good = chebyshev(a, b, lmin, 4.0, tol=1e-8, maxit=5000)
loose = chebyshev(a, b, lmin / 10, 4.0, tol=1e-8, maxit=5000)
assert good.converged and loose.converged
assert loose.iterations > good.iterations
def test_invalid_bounds_and_singular():
with pytest.raises(ValueError):
chebyshev([[1.0]], [1.0], 0.0, 1.0)
with pytest.raises(ValueError):
chebyshev([[1.0]], [1.0], -1.0, 1.0)
# singular A with a bound lmin > 0 that is wrong: no convergence
res = chebyshev([[1.0, 0.0], [0.0, 0.0]], [1.0, 1.0], 0.5, 1.0, maxit=300)
assert not res.converged
def test_hilbert_with_gershgorin_bound():
n = 4
h = [[1.0 / (i + j + 1) for j in range(n)] for i in range(n)]
xt = [1.0] * n
b = mv(h, xt)
lmax = max(sum(row) for row in h) # Gershgorin / row sums
lmin = 9.6e-5 # below lambda_min(H_4), which is about 9.670e-5
res = chebyshev(h, b, lmin, lmax, tol=1e-12, maxit=5000, check=10)
assert res.converged
r = [bi - t for bi, t in zip(b, mv(h, res.x))]
assert max(map(abs, r)) <= 1e-11 * max(map(abs, b))
What the next level needs
- the tests did not run in this build