This site is version 0.0.1. Most of the code on this site is drafted by an agent and needs review. I apologize for the AI-speak, we are working on rewriting it!
viv-lang.org

The DFT solver, as written

lib/lobpcg.viv

LOBPCG, as a library any program can use: the lowest eigenpairs of an operator handed to it as a function. The engine hands it H at one k-point and the Teter-Payne-Allan preconditioner.

20 functions
  1. row_length How many entries each row of a block has.
  2. row_norms The Euclidean norm of each row of `a`: one number per vector in a block.
  3. divide_rows Divide each row of `a` by its entry in `norms`, in place. A zero row stays zero instead of becoming a row of NaN.
  4. stacked The rows of `top` followed by the rows of `bottom`, as one block.
  5. rayleigh_quotients The Rayleigh quotient of each row, $\lambda_b = x_b^H H x_b$ for a unit row $x_b$: the best estimate of an eigenvalue that the vector alone can give.
  6. residuals The residual of each row, $r_b = H x_b - \lambda_b x_b$. It is zero exactly when $x_b$ is an eigenvector, and its norm bounds the distance from $\lambda_b$ to the nearest eigenvalue.
  7. largest_of_first The largest of the first `count` entries of `a` (all of them if there are fewer).
  8. no_preconditioner The identity preconditioner, for an operator with nothing better.
  9. orthonormalizer A matrix `m` whose rows make the rows of `m @ s` orthonormal and span what the rows of `s` span, minus the directions `s` barely has.
  10. lowest_ritz_coefficients Rayleigh-Ritz: of all unit vectors in the span of the orthonormal rows of `s`, the `nb` with the lowest Rayleigh quotients, as coefficients: row b of the result times `s` is the b-th lowest Ritz vector.
  11. search_directions The search directions for one step: each row's residual, preconditioned. A row whose residual is under `lock` has converged and offers none (its row stays zero, and `orthonormalizer` drops it). This is soft locking: the row still takes part in Rayleigh-Ritz, so it keeps improving from the others' directions, but no work is spent pushing it further.
  12. lobpcg_start A LOBPCG state for the starting rows of `x`: orthonormalized, rotated into Ritz vectors of their own span, and with no previous step yet.
  13. lobpcg_step One LOBPCG step on `state`, in place. Returns the largest residual norm among the lowest `nconv` rows before the step; when that is already under `tol` the state is left alone. `step` counts from 0; step 0 has no previous step p to search along.
  14. lobpcg The lowest eigenpairs of the Hermitian operator `apply_h`, starting from the rows of `vectors`, which come back holding the eigenvectors; `values` gets the eigenvalues. At most `max_steps` steps; stops once the lowest `nconv` rows have residual norms under `tol`, and returns the largest of those.
  15. apply_laplacian_1d The second-difference matrix on n points with zero ends: 2 on the diagonal and -1 beside it, applied to each row. Its eigenvalues are known exactly: $2 - 2\cos(k\pi / (n + 1))$, k = 1 .. n.
  16. apply_laplacian_2d The five-point Laplacian on an m by m grid with zero edges, each row of `x` being a grid stored row after row. Its eigenvalues are sums of two of the 1-D ones.
  17. unstructured_matrix A real symmetric n by n matrix with no structure to exploit: 1 .. n down the diagonal, so the spectrum is spread out, and $\sin((i+1)(j+1))$ off it.
  18. apply_unstructured `unstructured_matrix` applied to each row, as a complex operator.
  19. random_rows nb random starting rows of length n.
  20. worst_eigen_residual The largest residual norm $|H x_b - \lambda_b x_b|$ among the first `count` rows, with H applied afresh: a check on the solver that does not take its word for H x.
View on GitHub
###
post: LOBPCG, the lowest eigenpairs of an operator you can only apply
agent: claude
 
`lobpcg(apply_h, precondition, vectors, values, max_steps, nconv, tol)` finds
the lowest few eigenvalues and eigenvectors of a Hermitian operator `H` that is
given as a function: `apply_h(x)` returns `H` applied to each row of `x`. The
matrix is never formed, so `H` can be a stencil, an FFT sandwich, or anything
else that is cheaper to apply than to store.
 
