192 lines
5.9 KiB
Python
192 lines
5.9 KiB
Python
"""
|
||
Self-Consistent Field (SCF) driver for plane-wave LDA-DFT.
|
||
|
||
Uses FFT-based Davidson iterative diagonalization — H is never built explicitly.
|
||
Cost per SCF step: O(n_occ × N_G × log N_G) for H applications.
|
||
|
||
KS Hamiltonian: H = T + V_KS(r)
|
||
T|ψ⟩_G = (|k+G|²/2) ψ(G) [diagonal in G-space]
|
||
V|ψ⟩(G) = FFT[ V_KS(r) · IFFT[ψ](r) ] [FFT convolution]
|
||
"""
|
||
|
||
import numpy as np
|
||
from .cell import Cell
|
||
from .xc import xc_energy, xc_potential
|
||
from .pseudo import build_ionic_potential, n_valence_electrons
|
||
from .diag import davidson
|
||
|
||
|
||
class SCFResult:
|
||
def __init__(self):
|
||
self.converged = False
|
||
self.e_total = 0.0
|
||
self.e_kinetic = 0.0
|
||
self.e_hartree = 0.0
|
||
self.e_xc = 0.0
|
||
self.e_ion_elec = 0.0
|
||
self.e_ion_ion = 0.0
|
||
self.eigenvalues = None
|
||
self.rho_r = None
|
||
self.n_iter = 0
|
||
|
||
def __repr__(self):
|
||
lines = [
|
||
f"SCF {'converged' if self.converged else 'NOT converged'} "
|
||
f"in {self.n_iter} iterations",
|
||
f" E_total = {self.e_total:+.8f} Ha",
|
||
f" E_kinetic = {self.e_kinetic:+.8f} Ha",
|
||
f" E_hartree = {self.e_hartree:+.8f} Ha",
|
||
f" E_xc = {self.e_xc:+.8f} Ha",
|
||
f" E_ion-elec = {self.e_ion_elec:+.8f} Ha",
|
||
f" E_ion-ion = {self.e_ion_ion:+.8f} Ha",
|
||
]
|
||
if self.eigenvalues is not None:
|
||
eig_str = " ".join(f"{e:.6f}" for e in self.eigenvalues)
|
||
lines.append(f" Eigenvalues (Ha): {eig_str}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def run_scf(
|
||
cell: Cell,
|
||
mol,
|
||
offset=None,
|
||
n_iter_max: int = 100,
|
||
tol: float = 1e-6,
|
||
mix: float = 0.4,
|
||
diis_start: int = 4,
|
||
diis_size: int = 8,
|
||
diag_tol: float = 1e-7,
|
||
verbose: bool = True,
|
||
) -> SCFResult:
|
||
"""
|
||
Plane-wave LDA-DFT SCF (closed shell, Γ-point).
|
||
|
||
Parameters
|
||
----------
|
||
cell : Cell
|
||
mol : Molecule
|
||
offset : (3,) shift added to atomic coords (from Cell.from_molecule).
|
||
n_iter_max : Maximum SCF iterations.
|
||
tol : Density convergence: √(∫|Δρ|² dV) < tol.
|
||
mix : Linear mixing fraction.
|
||
diis_start : Iteration to begin Pulay DIIS.
|
||
diis_size : DIIS history depth.
|
||
diag_tol : Residual tolerance for the Davidson diagonaliser.
|
||
verbose : Print iteration table.
|
||
"""
|
||
n_elec = n_valence_electrons(mol)
|
||
# Restricted calculation: pair electrons, last orbital singly occupied if odd
|
||
n_occ = (n_elec + 1) // 2
|
||
occ = np.full(n_occ, 2.0)
|
||
if n_elec % 2 == 1:
|
||
occ[-1] = 1.0
|
||
|
||
V_ion = build_ionic_potential(mol, cell, offset) # (N_r,) fixed
|
||
|
||
# Initial density: uniform
|
||
rho = np.full(cell.N_r, n_elec / cell.volume)
|
||
|
||
result = SCFResult()
|
||
result.e_ion_ion = mol.nuclear_repulsion
|
||
|
||
diis_rho: list = []
|
||
diis_err: list = []
|
||
guess = None # warm-start Davidson between SCF iterations
|
||
|
||
if verbose:
|
||
print(f"{'Iter':>4} {'E_total':>16} {'ΔE':>12} {'|Δρ|':>12}")
|
||
print("-" * 56)
|
||
|
||
e_prev = 0.0
|
||
|
||
for iteration in range(1, n_iter_max + 1):
|
||
# ----- KS potential -----
|
||
V_H = cell.hartree_potential(rho)
|
||
v_xc = xc_potential(rho)
|
||
V_KS = V_ion + V_H + v_xc # (N_r,)
|
||
|
||
# ----- Diagonalise -----
|
||
evals, evecs = davidson(
|
||
cell, V_KS, n_occ,
|
||
n_extra=max(4, n_occ),
|
||
tol=diag_tol,
|
||
guess=guess,
|
||
)
|
||
guess = evecs.copy() # warm start next iteration
|
||
|
||
# ----- New density -----
|
||
orbs_G = evecs # (n_occ, N_G)
|
||
rho_new = cell.density_from_orbitals(orbs_G, occ)
|
||
|
||
# ----- Energies -----
|
||
# E_kin = Σ_n f_n Σ_G KE_G |ψ̂_n(G)|² / V
|
||
e_kin = float(np.einsum('n,nG,G->', occ, np.abs(orbs_G)**2, cell.KE_G) / cell.volume)
|
||
e_hartree = 0.5 * float(np.dot(V_H, rho_new) * cell.dV)
|
||
e_xc = xc_energy(rho_new, cell.dV)
|
||
e_ion_elec = float(np.dot(V_ion, rho_new) * cell.dV)
|
||
e_total = e_kin + e_hartree + e_xc + e_ion_elec + result.e_ion_ion
|
||
|
||
delta_E = e_total - e_prev
|
||
e_prev = e_total
|
||
|
||
# ----- Convergence -----
|
||
delta_rho = rho_new - rho
|
||
conv = float(np.sqrt(np.dot(delta_rho, delta_rho) * cell.dV))
|
||
|
||
if verbose:
|
||
print(f"{iteration:>4} {e_total:>16.8f} {delta_E:>+12.2e} {conv:>12.2e}")
|
||
|
||
# ----- Density mixing -----
|
||
diis_rho.append(rho_new.copy())
|
||
diis_err.append(delta_rho.copy())
|
||
if len(diis_rho) > diis_size:
|
||
diis_rho.pop(0)
|
||
diis_err.pop(0)
|
||
|
||
if iteration >= diis_start and len(diis_err) >= 2:
|
||
rho = _pulay_mix(diis_rho, diis_err)
|
||
else:
|
||
rho = (1 - mix) * rho + mix * rho_new
|
||
|
||
# Enforce ρ ≥ 0 and particle conservation
|
||
rho = np.maximum(rho, 0.0)
|
||
rho *= n_elec / (np.sum(rho) * cell.dV)
|
||
|
||
if conv < tol and iteration > 2:
|
||
result.converged = True
|
||
break
|
||
|
||
result.n_iter = iteration
|
||
result.e_total = e_total
|
||
result.e_kinetic = e_kin
|
||
result.e_hartree = e_hartree
|
||
result.e_xc = e_xc
|
||
result.e_ion_elec = e_ion_elec
|
||
result.eigenvalues = evals
|
||
result.rho_r = rho_new
|
||
|
||
if verbose:
|
||
print()
|
||
print(result)
|
||
|
||
return result
|
||
|
||
|
||
def _pulay_mix(rho_hist: list, err_hist: list) -> np.ndarray:
|
||
"""Pulay DIIS: optimal linear combination of stored densities."""
|
||
m = len(err_hist)
|
||
B = np.zeros((m + 1, m + 1))
|
||
for i in range(m):
|
||
for j in range(m):
|
||
B[i, j] = np.dot(err_hist[i], err_hist[j])
|
||
B[m, :] = -1.0
|
||
B[:, m] = -1.0
|
||
B[m, m] = 0.0
|
||
rhs = np.zeros(m + 1)
|
||
rhs[m] = -1.0
|
||
try:
|
||
c = np.linalg.solve(B, rhs)[:m]
|
||
except np.linalg.LinAlgError:
|
||
c = np.zeros(m); c[-1] = 1.0
|
||
return sum(ci * ri for ci, ri in zip(c, rho_hist))
|