Coverage for eminus/dft.py: 94.81%

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Main DFT functions based on the DFT++ formulation.""" 

4 

5import numpy as np 

6from numpy.linalg import multi_dot 

7from numpy.random import Generator, SFC64 

8from scipy.linalg import eig, eigh, eigvalsh, inv, sqrtm 

9 

10from .gga import calc_Vtau, get_grad_field, get_tau, gradient_correction 

11from .gth import calc_Vnonloc 

12from .logger import log 

13from .utils import handle_k, handle_spin, pseudo_uniform 

14from .xc import get_vxc 

15 

16 

17def get_phi(atoms, n): 

18 """Solve the Poisson equation. 

19 

20 Reference: Comput. Phys. Commun. 128, 1. 

21 

22 Args: 

23 atoms: Atoms object. 

24 n: Real-space electronic density. 

25 

26 Returns: 

27 Hartree field. 

28 """ 

29 # phi = -4 pi Linv(O(J(n))) 

30 return -4 * np.pi * atoms.Linv(atoms.O(atoms.J(n))) 

31 

32 

33@handle_k 

34@handle_spin 

35def orth(atoms, W): 

36 """Orthogonalize coefficient matrix W. 

37 

38 Reference: Comput. Phys. Commun. 128, 1. 

39 

40 Args: 

41 atoms: Atoms object. 

42 W: Expansion coefficients of unconstrained wave functions in reciprocal space. 

43 

44 Returns: 

45 Orthogonalized wave functions. 

46 """ 

47 # Y = W (Wdag O(W))^-0.5 

48 return W @ inv(sqrtm(W.conj().T @ atoms.O(W))) 

49 

50 

51def orth_unocc(atoms, Y, Z): 

52 """Orthogonalize unoccupied matrix Z while maintaining orthogonality to Y. 

53 

54 Reference: Comput. Phys. Commun. 128, 1. 

55 

56 Args: 

57 atoms: Atoms object. 

58 Y: Expansion coefficients of orthogonal wave functions in reciprocal space. 

59 Z: Expansion coefficients of unconstrained wave functions in reciprocal space. 

60 

61 Returns: 

62 Orthogonalized wave functions. 

63 """ 

64 D = [np.empty_like(Zk) for Zk in Z] 

65 for ik in range(atoms.kpts.Nk): 

66 for spin in range(atoms.occ.Nspin): 

67 # rhoZ = (I - Y Ydag O) Z 

68 Yocc = Y[ik][spin][:, atoms.occ.f[ik][spin] > 0] 

69 rhoZ = Z[ik][spin] - Yocc @ Yocc.conj().T @ atoms.O(Z[ik][spin]) 

70 # D = rhoZ (rhoZdag O(rhoZ))^-0.5 

71 D[ik][spin] = rhoZ @ inv(sqrtm(rhoZ.conj().T @ atoms.O(rhoZ))) 

72 return D 

73 

74 

75def get_n_total(atoms, Y, n_spin=None): 

76 """Calculate the total electronic density. 

77 

78 Reference: Comput. Phys. Commun. 128, 1. 

79 

80 Args: 

81 atoms: Atoms object. 

82 Y: Expansion coefficients of orthogonal wave functions in reciprocal space. 

83 

84 Keyword Args: 

85 n_spin: Real-space electronic densities per spin channel. 

86 

87 Returns: 

88 Electronic density. 

89 """ 

90 # Return the total density in the spin-paired case 

91 if n_spin is not None: 

92 return np.sum(n_spin, axis=0) 

93 

94 # n = (IW) F (IW)dag 

95 n = np.zeros(atoms.Ns) 

96 Yrs = atoms.I(Y) 

97 for ik in range(atoms.kpts.Nk): 

98 for spin in range(atoms.occ.Nspin): 

99 n += np.sum( 

100 atoms.occ.f[ik, spin] 

101 * atoms.kpts.wk[ik] 

102 * np.real(Yrs[ik][spin].conj() * Yrs[ik][spin]), 

103 axis=1, 

104 ) 

105 return n 

106 

107 

108@handle_k(mode="reduce") 

109def get_n_spin(atoms, Y, ik): 

110 """Calculate the electronic density per spin channel. 

111 

112 Reference: Comput. Phys. Commun. 128, 1. 

113 

114 Args: 

115 atoms: Atoms object. 

116 Y: Expansion coefficients of orthogonal wave functions in reciprocal space. 

117 ik: k-point index. 

118 

119 Returns: 

120 Electronic densities per spin channel. 

