Coverage for eminus/dft.py: 96.41%

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 = 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 # Extra parentheses are used to set the optimal ordering, reducing costs for large arrays 

181 tmp = OW @ (invU @ WHW) 

182 return atoms.kpts.wk[ik] * ((HW - tmp) @ (U12 @ (F @ U12))) + OW @ (U12 @ Q(Ht @ F - F @ Ht, U)) 

183 

184 

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

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

187 

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

189 

190 Args: 

191 scf: SCF object. 

192 ik: k-point index. 

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

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

195 

196 Keyword Args: 

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

198 phi: Hartree field. 

199 vxc: Exchange-correlation potential. 

200 vsigma: Contracted gradient potential derivative. 

201 vtau: Kinetic energy gradient potential derivative. 

202 

203 Returns: 

204 Hamiltonian applied on W. 

205 """ 

206 atoms = scf.atoms 

207 

208 # If vxc is None all other keyword arguments are None by design 

209 # In that case precompute values from the SCF class 

210 if vxc is None: 

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

212 

213 # This calculates the XC potential in the reciprocal space 

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

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

216 if "gga" in scf.xc_type: 

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

218 # Vkin = -0.5 L(W) 

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

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

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

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

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

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

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

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

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

228 return ( 

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

230 ) 

231 

232 

233def H_precompute(scf, W): 

234 """Create precomputed values as intermediate results. 

235 

236 Args: 

237 scf: SCF object. 

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

239 

240 Returns: 

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

242 """ 

243 # We are calculating everything here over both spin channels 

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

245 atoms = scf.atoms 

246 Y = orth(atoms, W) 

247 n_spin = get_n_spin(atoms, Y) 

248 n = get_n_total(atoms, Y, n_spin) 

249 if "gga" in scf.xc_type: 

250 dn_spin = get_grad_field(atoms, n_spin) 

251 else: 

252 dn_spin = None 

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

254 tau = get_tau(atoms, Y) 

255 else: 

256 tau = None 

257 phi = get_phi(atoms, n) 

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

259 return dn_spin, phi, vxc, vsigma, vtau 

260 

261 

262def Q(inp, U): 

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

264 

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

266 

267 Args: 

268 inp: Coefficients input array. 

269 U: Overlap of wave functions. 

270 

271 Returns: 

272 Q operator result. 

273 """ 

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

275 mu = mu[:, None] 

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

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

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

279 tmp = V.conj().T @ inp @ V 

280 return V @ (tmp / denom2) @ V.conj().T 

281 

282 

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

284 """Calculate eigenstates from H. 

285 

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

287 

288 Args: 

289 scf: SCF object. 

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

291 

292 Keyword Args: 

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

294 

295 Returns: 

296 Eigenstates in reciprocal space. 

297 """ 

298 if W is None: 

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

300 return None 

301 

302 atoms = scf.atoms 

303 Y = orth(atoms, W) 

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

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

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

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

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

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

310 return psi 

311 

312 

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

314 """Calculate eigenvalues from H. 

315 

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

317 

318 Args: 

319 scf: SCF object. 

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

321 

322 Keyword Args: 

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

324 

325 Returns: 

326 Eigenvalues. 

327 """ 

328 if W is None: 

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

330 return None 

331 

332 atoms = scf.atoms 

333 Y = orth(atoms, W) 

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

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

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

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

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

339 return epsilon 

340 

341 

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

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

344 

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

346 

347 Args: 

348 scf: SCF object. 

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

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

351 

352 Keyword Args: 

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

354 

355 Returns: 

356 Eigenvalues. 

357 """ 

358 if W is None or Z is None: 

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

360 return None 

361 

362 atoms = scf.atoms 

363 Y = orth(atoms, W) 

364 D = orth_unocc(atoms, Y, Z) 

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

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

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

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

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

370 return epsilon 

371 

372 

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

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

375 

376 Args: 

377 scf: SCF object. 

378 

379 Keyword Args: 

380 Nstate: Number of states. 

381 seed: Seed to initialize the random number generator. 

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

383 

384 Returns: 

385 Initial-guess orthogonal wave functions in reciprocal space. 

386 """ 

387 atoms = scf.atoms 

388 if Nstate is None: 

389 Nstate = atoms.occ.Nstate 

390 

391 rng = Generator(SFC64(seed)) 

392 W = [] 

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

394 if symmetric: 

395 W_ik = xp.asarray( 

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

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

398 ) 

399 W.append(xp.stack([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(xp.asarray(W_ik)) 

405 return orth(atoms, W) 

406 

407 

408def guess_filter(scf, Nstate=None, seed=42, symmetric=False): 

409 """Generate filtered random initial-guess coefficients as starting values. 

410 

411 A weighting corresponding to a Gaussian distribution in G^2 is applied to the random values, 

412 which acts similar to a low-pass filter. 

413 

414 Reference: https://doi.org/10.1017/CBO9780511609633 

415 

416 Args: 

417 scf: SCF object. 

418 

419 Keyword Args: 

420 Nstate: Number of states. 

421 seed: Seed to initialize the random number generator. 

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

423 

424 Returns: 

425 Initial-guess orthogonal wave functions in reciprocal space. 

426 """ 

427 atoms = scf.atoms 

428 if Nstate is None: 

429 Nstate = atoms.occ.Nstate 

430 

431 rng = Generator(SFC64(seed)) 

432 W = [] 

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

434 weight = 1 / math.sqrt(2 * math.pi) * xp.exp(-atoms.Gk2c[ik] / 2) 

435 if symmetric: 

436 W_ik = xp.asarray( 

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

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

439 ) 

440 W.append(weight[:, None] * xp.stack([W_ik] * atoms.occ.Nspin)) 

441 else: 

442 W_ik = rng.standard_normal( 

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

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

445 W.append(weight[:, None] * xp.asarray(W_ik)) 

446 return orth(atoms, W) 

447 

448 

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

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

451 

452 Args: 

453 scf: SCF object. 

454 

455 Keyword Args: 

456 Nstate: Number of states. 

457 seed: Seed to initialize the random number generator. 

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

459 

460 Returns: 

461 Initial-guess orthogonal wave functions in reciprocal space. 

462 """ 

463 atoms = scf.atoms 

464 if Nstate is None: 

465 Nstate = atoms.occ.Nstate 

466 

467 W = [] 

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

469 if symmetric: 

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

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

472 else: 

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

474 W.append(W_ik) 

475 return orth(atoms, W)