187 lines
5.9 KiB
Python
187 lines
5.9 KiB
Python
"""
|
||
Iterative diagonalization for the plane-wave KS Hamiltonian.
|
||
|
||
Lanczos algorithm with full reorthogonalisation and restart.
|
||
|
||
Physical normalisation convention: Σ_G |c(G)|² / V = 1 (matches the SCF density).
|
||
All Lanczos vectors satisfy np.vdot(q, q).real / V = 1.
|
||
|
||
The α and β values form a real symmetric tridiagonal T. The k lowest
|
||
eigenpairs of T approximate those of H at cost O(k × N_G log N_G) per step.
|
||
|
||
Reference: Cullum & Willoughby (1985); Payne et al., RMP 64, 1045 (1992).
|
||
"""
|
||
|
||
import numpy as np
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Hamiltonian application
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def apply_H(psi_G: np.ndarray, cell, V_r: np.ndarray) -> np.ndarray:
|
||
"""H|ψ⟩ = T|ψ⟩ + V|ψ⟩. O(N_G log N_G) via FFT. Preserves normalisation."""
|
||
return cell.KE_G * psi_G + cell.r_to_G(V_r * cell.G_to_r(psi_G))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Physical-norm helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _norm(v, V):
|
||
"""Physical norm: √(Σ|c|²/V)."""
|
||
return np.sqrt(np.vdot(v, v).real / V)
|
||
|
||
|
||
def _inner(a, b, V):
|
||
"""Physical inner product ⟨a|b⟩ = Σ a*(G) b(G) / V."""
|
||
return np.vdot(a, b) / V
|
||
|
||
|
||
def _normalise(v, V):
|
||
"""Return v / ‖v‖_phys. Physical norm of result = 1."""
|
||
n = _norm(v, V)
|
||
return v / n if n > 1e-14 else v
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Lanczos with restart
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _lanczos(cell, V_r, n_eig, v0, tol, max_steps, max_sub):
|
||
"""
|
||
Run Lanczos from starting vector v0, returning the n_eig lowest eigenpairs.
|
||
|
||
All Lanczos vectors are physically normalised: ‖q‖² / V = 1.
|
||
The tridiagonal T has α on the diagonal and β on the sub/super-diagonal.
|
||
Ritz vectors: Q_mat @ y (physically normalised if y is unit-normalised).
|
||
"""
|
||
N_G = cell.N_G
|
||
V = cell.volume
|
||
|
||
q = _normalise(v0.astype(complex), V)
|
||
|
||
Q = [q] # physically-normalised Lanczos vectors
|
||
alphas = []
|
||
betas = []
|
||
|
||
Hq = apply_H(q, cell, V_r)
|
||
alpha = _inner(q, Hq, V).real
|
||
alphas.append(alpha)
|
||
r = Hq - alpha * q
|
||
|
||
best_evals = None
|
||
best_Y = None # (m, n_eig) — Ritz vectors in tridiagonal basis
|
||
|
||
for j in range(1, max_steps + 1):
|
||
beta = _norm(r, V)
|
||
betas.append(beta)
|
||
|
||
if beta < 1e-12:
|
||
break
|
||
|
||
# Next Lanczos vector: r / ‖r‖_phys → physically normalised
|
||
q = _normalise(r, V)
|
||
|
||
# Full reorthogonalisation (Gram-Schmidt in physical inner product)
|
||
for qp in Q:
|
||
q -= _inner(qp, q, V) * qp
|
||
q = _normalise(q, V)
|
||
|
||
Q.append(q)
|
||
|
||
Hq = apply_H(q, cell, V_r)
|
||
alpha = _inner(q, Hq, V).real
|
||
alphas.append(alpha)
|
||
# Lanczos three-term recurrence: r = Hq - α q - β q_{j-1}
|
||
r = Hq - alpha * q - beta * Q[-2]
|
||
|
||
# Diagonalise tridiagonal T (real symmetric)
|
||
m = len(alphas)
|
||
T = (np.diag(alphas) +
|
||
np.diag(betas, -1) +
|
||
np.diag(betas, 1))
|
||
theta, Y = np.linalg.eigh(T) # (m,), (m,m)
|
||
|
||
best_evals = theta[:n_eig]
|
||
best_Y = Y[:, :n_eig]
|
||
|
||
# Residual bound: |β_m Y[m-1, n]| (standard Lanczos error estimate)
|
||
err = np.abs(betas[-1] * Y[m - 1, :n_eig])
|
||
if np.all(err < tol):
|
||
break
|
||
|
||
# Restart: compress subspace when it exceeds max_sub
|
||
if m >= max_sub:
|
||
# New starting vector = first Ritz vector
|
||
Q_mat = np.column_stack(Q) # (N_G, m)
|
||
ritz0 = Q_mat @ best_Y[:, 0] # (N_G,)
|
||
q = _normalise(ritz0, V)
|
||
|
||
Q = [q]
|
||
alphas = []
|
||
betas = []
|
||
Hq = apply_H(q, cell, V_r)
|
||
alpha = _inner(q, Hq, V).real
|
||
alphas.append(alpha)
|
||
r = Hq - alpha * q
|
||
|
||
# Final Ritz vectors
|
||
Q_mat = np.column_stack(Q) # (N_G, m)
|
||
ritz = Q_mat @ best_Y # (N_G, n_eig)
|
||
|
||
# Normalise (should already be near 1 if Q is orthonormal)
|
||
for n in range(n_eig):
|
||
nrm = _norm(ritz[:, n], V)
|
||
if nrm > 1e-14:
|
||
ritz[:, n] /= nrm
|
||
|
||
return best_evals, ritz.T # (n_eig,), (n_eig, N_G)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Public API (named "davidson" for backward compatibility with scf.py)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def davidson(
|
||
cell,
|
||
V_r: np.ndarray,
|
||
n_bands: int,
|
||
n_extra: int = 8,
|
||
tol: float = 1e-7,
|
||
max_iter: int = 200,
|
||
max_sub: int = None,
|
||
guess: np.ndarray = None,
|
||
) -> tuple:
|
||
"""
|
||
Find the lowest `n_bands` eigenpairs of H = T + V(r) via Lanczos.
|
||
|
||
Parameters
|
||
----------
|
||
cell : Cell
|
||
V_r : (N_r,) float KS potential.
|
||
n_bands : int
|
||
n_extra : int Extra Lanczos steps after convergence of lowest band.
|
||
tol : float Lanczos residual bound per eigenvalue.
|
||
max_iter: int Maximum Lanczos steps per restart window.
|
||
max_sub : int Restart threshold (default max(4*n_bands+n_extra, 20)).
|
||
guess : (n_bands, N_G) complex Warm-start vectors.
|
||
|
||
Returns
|
||
-------
|
||
evals : (n_bands,) float
|
||
evecs : (n_bands, N_G) complex Physically normalised (Σ|c|²/V = 1).
|
||
"""
|
||
V = cell.volume
|
||
if max_sub is None:
|
||
max_sub = max(4 * n_bands + n_extra, 20)
|
||
|
||
# Starting vector
|
||
if guess is not None and guess.shape[0] > 0:
|
||
v0 = guess[0].copy().astype(complex)
|
||
else:
|
||
v0 = np.zeros(cell.N_G, dtype=complex)
|
||
v0[np.argmin(cell.KE_G)] = np.sqrt(V) # unit vector at lowest KE point
|
||
|
||
return _lanczos(cell, V_r, n_bands, v0, tol, max_iter, max_sub)
|