1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
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
|