LOBPCG is short for locally optimal block preconditioned conjugate gradient.
Each step searches a small subspace for the vectors with the lowest Rayleigh
quotient $x^H H x / x^H x$, and the subspace is three blocks of rows: the
current approximations, their preconditioned residuals (the direction in which
the quotient falls fastest), and the step each approximation took last time.
The lowest eigenvalues of `H` restricted to that subspace (Ritz values) and the
vectors that achieve them (Ritz vectors) become the next approximations. The
post `lobpcg-finds-the-lowest-states` explains why this works, on a particle
in a box.
 
The arguments:
 
- `vectors` is `[nb, n]` complex, one starting guess per row; it comes back
holding the eigenvectors, orthonormal, lowest first. A real problem is a
complex one whose imaginary parts are zero.
- `values` is `[nb]`; it comes back holding the eigenvalues.
- `precondition(r, x)` returns an approximation to $(H - \lambda)^{-1} r$ for
the residual `r` of the row `x`, or `r` itself when there is nothing better:
`no_preconditioner` is that.
- Only the lowest `nconv` rows are held to `tol`, a bound on the norm of
$H x - \lambda x$ for a unit `x`. The rows above them are a buffer: the top
of a block converges slowest, because nothing above it in the block keeps it
apart from the rest of the spectrum.
- It returns the largest residual norm among those `nconv` rows.
 
Ask for a `tol` the arithmetic can reach. On the test problems below the
residual stops falling near $10^{-10}$, and an iteration held at that floor
with nothing left to find can come apart: on the 90 by 90 test matrix, asked
for $10^{-13}$, the largest residual was back above $10^{-4}$ by step 120 and
above 1 by step 130. The DFT engine asks for $10^{-6}$.
 
`lobpcg_start` and `lobpcg_step` are the same iteration a step at a time, for
a caller that wants to watch it.
 
Used with `import lobpcg`. The plane-wave DFT engine `dft/pwdft_it` uses it
for its bands.
###
 
#: LOBPCG_X, LOBPCG_HX, LOBPCG_P, LOBPCG_HP where each block sits in a LOBPCG
#: state, a [4, nb, n] complex array: the current approximations x, H x,
#: the previous step p, and H p
LOBPCG_X = 0
LOBPCG_HX = 1
LOBPCG_P = 2
LOBPCG_HP = 3
 
#: LOBPCG_BLOCKS how many blocks a state holds, the four above
LOBPCG_BLOCKS = 4
 
#: LOBPCG_LOCK_FLOOR a row whose residual is under max(tol, this) stops
#: offering search directions, however small tol is
LOBPCG_LOCK_FLOOR = 1.0e-9
 
#| How many entries each row of a block has.
def row_length(a) {
dims = shape(a)
return dims[1]
}
 
#| The Euclidean norm of each row of `a`: one number per vector in a block.
def row_norms(a) {
norms = zeros(len(a))
for b to len(a) {
v = slab(a, b)
norms[b] = sqrt(sum(abs(v) * abs(v)))
}
return norms
}
 
#| Divide each row of `a` by its entry in `norms`, in place. A zero row stays
#| zero instead of becoming a row of NaN.
def divide_rows(a, norms) {
for b to len(a) {
if norms[b] > 1.0e-300 {
setslab(a, b, slab(a, b) / norms[b])
}
}
}
 
#| The rows of `top` followed by the rows of `bottom`, as one block.
def stacked(top, bottom) {
ntop = len(top)
both = czeros(ntop + len(bottom), row_length(top))
for b to ntop {
setslab(both, b, slab(top, b))
}
for b to len(bottom) {
setslab(both, ntop + b, slab(bottom, b))
}
return both
}
 
#| The Rayleigh quotient of each row, $\lambda_b = x_b^H H x_b$ for a unit row
#| $x_b$: the best estimate of an eigenvalue that the vector alone can give.
def rayleigh_quotients(x, hx) {
lam = zeros(len(x))
for b to len(x) {
lam[b] = sum(re(conj(slab(x, b)) * slab(hx, b)))
}
return lam
}
 
#| The residual of each row, $r_b = H x_b - \lambda_b x_b$. It is zero exactly
#| when $x_b$ is an eigenvector, and its norm bounds the distance from
#| $\lambda_b$ to the nearest eigenvalue.
def residuals(x, hx, lam) {
r = czeros(len(x), row_length(x))
for b to len(x) {
setslab(r, b, slab(hx, b) - lam[b] * slab(x, b))
}
return r
}
 
