diff options
Diffstat (limited to 'eigen.py')
| -rw-r--r-- | eigen.py | 187 |
1 files changed, 187 insertions, 0 deletions
diff --git a/eigen.py b/eigen.py new file mode 100644 index 0000000..1dbb2a9 --- /dev/null +++ b/eigen.py @@ -0,0 +1,187 @@ +from math import * + +import numpy as np +import numpy.fft + +import visualize + + +# Normalize wavefunction with respect to volume element dV +def normalize(state, dV): + norm2 = np.vdot(state, state) * dV + return state / sqrt(abs(norm2)) + + +# Gram-Schmidt orthonormalization with respect to dV +def orthonormalize(basis, dV): + for n in range(len(basis)): + psi = basis[n] + + for i in range(n): + psi -= basis[i] * np.vdot(basis[i], psi) * dV + + basis[n] = normalize(psi, dV) + + return basis + + +# Zero all components with kinetic energy less than Ecut +def enforce_cutoff(s, psis_q, kq2): + for i in range(s.Nres): + for j in range(s.Nres): + for k in range(s.Nres): + if kq2[i, j, k] > 2 * s.Ecut: + for psi_q in psis_q: + psi_q[i, j, k] = 0 + return orthonormalize(psis_q, s.dq) + + +# Apply the Kohn-Sham Hamiltonian to a normalized state psi +def apply_hamiltonian(s, psi_q, Veff_r, kq2): + # Apply kinetic energy in q-domain + Tpsi_q = 0.5 * kq2 * psi_q + + # Apply local+nonlocal potential in r-domain + psi_r = np.fft.ifftn(psi_q) + # We don't normalize: the iFFT ensures a correctly normalized result + Vpsi_r = Veff_r * psi_r # local part + for a in s.atoms: # nonlocal part + Vnon_r = a.apply_nonlocal_potential(psi_r) * s.dr + Vpsi_r += Vnon_r + Vpsi_q = np.fft.fftn(Vpsi_r) + + return Tpsi_q + Vpsi_q + + +# Resisual vector R, which the eigensolver tries to minimize +def get_residual(s, psi_q, Hpsi_q, kq2): + ene = np.vdot(psi_q, Hpsi_q) * s.dq # assuming normalization + R_q = Hpsi_q - ene * psi_q + + # Calculate preconditioned residual PR for faster convergence: + # The ideal PR satisfies PR = (H - ene)^-1 R, such that (H - ene) PR = R. + # This operator (H - ene) is hard to invert, so we approximate it using + # the kinetic energy only: -0.5 \nabla^2 PR = R (the Poisson equation). + PR_q = 2 * R_q / (kq2 + 1) + # This is an even rougher approximation, but it works well enough. + + return PR_q + + +# One iteration of the block Davidson eigensolver algorithm +def diagonalize_subspace(s, psis, V, kq2): + # Apply Hamiltonian to psis and cache results + Hbasis = [] + for n in range(len(psis)): + Hbasis.append(apply_hamiltonian(s, psis[n], V, kq2)) + + # Calculate corresponding residuals + resids = [] + for n in range(len(psis)): + R = get_residual(s, psis[n], Hbasis[n], kq2) + resids.append(R) + + # Augment the input basis (psis) with the residuals + basis = psis + resids + basis = enforce_cutoff(s, basis, kq2) + + Ndim = len(basis) + + # Apply Hamiltonian to orthonormalized residuals and cache results + for n in range(len(psis), Ndim): + Hbasis.append(apply_hamiltonian(s, basis[n], V, kq2)) + + # Subspace Hamiltonian matrix with elements <basis[i] | H | basis[j]> + Hsub = np.zeros((Ndim, Ndim), dtype="complex128") + # Exploit Hermiticity: only calculate the diagonal and upper triangle, + # then copy the elements from the upper to the lower triangle. + for r in range(Ndim): + for c in range(r, Ndim): + Hsub[r, c] = np.vdot(basis[r], Hbasis[c]) * s.dq + for r in range(Ndim): + for c in range(0, r): + Hsub[r, c] = np.conj(Hsub[c, r]) + + # Find the eigenvectors and eigenvalues of Hsub + evals, evecs = np.linalg.eigh(Hsub) + + # Express the eigenvectors in the original basis (input psis) + result = [] + for n in range(Ndim // 2): + psi = np.zeros((s.Nres, s.Nres, s.Nres), dtype="complex128") + for b in range(Ndim): + psi += basis[b] * evecs[b, n] + result.append(psi) + + return evals, result + + +# Small optimization: cache |k+q|^2 ahead of time +def cache_kinetic_factor(s, kpoint): + kq2 = np.zeros((s.Nres, s.Nres, s.Nres)) + for i in range(s.Nres): + for j in range(s.Nres): + for k in range(s.Nres): + q = np.array((s.qx[i, j, k], s.qy[i, j, k], s.qz[i, j, k])) + kq = kpoint + q + kq2[i, j, k] = np.dot(kq, kq) + return kq2 + + +# Initial guess from atoms' individual (approximate) orbitals +def guess_orbitals(s): + psis_q = [] + for a in s.atoms: + psis_r = a.guess_orbitals(s.Norb) + + # Convert guesses to Fourier space + for psi_r in psis_r: + psi_q = np.fft.fftn(psi_r) + psis_q.append(psi_q) + + return orthonormalize(psis_q, s.dq) + + +# Get lowest eigenvalues and eigenstates of Kohn-Sham equation +def solve_kohnsham_equation(s, V, k, prefix="", thresh=1e-4): + kq2 = cache_kinetic_factor(s, k) + + # Initialize psis, which we will refine shortly + psis_q = guess_orbitals(s) + + print(prefix + "Eigensolver cycle 1 running...") + cycles = 0 + improv = 1 + energy_old = 0 + # Iterative eigensolver loop + while improv > thresh: + cycles += 1 + + # Eliminate high-frequency noise (little effect in practice) + psis_q = enforce_cutoff(s, psis_q, kq2) + + # Refine eigenstates using one step of the block Davidson method + energies, psis_q = diagonalize_subspace(s, psis_q, V, kq2) + + # To measure progress, we look at the change of (N+1)th eigenvalue, + # where N is the number of bands we are interested in converging. + energy_new = energies[s.Neig] + improv = abs(energy_new - energy_old) / abs(energy_new) + energy_old = energy_new + + # Clear line and (re)print progress information + print("\033[A\033[2K" + prefix, end="") + print("Eigensolver cycle {} done".format(cycles), end="") + if cycles > 1: + print(", improvement {:.2e}".format(improv), end="") + print("") + + # Convert psis to real space and normalize them there + psis_r = [] + for n in range(len(psis_q)): + psi_r = np.fft.ifftn(psis_q[n]) + psis_r.append(normalize(psi_r, s.dr)) + + # energies is a NumPy array right now + return energies.tolist(), psis_r, psis_q + |