Coverage for eminus/operators.py: 100.00%

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Basis set dependent operators for a plane wave basis. 

4 

5These operators act on discretized wave functions, i.e., the arrays W. 

6 

7These W are column vectors. This has been chosen to let theory and code coincide, e.g., 

8W^dagger W becomes :code:`W.conj().T @ W`. 

9 

10The downside is that the i-th state will be accessed with W[:, i] instead of W[i]. 

11Choosing the i-th state makes the array 1d. 

12 

13These operators can act on six different options, namely 

14 

151. the real-space 

162. the real-space (1d) 

173. the full reciprocal space 

184. the full reciprocal space (1d) 

195. the active reciprocal space 

206. the active reciprocal space (1d) 

21 

22The active space is the truncated reciprocal space by restricting it with a sphere given by ecut. 

23 

24Every spin dependence will be handled with handle_spin by calling the operators for each 

25spin individually. The same goes for the handling of k-points, while for k-points W is represented 

26as a list of arrays. This gives the final indexing for the k-point k, spin s, and state n of 

27W[ik][s, :, n]. 

28""" 

29 

30import copy 

31 

32import numpy as np 

33 

34from . import backend as xp 

35from .utils import handle_k, handle_spin 

36 

37 

38# Spin handling is trivial for this operator 

39@handle_k 

40def O(atoms, W): 

41 """Overlap operator. 

42 

43 This operator acts on the options 3, 4, 5, and 6. 

44 

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

46 

47 Args: 

48 atoms: Atoms object. 

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

50 

51 Returns: 

52 The operator applied on W. 

53 """ 

54 return atoms.Omega * W 

55 

56 

57@handle_spin 

58def L(atoms, W, ik=-1): 

59 """Laplacian operator with k-point dependency. 

60 

61 This operator acts on options 3 and 5. 

62 

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

64 

65 Args: 

66 atoms: Atoms object. 

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

68 

69 Keyword Args: 

70 ik: k-point index. 

71 

72 Returns: 

73 The operator applied on W. 

74 """ 

75 # Gk2 is a normal 1d row vector, reshape it so it can be applied to the column vector W 

76 if len(W) == len(atoms.Gk2c[ik]): 

77 Gk2 = atoms.Gk2c[ik][:, None] 

78 else: 

79 Gk2 = atoms.Gk2[ik][:, None] 

80 return -atoms.Omega * Gk2 * W 

81 

82 

83@handle_spin 

84def Linv(atoms, W): 

85 """Inverse Laplacian operator. 

86 

87 This operator acts on options 3 and 4. 

88 

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

90 

91 Args: 

92 atoms: Atoms object. 

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

94 

95 Returns: 

96 The operator applied on W. 

97 """ 

98 # Ignore the division by zero for the first elements 

99 with np.errstate(divide="ignore", invalid="ignore"): 

100 if W.ndim == 1: 

101 # One could do some proper indexing with [1:] but indexing is slow 

102 out = W / (atoms.G2 * -atoms.Omega) 

103 out[0] = 0 

104 else: 

105 # G2 is a normal 1d row vector, reshape it so it can be applied to the column vector W 

106 G2 = atoms.G2[:, None] 

107 out = W / (G2 * -atoms.Omega) 

108 out[0, :] = 0 

109 return out 

110 

111 

112@handle_k(mode="index") 

113def I(atoms, W, ik=-1, norm="forward"): 

114 """Backward transformation from reciprocal space to real-space. 

115 

116 This operator acts on the options 3, 4, 5, and 6. 

117 

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

119 

120 Args: 

121 atoms: Atoms object. 

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

123 

124 Keyword Args: 

125 ik: k-point index. 

126 norm: Normalization mode. 

127 

128 Returns: 

129 The operator applied on W. 

130 """ 

131 if W.ndim < 3: 

132 # If W is in the full space do nothing with W 

133 if len(W) == len(atoms._Gk2[ik]): 

134 Wfft = W 

135 # Fill with zeros if W is in the active space 

136 else: 

137 if W.ndim == 1: 

138 Wfft = xp.zeros(atoms.Ns, dtype=W.dtype) 

139 else: 

140 Wfft = xp.zeros((atoms.Ns, W.shape[-1]), dtype=W.dtype) 

141 Wfft[atoms._active[ik]] = W 

142 elif W.shape[1] == len(atoms._G2): 

143 Wfft = W 

144 else: 

145 Wfft = xp.zeros((atoms.occ.Nspin, atoms.Ns, W.shape[-1]), dtype=W.dtype) 

146 Wfft[:, atoms._active[ik][0]] = W 

147 

148 # Normally, we would have to multiply by n in the end for the correct normalization, but we can 

149 # ignore this step when properly setting the `norm` option for a faster operation 

150 if W.ndim == 1: 

151 Wfft = Wfft.reshape(list(atoms.s)) 

152 Finv = xp.ifftn(Wfft, norm=norm).ravel() 

153 # Here we reshape the input like in the 1d case but add an extra dimension in the end, 

154 # holding the number of states 

155 # Tell the function that the FFT only has to act on the respective 3 axes 

156 elif W.ndim == 2: 

157 Wfft = Wfft.reshape([*atoms.s, W.shape[-1]]) 

158 Finv = xp.ifftn(Wfft, norm=norm, axes=(0, 1, 2)).reshape((atoms.Ns, W.shape[-1])) 

159 else: 

160 Wfft = Wfft.reshape([atoms.occ.Nspin, *atoms.s, W.shape[-1]]) 

161 Finv = xp.ifftn(Wfft, norm=norm, axes=(1, 2, 3)).reshape( 

162 (atoms.occ.Nspin, atoms.Ns, W.shape[-1]) 

163 ) 

164 return Finv 

165 

166 

167@handle_k(mode="index") 

168def J(atoms, W, ik=-1, full=True, norm="forward"): 

169 """Forward transformation from real-space to reciprocal space. 

