Coverage for eminus/gga.py: 100.00%

1# SPDX-FileCopyrightText: 2023 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""DFT functions that are only needed for (meta-)GGA calculations. 

4 

5Most functions here have been optimized for performance and can be harder to understand than normal. 

6To mitigate this, easier (but slower) implementations have been added as comments. 

7""" 

8 

9from . import backend as xp 

10from .utils import handle_k 

11 

12 

13def get_grad_field(atoms, field, real=True): 

14 """Calculate the gradient of fields on the grid per spin channel. 

15 

16 Args: 

17 atoms: Atoms object. 

18 field: Real-space field per spin channel. 

19 

20 Keyword Args: 

21 real: Make the gradient a real array (the gradient of n_spin is real). 

22 

23 Returns: 

24 Gradients of field per spin channel. 

25 """ 

26 dfield = xp.empty((atoms.occ.Nspin, atoms.Ns, 3), dtype=complex) 

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

28 fieldG = atoms.J(field[spin]) 

29 for dim in range(3): 

30 dfield[spin, :, dim] = atoms.I(1j * atoms.G[:, dim] * fieldG) 

31 if real: 

32 return xp.real(dfield) 

33 return dfield 

34 

35 

36def gradient_correction(atoms, spin, dn_spin, vsigma): 

37 """Calculate the gradient correction for the exchange-correlation potential. 

38 

39 Reference: Chem. Phys. Lett. 199, 557. 

40 

41 Args: 

42 atoms: Atoms object. 

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

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

45 vsigma: Contracted gradient potential derivative. 

46 

47 Returns: 

48 Gradient correction in reciprocal space. 

49 """ 

50 # sigma is |dn|^2, while vsigma is n * d exc/d sigma 

51 h = xp.empty_like(dn_spin) 

52 if not atoms.unrestricted: 

53 # In the unpolarized case we have no spin mixing and only one spin density 

54 h[0] = 2 * vsigma[0, :, None] * dn_spin[0] 

55 else: 

56 # In the polarized case we would get for spin up (and similar for spin down) 

57 # Vxc_u = vxc_u - Nabla dot (2 vsigma_uu * dn_u + vsigma_ud * dn_d) 

58 # h is the expression in the brackets 

59 h[0] = 2 * vsigma[0, :, None] * dn_spin[0] + vsigma[1, :, None] * dn_spin[1] 

60 h[1] = 2 * vsigma[2, :, None] * dn_spin[1] + vsigma[1, :, None] * dn_spin[0] 

61 

62 # Calculate Nabla dot h 

63 # Normally we would calculate the correction for each spin, but we only need one at a time in H 

64 Gh = xp.empty((len(atoms.G2), 3), dtype=complex) 

65 for dim in range(3): 

66 Gh[:, dim] = atoms.J(h[spin, :, dim]) 

67 # return 1j * np.sum(atoms.G * Gh, axis=1) 

68 return xp.einsum("ir,ir->i", 1j * atoms.G, Gh) 

69 

70 

71@handle_k(mode="reduce") 

72def get_tau(atoms, Y, ik): 

73 """Calculate the positive-definite kinetic energy densities per spin. 

74 

75 Reference: J. Chem. Phys. 109, 2092. 

76 

77 Args: 

78 atoms: Atoms object. 

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

80 ik: k-point index. 

81 

82 Returns: 

83 Real-space positive-definite kinetic energy density. 

