Divide-and-conquer eigensolver in Python
Split a symmetric tridiagonal matrix into two halves plus a rank-one term, solve the halves recursively, and merge by solving a secular equation. Fast in practice because of deflation.
- File
divide_and_conquer_eigensolver.py, 179 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/divide-and-conquer-eigensolver/. The authors are listed in prov/J2346G.jsonl.
"""Cuppen's divide-and-conquer method for the symmetric tridiagonal eigenproblem,
with Dongarra-Sorensen deflation and the Gu-Eisenstat (Loewner) eigenvectors.
Python standard library only.
eigh_tridiagonal(a, b) returns (lam, q, stats) with lam[0] >= ... >= lam[n-1],
q[:, i] the unit eigenvector of lam[i], for the matrix with diagonal a and
off-diagonal b. The recursion depth is about log2(n).
"""
import math
U = 2.0 ** -53
def _secular_root(i, d, zsq, rho, stats):
"""Root of f(x) = 1 + rho sum zsq_j / (d_j - x) in (d_i, d_i+1), or above d_last.
d is strictly increasing, zsq > 0 and rho > 0. Returns (org, tau): the root is
d[org] + tau, and d[j] - root is computed as (d[j] - d[org]) - tau."""
k = len(d)
if i < k - 1:
mid = (d[i] + d[i + 1]) / 2.0
fmid = 1.0 + rho * sum(zsq[j] / (d[j] - mid) for j in range(k))
org = i if fmid >= 0 else i + 1
lo, hi = (0.0, mid - d[i]) if org == i else (mid - d[i + 1], 0.0)
else:
org, lo, hi = k - 1, 0.0, rho * sum(zsq)
delta = [d[j] - d[org] for j in range(k)]
tau = (lo + hi) / 2.0
for _ in range(200):
stats["secular_its"] += 1
psi = dpsi = phi = dphi = erretm = 0.0
for j in range(k):
t = zsq[j] / (delta[j] - tau)
if j <= i:
psi += t
dpsi += t / (delta[j] - tau)
else:
phi += t
dphi += t / (delta[j] - tau)
erretm += abs(t)
f = 1.0 + rho * (psi + phi)
if abs(f) <= 8 * k * U * (1.0 + rho * erretm) or hi - lo <= 4 * U * max(abs(lo), abs(hi)):
break
if f < 0:
lo = tau
else:
hi = tau
# rational model: rho psi ~ p + q / (delta_i - x), rho phi ~ r + s / (delta_i+1 - x)
q = rho * dpsi * (delta[i] - tau) ** 2
p = rho * psi - rho * dpsi * (delta[i] - tau)
if i < k - 1:
s = rho * dphi * (delta[i + 1] - tau) ** 2
r = rho * phi - rho * dphi * (delta[i + 1] - tau)
c0 = 1.0 + p + r
# c0 (a - x)(b - x) + q (b - x) + s (a - x) = 0 with a = delta_i, b = delta_i+1
a_, b_ = delta[i], delta[i + 1]
qa = c0
qb = -(c0 * (a_ + b_) + q + s)
qc = c0 * a_ * b_ + q * b_ + s * a_
x = None
if qa == 0.0:
x = -qc / qb if qb != 0 else None
else:
disc = qb * qb - 4 * qa * qc
if disc >= 0:
sq = math.sqrt(disc)
w = -(qb + math.copysign(sq, qb)) / 2.0
cands = [w / qa] + ([qc / w] if w != 0 else [])
inside = [y for y in cands if lo < y < hi]
x = inside[0] if inside else None
else:
c0 = 1.0 + p
x = delta[i] + q / c0 if c0 > 0 else None
tau = x if (x is not None and lo < x < hi) else (lo + hi) / 2.0
return org, tau
def _merge(dvals, qmat, z, rho, stats):
"""Eigen-decomposition of Q (D + rho z z^T) Q^T. Returns (lam, vecs) ascending."""
n = len(dvals)
nz = math.sqrt(sum(t * t for t in z))
if nz == 0.0 or rho == 0.0:
order = sorted(range(n), key=lambda i: dvals[i])
return [dvals[i] for i in order], [[row[i] for i in order] for row in qmat]
sign = 1.0 if rho > 0 else -1.0
rho = abs(rho) * nz * nz
z = [t / nz for t in z]
d = [sign * t for t in dvals] # solve with rho > 0; negate the eigenvalues at the end
order = sorted(range(n), key=lambda i: d[i])
d = [d[i] for i in order]
z = [z[i] for i in order]
q = [[row[i] for i in order] for row in qmat]
tol = 8 * U * max(max(abs(t) for t in d), rho)
deflated, kept = [], []
for i in range(n): # @line L5
if rho * abs(z[i]) <= tol:
deflated.append(i)
stats["deflated"] += 1
continue
if kept:
j = kept[-1]
r = math.hypot(z[j], z[i])
c, s = z[i] / r, z[j] / r
if abs(c * s * (d[i] - d[j])) <= tol:
dj = c * c * d[j] + s * s * d[i]
d[i] = s * s * d[j] + c * c * d[i]
d[j] = dj
z[j], z[i] = 0.0, r
for row in q:
qj, qi = row[j], row[i]
row[j], row[i] = c * qj - s * qi, s * qj + c * qi
kept.pop()
deflated.append(j)
stats["deflated"] += 1
kept.append(i)
kept.sort(key=lambda i: d[i])
dk = [d[i] for i in kept]
k = len(kept)
for t in range(k - 1):
if not dk[t] < dk[t + 1]: # equal poles after rotation: deflation missed; nudge
dk[t + 1] = dk[t] + 2 * U * max(abs(dk[t]), tol)
stats["roots"] += k
zk = [z[i] for i in kept]
zsq = [t * t for t in zk]
roots = [_secular_root(i, dk, zsq, rho, stats) for i in range(k)] # @line L6
diff = [[(dk[j] - dk[org]) - tau for (org, tau) in roots] for j in range(k)] # d_j - lam_i
zhat = []
for j in range(k): # @line L7
prod = -diff[j][j] / rho # (lam_j - d_j) / rho
for i in range(k):
if i != j:
prod *= -diff[j][i] / (dk[i] - dk[j])
zhat.append(math.copysign(math.sqrt(max(prod, 0.0)), zk[j]))
lam, vecs = [], []
for i in range(k):
s = [zhat[j] / diff[j][i] for j in range(k)] # @line L8
ns = math.sqrt(sum(t * t for t in s))
s = [t / ns for t in s]
lam.append(dk[roots[i][0]] + roots[i][1])
vecs.append([sum(q[r][kept[j]] * s[j] for j in range(k)) for r in range(n)]) # @line L9
for i in deflated:
lam.append(d[i])
vecs.append([q[r][i] for r in range(n)])
lam = [sign * t for t in lam]
order = sorted(range(n), key=lambda i: lam[i])
lam = [lam[i] for i in order]
cols = [vecs[i] for i in order]
return lam, [[cols[c][r] for c in range(n)] for r in range(n)]
def _dc(a, b, stats):
n = len(a)
if n == 1: # @line L1
return [a[0]], [[1.0]]
m = n // 2 # @line L2
rho = b[m - 1]
a1 = a[:m]
a2 = a[m:]
a1[-1] -= rho
a2[0] -= rho
d1, q1 = _dc(a1, b[:m - 1], stats) # @line L3
d2, q2 = _dc(a2, b[m:], stats)
stats["merges"] += 1
z = q1[-1][:] + q2[0][:] # @line L4
qmat = [row + [0.0] * (n - m) for row in q1] + [[0.0] * m + row for row in q2]
return _merge(d1 + d2, qmat, z, rho, stats)
def eigh_tridiagonal(a, b):
"""Eigenvalues in descending order, unit eigenvectors as columns, and counters."""
stats = {"merges": 0, "deflated": 0, "roots": 0, "secular_its": 0}
if not a:
return [], [], stats
lam, q = _dc([float(t) for t in a], [float(t) for t in b], stats)
lam = lam[::-1] # @line L10
q = [row[::-1] for row in q]
return lam, q, stats
test_divide_and_conquer_eigensolver.py 114 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/divide-and-conquer-eigensolver/. The authors are listed in prov/J2346G.jsonl.
import math
import random
from fractions import Fraction
from divide_and_conquer_eigensolver import U, eigh_tridiagonal
def check(a, b, c=10.0):
"""Assert residual and orthogonality bounds; return (lam, stats)."""
lam, q, stats = eigh_tridiagonal(a, b)
n = len(a)
scale = max([abs(t) for t in a] + [abs(t) for t in b] + [1e-300])
for col in range(n):
for i in range(n):
tq = a[i] * q[i][col]
if i > 0:
tq += b[i - 1] * q[i - 1][col]
if i < n - 1:
tq += b[i] * q[i + 1][col]
assert abs(tq - lam[col] * q[i][col]) <= c * n * U * scale
for i in range(n):
for j in range(n):
assert abs(sum(q[r][i] * q[r][j] for r in range(n)) - (i == j)) <= c * n * U
assert all(lam[i] >= lam[i + 1] for i in range(n - 1))
return lam, stats
def det_tridiagonal_exact(a, b):
p0, p1 = Fraction(1), Fraction(a[0])
for i in range(1, len(a)):
p0, p1 = p1, Fraction(a[i]) * p1 - Fraction(b[i - 1]) ** 2 * p0
return p1
def test_page_example():
lam, stats = check([2.0] * 4, [1.0] * 3)
for k in range(1, 5):
assert abs(lam[k - 1] - (2 + 2 * math.cos(k * math.pi / 5))) <= 4 * 4 * U
assert stats["merges"] == 3 and stats["deflated"] == 2 # equal poles at the top merge
# the reduced secular equation of the page has the roots 1.381966 and 3.618034
f = lambda x: 1 + 2 * (0.7236068 / (0.3819660 - x) + 0.2763932 / (2.6180340 - x))
assert abs(f(lam[0])) < 1e-6 and abs(f(lam[2])) < 1e-6
def test_toeplitz_spectrum_and_deflation():
n = 100
lam, stats = check([2.0] * n, [-1.0] * (n - 1))
exact = sorted((2 - 2 * 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(lam, exact)) <= 20 * n * U * 4
assert stats["deflated"] > 0
assert stats["secular_its"] <= 8 * stats["roots"] # observed about 4.3 evaluations per root
def test_random_seeded_trace_det():
rnd = random.Random(1)
n = 128
a = [rnd.gauss(0, 1) for _ in range(n)]
b = [rnd.gauss(0, 1) for _ in range(n - 1)]
lam, stats = check(a, b)
assert abs(sum(lam) - float(sum(Fraction(t) for t in a))) <= 10 * n * U * 4
assert stats["secular_its"] <= 8 * stats["roots"]
a = [float(rnd.randint(-4, 4)) for _ in range(8)]
b = [float(rnd.randint(1, 3)) for _ in range(7)]
lam, _ = check(a, b)
det = float(det_tridiagonal_exact(a, b))
assert abs(math.prod(lam) - det) <= 1e-12 * max(1.0, abs(det)) * 8
def test_wilkinson_close_pair_orthogonal():
n = 21
a = [abs(10 - i) * 1.0 for i in range(n)]
lam, _ = check(a, [1.0] * (n - 1))
assert 0 < lam[0] - lam[1] < 1e-12 # the top pair of W21+ agrees to about 13 digits
assert abs(lam[0] - 10.746194182903) < 1e-11
def test_trivial_and_structured():
assert check([3.0], [])[0] == [3.0]
lam, _ = check([1.0, 3.0], [1.0])
assert abs(lam[0] - (2 + math.sqrt(2))) < 4 * U * 4 and abs(lam[1] - (2 - math.sqrt(2))) < 4 * U * 4
lam, stats = check([5.0, -1.0, 2.0, 0.0], [0.0, 0.0, 0.0]) # diagonal
assert lam == [5.0, 2.0, 0.0, -1.0] and stats["secular_its"] == 0
lam, _ = check([0.0] * 5, [0.0] * 4) # zero matrix
assert lam == [0.0] * 5
lam, _ = check([1.0] * 6, [0.0, 0.0, 0.0, 0.0, 0.0]) # repeated eigenvalue
assert lam == [1.0] * 6
lam, _ = check([1.0, 1.0], [1e-300]) # nearly equal, deflated
assert lam == [1.0, 1.0]
lam, _ = check([1.0, 1.0, 1.0], [1e-10, 1e-10])
assert abs(lam[0] - 1 - math.sqrt(2) * 1e-10) < 4 * U and abs(lam[1] - 1) < 4 * U
def test_negative_coupling_and_equal_modulus():
n = 9
lam, _ = check([0.0] * n, [-1.0] * (n - 1)) # spectrum symmetric about 0: +-2cos(k pi/10)
for k in range(1, n + 1):
assert abs(lam[k - 1] - 2 * math.cos(k * math.pi / (n + 1))) <= 40 * U
assert abs(lam[n // 2]) <= 40 * U
rnd = random.Random(5)
a = [rnd.uniform(-1, 1) for _ in range(33)]
b = [-abs(rnd.uniform(0.1, 1)) for _ in range(32)]
check(a, b)
def test_glued_blocks_deflate_heavily():
block_a, block_b = [abs(5 - i) * 1.0 for i in range(11)], [1.0] * 10
a, b = [], []
for g in range(4):
a += block_a
b += block_b + ([1e-14] if g < 3 else [])
lam, stats = check(a, b)
assert stats["deflated"] >= 30
What the next level needs
- the tests did not run in this build