170 

171 This operator acts on options 1 and 2. 

172 

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

174 

175 Args: 

176 atoms: Atoms object. 

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

178 

179 Keyword Args: 

180 ik: k-point index. 

181 full: Whether to transform in the full or in the active space. 

182 norm: Normalization mode. 

183 

184 Returns: 

185 The operator applied on W. 

186 """ 

187 # Normally, we would have to divide by n in the end for the correct normalization, but we can 

188 # ignore this step when properly setting the `norm` option for a faster operation 

189 if W.ndim == 1: 

190 Wfft = W.reshape(list(atoms.s)) 

191 F = xp.fftn(Wfft, norm=norm).ravel() 

192 elif W.ndim == 2: 

193 Wfft = W.reshape([*atoms.s, W.shape[-1]]) 

194 F = xp.fftn(Wfft, norm=norm, axes=(0, 1, 2)).reshape((atoms.Ns, W.shape[-1])) 

195 else: 

196 Wfft = W.reshape([atoms.occ.Nspin, *atoms.s, W.shape[-1]]) 

197 F = xp.fftn(Wfft, norm=norm, axes=(1, 2, 3)).reshape( 

198 (atoms.occ.Nspin, atoms.Ns, W.shape[-1]) 

199 ) 

200 # There is no way to know if J has to transform to the full or the active space 

201 # but normally it transforms to the full space 

202 if not full: 

203 if F.ndim < 3: 

204 return F[atoms._active[ik]] 

205 return F[:, atoms._active[ik][0]] 

206 return F 

207 

208 

209@handle_k(mode="index") 

210def Idag(atoms, W, ik=-1, full=False): 

211 """Conjugated backward transformation from real-space to reciprocal space. 

212 

213 This operator acts on options 1 and 2. 

214 

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

216 

217 Args: 

218 atoms: Atoms object. 

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

220 

221 Keyword Args: 

222 ik: k-point index. 

223 full: Whether to transform in the full or in the active space. 

224 

225 Returns: 

226 The operator applied on W. 

227 """ 

228 # This is just the J operator with the appropriate norm 

229 # If we would not chnage the norm, we would have to multiply the result with n 

230 return J(atoms, W, ik, full=full, norm="backward") 

231 

232 

233@handle_k(mode="index") 

234def Jdag(atoms, W, ik=-1): 

235 """Conjugated forward transformation from reciprocal space to real-space. 

236 

237 This operator acts on the options 3, 4, 5, and 6. 

238 

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

240 

241 Args: 

242 atoms: Atoms object. 

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

244 

245 Keyword Args: 

246 ik: k-point index. 

247 

248 Returns: 

249 The operator applied on W. 

250 """ 

251 # This is just the J operator with the appropriate norm 

252 # If we would not chnage the norm, we would have to divide the result by n 

253 return I(atoms, W, ik, norm="backward") 

254 

255 

256@handle_spin 

257def K(atoms, W, ik): 

258 """Preconditioning operator with k-point dependency. 

259 

260 This operator acts on options 3 and 5. 

261 

262 Reference: Comput. Mater. Sci. 14, 4. 

263 

264 Args: 

265 atoms: Atoms object. 

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

267 ik: k-point index. 

268 

269 Returns: 

270 The operator applied on W. 

271 """ 

272 # Gk2c is a normal 1d row vector, reshape it so it can be applied to the column vector W 

273 return W / (1 + atoms.Gk2c[ik][:, None]) 

274 

275 

276def T(atoms, W, dr): 

277 """Translation operator. 

278 

279 This operator acts on options 5 and 6. 

280 

281 Reference: https://ccrma.stanford.edu/~jos/st/Shift_Theorem.html 

282 

283 Args: 

284 atoms: Atoms object. 

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

286 dr: Real-space shifting vector. 

287 

288 Returns: 

289 The operator applied on W. 

290 """ 

291 # We can not use a fancy decorator for this operator, so handle it here 

292 if xp.is_array(W) and W.ndim == 3: 

293 return xp.stack([T(atoms, Wspin, dr) for Wspin in W]) 

294 

295 if xp.is_array(W): 

296 atoms.kpts._assert_gamma_only() 

297 if len(W) == len(atoms.Gk2c[0]): 

298 G = atoms.G[atoms.active[0]] 

299 else: 

300 G = atoms.G 

301 factor = xp.exp(-1j * (G @ dr)) 

302 if W.ndim == 2: 

303 factor = factor[:, None] 

304 return factor * W 

305 

306 # If W is a list we have to account for k-points 

307 Wshift = copy.deepcopy(W) 

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

309 # Do the shift by multiplying a phase factor, given by the shift theorem 

310 if W[ik].shape[1] == len(atoms.Gk2c[ik]): 

311 Gk = atoms.G[atoms.active[ik]] + atoms.kpts.k[ik] 

312 else: 

313 Gk = atoms.G + atoms.kpts.k[ik] 

314 Wshift[ik] = xp.exp(-1j * (Gk @ dr))[:, None] * W[ik] 

315 return Wshift