Coverage for eminus/operators.py: 100.00%
88 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-10-21 12:19 +0000
« prev ^ index » next coverage.py v7.11.0, created at 2025-10-21 12:19 +0000
1# SPDX-FileCopyrightText: 2021 The eminus developers
2# SPDX-License-Identifier: Apache-2.0
3"""Basis set dependent operators for a plane wave basis.
5These operators act on discretized wave functions, i.e., the arrays W.
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`.
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.
13These operators can act on six different options, namely
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)
22The active space is the truncated reciprocal space by restricting it with a sphere given by ecut.
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"""
30import copy
32import numpy as np
34from . import backend as xp
35from .utils import handle_k, handle_spin
38# Spin handling is trivial for this operator
39@handle_k
40def O(atoms, W):
41 """Overlap operator.
43 This operator acts on the options 3, 4, 5, and 6.
45 Reference: Comput. Phys. Commun. 128, 1.
47 Args:
48 atoms: Atoms object.
49 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
51 Returns:
52 The operator applied on W.
53 """
54 return atoms.Omega * W
57@handle_spin
58def L(atoms, W, ik=-1):
59 """Laplacian operator with k-point dependency.
61 This operator acts on options 3 and 5.
63 Reference: Comput. Phys. Commun. 128, 1.
65 Args:
66 atoms: Atoms object.
67 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
69 Keyword Args:
70 ik: k-point index.
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
83@handle_spin
84def Linv(atoms, W):
85 """Inverse Laplacian operator.
87 This operator acts on options 3 and 4.
89 Reference: Comput. Phys. Commun. 128, 1.
91 Args:
92 atoms: Atoms object.
93 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
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
112@handle_k(mode="index")
113def I(atoms, W, ik=-1, norm="forward"):
114 """Backward transformation from reciprocal space to real-space.
116 This operator acts on the options 3, 4, 5, and 6.
118 Reference: Comput. Phys. Commun. 128, 1.
120 Args:
121 atoms: Atoms object.
122 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
124 Keyword Args:
125 ik: k-point index.
126 norm: Normalization mode.
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
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
167@handle_k(mode="index")
168def J(atoms, W, ik=-1, full=True, norm="forward"):
169 """Forward transformation from real-space to reciprocal space.
171 This operator acts on options 1 and 2.
173 Reference: Comput. Phys. Commun. 128, 1.
175 Args:
176 atoms: Atoms object.
177 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
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.
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
209@handle_k(mode="index")
210def Idag(atoms, W, ik=-1, full=False):
211 """Conjugated backward transformation from real-space to reciprocal space.
213 This operator acts on options 1 and 2.
215 Reference: Comput. Phys. Commun. 128, 1.
217 Args:
218 atoms: Atoms object.
219 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
221 Keyword Args:
222 ik: k-point index.
223 full: Whether to transform in the full or in the active space.
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")
233@handle_k(mode="index")
234def Jdag(atoms, W, ik=-1):
235 """Conjugated forward transformation from reciprocal space to real-space.
237 This operator acts on the options 3, 4, 5, and 6.
239 Reference: Comput. Phys. Commun. 128, 1.
241 Args:
242 atoms: Atoms object.
243 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
245 Keyword Args:
246 ik: k-point index.
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")
256@handle_spin
257def K(atoms, W, ik):
258 """Preconditioning operator with k-point dependency.
260 This operator acts on options 3 and 5.
262 Reference: Comput. Mater. Sci. 14, 4.
264 Args:
265 atoms: Atoms object.
266 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
267 ik: k-point index.
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])
276def T(atoms, W, dr):
277 """Translation operator.
279 This operator acts on options 5 and 6.
281 Reference: https://ccrma.stanford.edu/~jos/st/Shift_Theorem.html
283 Args:
284 atoms: Atoms object.
285 W: Expansion coefficients of unconstrained wave functions in reciprocal space.
286 dr: Real-space shifting vector.
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])
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
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