Lanczos algorithm in Python
A three-term recurrence building an orthonormal basis of the Krylov space K_k(A,b) for a symmetric A and the tridiagonal T_k = Q_k* A Q_k, whose eigenvalues approximate the extreme eigenvalues of A.
- File
lanczos_algorithm.py, 104 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/lanczos-algorithm/. The authors are listed in prov/GKZVXW.jsonl.
"""Symmetric Lanczos algorithm with optional full reorthogonalization.
lanczos(matvec, b, k) returns (alpha, beta, q, breakdown): the diagonal alpha_1..alpha_m and
off-diagonal beta_1..beta_m of the tridiagonal T_m (beta_m couples to the next vector, so the
residual of a Ritz pair is beta_m |s_mi|), the orthonormal Lanczos vectors q_1..q_m, and whether
an invariant subspace was found (m < k). ritz(alpha, beta, q) returns the Ritz pairs with their
residual estimates. Standard library only; vectors are lists of floats.
"""
import math
U = 2.0 ** -53 # unit roundoff of IEEE double precision
def dot(x, y):
return sum(a * b for a, b in zip(x, y))
def norm2(x):
return math.sqrt(dot(x, x))
def lanczos(matvec, b, k, reorth=True):
"""Run k steps of the Lanczos recurrence on the symmetric operator matvec, start vector b."""
nb = norm2(b)
if nb == 0.0:
raise ValueError("start vector is zero")
q = [[t / nb for t in b]] # @line L1
alpha, beta = [], []
q_prev, b_prev = [0.0] * len(b), 0.0
for j in range(k): # @line L2
w = matvec(q[j])
w = [wi - b_prev * pi for wi, pi in zip(w, q_prev)] # @line L3
a = dot(q[j], w) # @line L4
w = [wi - a * qi for wi, qi in zip(w, q[j])]
if reorth: # @line L5: classical Gram-Schmidt, applied twice
for _ in range(2):
for qi in q:
h = dot(qi, w)
w = [wi - h * t for wi, t in zip(w, qi)]
alpha.append(a)
bj = norm2(w) # @line L6
beta.append(bj)
scale = max(abs(a), b_prev, 1.0)
if bj <= 10.0 * len(b) * U * scale:
return alpha, beta, q, True
if j + 1 < k:
q_prev, b_prev = q[j], bj
q.append([t / bj for t in w]) # @line L7
return alpha, beta, q, False
def tridiag_eig(alpha, beta):
"""Eigenvalues and eigenvectors of the symmetric tridiagonal T (diagonal alpha, off-diagonal beta[:-1])
by cyclic Jacobi on the dense matrix; adequate for the small T of Lanczos."""
m = len(alpha)
t = [[0.0] * m for _ in range(m)]
for i in range(m):
t[i][i] = alpha[i]
if i + 1 < m:
t[i][i + 1] = t[i + 1][i] = beta[i]
s = [[1.0 if i == j else 0.0 for j in range(m)] for i in range(m)]
for _ in range(60):
rotated = False
for p in range(m - 1):
for r in range(p + 1, m):
apq = t[p][r]
if apq == 0.0 or abs(apq) <= U * math.sqrt(abs(t[p][p] * t[r][r])) * 0.5:
continue
rotated = True
tau = (t[r][r] - t[p][p]) / (2.0 * apq)
tt = (1.0 if tau >= 0 else -1.0) / (abs(tau) + math.hypot(1.0, tau))
c = 1.0 / math.hypot(1.0, tt)
sn = tt * c
t[p][p] -= tt * apq
t[r][r] += tt * apq
t[p][r] = t[r][p] = 0.0
for i in range(m):
if i != p and i != r:
x, y = t[i][p], t[i][r]
t[i][p] = t[p][i] = c * x - sn * y
t[i][r] = t[r][i] = sn * x + c * y
for i in range(m):
x, y = s[i][p], s[i][r]
s[i][p] = c * x - sn * y
s[i][r] = sn * x + c * y
if not rotated:
break
order = sorted(range(m), key=lambda i: -t[i][i])
return [t[i][i] for i in order], [[s[r][i] for r in range(m)] for i in order]
def ritz(alpha, beta, q):
"""Ritz values theta_1 >= ... >= theta_m, Ritz vectors y_i = Q_m s_i, and the residual
estimates beta_m |s_mi| = ||A y_i - theta_i y_i||_2 (exact in exact arithmetic)."""
m = len(alpha)
theta, svecs = tridiag_eig(alpha, beta) # @line L8
n = len(q[0])
ys = [[sum(q[j][r] * s[j] for j in range(m)) for r in range(n)] for s in svecs]
est = [beta[m - 1] * abs(s[m - 1]) for s in svecs]
return theta, ys, est
test_lanczos_algorithm.py 135 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/lanczos-algorithm/. The authors are listed in prov/GKZVXW.jsonl.
import math
import random
from lanczos_algorithm import U, dot, lanczos, norm2, ritz
def diag_op(d):
return lambda x: [di * xi for di, xi in zip(d, x)]
def dense_op(a):
return lambda x: [sum(r[j] * x[j] for j in range(len(x))) for r in a]
def orth_error(q):
return max(abs(dot(q[i], q[j]) - (1.0 if i == j else 0.0)) for i in range(len(q)) for j in range(len(q)))
def chebyshev(m, x):
return math.cosh(m * math.acosh(x)) if x >= 1 else math.cos(m * math.acos(x))
def test_page_example():
n = 100
d = [float(i) for i in range(1, n + 1)]
want = {5: 4.205216, 10: 0.862128, 20: 0.050398, 30: 6.422876e-4, 40: 7.154212e-7}
for k, err in want.items():
al, be, q, br = lanczos(diag_op(d), [1.0] * n, k)
th, ys, est = ritz(al, be, q)
assert abs((100.0 - th[0]) - err) <= 1e-5 * err + 1e-12
assert abs((th[-1] - 1.0) - err) <= 1e-5 * err + 1e-12 # the spectrum is symmetric about 50.5
def test_kaniel_paige_saad_bound():
rng = random.Random(5)
for trial in range(6):
n = 60
d = sorted((rng.uniform(0, 10) for _ in range(n)), reverse=True)
b = [rng.gauss(0, 1) for _ in range(n)]
cos2 = b[0] ** 2 / dot(b, b)
tan2 = (1 - cos2) / cos2
gamma = (d[0] - d[1]) / (d[1] - d[-1])
for k in (5, 10, 15, 20):
al, be, q, br = lanczos(diag_op(d), b, k)
th = ritz(al, be, q)[0]
bound = (d[0] - d[-1]) * tan2 / chebyshev(k - 1, 1 + 2 * gamma) ** 2
err = d[0] - th[0]
assert -1e-12 <= err <= bound * (1 + 1e-8) + 1e-12
def test_three_term_relation_and_residuals():
rng = random.Random(9)
n, k = 40, 15
m = [[rng.gauss(0, 1) for _ in range(n)] for _ in range(n)]
a = [[m[i][j] + m[j][i] for j in range(n)] for i in range(n)]
op = dense_op(a)
al, be, q, br = lanczos(op, [rng.gauss(0, 1) for _ in range(n)], k + 1)
assert not br and orth_error(q) <= 10 * k * U * 10
for j in range(k): # A q_j = beta_{j-1} q_{j-1} + alpha_j q_j + beta_j q_{j+1}
aq = op(q[j])
rhs = [al[j] * q[j][i] + be[j] * q[j + 1][i] + (be[j - 1] * q[j - 1][i] if j else 0.0) for i in range(n)]
assert norm2([x - y for x, y in zip(aq, rhs)]) <= 1e3 * n * U * 10
al, be = al[:k], be[:k]
th, ys, est = ritz(al, be, q[:k])
for t, y, e in zip(th, ys, est):
r = [x - t * yy for x, yy in zip(op(y), y)]
assert abs(norm2(r) - e) <= 1e-10 * max(1.0, e) * 1e2
assert abs(norm2(y) - 1.0) <= 1e-12
def test_ritz_values_inside_spectrum():
rng = random.Random(13)
d = [rng.uniform(-3, 7) for _ in range(50)]
al, be, q, br = lanczos(diag_op(d), [1.0] * 50, 20)
th = ritz(al, be, q)[0]
assert min(d) - 1e-12 <= min(th) and max(th) <= max(d) + 1e-12
assert all(th[i] >= th[i + 1] for i in range(len(th) - 1))
def test_exact_breakdown_with_few_distinct_eigenvalues():
d = [1.0] * 5 + [2.0] * 7 + [5.0] * 3
al, be, q, br = lanczos(diag_op(d), [1.0] * 15, 10)
assert br and len(al) == 3
th = ritz(al, be, q)[0]
assert max(abs(x - y) for x, y in zip(th, [5.0, 2.0, 1.0])) < 1e-13
def test_full_n_steps_keep_trace_and_spectrum():
n = 12
a = [[2.0 if i == j else (-1.0 if abs(i - j) == 1 else 0.0) for j in range(n)] for i in range(n)]
al, be, q, br = lanczos(dense_op(a), [float(i + 1) for i in range(n)], n)
assert len(al) == n
assert abs(sum(al) - 2.0 * n) < 1e-12
th = ritz(al, be, q)[0]
want = sorted((2.0 - 2.0 * math.cos(k * math.pi / (n + 1)) for k in range(1, n + 1)), reverse=True)
assert max(abs(x - y) for x, y in zip(th, want)) < 1e-12
assert orth_error(q) <= 10 * n * U * 10
def test_start_vector_orthogonal_to_wanted_eigenvector():
"""The all-ones vector is orthogonal to the top eigenvector of the even-order Toeplitz matrix."""
n = 10
a = [[2.0 if i == j else (-1.0 if abs(i - j) == 1 else 0.0) for j in range(n)] for i in range(n)]
al, be, q, br = lanczos(dense_op(a), [1.0] * n, n)
th = ritz(al, be, q)[0]
top = 2.0 - 2.0 * math.cos(n * math.pi / (n + 1))
assert br and len(al) == n // 2
assert abs(th[0] - top) > 1e-3
def test_loss_of_orthogonality_and_ghosts():
"""Without reorthogonalization, a converged Ritz value is copied again and again (Paige)."""
n = 30
d = [float(i) for i in range(1, n)] + [60.0]
counts = []
for k in (30, 60, 80):
al, be, q, br = lanczos(diag_op(d), [1.0] * n, k, reorth=False)
th = ritz(al, be, q)[0]
counts.append(sum(1 for t in th if abs(t - 60.0) < 1e-8))
if k == 60:
assert orth_error(q) > 0.1
assert counts[0] >= 2 and counts[-1] > counts[0]
al, be, q, br = lanczos(diag_op(d), [1.0] * n, 30, reorth=True)
th = ritz(al, be, q)[0]
assert sum(1 for t in th if abs(t - 60.0) < 1e-8) == 1
def test_one_by_one_and_zero_matrix():
al, be, q, br = lanczos(diag_op([3.0]), [2.0], 3)
assert br and al == [3.0] and q == [[1.0]]
al, be, q, br = lanczos(diag_op([0.0, 0.0, 0.0]), [1.0, 2.0, 2.0], 3)
assert br and al == [0.0]
What the next level needs
- the tests did not run in this build