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(atoms.occ.Nspin, atoms.Ns, W.shape[-1]) 

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

199 # but normally it transforms to the full space 

200 if not full: 

201 if F.ndim < 3: 

202 return F[atoms._active[ik]] 

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

204 return F 

205 

206 

207@handle_k(mode="index") 

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

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

210 

211 This operator acts on options 1 and 2. 

212 

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

214 

215 Args: 

216 atoms: Atoms object. 

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

218 

219 Keyword Args: 

220 ik: k-point index. 

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

222 

223 Returns: 

224 The operator applied on W. 

225 """ 

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

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

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

229 

230 

231@handle_k(mode="index") 

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

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

234 

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

236 

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

238 

239 Args: 

240 atoms: Atoms object. 

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

242 

243 Keyword Args: 

244 ik: k-point index. 

245 

246 Returns: 

247 The operator applied on W. 

248 """ 

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

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

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

252 

253 

254@handle_spin 

255def K(atoms, W, ik): 

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

257 

258 This operator acts on options 3 and 5. 

259 

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

261 

262 Args: 

263 atoms: Atoms object. 

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

265 ik: k-point index. 

266 

267 Returns: 

268 The operator applied on W. 

269 """ 

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

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

272 

273 

274def T(atoms, W, dr): 

275 """Translation operator. 

276 

277 This operator acts on options 5 and 6. 

278 

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

280 

281 Args: 

282 atoms: Atoms object. 

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

284 dr: Real-space shifting vector. 

285 

286 Returns: 

287 The operator applied on W. 

288 """ 

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

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

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

292 

293 if xp.is_array(W): 

294 atoms.kpts._assert_gamma_only() 

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

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

297 else: 

298 G = atoms.G 

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

300 if W.ndim == 2: 

301 factor = factor[:, None] 

302 return factor * W 

303 

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

305 Wshift = copy.deepcopy(W) 

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

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

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

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

310 else: 

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

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

313 return Wshift