Coverage for eminus/energies.py: 95.83%
168 statements
« prev ^ index » next coverage.py v7.6.4, created at 2024-11-08 12:59 +0000
1# SPDX-FileCopyrightText: 2021 The eminus developers
2# SPDX-License-Identifier: Apache-2.0
3"""Calculate different energy contributions."""
5import dataclasses
7import numpy as np
8from scipy.linalg import inv, norm
9from scipy.special import erfc
11from .dft import get_n_single, get_phi, H
12from .extras import dispersion
13from .gga import get_grad_field, get_tau
14from .logger import log
15from .tools import electronic_entropy
16from .utils import handle_k
17from .xc import get_exc
20@dataclasses.dataclass
21class Energy:
22 """Energy class to save energy contributions in one place."""
24 Ekin: float = 0 #: Kinetic energy.
25 Ecoul: float = 0 #: Coulomb energy.
26 Exc: float = 0 #: Exchange-correlation energy.
27 Eloc: float = 0 #: Local energy.
28 Enonloc: float = 0 #: Non-local energy.
29 Eewald: float = 0 #: Ewald energy.
30 Esic: float = 0 #: Self-interaction correction energy.
31 Edisp: float = 0 #: Dispersion correction energy.
32 Eentropy: float = 0 #: Fillings entropic energy.
34 @property
35 def Etot(self):
36 """Total energy is the sum of all energy contributions."""
37 return sum(getattr(self, ie.name) for ie in dataclasses.fields(self))
39 def extrapolate(self):
40 """Calculate the total energy at T=0.
42 Reference: J. Phys.: Condens. Matter 1, 689.
44 Returns:
45 Total energy extrapolated to T=0.
46 """
47 return self.Etot - 0.5 * self.Eentropy
49 def __repr__(self):
50 """Print the energies stored in the Energy object."""
51 out = ""
52 for ie in dataclasses.fields(self):
53 energy = getattr(self, ie.name)
54 if energy != 0:
55 out += f"{ie.name:<9}: {energy:+.9f} Eh\n"
56 return f"{out}{'-' * 26}\nEtot : {self.Etot:+.9f} Eh"
59def get_E(scf):
60 """Calculate energy contributions and update energies needed in one SCF step.
62 Args:
63 scf: SCF object.
65 Returns:
66 Total energy.
67 """
68 scf.energies.Ekin = get_Ekin(scf.atoms, scf.Y)
69 scf.energies.Ecoul = get_Ecoul(scf.atoms, scf.n, scf.phi)
70 scf.energies.Exc = get_Exc(scf, scf.n, scf.exc, Nspin=scf.atoms.occ.Nspin)
71 scf.energies.Eloc = get_Eloc(scf, scf.n)
72 scf.energies.Enonloc = get_Enonloc(scf, scf.Y)
73 return scf.energies.Etot
76@handle_k(mode="reduce")
77def get_Ekin(atoms, Y, ik):
78 """Calculate the kinetic energy.
80 Reference: Comput. Phys. Commun. 128, 1.
82 Args:
83 atoms: Atoms object.
84 Y: Expansion coefficients of orthogonal wave functions in reciprocal space.
85 ik: k-point index.
87 Returns:
88 Kinetic energy in Hartree.
89 """
90 # Ekin = -0.5 Tr(F Wdag L(W))
91 Ekin = 0
92 for spin in range(atoms.occ.Nspin):
93 Ekin += (
94 -0.5
95 * atoms.kpts.wk[ik]
96 * np.trace(atoms.occ.F[ik][spin] @ Y[spin].conj().T @ atoms.L(Y[spin], ik))
97 )
98 return np.real(Ekin)
101def get_Ecoul(atoms, n, phi=None):
102 """Calculate the Coulomb energy.
104 Reference: Comput. Phys. Commun. 128, 1.
106 Args:
107 atoms: Atoms object.
108 n: Real-space electronic density.
110 Keyword Args:
111 phi: Hartree field.
113 Returns:
114 Coulomb energy in Hartree.
115 """
116 if phi is None:
117 phi = get_phi(atoms, n)
118 # Ecoul = 0.5 (J(n))dag O(phi)
119 return np.real(0.5 * n @ atoms.Jdag(atoms.O(phi)))
122def get_Exc(scf, n, exc=None, n_spin=None, dn_spin=None, tau=None, Nspin=2):
123 """Calculate the exchange-correlation energy.
125 Reference: Comput. Phys. Commun. 128, 1.
127 Args:
128 scf: SCF object.
129 n: Real-space electronic density.
131 Keyword Args:
132 exc: Exchange-correlation energy density.
133 n_spin: Real-space electronic densities per spin channel.
134 dn_spin: Real-space gradient of densities per spin channel.
135 tau: Real-space kinetic energy densities per spin channel.
136 Nspin: Number of spin states.
138 Returns:
139 Exchange-correlation energy in Hartree.
140 """
141 atoms = scf.atoms
142 if exc is None:
143 exc = get_exc(scf.xc, n_spin, Nspin, dn_spin, tau, scf.xc_params)
144 # Exc = (J(n))dag O(J(exc))
145 return np.real(n @ atoms.Jdag(atoms.O(atoms.J(exc))))
148def get_Eloc(scf, n):
149 """Calculate the local energy contribution.
151 Reference: Comput. Phys. Commun. 128, 1.
153 Args:
154 scf: SCF object.
155 n: Real-space electronic density.
157 Returns:
158 Local energy contribution in Hartree.
159 """
160 return np.real(np.vdot(scf.Vloc, n))
163@handle_k(mode="reduce")
164def get_Enonloc(scf, Y, ik):
165 """Calculate the non-local GTH energy contribution.
167 Adapted from https://github.com/f-fathurrahman/PWDFT.jl/blob/master/src/calc_energies.jl
169 Reference: Phys. Rev. B 54, 1703.
171 Args:
172 scf: SCF object.
173 Y: Expansion coefficients of orthogonal wave functions in reciprocal space.
174 ik: k-point index.
176 Returns:
177 Non-local GTH energy contribution in Hartree.
178 """
179 atoms = scf.atoms
181 Enonloc = 0
182 if scf.pot != "gth" or scf.gth.NbetaNL == 0: # Only calculate the non-local part if necessary
183 return Enonloc
185 for spin in range(atoms.occ.Nspin):
186 betaNL_psi = (Y[spin].conj().T @ scf.gth.betaNL[ik]).conj()
188 enl = np.zeros(Y.shape[-1], dtype=complex)
189 for ia in range(atoms.Natoms):
190 psp = scf.gth[atoms.atom[ia]]
191 for l in range(psp["lmax"]):
192 for m in range(-l, l + 1):
193 for iprj in range(psp["Nproj_l"][l]):
194 ibeta = scf.gth.prj2beta[iprj, ia, l, m + psp["lmax"] - 1] - 1
195 for jprj in range(psp["Nproj_l"][l]):
196 jbeta = scf.gth.prj2beta[jprj, ia, l, m + psp["lmax"] - 1] - 1
197 hij = psp["h"][l, iprj, jprj]
198 enl += hij * betaNL_psi[:, ibeta].conj() * betaNL_psi[:, jbeta]
199 Enonloc += np.sum(atoms.occ.f[ik, spin] * atoms.kpts.wk[ik] * enl)
200 # We have to multiply with the cell volume, because of different orthogonalization methods
201 return np.real(Enonloc * atoms.Omega)
204def get_Eewald(atoms, gcut=2, gamma=1e-8):
205 """Calculate the Ewald energy.
207 Reference: J. Chem. Theory Comput. 10, 381.
209 Args:
210 atoms: Atoms object.
212 Keyword Args:
213 gcut: G-vector cut-off.
214 gamma: Error tolerance.
216 Returns:
217 Ewald energy in Hartree.
218 """
220 # For a plane wave code we have multiple contributions for the Ewald energy
221 # Namely, a sum from contributions from real-space, reciprocal space,
222 # the self energy, (the dipole term [neglected]), and an additional electroneutrality term
223 def get_index_vectors(s):
224 """Create index vectors of s periodic images."""
225 m1 = np.arange(-s[0], s[0] + 1)
226 m2 = np.arange(-s[1], s[1] + 1)
227 m3 = np.arange(-s[2], s[2] + 1)
228 M = np.transpose(np.meshgrid(m1, m2, m3)).reshape(-1, 3)
229 # Delete the [0, 0, 0] element, to prevent checking for it in every loop
230 return M[~(M == 0).all(axis=1)]
232 # Calculate the Ewald splitting parameter nu
233 gexp = -np.log(gamma)
234 nu = 0.5 * np.sqrt(gcut**2 / gexp)
236 # Start by calculating the self-energy
237 Eewald = -nu / np.sqrt(np.pi) * np.sum(atoms.Z**2)
238 # Add the electroneutrality term
239 Eewald += -np.pi * np.sum(atoms.Z) ** 2 / (2 * nu**2 * atoms.Omega)
241 # Calculate the real-space contribution
242 # Calculate the amount of images that have to be considered per axis
243 Rm = norm(atoms.a, axis=1)
244 tmax = np.sqrt(0.5 * gexp) / nu
245 s = np.rint(tmax / Rm + 1.5)
246 # Collect all box index vectors in a matrix
247 M = get_index_vectors(s)
248 # Calculate the translation vectors
249 T = M @ atoms.a
251 # Calculate the reciprocal space contribution
252 # Calculate the amount of reciprocal images that have to be considered per axis
253 g = 2 * np.pi * inv(atoms.a.T)
254 gm = norm(g, axis=1)
255 s = np.rint(gcut / gm + 1.5)
256 # Collect all box index vectors in a matrix
257 M = get_index_vectors(s)
258 # Calculate the reciprocal translation vectors and precompute the prefactor
259 G = M @ g
260 G2 = norm(G, axis=1) ** 2
261 prefactor = 2 * np.pi / atoms.Omega * np.exp(-0.25 * G2 / nu**2) / G2
263 for ia in range(atoms.Natoms):
264 for ja in range(atoms.Natoms):
265 dpos = atoms.pos[ia] - atoms.pos[ja]
266 ZiZj = atoms.Z[ia] * atoms.Z[ja]
268 # Add the real-space contribution
269 rmag = np.sqrt(norm(dpos - T, axis=1) ** 2)
270 Eewald += 0.5 * ZiZj * np.sum(erfc(rmag * nu) / rmag)
271 # The T=[0, 0, 0] element is omitted in M but needed if ia!=ja, so add it manually
272 if ia != ja:
273 rmag = np.sqrt(norm(dpos) ** 2)
274 Eewald += 0.5 * ZiZj * erfc(rmag * nu) / rmag
276 # Add the reciprocal space contribution
277 Gpos = np.sum(G * dpos, axis=1)
278 Eewald += ZiZj * np.sum(prefactor * np.cos(Gpos))
279 return Eewald
282def get_Esic(scf, Y, n_single=None):
283 """Calculate the Perdew-Zunger self-interaction energy.
285 Reference: Phys. Rev. B 23, 5048.
287 Args:
288 scf: SCF object.
289 Y: Expansion coefficients of orthogonal wave functions in reciprocal space.
291 Keyword Args:
292 n_single: Single-electron densities.
294 Returns:
295 PZ self-interaction energy.
296 """
297 if Y is None:
298 log.warning('The provided wave function is "None".')
299 return None
301 atoms = scf.atoms
302 # E_PZ-SIC = \sum_i Ecoul[n_i] + Exc[n_i, 0]
303 if n_single is None:
304 n_single = get_n_single(atoms, Y)
306 Esic = 0
307 for i in range(atoms.occ.Nstate):
308 for spin in range(atoms.occ.Nspin):
309 if np.sum(atoms.occ.f[:, spin, i] * atoms.kpts.wk) > 0:
310 # Create normalized single-particle densities in the form of electronic densities
311 # per spin channel, since spin-polarized functionals expect this form
312 ni = np.zeros((2, atoms.Ns))
313 # Normalize single-particle densities to 1
314 ni[0] = n_single[spin, :, i] / np.sum(atoms.occ.f[:, spin, i] * atoms.kpts.wk)
316 # Get the gradient of the single-particle density
317 if "gga" in scf.xc_type:
318 dni = np.zeros((2, atoms.Ns, 3))
319 dni[0] = get_grad_field(atoms, ni)[0]
320 else:
321 dni = None
323 # Get the kinetic energy density of the corresponding orbital
324 if scf.xc_type == "meta-gga":
325 # Use only one orbital for the calculation
326 Ytmp = []
327 for ik in range(atoms.kpts.Nk):
328 Ytmp.append(np.zeros_like(Y[ik]))
329 Ytmp[ik][0, :, 0] = Y[ik][spin, :, i]
330 taui = np.zeros_like(ni)
331 # We also have to normalize to one again
332 taui[0] = get_tau(atoms, Ytmp)[0] / np.sum(
333 atoms.occ.f[:, spin, i] * atoms.kpts.wk
334 )
335 else:
336 taui = None
338 coul = get_Ecoul(atoms, ni[0])
339 # The exchange part for a SIC correction has to be spin-polarized
340 xc = get_Exc(scf, ni[0], n_spin=ni, dn_spin=dni, tau=taui, Nspin=2)
341 # SIC energy is scaled by the occupation number
342 Esic += (coul + xc) * np.sum(atoms.occ.f[:, spin, i] * atoms.kpts.wk)
343 scf.energies.Esic = Esic
344 return Esic
347def get_Edisp(scf, version="d3bj", atm=True, xc=None): # noqa: D103
348 try:
349 return dispersion.get_Edisp(scf, version, atm, xc)
350 except ImportError:
351 log.warning("You have to install the dispersion extra to use this function.")
352 return 0
355get_Edisp.__doc__ = dispersion.get_Edisp.__doc__
358def get_Eband(scf, Y, **kwargs):
359 """Calculate the band energy for occupied or unoccupied bands.
361 Reference: Comput. Phys. Commun. 128, 1.
363 Args:
364 scf: SCF object.
365 Y: Expansion coefficients of orthogonal wave functions in reciprocal space.
367 Keyword Args:
368 **kwargs: See :func:`eminus.dft.H`.
370 Returns:
371 Band energy in Hartree.
372 """
373 atoms = scf.atoms
374 # Eband = Tr(Wdag H(W))
375 Eband = 0
376 for ik in range(atoms.kpts.Nk):
377 for spin in range(atoms.occ.Nspin):
378 Eband += atoms.kpts.wk[ik] * np.trace(
379 Y[ik][spin].conj().T @ H(scf, ik, spin, Y, **kwargs)
380 )
381 return np.real(Eband)
384def get_Eentropy(scf, epsilon, Efermi):
385 """Calculate the fillings entropic energy.
387 Reference: J. Phys. Condens. Matter 1, 689.
389 Args:
390 scf: SCF object.
391 epsilon: Eigenenergies.
392 Efermi: Fermi energy.
394 Returns:
395 Entropic energy in Hartree.
396 """
397 occ = scf.atoms.occ
399 Eentropy = 0
400 for ik in range(scf.kpts.Nk):
401 for spin in range(occ.Nspin):
402 for i in range(occ.Nstate):
403 # Beware the sign change, it is handled in the electronic_entropy function
404 Eentropy += (
405 occ.wk[ik]
406 * occ.smearing
407 * electronic_entropy(epsilon[ik, spin, i], Efermi, occ.smearing)
408 )
410 Eentropy *= 2 / occ.Nspin
411 scf.energies.Eentropy = Eentropy
412 return Eentropy