Coverage for eminus/gga.py: 100.00%

1# SPDX-FileCopyrightText: 2021 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 

9import numpy as np 

10 

11from .utils import handle_k 

12 

13 

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

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

16 

17 Args: 

18 atoms: Atoms object. 

19 field: Real-space field per spin channel. 

20 

21 Keyword Args: 

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

23 

24 Returns: 

25 Gradients of field per spin channel. 

26 """ 

27 dfield = np.empty((atoms.occ.Nspin, atoms.Ns, 3), dtype=complex) 

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

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

30 for dim in range(3): 

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

32 if real: 

33 return np.real(dfield) 

34 return dfield 

35 

36 

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

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

39 

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

41 

42 Args: 

43 atoms: Atoms object. 

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

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

46 vsigma: Contracted gradient potential derivative. 

47 

48 Returns: 

49 Gradient correction in reciprocal space. 

50 """ 

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

52 h = np.empty_like(dn_spin) 

53 if not atoms.unrestricted: 

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

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

56 else: 

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

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

59 # h is the expression in the brackets 

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

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

62 

63 # Calculate Nabla dot h 

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

65 Gh = np.empty((len(atoms.G2), 3), dtype=complex) 

66 for dim in range(3): 

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

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

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

70 

71 

72@handle_k(mode="reduce") 

73def get_tau(atoms, Y, ik): 

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

75 

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

77 

78 Args: 

79 atoms: Atoms object. 

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

81 ik: k-point index. 

82 

83 Returns: 

84 Real-space positive-definite kinetic energy density. 

85 """ 

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

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

88 

89 # Yrs = atoms.I(Y) 

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

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

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

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

94 # return np.real(tau) 

95 

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

97 # Calculate the active G vectors and broadcast to a desired shape 

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

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

100 for dim in range(3): 

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

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

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

104 # Sum over all states 

105 # Use the definition with a division by two 

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

107 # Or in compressed Einstein notation: 

108 return ( 

109 0.5 

110 * atoms.kpts.wk[ik] 

111 * np.real( 

112 np.einsum("sj,sijr,sijr->si", atoms.occ.f[ik], dYrs.conj(), dYrs, optimize="greedy") 

113 ) 

114 ) 

115 

116 

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

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

119 

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

121 

122 Args: 

123 scf: SCF object. 

124 ik: k-point index. 

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

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

127 vtau: Kinetic energy density potential derivative. 

128 

129 Returns: 

130 Tau-dependent potential contribution in reciprocal space. 

131 """ 

132 atoms = scf.atoms 

133 

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

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

136 return Vpsi 

137 

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

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

140 

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

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

143 # Wrs = atoms.I(W) 

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

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

146 # for r in range(3): 

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

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

149 

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

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

152 # Calculate the active G vectors and broadcast to a desired shape 

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

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

155 vtau_spin = vtau[spin, :, None] 

156 W_spin = W[ik][spin] 

157 for dim in range(3): 

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

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

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

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

162 # Sum over dimensions 

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

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

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

166 return atoms.O(Vpsi)