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( 

111 xp.einsum("sj,sijr,sijr->si", xp.astype(atoms.occ.f[ik], complex), dYrs.conj(), dYrs) 

112 ) 

113 ) 

114 

115 

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

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

118 

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

120 

121 Args: 

122 scf: SCF object. 

123 ik: k-point index. 

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

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

126 vtau: Kinetic energy density potential derivative. 

127 

128 Returns: 

129 Tau-dependent potential contribution in reciprocal space. 

130 """ 

131 atoms = scf.atoms 

132 

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

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

135 

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

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

138 

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

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

141 # Wrs = atoms.I(W) 

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

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

144 # for r in range(3): 

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

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

147 

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

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

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

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

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

153 vtau_spin = vtau[spin, :, None] 

154 W_spin = W[ik][spin] 

155 for dim in range(3): 

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

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

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

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

160 # Sum over dimensions 

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

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

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

164 return atoms.O(Vpsi)