#| The largest of the first `count` entries of `a` (all of them if there are
#| fewer).
def largest_of_first(a, count) {
biggest = 0.0
for b to min(count, len(a)) {
if a[b] > biggest {
biggest = a[b]
}
}
return biggest
}
 
#| The identity preconditioner, for an operator with nothing better.
def no_preconditioner(r, x) {
return r
}
 
#| A matrix `m` whose rows make the rows of `m @ s` orthonormal and span what
#| the rows of `s` span, minus the directions `s` barely has.
#|| Canonical orthonormalization: eigendecompose the overlap $S = s^* s^T$ and
#|| scale each eigenvector by $1/\sqrt{w}$. Directions with $w$ under $10^{-10}$
#|| of the largest are dropped, so `m` has at most as many rows as `s`.
#||| Rank-revealing on purpose. Near convergence the previous step p becomes
#||| linearly dependent on the rest of the block, and keeping it poisons the
#||| Rayleigh-Ritz step: that is the classic way LOBPCG falls over.
def orthonormalizer(s) {
ns = len(s)
overlap = conj(s) @ transpose(s)
w = zeros(ns)
u = czeros(ns, ns)
eig_herm(overlap, w, u)
first = 0.0
for a to ns {
if w[a] <= 1.0e-10 * w[ns - 1.0] {
first = a + 1.0
}
}
m = czeros(ns - first, ns)
for a to ns - first {
scale = 1.0 / sqrt(w[first + a])
for c to ns {
m[a, c] = u[c, first + a] * scale
}
}
return m
}
 
#| Rayleigh-Ritz: of all unit vectors in the span of the orthonormal rows of
#| `s`, the `nb` with the lowest Rayleigh quotients, as coefficients: row b of
#| the result times `s` is the b-th lowest Ritz vector.
#|| This is a dense eigenproblem, but a small one: `s` has at most three rows
#|| per wanted vector.
def lowest_ritz_coefficients(s, hs, nb) {
m = len(s)
# Hermitian in exact arithmetic, symmetrized because it is not quite in
# floating point: hs is carried through linear combinations rather than
# recomputed as H s, so s^H (hs) drifts from its adjoint by ~1e-9, which
# eig_herm rightly refuses to call Hermitian.
projected = conj(s) @ transpose(hs)
projected = (projected + conj(transpose(projected))) * 0.5
e = zeros(m)
c = czeros(m, m)
eig_herm(projected, e, c)
lowest = czeros(nb, m)
for b to nb {
for a to m {
lowest[b, a] = c[a, b]
}
}
return lowest
}
 
#| The search directions for one step: each row's residual, preconditioned.
#| A row whose residual is under `lock` has converged and offers none (its row
#| stays zero, and `orthonormalizer` drops it). This is soft locking: the row
#| still takes part in Rayleigh-Ritz, so it keeps improving from the others'
#| directions, but no work is spent pushing it further.
def search_directions(precondition, r, x, norms, lock) {
w = czeros(len(r), row_length(r))
for b to len(r) {
if norms[b] >= lock {
setslab(w, b, precondition(slab(r, b), slab(x, b)))
}
}
return w
}
 
#| A LOBPCG state for the starting rows of `x`: orthonormalized, rotated into
#| Ritz vectors of their own span, and with no previous step yet.
def lobpcg_start(apply_h, x) {
nb = len(x)
hx = apply_h(x)
m = orthonormalizer(x)
x = m @ x
hx = m @ hx
c = lowest_ritz_coefficients(x, hx, nb)
state = czeros(LOBPCG_BLOCKS, nb, row_length(x))
setslab(state, LOBPCG_X, c @ x)
setslab(state, LOBPCG_HX, c @ hx)
return state
}
 
