Coverage for eminus/energies.py: 95.35%

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Calculate different energy contributions.""" 

4 

5import dataclasses 

6import math 

7 

8from . import backend as xp 

9from . import config 

10from .dft import get_n_single, get_phi, H 

11from .extras import d3 

12from .gga import get_grad_field, get_tau 

13from .logger import log 

14from .tools import electronic_entropy 

15from .utils import handle_k 

16from .xc import get_exc 

17 

18 

19@dataclasses.dataclass 

20class Energy: 

21 """Energy class to save energy contributions in one place.""" 

22 

23 Ekin: float = 0 #: Kinetic energy. 

24 Ecoul: float = 0 #: Coulomb energy. 

25 Exc: float = 0 #: Exchange-correlation energy. 

26 Eloc: float = 0 #: Local energy. 

27 Enonloc: float = 0 #: Non-local energy. 

28 Eewald: float = 0 #: Ewald energy. 

29 Esic: float = 0 #: Self-interaction correction energy. 

30 Edisp: float = 0 #: Dispersion correction energy. 

31 Eentropy: float = 0 #: Fillings entropic energy. 

32 

33 @property 

34 def Etot(self): 

35 """Total energy is the sum of all energy contributions.""" 

36 return sum(getattr(self, ie.name) for ie in dataclasses.fields(self)) 

37 

38 def extrapolate(self): 

39 """Calculate the total energy at T=0. 

40 

41 Reference: J. Phys.: Condens. Matter 1, 689. 

42 

43 Returns: 

44 Total energy extrapolated to T=0. 

45 """ 

46 return self.Etot - 0.5 * self.Eentropy 

47 

48 def __repr__(self): 

49 """Print the energies stored in the Energy object.""" 

50 out = "" 

51 for ie in dataclasses.fields(self): 

52 energy = getattr(self, ie.name) 

53 if energy != 0: 

54 out += f"{ie.name:<9}: {energy:+.9f} Eh\n" 

55 return f"{out}{'-' * 26}\nEtot : {self.Etot:+.9f} Eh" 

56 

57 

58def get_E(scf): 

59 """Calculate energy contributions and update energies needed in one SCF step. 

60 

61 Args: 

62 scf: SCF object. 

63 

64 Returns: 

65 Total energy. 

66 """ 

67 scf.energies.Ekin = get_Ekin(scf.atoms, scf.Y) 

68 scf.energies.Ecoul = get_Ecoul(scf.atoms, scf.n, scf.phi) 

69 scf.energies.Exc = get_Exc(scf, scf.n, scf.exc, Nspin=scf.atoms.occ.Nspin) 

70 scf.energies.Eloc = get_Eloc(scf, scf.n) 

71 scf.energies.Enonloc = get_Enonloc(scf, scf.Y) 

72 return scf.energies.Etot 

73 

74 

75@handle_k(mode="reduce") 

76def get_Ekin(atoms, Y, ik): 

77 """Calculate the kinetic energy. 

78 

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

80 

81 Args: 

82 atoms: Atoms object. 

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

84 ik: k-point index. 

85 

86 Returns: 

87 Kinetic energy in Hartree. 

88 """ 

89 # Ekin = -0.5 Tr(F Wdag L(W)) 

90 Ekin = 0 

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

92 Ekin += ( 

93 -0.5 

94 * atoms.kpts.wk[ik] 

95 * xp.trace( 

96 xp.astype(atoms.occ.F[ik][spin], complex) @ Y[spin].conj().T @ atoms.L(Y[spin], ik) 

97 ) 

98 ) 

99 return float(xp.real(Ekin)) 

100 

101 

102def get_Ecoul(atoms, n, phi=None): 

103 """Calculate the Coulomb energy. 

104 

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

106 

107 Args: 

108 atoms: Atoms object. 

109 n: Real-space electronic density. 

110 

111 Keyword Args: 

112 phi: Hartree field. 

113 

114 Returns: 

115 Coulomb energy in Hartree. 