84 """ 

85 # The "intuitive" way is the one commented out below (without k-point dependency) 

86 # Sadly, this implementation is really slow for various reasons so use the faster one below 

87 

88 # Yrs = atoms.I(Y) 

89 # tau = np.zeros((atoms.occ.Nspin, atoms.Ns), dtype=complex) 

90 # for i in range(atoms.occ.Nstate): 

91 # dYrs = get_grad_field(atoms, Yrs[..., i], real=False) 

92 # tau += 0.5 * atoms.occ.f[:, i, None] * np.sum(dYrs.conj() * dYrs, axis=2) 

93 # return np.real(tau) 

94 

95 dYrs = xp.empty((atoms.occ.Nspin, atoms.Ns, Y.shape[-1], 3), dtype=complex) 

96 # Calculate the active G vectors and broadcast to the needed shape 

97 Gkc = atoms.G[atoms.active[ik]][:, None, :] + atoms.kpts.k[ik] 

98 # Calculate the gradients of Y in the active(!) reciprocal space and transform to real-space 

99 for dim in range(3): 

100 dYrs[..., dim] = atoms.I(1j * Gkc[..., dim] * Y, ik) 

101 # Sum over dimensions (dYx* dYx + dYy* dYy + dYz* dYz) 

102 # sumdYrs = np.real(np.sum(dYrs.conj() * dYrs, axis=3)) 

103 # Sum over all states 

104 # Use the definition with a division by two 

105 # return 0.5 * np.sum(atoms.occ.f[:, None, :] * sumdYrs, axis=2) 

106 # Or in compressed Einstein notation: 

107 return ( 

108 0.5 

109 * atoms.kpts.wk[ik] 

110 * xp.real(xp.einsum("sj,sijr,sijr->si", atoms.occ.f[ik], dYrs.conj(), dYrs)) 

111 ) 

112 

113 

114def calc_Vtau(scf, ik, spin, W, vtau): 

115 """Calculate the tau-dependent potential contribution for meta-GGAs. 

116 

117 Reference: J. Chem. Phys. 145, 204114. 

118 

119 Args: 

120 scf: SCF object. 

121 ik: k-point index. 

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

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

124 vtau: Kinetic energy density potential derivative. 

125 

126 Returns: 

127 Tau-dependent potential contribution in reciprocal space. 

128 """ 

129 atoms = scf.atoms 

130 

131 if scf.xc_type != "meta-gga": # Only calculate the contribution for meta-GGAs 

132 return xp.zeros((len(atoms.Gk2c[ik]), W[ik].shape[-1]), dtype=complex) 

133 

134 # The "intuitive" way is the one commented out below (without k-point dependency) 

135 # Sadly, this implementation is really slow for various reasons so use the faster one below 

136 

137 # GVpsi = np.empty((len(atoms.G2c), 3), dtype=complex) 

138 # Gc = atoms.G[atoms.active] 

139 # Wrs = atoms.I(W) 

140 # for i in range(atoms.occ.Nstate): 

141 # dWrs = get_grad_field(atoms, Wrs[..., i], real=False) 

142 # for r in range(3): 

143 # GVpsi[:, r] = atoms.J(vtau[spin] * dWrs[spin, :, r], full=False) 

144 # Vpsi[:, i] = -0.5 * 1j * np.sum(Gc * GVpsi, axis=1) 

145 

146 GVpsi = xp.empty((len(atoms.Gk2c[ik]), W[ik].shape[-1], 3), dtype=complex) 

147 dWrs = xp.empty((atoms.Ns, W[ik].shape[-1], 3), dtype=complex) 

148 # Calculate the active G vectors and broadcast to the needed shape 

149 Gkc = atoms.G[atoms.active[ik]][:, None, :] + scf.kpts.k[ik] 

150 # We only calculate Vtau for one spin channel, index, and reshape before the loop 

151 vtau_spin = vtau[spin, :, None] 

152 W_spin = W[ik][spin] 

153 for dim in range(3): 

154 # Calculate the gradients of W in the active(!) space and transform to real-space 

155 dWrs[..., dim] = atoms.I(1j * Gkc[..., dim] * W_spin, ik) 

156 # Calculate dexc/dtau * gradpsi and transform to the active reciprocal space 

157 GVpsi[..., dim] = atoms.J(vtau_spin * dWrs[..., dim], ik, full=False) 

158 # Sum over dimensions 

159 # Calculate -0.5 Nabla dot Gvpsi (compare with gradient_correction) 

160 # Vpsi = -0.5 * 1j * np.sum(Gkc * GVpsi, axis=2) 

161 Vpsi = -0.5 * xp.einsum("ior,ijr->ij", 1j * Gkc, GVpsi) 

162 return atoms.O(Vpsi)