121 """ 

122 Yrs = atoms.I(Y, ik) 

123 n = np.empty((atoms.occ.Nspin, atoms.Ns)) 

124 for spin in range(atoms.occ.Nspin): 

125 n[spin] = np.sum( 

126 atoms.occ.f[ik, spin] * atoms.kpts.wk[ik] * np.real(Yrs[spin].conj() * Yrs[spin]), 

127 axis=1, 

128 ) 

129 return n 

130 

131 

132@handle_k(mode="reduce") 

133def get_n_single(atoms, Y, ik): 

134 """Calculate the single-electron densities. 

135 

136 Args: 

137 atoms: Atoms object. 

138 Y: Expansion coefficients of orthogonal wave functions in reciprocal space. 

139 ik: k-point index. 

140 

141 Returns: 

142 Single-electron densities. 

143 """ 

144 Yrs = atoms.I(Y, ik) 

145 n = np.empty((atoms.occ.Nspin, atoms.Ns, atoms.occ.Nstate)) 

146 for spin in range(atoms.occ.Nspin): 

147 n[spin] = atoms.occ.f[ik, spin] * atoms.kpts.wk[ik] * np.real(Yrs[spin].conj() * Yrs[spin]) 

148 return n 

149 

150 

151def get_grad(scf, ik, spin, W, **kwargs): 

152 """Calculate the energy gradient with respect to W. 

153 

154 Reference: Comput. Phys. Commun. 128, 1. 

155 

156 Args: 

157 scf: SCF object. 

158 ik: k-point index. 

159 spin: Spin variable to track whether to do the calculation for spin up or down. 

160 W: Expansion coefficients of unconstrained wave functions in reciprocal space. 

161 

162 Keyword Args: 

163 **kwargs: See :func:`H`. 

164 

165 Returns: 

166 Gradient. 

167 """ 

168 atoms = scf.atoms 

169 F = atoms.occ.F[ik][spin] 

170 HW = H(scf, ik, spin, W, **kwargs) 

171 WHW = W[ik][spin].conj().T @ HW 

172 # U = Wdag O(W) 

173 OW = atoms.O(W[ik][spin]) 

174 U = W[ik][spin].conj().T @ OW 

175 invU = inv(U) 

176 U12 = sqrtm(invU) 

177 # Htilde = U^-0.5 Wdag H(W) U^-0.5 

178 Ht = multi_dot([U12, WHW, U12]) 

179 # grad E = H(W) - O(W) U^-1 (Wdag H(W)) (U^-0.5 F U^-0.5) + O(W) (U^-0.5 Q(Htilde F - F Htilde)) 

180 tmp = multi_dot([OW, invU, WHW]) 

181 return atoms.kpts.wk[ik] * ( 

182 multi_dot([HW - tmp, U12, F, U12]) + multi_dot([OW, U12, Q(Ht @ F - F @ Ht, U)]) 

183 ) 

184 

185 

186def H(scf, ik, spin, W, dn_spin=None, phi=None, vxc=None, vsigma=None, vtau=None): 

187 """Left-hand side of the eigenvalue equation. 

188 

189 Reference: Comput. Phys. Commun. 128, 1. 

190 

191 Args: 

192 scf: SCF object. 

193 ik: k-point index. 

194 spin: Spin variable to track whether to do the calculation for spin up or down. 

195 W: Expansion coefficients of unconstrained wave functions in reciprocal space. 

196 

197 Keyword Args: 

198 dn_spin: Real-space gradient of densities per spin channel. 

199 phi: Hartree field. 

200 vxc: Exchange-correlation potential. 

201 vsigma: Contracted gradient potential derivative. 

202 vtau: Kinetic energy gradient potential derivative. 

203 

204 Returns: 

205 Hamiltonian applied on W. 

206 """ 

207 atoms = scf.atoms 

208 

209 # If dn_spin is None all other keyword arguments are None by design 

210 # In that case precompute values from the SCF class 

211 if phi is None: 

212 dn_spin, phi, vxc, vsigma, vtau = H_precompute(scf, W) 

213 

214 # This calculates the XC potential in the reciprocal space 

215 Gvxc = atoms.J(vxc[spin]) 

216 # Calculate the gradient correction to the potential if a (meta-)GGA functional is used 

217 if "gga" in scf.xc_type: 

218 Gvxc = Gvxc - gradient_correction(atoms, spin, dn_spin, vsigma) 

219 # Vkin = -0.5 L(W) 

220 Vkin_psi = -0.5 * atoms.L(W[ik][spin], ik) 

221 # Veff = Jdag(Vion) + Jdag(O(J(vxc))) + Jdag(O(phi)) 

222 # We get the full potential in the functional definition (different to the DFT++ notation) 

223 # Normally Vxc = Jdag(O(J(exc))) + diag(exc') Jdag(O(J(n))) (for LDA functionals) 

224 Veff = scf.Vloc + atoms.Jdag(atoms.O(Gvxc + phi)) 

225 Vnonloc_psi = calc_Vnonloc(scf, ik, spin, W) 

226 Vtau_psi = calc_Vtau(scf, ik, spin, W, vtau) 

227 # H = Vkin + Idag(diag(Veff))I + Vnonloc (+ Vtau) 

228 # Diag(a) * B can be written as a * B if a is a column vector 

229 return ( 

230 Vkin_psi + atoms.Idag(Veff[:, None] * atoms.I(W[ik][spin], ik), ik) + Vnonloc_psi + Vtau_psi 

231 ) 

232 

233 

234def H_precompute(scf, W): 

235 """Create precomputed values as intermediate results. 

236 

237 Args: 

238 scf: SCF object. 

239 W: Expansion coefficients of unconstrained wave functions in reciprocal space. 

240 

241 Returns: 

242 dn_spin, phi, vxc, vsigma, and vtau. 

243 """ 

244 # We are calculating everything here over both spin channels 

245 # We would be fine/could improve performance by only calculating the currently used spin channel 

246 atoms = scf.atoms 

247 Y = orth(atoms, W) 

248 n_spin = get_n_spin(atoms, Y) 

249 n = get_n_total(atoms, Y, n_spin) 

250 if "gga" in scf.xc_type: 

251 dn_spin = get_grad_field(atoms, n_spin) 

252 else: 

253 dn_spin = None 

254 if scf.xc_type == "meta-gga": 

255 tau = get_tau(atoms, Y) 

256 else: 

257 tau = None 

258 phi = get_phi(atoms, n) 

259 vxc, vsigma, vtau = get_vxc(scf.xc, n_spin, atoms.occ.Nspin, dn_spin, tau, scf.xc_params) 

260 return dn_spin, phi, vxc, vsigma, vtau 

261 

262 

263def Q(inp, U): 

264 """Operator needed to calculate gradients with non-constant occupations. 

265 

266 Reference: Comput. Phys. Commun. 128, 1. 

267 

268 Args: 

269 inp: Coefficients input array. 

270 U: Overlap of wave functions. 

271 

272 Returns: 

273 Q operator result. 

274 """ 

275 mu, V = eig(U) 

276 mu = mu[:, None] 

277 denom = np.sqrt(mu) @ np.ones((1, len(mu))) 

278 denom2 = denom + denom.conj().T 

279 # return V @ ((V.conj().T @ inp @ V) / denom2) @ V.conj().T 

280 tmp = multi_dot([V.conj().T, inp, V]) 

281 return multi_dot([V, tmp / denom2, V.conj().T]) 

282 

283 

284def get_psi(scf, W, **kwargs): 

285 """Calculate eigenstates from H. 

286 

287 Reference: Comput. Phys. Commun. 128, 1. 

288 

289 Args: 

290 scf: SCF object. 

291 W: Expansion coefficients of unconstrained wave functions in reciprocal space. 

292 

293 Keyword Args: 

294 **kwargs: See :func:`H`. 

295 

296 Returns: 

297 Eigenstates in reciprocal space. 

298 """ 

299 if W is None: 

300 log.warning('The provided wave function is "None".') 

301 return None 

302 

303 atoms = scf.atoms 

304 Y = orth(atoms, W) 

305 psi = [np.empty_like(Yk) for Yk in Y] 

306 for ik in range(atoms.kpts.Nk): 

307 for spin in range(atoms.occ.Nspin): 

308 mu = Y[ik][spin].conj().T @ H(scf, ik, spin, Y, **kwargs) 

309 _, D = eigh(mu) 

310 psi[ik][spin] = Y[ik][spin] @ D 

311 return psi 

312 

313 

314def get_epsilon(scf, W, **kwargs): 

315 """Calculate eigenvalues from H. 

316 

317 Reference: Comput. Phys. Commun. 128, 1. 

318 

319 Args: 

320 scf: SCF object. 

321 W: Expansion coefficients of unconstrained wave functions in reciprocal space. 

322 

323 Keyword Args: 

324 **kwargs: See :func:`H`. 

325 

326 Returns: 

327 Eigenvalues. 

328 """ 

329 if W is None: 

330 log.warning('The provided wave function is "None".') 

331 return None 

332 

333 atoms = scf.atoms 

334 Y = orth(atoms, W) 

335 epsilon = np.empty((atoms.kpts.Nk, atoms.occ.Nspin, Y[0].shape[-1])) 

336 for ik in range(atoms.kpts.Nk): 

337 for spin in range(atoms.occ.Nspin): 

338 mu = Y[ik][spin].conj().T @ H(scf, ik, spin, Y, **kwargs) 

339 epsilon[ik][spin] = np.sort(eigvalsh(mu)) 

340 return epsilon 

341 

342 

343def get_epsilon_unocc(scf, W, Z, **kwargs): 

344 """Calculate eigenvalues from H of unoccupied states. 

345 

346 Reference: Comput. Phys. Commun. 128, 1. 

347 

348 Args: 

349 scf: SCF object. 

350 W: Expansion coefficients of unconstrained wave functions in reciprocal space. 

351 Z: Expansion coefficients of unconstrained wave functions in reciprocal space. 

352 

353 Keyword Args: 

354 **kwargs: See :func:`H`. 

355 

356 Returns: 

357 Eigenvalues. 

358 """ 

359 if W is None or Z is None: 

360 log.warning('The provided wave function is "None".') 

361 return None 

362 

363 atoms = scf.atoms 

364 Y = orth(atoms, W) 

365 D = orth_unocc(atoms, Y, Z) 

366 epsilon = np.empty((atoms.kpts.Nk, atoms.occ.Nspin, D[0].shape[-1])) 

367 for ik in range(atoms.kpts.Nk): 

368 for spin in range(atoms.occ.Nspin): 

369 mu = D[ik][spin].conj().T @ H(scf, ik, spin, D, **kwargs) 

370 epsilon[ik][spin] = np.sort(eigvalsh(mu)) 

371 return epsilon 

372 

373 

374def guess_random(scf, Nstate=None, seed=42, symmetric=False): 

375 """Generate random initial-guess coefficients as starting values. 

376 

377 Args: 

378 scf: SCF object. 

379 

380 Keyword Args: 

381 Nstate: Number of states. 

382 seed: Seed to initialize the random number generator. 

383 symmetric: Whether to use the same guess for both spin channels. 

384 

385 Returns: 

386 Initial-guess orthogonal wave functions in reciprocal space. 

387 """ 

388 atoms = scf.atoms 

389 if Nstate is None: 

390 Nstate = atoms.occ.Nstate 

391 

392 rng = Generator(SFC64(seed)) 

393 W = [] 

394 for ik in range(atoms.kpts.Nk): 

395 if symmetric: 

396 W_ik = rng.standard_normal((len(atoms.Gk2c[ik]), Nstate)) + 1j * rng.standard_normal( 

397 (len(atoms.Gk2c[ik]), Nstate) 

398 ) 

399 W.append(np.array([W_ik] * atoms.occ.Nspin)) 

400 else: 

401 W_ik = rng.standard_normal( 

402 (atoms.occ.Nspin, len(atoms.Gk2c[ik]), Nstate) 

403 ) + 1j * rng.standard_normal((atoms.occ.Nspin, len(atoms.Gk2c[ik]), Nstate)) 

404 W.append(W_ik) 

405 return orth(atoms, W) 

406 

407 

408def guess_pseudo(scf, Nstate=None, seed=1234, symmetric=False): 

409 """Generate initial-guess coefficients using pseudo-random starting values. 

410 

411 Args: 

412 scf: SCF object. 

413 

414 Keyword Args: 

415 Nstate: Number of states. 

416 seed: Seed to initialize the random number generator. 

417 symmetric: Whether to use the same guess for both spin channels. 

418 

419 Returns: 

420 Initial-guess orthogonal wave functions in reciprocal space. 

421 """ 

422 atoms = scf.atoms 

423 if Nstate is None: 

424 Nstate = atoms.occ.Nstate 

425 

426 W = [] 

427 for ik in range(atoms.kpts.Nk): 

428 if symmetric: 

429 W_ik = pseudo_uniform((1, len(atoms.Gk2c[ik]), Nstate), seed=seed) 

430 W.append(np.array([W_ik[0]] * atoms.occ.Nspin)) 

431 else: 

432 W_ik = pseudo_uniform((atoms.occ.Nspin, len(atoms.Gk2c[ik]), Nstate), seed=seed) 

433 W.append(W_ik) 

434 return orth(atoms, W)