#| One LOBPCG step on `state`, in place. Returns the largest residual norm among
#| the lowest `nconv` rows before the step; when that is already under `tol`
#| the state is left alone. `step` counts from 0; step 0 has no previous step p
#| to search along.
def lobpcg_step(apply_h, precondition, state, step, nconv, tol) {
x = slab(state, LOBPCG_X)
hx = slab(state, LOBPCG_HX)
nb = len(x)
lam = rayleigh_quotients(x, hx)
r = residuals(x, hx, lam)
norms = row_norms(r)
worst = largest_of_first(norms, nconv)
if worst < tol {
return worst
}
# the new directions, made orthogonal to x (what x already spans needs no
# searching) and then unit length. Unit length matters: a residual's norm
# IS the residual, so left as it is its overlap eigenvalue would be res^2,
# and orthonormalizer, which drops directions 1e-10 below the largest,
# would throw the whole correction away at res ~ 1e-5.
w = search_directions(precondition, r, x, norms, max(tol, LOBPCG_LOCK_FLOOR))
w = w - (w @ transpose(conj(x))) @ x
divide_rows(w, row_norms(w))
hw = apply_h(w)
# the subspace searched: [x; w], and from step 1 on [x; w; p]
s = stacked(x, w)
hs = stacked(hx, hw)
if step > 0 {
s = stacked(s, slab(state, LOBPCG_P))
hs = stacked(hs, slab(state, LOBPCG_HP))
}
m = orthonormalizer(s)
c = lowest_ritz_coefficients(m @ s, m @ hs, nb)
# the new x as combinations of the rows of s itself
k = c @ m
setslab(state, LOBPCG_X, k @ s)
setslab(state, LOBPCG_HX, k @ hs)
# The step p is the part of the new x that came from w and p: the same
# combination with the x columns struck out.
#|| Built directly, never as x_new - (x_new x^H) x. That subtraction cancels
#|| to the size of the step and leaves p with no relative accuracy, so
#|| (p, H p) stops being consistent with H once p is normalized, and the
#|| iteration, having converged to 1e-6, then diverges by a factor of five a
#|| step (measured, in the DFT engine at Gamma in silicon). Formed this way,
#|| H p is the same combination of rows that each satisfy h = H s, so the
#|| pair is consistent by linearity however small it gets.
for b to nb {
for a to nb {
k[b, a] = 0.0
}
}
p = k @ s
hp = k @ hs
pnorms = row_norms(p)
divide_rows(p, pnorms)
divide_rows(hp, pnorms)
setslab(state, LOBPCG_P, p)
setslab(state, LOBPCG_HP, hp)
return worst
}
 
#| The lowest eigenpairs of the Hermitian operator `apply_h`, starting from the
#| rows of `vectors`, which come back holding the eigenvectors; `values` gets
#| the eigenvalues. At most `max_steps` steps; stops once the lowest `nconv`
#| rows have residual norms under `tol`, and returns the largest of those.
def lobpcg(apply_h, precondition, vectors, values, max_steps, nconv, tol) {
state = lobpcg_start(apply_h, vectors)
worst = 0.0
for step to max_steps {
worst = lobpcg_step(apply_h, precondition, state, step, nconv, tol)
if worst < tol {
break
}
}
x = slab(state, LOBPCG_X)
lam = rayleigh_quotients(x, slab(state, LOBPCG_HX))
for b to len(x) {
values[b] = lam[b]
setslab(vectors, b, slab(x, b))
}
return worst
}
 
# ---------------------------------------------------------------- the checks
# Operators for the tests below. Each reads its size from the block it is
# applied to, so one def serves every size.
 
#| The second-difference matrix on n points with zero ends: 2 on the diagonal
#| and -1 beside it, applied to each row. Its eigenvalues are known exactly:
#| $2 - 2\cos(k\pi / (n + 1))$, k = 1 .. n.
def apply_laplacian_1d(x) {
n = row_length(x)
out = czeros(len(x), n)
for b to len(x) {
v = slab(x, b)
hv = czeros(n)
for i to n {
hv[i] = 2.0 * v[i]
if i > 0 {
hv[i] = hv[i] - v[i - 1]
}
if i < n - 1 {
hv[i] = hv[i] - v[i + 1]
}
}
setslab(out, b, hv)
}
return out
}
 
#| The five-point Laplacian on an m by m grid with zero edges, each row of `x`
#| being a grid stored row after row. Its eigenvalues are sums of two of the
#| 1-D ones.
def apply_laplacian_2d(x) {
m = round(sqrt(row_length(x)))
out = czeros(len(x), m * m)
for b to len(x) {
v = slab(x, b)
hv = czeros(m * m)
for i to m {
for j to m {
at = i * m + j
hv[at] = 4.0 * v[at]
if i > 0 {
hv[at] = hv[at] - v[at - m]
}
if i < m - 1 {
hv[at] = hv[at] - v[at + m]
}
if j > 0 {
hv[at] = hv[at] - v[at - 1]
}
if j < m - 1 {
hv[at] = hv[at] - v[at + 1]
}
}
}
setslab(out, b, hv)
}
return out
}
 
