Coverage for gpaw/pw/lfc.py: 98%
281 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1from math import pi
3import gpaw.cgpaw as cgpaw
4import numpy as np
5from gpaw.lfc import BaseLFC
6from gpaw.spherical_harmonics import Y, nablarlYL
7from gpaw.ffbt import rescaled_fourier_bessel_transform
8from gpaw.utilities.blas import mmm
11class PWLFC(BaseLFC):
12 def __init__(self, spline_aj, pd, blocksize=5000, comm=None):
13 """Reciprocal-space plane-wave localized function collection.
15 spline_aj: list of list of spline objects
16 Splines.
17 pd: PWDescriptor
18 Plane-wave descriptor object.
19 blocksize: int
20 Block-size to use when looping over G-vectors. Use None for
21 doing all G-vectors in one big block.
22 comm: communicator
23 Communicator for operations that support parallelization
24 over planewaves (only integrate so far)."""
26 self.pd = pd
27 self.spline_aj = spline_aj
29 self.dtype = pd.dtype
31 self.initialized = False
33 # These will be filled in later:
34 self.Y_qGL = []
35 self.emiGR_qGa = []
36 self.f_qGs = []
37 self.l_s = None
38 self.a_J = None
39 self.s_J = None
40 self.lmax = None
42 if blocksize is not None:
43 if pd.ngmax <= blocksize:
44 # No need to block G-vectors
45 blocksize = None
46 self.blocksize = blocksize
48 # These are set later in set_potitions():
49 self.eikR_qa = None
50 self.my_atom_indices = None
51 self.my_indices = None
52 self.pos_av = None
53 self.nI = None
55 if comm is None:
56 comm = pd.gd.comm
57 else:
58 assert False
59 self.comm = comm
61 def initialize(self):
62 """Initialize position-independent stuff."""
63 if self.initialized:
64 return
66 splines = {} # Dict[Spline, int]
67 for spline_j in self.spline_aj:
68 for spline in spline_j:
69 if spline not in splines:
70 splines[spline] = len(splines)
71 nsplines = len(splines)
73 nJ = sum(len(spline_j) for spline_j in self.spline_aj)
75 self.f_qGs = [np.empty((mynG, nsplines)) for mynG in self.pd.myng_q]
76 self.l_s = np.empty(nsplines, np.int32)
77 self.a_J = np.empty(nJ, np.int32)
78 self.s_J = np.empty(nJ, np.int32)
80 # Fourier transform radial functions:
81 J = 0
82 done = set() # Set[Spline]
83 for a, spline_j in enumerate(self.spline_aj):
84 for spline in spline_j:
85 s = splines[spline] # get spline index
86 if spline not in done:
87 f = rescaled_fourier_bessel_transform(spline)
88 for f_Gs, G2_G in zip(self.f_qGs, self.pd.G2_qG):
89 G_G = G2_G**0.5
90 f_Gs[:, s] = f.map(G_G)
91 self.l_s[s] = spline.get_angular_momentum_number()
92 done.add(spline)
93 self.a_J[J] = a
94 self.s_J[J] = s
95 J += 1
97 self.lmax = max(self.l_s, default=-1)
99 # Spherical harmonics:
100 for q, K_v in enumerate(self.pd.K_qv):
101 G_Gv = self.pd.get_reciprocal_vectors(q=q)
102 Y_GL = np.empty((len(G_Gv), (self.lmax + 1)**2))
103 for L in range((self.lmax + 1)**2):
104 Y_GL[:, L] = Y(L, *G_Gv.T)
105 self.Y_qGL.append(Y_GL)
107 self.initialized = True
109 def estimate_memory(self, mem):
110 splines = set()
111 lmax = -1
112 for spline_j in self.spline_aj:
113 for spline in spline_j:
114 splines.add(spline)
115 l = spline.get_angular_momentum_number()
116 lmax = max(lmax, l)
117 nbytes = ((len(splines) + (lmax + 1)**2) *
118 sum(G2_G.nbytes for G2_G in self.pd.G2_qG))
119 mem.subnode('Arrays', nbytes)
121 def get_function_count(self, a):
122 return sum(2 * spline.get_angular_momentum_number() + 1
123 for spline in self.spline_aj[a])
125 def set_positions(self, spos_ac, atom_partition=None):
126 self.initialize()
127 kd = self.pd.kd
128 if kd is None or kd.gamma:
129 self.eikR_qa = np.ones((1, len(spos_ac)))
130 else:
131 self.eikR_qa = np.exp(2j * pi * np.dot(kd.ibzk_qc, spos_ac.T))
133 self.pos_av = np.dot(spos_ac, self.pd.gd.cell_cv)
135 del self.emiGR_qGa[:]
136 G_Qv = self.pd.G_Qv
137 for Q_G in self.pd.myQ_qG:
138 GR_Ga = np.dot(G_Qv[Q_G], self.pos_av.T)
139 self.emiGR_qGa.append(np.exp(-1j * GR_Ga))
141 if atom_partition is None:
142 assert self.comm.size == 1
143 rank_a = np.zeros(len(spos_ac), int)
144 else:
145 rank_a = atom_partition.rank_a
147 self.my_atom_indices = []
148 self.my_indices = []
149 I1 = 0
150 for a, rank in enumerate(rank_a):
151 I2 = I1 + self.get_function_count(a)
152 if rank == self.comm.rank:
153 self.my_atom_indices.append(a)
154 self.my_indices.append((a, I1, I2))
155 I1 = I2
156 self.nI = I1
158 def expand(self, q=-1, G1=0, G2=None, cc=False):
159 """Expand functions in plane-waves.
161 q: int
162 k-point index.
163 G1: int
164 Start G-vector index.
165 G2: int
166 End G-vector index.
167 cc: bool
168 Complex conjugate.
169 """
170 if G2 is None:
171 G2 = self.Y_qGL[q].shape[0]
173 emiGR_Ga = self.emiGR_qGa[q][G1:G2]
174 f_Gs = self.f_qGs[q][G1:G2]
175 Y_GL = self.Y_qGL[q][G1:G2]
177 if self.pd.dtype == complex:
178 f_GI = np.empty((G2 - G1, self.nI), complex)
179 else:
180 # Special layout because BLAS does not have real-complex
181 # multiplications. f_GI(G,I) layout:
182 #
183 # real(G1, 0), real(G1, 1), ...
184 # imag(G1, 0), imag(G1, 1), ...
185 # real(G1+1, 0), real(G1+1, 1), ...
186 # imag(G1+1, 0), imag(G1+1, 1), ...
187 # ...
189 f_GI = np.empty((2 * (G2 - G1), self.nI))
191 if True:
192 # Fast C-code:
193 cgpaw.pwlfc_expand_old(f_Gs, emiGR_Ga, Y_GL,
194 self.l_s, self.a_J, self.s_J,
195 cc, f_GI)
196 return f_GI
198 # Equivalent slow Python code:
199 f_GI = np.empty((G2 - G1, self.nI), complex)
200 I1 = 0
201 for J, (a, s) in enumerate(zip(self.a_J, self.s_J)):
202 l = self.l_s[s]
203 I2 = I1 + 2 * l + 1
204 f_GI[:, I1:I2] = (f_Gs[:, s] *
205 emiGR_Ga[:, a] *
206 Y_GL[:, l**2:(l + 1)**2].T *
207 (-1.0j)**l).T
208 I1 = I2
209 if cc:
210 f_GI = f_GI.conj()
211 if self.pd.dtype == float:
212 f_GI = f_GI.T.copy().view(float).T.copy()
214 return f_GI
216 def block(self, q=-1, ensure_same_number_of_blocks=False):
217 nG = self.Y_qGL[q].shape[0]
218 B = self.blocksize
219 if B:
220 G1 = 0
221 while G1 < nG:
222 G2 = min(G1 + B, nG)
223 yield G1, G2
224 G1 = G2
225 if ensure_same_number_of_blocks:
226 # Make sure we yield the same number of times:
227 nb = (self.pd.maxmyng + B - 1) // B
228 mynb = (nG + B - 1) // B
229 if mynb < nb:
230 yield nG, nG # empty block
231 else:
232 yield 0, nG
234 def add(self, a_xG, c_axi=1.0, q=-1, f0_IG=None):
235 c_xI = np.empty(a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
237 if isinstance(c_axi, float):
238 assert q == -1 and a_xG.ndim == 1
239 c_xI[:] = c_axi
240 else:
241 assert q != -1 or self.pd.only_one_k_point
242 if self.comm.size != 1:
243 c_xI[:] = 0.0
244 for a, I1, I2 in self.my_indices:
245 c_xI[..., I1:I2] = c_axi[a] * self.eikR_qa[q][a].conj()
246 if self.comm.size != 1:
247 self.comm.sum(c_xI)
249 nx = np.prod(c_xI.shape[:-1], dtype=int)
250 c_xI = c_xI.reshape((nx, self.nI))
251 a_xG = a_xG.reshape((nx, a_xG.shape[-1])).view(self.pd.dtype)
253 for G1, G2 in self.block(q):
254 if f0_IG is None:
255 f_GI = self.expand(q, G1, G2, cc=False)
256 else:
257 1 / 0
258 # f_IG = f0_IG
260 if self.pd.dtype == float:
261 # f_IG = f_IG.view(float)
262 G1 *= 2
263 G2 *= 2
265 mmm(1.0 / self.pd.gd.dv, c_xI, 'N', f_GI, 'T',
266 1.0, a_xG[:, G1:G2])
268 def integrate(self, a_xG, c_axi=None, q=-1):
269 c_xI = np.zeros(a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
271 nx = np.prod(c_xI.shape[:-1], dtype=int)
272 b_xI = c_xI.reshape((nx, self.nI))
273 a_xG = a_xG.reshape((nx, a_xG.shape[-1]))
275 alpha = 1.0 / self.pd.gd.N_c.prod()
276 if self.pd.dtype == float:
277 alpha *= 2
278 a_xG = a_xG.view(float)
280 if c_axi is None:
281 c_axi = self.dict(a_xG.shape[:-1])
283 x = 0.0
284 for G1, G2 in self.block(q):
285 f_GI = self.expand(q, G1, G2, cc=self.pd.dtype == complex)
286 if self.pd.dtype == float:
287 if G1 == 0 and self.comm.rank == 0:
288 f_GI[0] *= 0.5
289 G1 *= 2
290 G2 *= 2
291 mmm(alpha, a_xG[:, G1:G2], 'N', f_GI, 'N', x, b_xI)
292 x = 1.0
294 self.comm.sum(b_xI)
295 for a, I1, I2 in self.my_indices:
296 c_axi[a][:] = self.eikR_qa[q][a] * c_xI[..., I1:I2]
298 return c_axi
300 def matrix_elements(self, psit, out):
301 P_ani = {a: P_in.T for a, P_in in out.items()}
302 self.integrate(psit.array, P_ani, psit.kpt)
304 def derivative(self, a_xG, c_axiv=None, q=-1):
305 c_vxI = np.zeros((3,) + a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
306 nx = np.prod(c_vxI.shape[1:-1], dtype=int)
307 b_vxI = c_vxI.reshape((3, nx, self.nI))
308 a_xG = a_xG.reshape((nx, a_xG.shape[-1])).view(self.pd.dtype)
310 alpha = 1.0 / self.pd.gd.N_c.prod()
312 if c_axiv is None:
313 c_axiv = self.dict(a_xG.shape[:-1], derivative=True)
315 K_v = self.pd.K_qv[q]
317 x = 0.0
318 for G1, G2 in self.block(q):
319 f_GI = self.expand(q, G1, G2, cc=True)
320 G_Gv = self.pd.G_Qv[self.pd.myQ_qG[q][G1:G2]]
321 if self.pd.dtype == float:
322 d_GI = np.empty_like(f_GI)
323 for v in range(3):
324 d_GI[::2] = f_GI[1::2] * G_Gv[:, v, np.newaxis]
325 d_GI[1::2] = f_GI[::2] * G_Gv[:, v, np.newaxis]
326 mmm(2 * alpha,
327 a_xG[:, 2 * G1:2 * G2], 'N',
328 d_GI, 'N',
329 x, b_vxI[v])
330 else:
331 for v in range(3):
332 mmm(-alpha,
333 a_xG[:, G1:G2], 'N',
334 f_GI * (G_Gv[:, v] + K_v[v])[:, np.newaxis], 'N',
335 x, b_vxI[v])
336 x = 1.0
338 self.comm.sum(c_vxI)
340 for v in range(3):
341 if self.pd.dtype == float:
342 for a, I1, I2 in self.my_indices:
343 c_axiv[a][..., v] = c_vxI[v, ..., I1:I2]
344 else:
345 for a, I1, I2 in self.my_indices:
346 c_axiv[a][..., v] = (1.0j * self.eikR_qa[q][a] *
347 c_vxI[v, ..., I1:I2])
349 return c_axiv
351 def stress_tensor_contribution(self, a_xG, c_axi=1.0, q=-1):
352 cache = {}
353 things = []
354 I1 = 0
355 lmax = 0
356 for a, spline_j in enumerate(self.spline_aj):
357 for spline in spline_j:
358 if spline not in cache:
359 s = rescaled_fourier_bessel_transform(spline)
360 G_G = self.pd.G2_qG[q]**0.5
361 f_G = []
362 dfdGoG_G = []
363 for G in G_G:
364 f, dfdG = s.get_value_and_derivative(G)
365 if G < 1e-10:
366 G = 1.0
367 f_G.append(f)
368 dfdGoG_G.append(dfdG / G)
369 f_G = np.array(f_G)
370 dfdGoG_G = np.array(dfdGoG_G)
371 cache[spline] = (f_G, dfdGoG_G)
372 else:
373 f_G, dfdGoG_G = cache[spline]
374 l = spline.l
375 lmax = max(l, lmax)
376 I2 = I1 + 2 * l + 1
377 things.append((a, l, I1, I2, f_G, dfdGoG_G))
378 I1 = I2
380 if isinstance(c_axi, float):
381 c_axi = {a: c_axi for a in range(len(self.pos_av))}
383 G0_Gv = self.pd.get_reciprocal_vectors(q=q)
385 stress_vv = np.zeros((3, 3))
386 for G1, G2 in self.block(q, ensure_same_number_of_blocks=True):
387 G_Gv = G0_Gv[G1:G2]
388 Z_LvG = np.array([nablarlYL(L, G_Gv.T)
389 for L in range((lmax + 1)**2)])
390 aa_xG = a_xG[..., G1:G2]
391 for v1 in range(3):
392 for v2 in range(3):
393 stress_vv[v1, v2] += self._stress_tensor_contribution(
394 v1, v2, things, G1, G2, G_Gv, aa_xG, c_axi, q, Z_LvG)
396 self.comm.sum(stress_vv)
398 return stress_vv
400 def _stress_tensor_contribution(self, v1, v2, things, G1, G2,
401 G_Gv, a_xG, c_axi, q, Z_LvG):
402 f_IG = np.empty((self.nI, G2 - G1), complex)
403 emiGR_Ga = self.emiGR_qGa[q][G1:G2]
404 Y_LG = self.Y_qGL[q].T
405 for a, l, I1, I2, f_G, dfdGoG_G in things:
406 L1 = l**2
407 L2 = (l + 1)**2
408 f_IG[I1:I2] = (emiGR_Ga[:, a] * (-1.0j)**l *
409 (dfdGoG_G[G1:G2] * G_Gv[:, v1] * G_Gv[:, v2] *
410 Y_LG[L1:L2, G1:G2] +
411 f_G[G1:G2] * G_Gv[:, v1] * Z_LvG[L1:L2, v2]))
413 c_xI = np.zeros(a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
415 x = np.prod(c_xI.shape[:-1], dtype=int)
416 b_xI = c_xI.reshape((x, self.nI))
417 a_xG = a_xG.reshape((x, a_xG.shape[-1]))
419 alpha = 1.0 / self.pd.gd.N_c.prod()
420 if self.pd.dtype == float:
421 alpha *= 2
422 if G1 == 0 and self.pd.gd.comm.rank == 0:
423 f_IG[:, 0] *= 0.5
424 f_IG = f_IG.view(float)
425 a_xG = a_xG.copy().view(float)
427 mmm(alpha, a_xG, 'N', f_IG, 'C', 0.0, b_xI)
428 self.comm.sum(b_xI)
430 stress = 0.0
431 for a, I1, I2 in self.my_indices:
432 stress -= self.eikR_qa[q][a] * (c_axi[a] * c_xI[..., I1:I2]).sum()
433 return stress.real