Coverage for gpaw/response/paw.py: 99%
210 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1from __future__ import annotations
3import numpy as np
4from scipy.special import spherical_jn
5from dataclasses import dataclass
7from gpaw.spline import Spline
8from gpaw.ffbt import rescaled_fourier_bessel_transform
9from gpaw.gaunt import gaunt, super_gaunt
10from gpaw.spherical_harmonics import Y
11from gpaw.atom.radialgd import RadialGridDescriptor
12from gpaw.sphere.rshe import RealSphericalHarmonicsExpansion
13from gpaw.response.pw_parallelization import Blocks1D
16# Important note: The test suite monkeypatches this value to 2**10 so
17# you may get different results in tests and production until we
18# implement a better solution (e.g. generating setup-dependent radial_points).
19#
20# The motivation for lowering to 2**10 in tests is that many tests
21# take 3-4 times longer if we do not.
22#
23# See https://gitlab.com/gpaw/gpaw/-/issues/984
24DEFAULT_RADIAL_POINTS = 2**12
27@dataclass
28class LeanPAWDataset:
29 rgd: RadialGridDescriptor
30 l_j: np.ndarray
31 rcut_j: np.ndarray
32 phit_jg: np.ndarray
33 phi_jg: np.ndarray
34 # Number of radial points in spline interpolation
35 radial_points: int | None = None
37 def __post_init__(self):
38 if self.radial_points is None:
39 # We assign this late due to monkeypatch in testing
40 self.radial_points = DEFAULT_RADIAL_POINTS
42 # Number of basis functions
43 self.ni = np.sum([2 * l + 1 for l in self.l_j])
44 # Maximum angular momentum index l
45 self.lmax = np.max(self.l_j)
46 # Grid cutoff to create spline representation
47 self.gcut2 = self.rgd.ceil(2 * max(self.rcut_j))
49 # Set up cache
50 self.dn_g_cache = {}
51 self.dn_kspline_cache = {}
53 def dn_kspline(self, j1: int, j2: int, l: int):
54 """Get spline representation of Δn_jj'l(k)."""
55 if (j1, j2, l) not in self.dn_kspline_cache:
56 dn_g = self.get_pair_density_correction(j1, j2)
57 self.dn_kspline_cache[j1, j2, l] = \
58 self.rescaled_fourier_bessel_transform(dn_g, l)
59 return self.dn_kspline_cache[j1, j2, l]
61 def get_pair_density_correction(self, j1: int, j2: int):
62 """Get Δn_jj'(r), while keeping the newest correction cached."""
63 if (j1, j2) not in self.dn_g_cache:
64 self.dn_g_cache = {} # keep only one density in cache
65 self.dn_g_cache[j1, j2] = self.calculate_pair_density_correction(
66 j1, j2)
67 return self.dn_g_cache[j1, j2]
69 def calculate_pair_density_correction(self, j1, j2):
70 """Calculate the pair density PAW correction for two partial waves.
71 ˷ ˷
72 Δn (r) = φ (r) φ (r) - φ (r) φ (r)
73 jj' j j' j j'
74 """
75 # (Real) radial functions for the partial waves
76 phi_jg = self.phi_jg
77 phit_jg = self.phit_jg
78 return phi_jg[j1] * phi_jg[j2] - phit_jg[j1] * phit_jg[j2]
80 def rescaled_fourier_bessel_transform(self, f_g, l):
81 """Calculate the rescaled Fourier Bessel transform f_l(k)
83 rc
84 4π ⌠ 2
85 f (k) = ‾‾‾ ⎪ r dr j (kr) f(r)
86 l k^l ⌡ l
87 0
88 """
89 # First, we make a spline representation of the radial function f(r),
90 # which is rescaled by a factor of r^-l by the radial grid descriptor:
91 spline = self.rgd.spline(
92 f_g[:self.gcut2], l=l, points=self.radial_points)
93 # Once rescaled by r^-l, we can use the gpaw.ffbt module to Fast
94 # Fourier Bessel Transform the radial spline r^-l f(r) and obtain
95 # f_l(k) in a reciprocal spline representation
96 kspline = rescaled_fourier_bessel_transform(
97 spline, N=4 * self.radial_points)
98 # Since this procedures hinges on a series of hardcoded parameters, we
99 # return a self-testing version of the spline. If someone wants to run
100 # calculations without dynamic testing of the FFBT methodology (i.e.
101 # in an "unprotected" mode), simply return the bare `kspline` here.
102 return SelfTestingKSpline(self.rgd, f_g, kspline)
105class SelfTestingKSpline(Spline):
106 """Self-testing reciprocal spline representation, f_l(k)."""
108 def __init__(self, rgd, f_g, spline: Spline):
109 # Store original real-space representation of the radial function f(r)
110 self.rgd = rgd
111 self.f_g = f_g
112 super().__init__(spline.spline)
114 def map(self, k_G):
115 self.test_spline_representation(k_G)
116 return super().map(k_G)
118 def test_spline_representation(self, k_G):
119 """Test validity of the FFBT implementation on input domain.
121 At present, LeanPAWDataset's FFBT implementation relies on a range of
122 hardcoded parameters, which are not guaranteed to work for all cases.
124 In particular, the uniform radial grid used for the FFBT is defined
125 through the `rcut` and `N` parameters in
126 `rescaled_fourier_bessel_transform()`
127 where the former is hardcoded inside the function itself.
129 Furthermore, the `points` parameter to `rgd.spline()` controls the
130 fidelity of the interpolation between nonlinear and equidistant radial
131 grids needed to make use of the FFBT algorithm.
133 To make a generally reliable implementation, one would need to control
134 all of these parameters based on the setup, e.g. the nonlinear radial
135 grid spacing. In doing so, one should be mindful that the `rcut`
136 parameter defines the reciprocal grid spacing of the kspline and that
137 `N` controls the range of the reciprocal space domain.
139 For now, we simply check that the requested plane waves are within the
140 computed k-range of the FFBT and check that the resulting transforms
141 match a manual calculation at a few selected K-vectors.
142 """
143 kmax = np.max(k_G)
144 assert kmax <= self.get_cutoff()
146 # Manual calculation at finite k
147 k_k = np.array([kmax, np.average(k_G)])
148 f_k = 4 * np.pi * fourier_bessel_transform(
149 np.array(k_k), self.l, self.rgd, self.f_g)
150 # Manual calculation at k=0
151 if self.l == 0: # Vanishes for l>0
152 k_k = np.append(k_k, [0.])
153 f_k = np.append(f_k, [self.rgd.integrate(self.f_g)])
154 # FFBT calculation
155 myf_k = k_k**self.l * super().map(k_k)
156 assert np.allclose(myf_k, f_k, rtol=1e-2, atol=1e-3), \
157 f'FFBT mismatch: {myf_k}, {f_k}'
160def calculate_pair_density_correction(qG_Gv: np.ndarray, *,
161 pawdata: LeanPAWDataset):
162 r"""Calculate the atom-centered PAW correction to the pair density.
163 ˍ
164 The atom-centered pair density correction tensor, Q_aii', is defined as the
165 atom-centered Fourier transform
167 ˍ / ˷ ˷
168 Q_aii'(G+q) = | dr e^-i(G+q)r [φ_ai^*(r) φ_ai'(r) - φ_ai^*(r) φ_ai'(r)]
169 /
171 evaluated with the augmentation sphere center at the origin. The full pair
172 density correction tensor is then given by
173 ˍ
174 Q_aii'(G+q) = e^(-i[G+q].R_a) Q_aii'(G+q)
176 Expanding the plane wave coefficient into real spherical harmonics and
177 spherical Bessel functions, the correction can split into angular and
178 radial contributions
180 l
181 __ __
182 \ \ l m ˰ m,m_i,m_i'
183 Q_aii'(K) = / / (-i) Y (K) g
184 ‾‾ ‾‾ l l,l_i,l_i'
185 l m=-l
186 rc
187 / a a ˷a ˷a
188 × 4π | r^2 dr j_l(|K|r) [φ (r) φ (r) - φ (r) φ (r)]
189 / j_i j_i' j_i j_i'
190 0
192 where K = G+q and g denotes the Gaunt coefficients.
194 For more information, see [PRB 103, 245110 (2021)]. In particular, it
195 should be noted that the partial waves themselves are defined via real
196 spherical harmonics and radial functions φ_j from the PAW setups:
198 a m_i ˰ a
199 φ (r) = Y (r) φ (r)
200 i l_i j_i
201 """
202 ni = pawdata.ni # Number of partial waves
203 l_j = pawdata.l_j # l-index for each radial function index j
204 G_LLL = gaunt(pawdata.lmax)
206 # Initialize correction tensor
207 npw = qG_Gv.shape[0]
208 Qbar_Gii = np.zeros((npw, ni, ni), dtype=complex)
210 # K-vector norm
211 k_G = np.linalg.norm(qG_Gv, axis=1)
213 # Loop of radial function indices for partial waves i and i'
214 i1_counter = 0
215 for j1, l1 in enumerate(l_j):
216 i2_counter = 0
217 for j2, l2 in enumerate(l_j):
218 # Sample l according to the Gaunt coefficient selection rules, see
219 # e.g. gpaw.test.test_gaunt
220 for l in range(abs(l1 - l2), l1 + l2 + 1, 2):
221 # Calculate the spherical Fourier-Bessel transform
222 # rc
223 # 4π /
224 # Δn_jj'l(k) = ‾‾‾ | r^2 dr j_l(kr) Δn_jj'(r)
225 # k^l /
226 # 0
227 # To evaluate the radial integral efficiently, we rely on the
228 # Fast Fourier Bessel Transform (FFBT) algorithm, see gpaw.ffbt
229 # The FFBT produces a radial spline representation of
230 # Δn_jj'l(k), which we map to the K-vector norms in question:
231 dn_G = pawdata.dn_kspline(j1, j2, l).map(k_G)
233 # Angular part of the integral
234 f_G = (-1j)**l * dn_G
235 for m in range(l**2, (l + 1)**2):
236 # Calculate the solid harmonic
237 # m ˰
238 # |K|^l Y (K)
239 # l
240 klY_G = Y(m, *qG_Gv.T)
242 # Generate m-indices for each radial function
243 for m1 in range(2 * l1 + 1):
244 for m2 in range(2 * l2 + 1):
245 # Set up the i=(l,m) index for each partial wave
246 i1 = i1_counter + m1
247 i2 = i2_counter + m2
248 # Extract Gaunt coefficients
249 gaunt_coeff = G_LLL[l1**2 + m1, l2**2 + m2, m]
250 if (gaunt_coeff == 0):
251 continue
253 # Add contribution to the PAW correction
254 Qbar_Gii[:, i1, i2] += gaunt_coeff * klY_G * f_G
256 # Add to i and i' counters
257 i2_counter += 2 * l2 + 1
258 i1_counter += 2 * l1 + 1
259 return Qbar_Gii
262def calculate_matrix_element_correction(qG_Gv, pawdata,
263 rshe: RealSphericalHarmonicsExpansion):
264 r"""Calculate the atom-centered correction to a generalized matrix element.
266 For matrix elements corresponding to the expectation value of a plane wave
267 coefficient e^-i(G+q)r and a known functional of the (spin-)density
268 f[n](r), the PAW correction tensor is given by
270 F_aii'(G+q) = <φ_ai| e^-i(G+q)r f[n](r) |φ_ai'>
271 ˷ ˷
272 - <φ_ai| e^-i(G+q)r f[n](r) |φ_ai'>
273 ˍ
274 = e^(-i[G+q].R_a) F_aii'(G+q)
275 ˍ
276 where F_aii'(G+q) is the atom-centered PAW correction tensor.
278 Expanding the functional f[n](r) in the atom-centered frame in real
279 spherical harmonics (corresponding to the input rshe),
281 l
282 __ __
283 → \ \ m ˰ m
284 f[n](r) = / / Y (r) f (r)
285 ‾‾ ‾‾ l l
286 l m=-l
288 expansion of the plane-wave coefficient in real spherical harmonics and
289 spherical Bessel functions j_l(Kr) yields the following expression for the
290 atom-centered correction tensor [publication in preparation]:
292 l l'
293 __ __ __ __
294 ˍ \ \ \ \ l' m'˰ m_i,m_i',m,m'
295 F_aii'(K) = 4π / / / / (-i) Y (K) G
296 ‾‾ ‾‾ ‾‾ ‾‾ l' l_i,l_i',l,l'
297 l m=-l l' m'=-l'
299 rc
300 / 2 a a ˷a ˷a m
301 × | r dr j (Kr) [φ (r) φ (r) - φ (r) φ (r)] f (r)
302 / l' j_i j_i' j_i j_i' l
303 0
305 where K=G+q and G_LLLL denotes the super Gaunt coefficients, which yield
306 the integrals over four spherical harmonics.
307 """
308 rgd = rshe.rgd
309 assert rgd is pawdata.xc_correction.rgd
310 ni = pawdata.ni # Number of partial waves
311 l_j = pawdata.l_j # l-index for each radial function index j
312 lmax = max(l_j)
313 assert max(rshe.l_M) <= 2 * lmax
314 G_LLLL = super_gaunt(lmax)
315 # (Real) radial functions for the partial waves
316 phi_jg = pawdata.phi_jg
317 phit_jg = pawdata.phit_jg
318 # Truncate the radial functions to span only the radial grid coordinates
319 # which need correction
320 assert np.allclose(rgd.r_g, pawdata.rgd.r_g[:rgd.N])
321 phi_jg = np.array(phi_jg)[:, :rgd.N]
322 phit_jg = np.array(phit_jg)[:, :rgd.N]
324 # Initialize correction tensor
325 npw = qG_Gv.shape[0]
326 Fbar_Gii = np.zeros((npw, ni, ni), dtype=complex)
328 # K-vector norm and direction
329 k_G = np.linalg.norm(qG_Gv, axis=1)
330 Kd_Gv = qG_Gv.copy()
331 Kd_Gv[k_G > 1e-10] /= k_G[k_G > 1e-10, np.newaxis]
333 # Loop of radial function indices for partial waves i and i'
334 i1_counter = 0
335 for j1, l1 in enumerate(l_j):
336 i2_counter = 0
337 for j2, l2 in enumerate(l_j):
338 # Calculate the radial partial wave correction
339 # ˷ ˷
340 # Δn_jj'(r) = φ_j(r) φ_j'(r) - φ_j(r) φ_j'(r)
341 dn_g = phi_jg[j1] * phi_jg[j2] - phit_jg[j1] * phit_jg[j2]
343 # Loop through the angular components in the real spherical
344 # harmonics expansion of f[n](r)
345 for l, L, f_g in zip(rshe.l_M, rshe.L_M, rshe.f_gM.T):
346 dnf_g = dn_g * f_g
347 # Apply Gaunt coefficient selection rules to loop through
348 # the l' coefficients of the plane-wave expansion
349 lpmin = np.min(abs(
350 np.arange(abs(l1 - l2), l1 + l2 + 1) - l))
351 for lp in range(lpmin, l1 + l2 + l + 1):
352 if not (l1 + l2 + l + lp) % 2 == 0:
353 continue
354 # Calculate radial part of the correction
355 dnf_G = parallel_fourier_bessel_transform(
356 k_G, lp, rgd, dnf_g)
358 # Calculate angular part of the correction
359 x_G = 4 * np.pi * (-1j)**lp * dnf_G
360 # Loop through available m-indices for the partial waves
361 # and generate the composite L=(l,m) index as well as the
362 # partial wave index i
363 for m1 in range(2 * l1 + 1):
364 L1 = l1**2 + m1
365 i1 = i1_counter + m1
366 for m2 in range(2 * l2 + 1):
367 L2 = l2**2 + m2
368 i2 = i2_counter + m2
369 # Loop through m' indices of the plane-wave
370 # expansion and generate the L' composite index
371 for mp in range(2 * lp + 1):
372 Lp = lp**2 + mp
373 # If the angular integral (super gaunt
374 # coefficient) is finite,
375 coeff = G_LLLL[L1, L2, L, Lp]
376 if abs(coeff) > 1e-10:
377 # Calculate spherical harmonic and add
378 # contribution to the PAW correction
379 Y_G = Y(Lp, *Kd_Gv.T)
380 Fbar_Gii[:, i1, i2] += coeff * Y_G * x_G
382 # Add to i and i' counters
383 i2_counter += 2 * l2 + 1
384 i1_counter += 2 * l1 + 1
385 return Fbar_Gii
388def parallel_fourier_bessel_transform(k_G, *args, comm=None):
389 """Distribute FBT plane-wave components over a given communicator."""
390 # NB: If we need to do something similar elsewhere, we can generalize this
391 # function to a decorator!
392 if comm is None:
393 from gpaw.mpi import world as comm
394 Gblocks = Blocks1D(comm, len(k_G))
395 f_myG = fourier_bessel_transform(k_G[Gblocks.myslice], *args)
396 return Gblocks.all_gather(f_myG)
399def fourier_bessel_transform(k_G, l, rgd, f_g):
400 """Perform a spherical Fourier-Bessel transform of a radial function f(r).
402 Computes the transform
404 max
405 r
406 ⌠ 2
407 f(k) = ⎪ r dr j (kr) f(r)
408 ⌡ l
409 0
411 on the supplied radial grid.
412 """
413 # Vectorize calculation of spherical Bessel functions
414 l_Gg = l * np.ones((len(k_G), rgd.N), dtype=int)
415 kr_Gg = k_G[:, np.newaxis] * rgd.r_g[np.newaxis]
416 jl_Gg = spherical_jn(l_Gg, kr_Gg) # so slow...
417 # Integrate the radial grid using linear interpolation
418 f_G = rgd.integrate_trapz(jl_Gg * f_g[np.newaxis])
419 return f_G
422class PWPAWCorrectionData:
423 def __init__(self, Q_aGii, qpd, pawdatasets, pos_av, atomrotations):
424 # Sometimes we loop over these in ways that are very dangerous.
425 # It must be list, not dictionary.
426 assert isinstance(Q_aGii, list)
427 assert len(Q_aGii) == len(pos_av) == len(pawdatasets.by_atom)
429 self.Q_aGii = Q_aGii
431 self.qpd = qpd
432 self.pawdatasets = pawdatasets
433 self.pos_av = pos_av
434 self.atomrotations = atomrotations
436 def _new(self, Q_aGii):
437 return PWPAWCorrectionData(Q_aGii, qpd=self.qpd,
438 pawdatasets=self.pawdatasets,
439 pos_av=self.pos_av,
440 atomrotations=self.atomrotations)
442 def remap(self, M_vv, G_Gv, sym, sign):
443 Q_aGii = []
444 for a, Q_Gii in enumerate(self.Q_aGii):
445 x_G = self._get_x_G(G_Gv, M_vv, self.pos_av[a])
446 U_ii = self.atomrotations.get_by_a(a).R_sii[sym]
448 Q_Gii = np.einsum('ij,kjl,ml->kim',
449 U_ii,
450 Q_Gii * x_G[:, None, None],
451 U_ii,
452 optimize='optimal')
453 if sign == -1:
454 Q_Gii = Q_Gii.conj()
455 Q_aGii.append(Q_Gii)
457 return self._new(Q_aGii)
459 def _get_x_G(self, G_Gv, M_vv, pos_v):
460 # This doesn't really belong here. Or does it? Maybe this formula
461 # is only used with PAW corrections.
462 return np.exp(1j * (G_Gv @ (pos_v - M_vv @ pos_v)))
464 def remap_by_symop(self, symop, G_Gv, M_vv):
465 return self.remap(M_vv, G_Gv, symop.symno, symop.sign)
467 def multiply(self, P_ani, band):
468 assert isinstance(P_ani, list)
469 assert len(P_ani) == len(self.Q_aGii)
471 C1_aGi = [Qa_Gii @ P1_ni[band].conj()
472 for Qa_Gii, P1_ni in zip(self.Q_aGii, P_ani)]
473 return C1_aGi
475 def reduce_ecut(self, G2G):
476 # XXX actually we should return this with another PW descriptor.
477 return self._new([Q_Gii.take(G2G, axis=0) for Q_Gii in self.Q_aGii])
479 def almost_equal(self, otherpawcorr, G_G):
480 for a, Q_Gii in enumerate(otherpawcorr.Q_aGii):
481 e = abs(self.Q_aGii[a] - Q_Gii[G_G]).max()
482 if e > 1e-12:
483 return False
484 return True
487def get_pair_density_paw_corrections(pawdatasets, qpd, spos_ac, atomrotations):
488 r"""Calculate and bundle paw corrections to the pair densities as a
489 PWPAWCorrectionData object.
491 The pair density PAW correction tensor is given by:
493 /
494 Q_aii'(G+q) = | dr e^(-i[G+q].r) [φ_ai^*(r-R_a) φ_ai'(r-R_a)
495 / ˷ ˷
496 - φ_ai^*(r-R_a) φ_ai'(r-R_a)]
497 """
498 qG_Gv = qpd.get_reciprocal_vectors(add_q=True)
499 pos_av = spos_ac @ qpd.gd.cell_cv
501 # Calculate pair density PAW correction tensor
502 Qbar_xGii = {}
503 for species_index, pawdata in pawdatasets.by_species.items():
504 # Calculate atom-centered correction tensor
505 Qbar_Gii = calculate_pair_density_correction(qG_Gv, pawdata=pawdata)
506 # Add dependency on the atomic position (phase factor)
507 Qbar_xGii[species_index] = Qbar_Gii
509 Q_aGii = []
510 for a, (pos_v, pawdata) in enumerate(zip(pos_av, pawdatasets.by_atom)):
511 x_G = np.exp(-1j * (qG_Gv @ pos_v))
512 species_index = pawdatasets.id_by_atom[a]
513 Qbar_Gii = Qbar_xGii[species_index]
514 Q_aGii.append(x_G[:, np.newaxis, np.newaxis] * Qbar_Gii)
516 return PWPAWCorrectionData(Q_aGii, qpd=qpd,
517 pawdatasets=pawdatasets,
518 pos_av=pos_av,
519 atomrotations=atomrotations)
522def get_matrix_element_paw_corrections(qpd, pawdata_a, rshe_a, spos_ac):
523 r"""Calculate the PAW correction to a generalized matrix element.
525 For a given functional of the electron (spin-)density f[n](r), the PAW
526 correction is given by
527 ˍ
528 F_aii'(G+q) = e^(-i[G+q].R_a) F_aii'(G+q)
529 ˍ
530 where F_aii'(G+q) is the atom-centered correction (see above).
531 """
532 qG_Gv = qpd.get_reciprocal_vectors(add_q=True)
534 F_aGii = []
535 for pawdata, rshe, spos_c in zip(pawdata_a.by_atom, rshe_a, spos_ac):
536 # Calculate atom-centered PAW correction
537 Fbar_Gii = calculate_matrix_element_correction(
538 qG_Gv, pawdata, rshe)
540 # XXX Can time be saved by doing some of the processing per species
541 # rather than per atom?
543 # Add dependency on the atomic position (phase factor)
544 pos_v = spos_c @ qpd.gd.cell_cv
545 x_G = np.exp(-1j * (qG_Gv @ pos_v))
546 F_aGii.append(x_G[:, np.newaxis, np.newaxis] * Fbar_Gii)
548 return F_aGii