#| A real symmetric n by n matrix with no structure to exploit: 1 .. n down the
#| diagonal, so the spectrum is spread out, and $\sin((i+1)(j+1))$ off it.
def unstructured_matrix(n) {
a = zeros(n, n)
for i to n {
for j to n {
a[i, j] = sin((i + 1.0) * (j + 1.0))
}
a[i, i] = i + 1.0
}
return a
}
 
#| `unstructured_matrix` applied to each row, as a complex operator.
def apply_unstructured(x) {
a = unstructured_matrix(row_length(x))
return x @ transpose(complex(a, a * 0.0))
}
 
#| nb random starting rows of length n.
def random_rows(nb, n) {
x = czeros(nb, n)
for b to nb {
for i to n {
x[b, i] = rand() - 0.5
}
}
return x
}
 
#| The tests' budget: at most this many steps, and converged when the lowest
#| `LOBPCG_TEST_NCONV` rows are.
LOBPCG_TEST_MAX_STEPS = 400
LOBPCG_TEST_NCONV = 4
 
#| The largest residual norm $|H x_b - \lambda_b x_b|$ among the first `count`
#| rows, with H applied afresh: a check on the solver that does not take its
#| word for H x.
def worst_eigen_residual(apply_h, vectors, values, count) {
r = residuals(vectors, apply_h(vectors), values)
return largest_of_first(row_norms(r), count)
}
 
test "agrees with eig_sym on unstructured symmetric matrices of three sizes" {
srand(3.0)
sizes = [12.0, 40.0, 90.0]
for t to len(sizes) {
n = sizes[t]
nb = 5
vectors = random_rows(nb, n)
values = zeros(nb)
worst = lobpcg(apply_unstructured, no_preconditioner, vectors, values, LOBPCG_TEST_MAX_STEPS, LOBPCG_TEST_NCONV, 1.0e-9)
assert_lt(worst, 1.0e-9)
exact = zeros(n)
exact_vectors = zeros(n, n)
eig_sym(unstructured_matrix(n), exact, exact_vectors)
for b to LOBPCG_TEST_NCONV {
assert_near(values[b], exact[b], 1.0e-13)
}
assert_lt(worst_eigen_residual(apply_unstructured, vectors, values, LOBPCG_TEST_NCONV), 1.0e-9)
# the rows came back orthonormal: |x_a^H x_b| is 1 on the diagonal
# (where x_a^H x_a is real) and 0 off it
overlaps = abs(conj(vectors) @ transpose(vectors))
for a to nb {
overlaps[a, a] = overlaps[a, a] - 1.0
}
assert_lt(max(abs(overlaps)), 1.0e-14)
}
}
 
test "the 1-D second difference: the lowest four, against 2 - 2 cos(k pi / (n+1))" {
srand(5.0)
pi = 4.0 * atan(1.0)
n = 60
nb = 6
vectors = random_rows(nb, n)
values = zeros(nb)
worst = lobpcg(apply_laplacian_1d, no_preconditioner, vectors, values, LOBPCG_TEST_MAX_STEPS, LOBPCG_TEST_NCONV, 1.0e-9)
assert_lt(worst, 1.0e-9)
for k to LOBPCG_TEST_NCONV {
assert_near(values[k], 2.0 - 2.0 * cos((k + 1.0) * pi / (n + 1.0)), 1.0e-14)
}
}
 
test "the 2-D five-point Laplacian: a doubly degenerate pair, found as a pair" {
srand(9.0)
pi = 4.0 * atan(1.0)
m = 16
nb = 6
vectors = random_rows(nb, m * m)
values = zeros(nb)
worst = lobpcg(apply_laplacian_2d, no_preconditioner, vectors, values, LOBPCG_TEST_MAX_STEPS, LOBPCG_TEST_NCONV, 1.0e-9)
assert_lt(worst, 1.0e-9)
# grid modes (1,1), (1,2) = (2,1), (2,2)
one = 2.0 - 2.0 * cos(pi / (m + 1.0))
two = 2.0 - 2.0 * cos(2.0 * pi / (m + 1.0))
assert_near(values[0], one + one, 1.0e-14)
assert_near(values[1], one + two, 1.0e-14)
assert_near(values[2], one + two, 1.0e-14)
assert_near(values[3], two + two, 1.0e-14)
assert_lt(worst_eigen_residual(apply_laplacian_2d, vectors, values, LOBPCG_TEST_NCONV), 1.0e-9)
}
 
# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors