Coverage for eminus/dft.py: 96.08%

1# SPDX-FileCopyrightText: 2022 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

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

4 

5import math 

6 

7from numpy.random import Generator, SFC64 

8 

9from . import backend as xp 

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 * math.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 @ xp.linalg.inv(xp.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 = [xp.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 @ xp.linalg.inv(xp.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 xp.sum(n_spin, axis=0) 

93 

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

95 n = xp.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 += xp.sum( 

100 atoms.occ.f[ik, spin] 

101 * atoms.kpts.wk[ik] 

102 * xp.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 = xp.empty((atoms.occ.Nspin, atoms.Ns)) 

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

125 n[spin] = xp.sum( 

126 atoms.occ.f[ik, spin] * atoms.kpts.wk[ik] * xp.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 = xp.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] * xp.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 = xp.asarray(atoms.occ.F[ik][spin], dtype=complex) 

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 = xp.linalg.inv(U) 

176 U12 = xp.sqrtm(invU) 

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

178 Ht = xp.linalg.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 = xp.linalg.multi_dot([OW, invU, WHW]) 

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

182 xp.linalg.multi_dot([HW - tmp, U12, F, U12]) 

183 + xp.linalg.multi_dot([OW, U12, Q(Ht @ F - F @ Ht, U)]) 

184 ) 

185 

186 

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

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

189 

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

191 

192 Args: 

193 scf: SCF object. 

194 ik: k-point index. 

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

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

197 

198 Keyword Args: 

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

200 phi: Hartree field. 

201 vxc: Exchange-correlation potential. 

202 vsigma: Contracted gradient potential derivative. 

203 vtau: Kinetic energy gradient potential derivative. 

204 

205 Returns: 

206 Hamiltonian applied on W. 

207 """ 

208 atoms = scf.atoms 

209 

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

211 # In that case precompute values from the SCF class 

212 if phi is None: 

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

214 

215 # This calculates the XC potential in the reciprocal space 

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

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

218 if "gga" in scf.xc_type: 

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

220 # Vkin = -0.5 L(W) 

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

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

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

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

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

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

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

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

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

230 return ( 

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

232 ) 

233 

234 

235def H_precompute(scf, W): 

236 """Create precomputed values as intermediate results. 

237 

238 Args: 

239 scf: SCF object. 

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

241 

242 Returns: 

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

244 """ 

245 # We are calculating everything here over both spin channels 

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

247 atoms = scf.atoms 

248 Y = orth(atoms, W) 

249 n_spin = get_n_spin(atoms, Y) 

250 n = get_n_total(atoms, Y, n_spin) 

251 if "gga" in scf.xc_type: 

252 dn_spin = get_grad_field(atoms, n_spin) 

253 else: 

254 dn_spin = None 

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

256 tau = get_tau(atoms, Y) 

257 else: 

258 tau = None 

259 phi = get_phi(atoms, n) 

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

261 return dn_spin, phi, vxc, vsigma, vtau 

262 

263 

264def Q(inp, U): 

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

266 

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

268 

269 Args: 

270 inp: Coefficients input array. 

271 U: Overlap of wave functions. 

272 

273 Returns: 

274 Q operator result. 

275 """ 

276 mu, V = xp.linalg.eig(U) 

277 mu = mu[:, None] 

278 denom = xp.sqrt(mu) @ xp.ones((1, len(mu)), dtype=complex) 

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

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

281 tmp = xp.linalg.multi_dot([V.conj().T, inp, V]) 

282 return xp.linalg.multi_dot([V, tmp / denom2, V.conj().T]) 

283 

284 

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

286 """Calculate eigenstates from H. 

287 

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

289 

290 Args: 

291 scf: SCF object. 

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

293 

294 Keyword Args: 

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

296 

297 Returns: 

298 Eigenstates in reciprocal space. 

299 """ 

300 if W is None: 

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

302 return None 

303 

304 atoms = scf.atoms 

305 Y = orth(atoms, W) 

306 psi = [xp.empty_like(Yk) for Yk in Y] 

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

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

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

310 _, D = xp.linalg.eigh(mu) 

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

312 return psi 

313 

314 

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

316 """Calculate eigenvalues from H. 

317 

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

319 

320 Args: 

321 scf: SCF object. 

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

323 

324 Keyword Args: 

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

326 

327 Returns: 

328 Eigenvalues. 

329 """ 

330 if W is None: 

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

332 return None 

333 

334 atoms = scf.atoms 

335 Y = orth(atoms, W) 

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

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

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

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

340 epsilon[ik][spin] = xp.sort(xp.linalg.eigvalsh(mu)) 

341 return epsilon 

342 

343 

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

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

346 

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

348 

349 Args: 

350 scf: SCF object. 

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

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

353 

354 Keyword Args: 

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

356 

357 Returns: 

358 Eigenvalues. 

359 """ 

360 if W is None or Z is None: 

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

362 return None 

363 

364 atoms = scf.atoms 

365 Y = orth(atoms, W) 

366 D = orth_unocc(atoms, Y, Z) 

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

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

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

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

371 epsilon[ik][spin] = xp.sort(xp.linalg.eigvalsh(mu)) 

372 return epsilon 

373 

374 

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

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

377 

378 Args: 

379 scf: SCF object. 

380 

381 Keyword Args: 

382 Nstate: Number of states. 

383 seed: Seed to initialize the random number generator. 

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

385 

386 Returns: 

387 Initial-guess orthogonal wave functions in reciprocal space. 

388 """ 

389 atoms = scf.atoms 

390 if Nstate is None: 

391 Nstate = atoms.occ.Nstate 

392 

393 rng = Generator(SFC64(seed)) 

394 W = [] 

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

396 if symmetric: 

397 W_ik = xp.asarray( 

398 rng.standard_normal((len(atoms.Gk2c[ik]), Nstate)) 

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

400 ) 

401 W.append(xp.stack([W_ik] * atoms.occ.Nspin)) 

402 else: 

403 W_ik = rng.standard_normal( 

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

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

406 W.append(xp.asarray(W_ik)) 

407 return orth(atoms, W) 

408 

409 

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

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

412 

413 Args: 

414 scf: SCF object. 

415 

416 Keyword Args: 

417 Nstate: Number of states. 

418 seed: Seed to initialize the random number generator. 

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

420 

421 Returns: 

422 Initial-guess orthogonal wave functions in reciprocal space. 

423 """ 

424 atoms = scf.atoms 

425 if Nstate is None: 

426 Nstate = atoms.occ.Nstate 

427 

428 W = [] 

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

430 if symmetric: 

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

432 W.append(xp.stack([W_ik[0]] * atoms.occ.Nspin)) 

433 else: 

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

435 W.append(W_ik) 

436 return orth(atoms, W)