116 """ 

117 if phi is None: 

118 phi = get_phi(atoms, n) 

119 # Ecoul = 0.5 (J(n))dag O(phi) 

120 return float(xp.real(0.5 * xp.astype(n, complex) @ atoms.Jdag(atoms.O(phi)))) 

121 

122 

123def get_Exc(scf, n, exc=None, n_spin=None, dn_spin=None, tau=None, Nspin=2): 

124 """Calculate the exchange-correlation energy. 

125 

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

127 

128 Args: 

129 scf: SCF object. 

130 n: Real-space electronic density. 

131 

132 Keyword Args: 

133 exc: Exchange-correlation energy density. 

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

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

136 tau: Real-space kinetic energy densities per spin channel. 

137 Nspin: Number of spin states. 

138 

139 Returns: 

140 Exchange-correlation energy in Hartree. 

141 """ 

142 atoms = scf.atoms 

143 if exc is None: 

144 exc = get_exc(scf.xc, n_spin, Nspin, dn_spin, tau, scf.xc_params) 

145 # Exc = (J(n))dag O(J(exc)) 

146 return float(xp.real(xp.astype(n, complex) @ atoms.Jdag(atoms.O(atoms.J(exc))))) 

147 

148 

149def get_Eloc(scf, n): 

150 """Calculate the local energy contribution. 

151 

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

153 

154 Args: 

155 scf: SCF object. 

156 n: Real-space electronic density. 

157 

158 Returns: 

159 Local energy contribution in Hartree. 

160 """ 

161 return float(xp.real(xp.vdot(scf.Vloc, xp.astype(n, complex)))) 

162 

163 

164@handle_k(mode="reduce") 

165def get_Enonloc(scf, Y, ik): 

166 """Calculate the non-local GTH energy contribution. 

167 

168 Adapted from https://github.com/f-fathurrahman/PWDFT.jl/blob/master/src/calc_energies.jl 

169 

170 Reference: Phys. Rev. B 54, 1703. 

171 

172 Args: 

173 scf: SCF object. 

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

175 ik: k-point index. 

176 

177 Returns: 

178 Non-local GTH energy contribution in Hartree. 

179 """ 

180 atoms = scf.atoms 

181 

182 Enonloc = 0 

183 if scf.pot != "gth" or scf.gth.NbetaNL == 0: # Only calculate the non-local part if necessary 

184 return Enonloc 

185 

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

187 betaNL_psi = (Y[spin].conj().T @ scf.gth.betaNL[ik]).conj() 

188 

189 enl = xp.zeros(Y.shape[-1], dtype=complex) 

190 for ia in range(atoms.Natoms): 

191 psp = scf.gth[atoms.atom[ia]] 

192 for l in range(psp["lmax"]): 

193 for m in range(-l, l + 1): 

194 for iprj in range(psp["Nproj_l"][l]): 

195 ibeta = scf.gth.prj2beta[iprj, ia, l, m + psp["lmax"] - 1] - 1 

196 for jprj in range(psp["Nproj_l"][l]): 

197 jbeta = scf.gth.prj2beta[jprj, ia, l, m + psp["lmax"] - 1] - 1 

198 hij = psp["h"][l, iprj, jprj] 

199 enl += hij * betaNL_psi[:, ibeta].conj() * betaNL_psi[:, jbeta] 

200 Enonloc += xp.sum(atoms.occ.f[ik, spin] * atoms.kpts.wk[ik] * enl) 

201 # We have to multiply with the cell volume, because of different orthogonalization methods 

202 return float(xp.real(Enonloc * atoms.Omega)) 

203 

204 

205def get_Eewald(atoms, gcut=2, gamma=1e-8): 

206 """Calculate the Ewald energy. 

207 

208 Reference: J. Chem. Theory Comput. 10, 381. 

209 

210 Args: 

211 atoms: Atoms object. 

212 

213 Keyword Args: 

214 gcut: G-vector cut-off. 

215 gamma: Error tolerance. 

216 

217 Returns: 

218 Ewald energy in Hartree. 

