Coverage for eminus/energies.py: 95.32%

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 m1 = xp.arange(-s[0], s[0] + 1) 

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

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

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

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

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

236 

237 # Calculate the Ewald splitting parameter nu 

238 gexp = -math.log(gamma) 

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

240 

241 # Start by calculating the self-energy 

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

243 # Add the electroneutrality term 

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

245 

246 # Calculate the real-space contribution 

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

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

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

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

251 # Collect all box index vectors in a matrix 

252 M = get_index_vectors(s) 

253 # Calculate the translation vectors 

254 T = M @ atoms.a 

255 

256 # Calculate the reciprocal space contribution 

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

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

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

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

261 # Collect all box index vectors in a matrix 

262 M = get_index_vectors(s) 

263 # Calculate the reciprocal translation vectors and precompute the prefactor 

264 G = M @ g 

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

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

267 

268 for ia in range(atoms.Natoms): 

269 for ja in range(atoms.Natoms): 

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

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

272 

273 # Add the real-space contribution 

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

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

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

277 if ia != ja: 

278 rmag = xp.linalg.norm(dpos) 

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

280 

281 # Add the reciprocal space contribution 

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

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

284 return float(Eewald) 

285 

286 

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

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

289 

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

291 

292 Args: 

293 scf: SCF object. 

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

295 

296 Keyword Args: 

297 n_single: Single-electron densities. 

298 

299 Returns: 

300 PZ self-interaction energy. 

301 """ 

302 if Y is None: 

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

304 return None 

305 

306 atoms = scf.atoms 

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

308 if n_single is None: 

309 n_single = get_n_single(atoms, Y) 

310 

311 Esic = 0 

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

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

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

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

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

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

318 # Normalize single-particle densities to 1 

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

320 

321 # Get the gradient of the single-particle density 

322 if "gga" in scf.xc_type: 

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

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

325 else: 

326 dni = None 

327 

328 # Get the kinetic energy density of the corresponding orbital 

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

330 # Use only one orbital for the calculation 

331 Ytmp = [] 

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

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

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

335 taui = xp.zeros_like(ni) 

336 # We also have to normalize to one again 

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

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

339 ) 

340 else: 

341 taui = None 

342 

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

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

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

346 # SIC energy is scaled by the occupation number 

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

348 scf.energies.Esic = Esic 

349 return float(Esic) 

350 

351 

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

353 try: 

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

355 except ImportError: 

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

357 return 0 

358 

359 

360get_Edisp.__doc__ = d3.get_Edisp.__doc__ 

361 

362 

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

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

365 

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

367 

368 Args: 

369 scf: SCF object. 

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

371 

372 Keyword Args: 

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

374 

375 Returns: 

376 Band energy in Hartree. 

377 """ 

378 atoms = scf.atoms 

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

380 Eband = 0 

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

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

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

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

385 ) 

386 return float(xp.real(Eband)) 

387 

388 

389def get_Eentropy(scf, epsilon, Efermi): 

390 """Calculate the fillings entropic energy. 

391 

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

393 

394 Args: 

395 scf: SCF object. 

396 epsilon: Eigenenergies. 

397 Efermi: Fermi energy. 

398 

399 Returns: 

400 Entropic energy in Hartree. 

401 """ 

402 occ = scf.atoms.occ 

403 

404 Eentropy = 0 

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

406 for spin in range(occ.Nspin): 

407 for i in range(occ.Nstate): 

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

409 Eentropy += ( 

410 occ.wk[ik] 

411 * occ.smearing 

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

413 ) 

414 

415 Eentropy *= 2 / occ.Nspin 

416 scf.energies.Eentropy = float(Eentropy) 

417 return float(Eentropy)