219 """ 

220 if config.backend == "torch": 

221 erfc = xp.special.erfc 

222 else: 

223 from scipy.special import erfc 

224 

225 # For a plane wave code we have multiple contributions for the Ewald energy 

226 # Namely, a sum from contributions from real-space, reciprocal space, 

227 # the self energy, (the dipole term [neglected]), and an additional electroneutrality term 

228 def get_index_vectors(s): 

229 """Create index vectors of s periodic images.""" 

230 s = xp.to_np(s) 

231 m1 = xp.arange(-s[0], s[0] + 1) 

232 m2 = xp.arange(-s[1], s[1] + 1) 

233 m3 = xp.arange(-s[2], s[2] + 1) 

234 M = xp.permute_dims(xp.stack(xp.meshgrid(m1, m2, m3)), (3, 2, 1, 0)).reshape(-1, 3) 

235 # Delete the [0, 0, 0] element, to prevent checking for it in every loop 

236 return M[~(M == 0).all(axis=1)] 

237 

238 # Calculate the Ewald splitting parameter nu 

239 gexp = -math.log(gamma) 

240 nu = 0.5 * math.sqrt(gcut**2 / gexp) 

241 

242 # Start by calculating the self-energy 

243 Eewald = -nu / math.sqrt(math.pi) * xp.sum(atoms.Z**2) 

244 # Add the electroneutrality term 

245 Eewald += -math.pi * xp.sum(atoms.Z) ** 2 / (2 * nu**2 * atoms.Omega) 

246 

247 # Calculate the real-space contribution 

248 # Calculate the amount of images that have to be considered per axis 

249 Rm = xp.linalg.norm(atoms.a, axis=1) 

250 tmax = math.sqrt(0.5 * gexp) / nu 

251 s = xp.round(tmax / Rm + 1.5) 

252 # Collect all box index vectors in a matrix 

253 M = get_index_vectors(s) 

254 # Calculate the translation vectors 

255 T = M @ atoms.a 

256 

257 # Calculate the reciprocal space contribution 

258 # Calculate the amount of reciprocal images that have to be considered per axis 

259 g = 2 * math.pi * xp.linalg.inv(atoms.a.T) 

260 gm = xp.linalg.norm(g, axis=1) 

261 s = xp.round(gcut / gm + 1.5) 

262 # Collect all box index vectors in a matrix 

263 M = get_index_vectors(s) 

264 # Calculate the reciprocal translation vectors and precompute the prefactor 

265 G = M @ g 

266 G2 = xp.linalg.norm(G, axis=1) ** 2 

267 prefactor = 2 * math.pi / float(atoms.Omega) * xp.exp(-0.25 * G2 / nu**2) / G2 

268 

269 for ia in range(atoms.Natoms): 

270 for ja in range(atoms.Natoms): 

271 dpos = atoms.pos[ia] - atoms.pos[ja] 

272 ZiZj = atoms.Z[ia] * atoms.Z[ja] 

273 

274 # Add the real-space contribution 

275 rmag = xp.linalg.norm(dpos - T, axis=1) 

276 Eewald += 0.5 * ZiZj * xp.sum(erfc(rmag * nu) / rmag) 

277 # The T=[0, 0, 0] element is omitted in M but needed if ia!=ja, so add it manually 

278 if ia != ja: 

279 rmag = xp.linalg.norm(dpos) 

280 Eewald += 0.5 * ZiZj * erfc(rmag * nu) / rmag 

281 

282 # Add the reciprocal space contribution 

283 Gpos = xp.sum(G * dpos, axis=1) 

284 Eewald += ZiZj * xp.sum(prefactor * xp.cos(Gpos)) 

285 return float(Eewald) 

286 

287 

288def get_Esic(scf, Y, n_single=None): 

289 """Calculate the Perdew-Zunger self-interaction energy. 

290 

291 Reference: Phys. Rev. B 23, 5048. 

292 

293 Args: 

294 scf: SCF object. 

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

296 

297 Keyword Args: 

298 n_single: Single-electron densities. 

299 

300 Returns: 

301 PZ self-interaction energy. 

302 """ 

303 if Y is None: 

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

305 return None 

306 

307 atoms = scf.atoms 

308 # E_PZ-SIC = \sum_i Ecoul[n_i] + Exc[n_i, 0] 

309 if n_single is None: 

310 n_single = get_n_single(atoms, Y) 

311 

312 Esic = 0 

313 for i in range(atoms.occ.Nstate): 

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

315 if xp.sum(atoms.occ.f[:, spin, i] * atoms.kpts.wk) > 0: 

316 # Create normalized single-particle densities in the form of electronic densities 

317 # per spin channel, since spin-polarized functionals expect this form 

318 ni = xp.zeros((2, atoms.Ns)) 

319 # Normalize single-particle densities to 1 

320 ni[0] = n_single[spin, :, i] / xp.sum(atoms.occ.f[:, spin, i] * atoms.kpts.wk) 

321 

322 # Get the gradient of the single-particle density 

323 if "gga" in scf.xc_type: 

324 dni = xp.zeros((2, atoms.Ns, 3)) 

325 dni[0] = get_grad_field(atoms, ni)[0] 

326 else: 

327 dni = None 

328 

329 # Get the kinetic energy density of the corresponding orbital 

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

331 # Use only one orbital for the calculation 

332 Ytmp = [] 

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

334 Ytmp.append(xp.zeros_like(Y[ik])) 

335 Ytmp[ik][0, :, 0] = Y[ik][spin, :, i] 

336 taui = xp.zeros_like(ni) 

337 # We also have to normalize to one again 

338 taui[0] = get_tau(atoms, Ytmp)[0] / xp.sum( 

339 atoms.occ.f[:, spin, i] * atoms.kpts.wk 

340 ) 

341 else: 

342 taui = None 

343 

344 coul = get_Ecoul(atoms, ni[0]) 

345 # The exchange part for a SIC correction has to be spin-polarized 

346 xc = get_Exc(scf, ni[0], n_spin=ni, dn_spin=dni, tau=taui, Nspin=2) 

347 # SIC energy is scaled by the occupation number 

348 Esic += (coul + xc) * xp.sum(atoms.occ.f[:, spin, i] * atoms.kpts.wk) 

349 scf.energies.Esic = Esic 

350 return float(Esic) 

351 

352 

353def get_Edisp(scf, version="d3bj", atm=True, xc=None): # noqa: D103 

354 try: 

355 return d3.get_Edisp(scf, version, atm, xc) 

356 except ImportError: 

357 log.warning("You have to install the d3 extra to use this function.") 

358 return 0 

359 

360 

361get_Edisp.__doc__ = d3.get_Edisp.__doc__ 

362 

363 

364def get_Eband(scf, Y, **kwargs): 

365 """Calculate the band energy for occupied or unoccupied bands. 

366 

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

368 

369 Args: 

370 scf: SCF object. 

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

372 

373 Keyword Args: 

374 **kwargs: See :func:`eminus.dft.H`. 

375 

376 Returns: 

377 Band energy in Hartree. 

378 """ 

379 atoms = scf.atoms 

380 # Eband = Tr(Wdag H(W)) 

381 Eband = 0 

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

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

384 Eband += atoms.kpts.wk[ik] * xp.trace( 

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

386 ) 

387 return float(xp.real(Eband)) 

388 

389 

390def get_Eentropy(scf, epsilon, Efermi): 

391 """Calculate the fillings entropic energy. 

392 

393 Reference: J. Phys. Condens. Matter 1, 689. 

394 

395 Args: 

396 scf: SCF object. 

397 epsilon: Eigenenergies. 

398 Efermi: Fermi energy. 

399 

400 Returns: 

401 Entropic energy in Hartree. 

402 """ 

403 occ = scf.atoms.occ 

404 

405 Eentropy = 0 

406 for ik in range(scf.kpts.Nk): 

407 for spin in range(occ.Nspin): 

408 for i in range(occ.Nstate): 

409 # Beware the sign change, it is handled in the electronic_entropy function 

410 Eentropy += ( 

411 occ.wk[ik] 

412 * occ.smearing 

413 * electronic_entropy(epsilon[ik, spin, i], Efermi, occ.smearing) 

414 ) 

415 

416 Eentropy *= 2 / occ.Nspin 

417 scf.energies.Eentropy = float(Eentropy) 

418 return float(Eentropy)