Coverage for gpaw/response/localft.py: 96%
248 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1"""Functionality to calculate the all-electron Fourier components of local
2functions of the electon (spin-)density."""
4from abc import ABC, abstractmethod
6import numpy as np
7from scipy.special import spherical_jn
9from ase.units import Bohr
11from gpaw.response import ResponseGroundStateAdapter, ResponseContext, timer
13from gpaw.spherical_harmonics import Yarr
14from gpaw.sphere.rshe import calculate_reduced_rshe
15from gpaw.xc import XC
16from gpaw.xc.libxc import LibXC
19class LocalFTCalculator(ABC):
20 r"""Calculator base class for calculators of all-electron plane-wave
21 components to arbitrary real-valued real-space functionals f[n](r) which
22 can be written as closed form functions of the local ground state
23 (spin-)density:
25 f[n](r) = f(n(r)).
27 Since n(r) is lattice periodic, so is f(r) and the plane-wave components
28 can be calculated as (see [PRB 103, 245110 (2021)] for definitions)
30 /
31 f(G) = |dr f(r) e^(-iG.r),
32 /
33 V0
35 where V0 is the unit-cell volume.
36 """
38 def __init__(self, gs, context, bg_density=None):
39 """Constructor for the LocalFTCalculator
41 Parameters
42 ----------
43 gs : ResponseGroundStateAdapter
44 Adapter containing relevant information about the underlying DFT
45 ground state
46 context : ResponseContext
47 bg_density : float
48 Spin-neutral background electron density (in Å^-3) to add to the
49 actual electron density in order to regularize functions which
50 diverge in vacuum.
51 """
52 assert isinstance(gs, ResponseGroundStateAdapter)
53 self.gs = gs
54 assert isinstance(context, ResponseContext)
55 self.context = context
57 if bg_density is None:
58 self.bg_density = None
59 else:
60 assert isinstance(bg_density, float)
61 # Convert to atomic units
62 self.bg_density = bg_density * Bohr**3. # Å^-3 -> Bohr^-3
64 @staticmethod
65 def from_rshe_parameters(gs, context, bg_density=None,
66 rshelmax=-1, rshewmin=None):
67 """Construct the LocalFTCalculator based on parameters for the
68 expansion of the PAW correction in real spherical harmonics
70 Parameters
71 ----------
72 rshelmax : int or None
73 Expand f(r) in real spherical harmonics inside the augmentation
74 spheres. If None, the plane-wave components will be calculated
75 without augmentation. The value of rshelmax indicates the maximum
76 index l to perform the expansion in (l < 6).
77 rshewmin : float or None
78 If None, the PAW correction will be fully expanded up to the chosen
79 lmax. Given as a float (0 < rshewmin < 1), rshewmin indicates what
80 coefficients to use in the expansion. If any (l,m) coefficient
81 contributes with less than a fraction of rshewmin on average, it
82 will not be included.
83 """
84 if rshelmax is None:
85 return LocalGridFTCalculator(gs, context, bg_density=bg_density)
86 else:
87 return LocalPAWFTCalculator(gs, context, bg_density=bg_density,
88 rshelmax=rshelmax, rshewmin=rshewmin)
90 @timer('LocalFT')
91 def __call__(self, qpd, add_f):
92 """Calculate the plane-wave components f(G).
94 Parameters
95 ----------
96 qpd : SingleQPWDescriptor
97 Defines the plane-wave basis to calculate the components in.
98 add_f : method
99 Defines the local function of the electron (spin-)density to
100 Fourier transform. Should take arguments gd (GridDescriptor),
101 n_sR (electron spin-density on the real space grid of gd) and
102 f_R (output array) and add the function f(R) to the output array.
103 Example:
104 >>> def add_total_density(gd, n_sR, f_R):
105 ... f_R += np.sum(n_sR, axis=0)
107 Returns
108 -------
109 f_G : np.array
110 Plane-wave components of the function f, indexes by the reciprocal
111 lattice vectors G.
112 """
113 self.context.print('Calculating f(G)')
114 f_G = self.calculate(qpd, add_f)
115 self.context.print('Finished calculating f(G)')
117 return f_G
119 @abstractmethod
120 def calculate(self, qpd, add_f):
121 pass
123 @staticmethod
124 def equivalent_real_space_grids(gd1, gd2):
125 assert gd1.comm.size == 1
126 assert gd2.comm.size == 1
127 return (gd1.N_c == gd2.N_c).all()
129 def get_electron_density(self, gd, pseudo=False):
130 """Get the electron density corresponding to a given grid descriptor.
131 """
132 gridrefinement = self.get_gridrefinement(gd)
134 if pseudo:
135 _get_electron_density = self.gs.get_pseudo_density
136 else:
137 _get_electron_density = self.gs.get_all_electron_density
139 n_sR, gdref = _get_electron_density(gridrefinement=gridrefinement)
141 assert self.equivalent_real_space_grids(gd, gdref)
143 if self.bg_density is not None:
144 # Add spin-neutral background electron density
145 self.context.print(' Adding a background a background electron '
146 f'density of {self.bg_density / Bohr**3.} Å^-3')
147 n_sR = n_sR.copy() # Make a copy in order not to modify gs
148 n_sR += self.bg_density / n_sR.shape[0]
150 return n_sR
152 def get_gridrefinement(self, gd):
153 if self.equivalent_real_space_grids(gd, self.gs.gd):
154 gridrefinement = 1
155 elif self.equivalent_real_space_grids(gd, self.gs.finegd):
156 gridrefinement = 2
157 else:
158 raise ValueError('The supplied gd is neither compatible with the '
159 'coarse nor the fine real-space grid of the '
160 'underlying ground state')
162 return gridrefinement
165class LocalGridFTCalculator(LocalFTCalculator):
167 def calculate(self, qpd, add_f):
168 """Calculate f(G) directly from the all-electron density on the cubic
169 real-space grid."""
170 n_sR = self.get_all_electron_density(qpd.gd)
171 f_G = self._calculate(qpd, n_sR, add_f)
173 return f_G
175 def _calculate(self, qpd, n_sR, add_f):
176 """In-place calculation of the plane-wave components."""
177 # Calculate f(r)
178 gd = qpd.gd
179 f_R = gd.zeros()
180 add_f(gd, n_sR, f_R)
182 # FFT to reciprocal space
183 f_G = fft_from_grid(f_R, qpd) # G = 1D grid of |G|^2/2 < ecut
185 return f_G
187 @timer('Calculate the all-electron density')
188 def get_all_electron_density(self, gd):
189 """Calculate the all-electron (spin-)density."""
190 self.context.print(' Calculating the all-electron density')
191 return self.get_electron_density(gd)
194class LocalPAWFTCalculator(LocalFTCalculator):
196 def __init__(self, gs, context, bg_density=None,
197 rshelmax=-1, rshewmin=None):
198 super().__init__(gs, context, bg_density=bg_density)
200 self.engine = LocalPAWFTEngine(self.context, rshelmax, rshewmin)
202 def calculate(self, qpd, add_f):
203 """Calculate f(G) with an expansion of f(r) in real spherical harmonics
204 inside the augmentation spheres."""
205 # Retrieve the pseudo (spin-)density on the real-space grid
206 nt_sR = self.get_pseudo_density(qpd.gd) # R = 3D real-space grid
208 # Retrieve the pseudo and all-electron atomic centered densities inside
209 # the augmentation spheres
210 R_av, micro_setups = self.extract_atom_centered_quantities()
212 # Let the engine perform the in-place calculation
213 f_G = self.engine.calculate(qpd, nt_sR, R_av, micro_setups, add_f)
215 return f_G
217 def get_pseudo_density(self, gd):
218 """Get the pseudo (spin-)density of the ground state."""
219 return self.get_electron_density(gd, pseudo=True)
221 def extract_atom_centered_quantities(self):
222 """Extract all relevant atom centered quantities that the engine needs
223 in order to calculate PAW corrections. Most of the information is
224 bundled as a list of MicroSetups for each atom."""
225 R_av = self.gs.atoms.positions / Bohr
226 micro_setups = self.gs.micro_setups
227 return R_av, micro_setups
230class MicroSetup:
232 def __init__(self, rgd, Y_nL, n_sLg, nt_sLg):
233 self.rgd = rgd
234 self.Y_nL = Y_nL
235 self.n_sLg = n_sLg
236 self.nt_sLg = nt_sLg
238 def evaluate_function(self, add_f):
239 """Evaluate a given function f(r) on the angular and radial grids."""
240 f_ng = np.array([self.rgd.zeros() for n in range(self.Y_nL.shape[0])])
241 for n, Y_L in enumerate(self.Y_nL):
242 n_sg = Y_L @ self.n_sLg
243 add_f(self.rgd, n_sg, f_ng[n])
244 return f_ng
246 def evaluate_paw_correction(self, add_f):
247 r"""Evaluate Δf_a[n_a,ñ_a](r) for a given function f(r).
249 Returns
250 -------
251 df_ng : nd.array
252 (f_ng - ft_ng) where (n=Lebedev index, g=radial grid index)
253 """
254 rgd = self.rgd
255 f_g = rgd.zeros()
256 ft_g = rgd.zeros()
257 df_ng = np.array([rgd.zeros() for n in range(self.Y_nL.shape[0])])
258 for n, Y_L in enumerate(self.Y_nL):
259 f_g[:] = 0.
260 n_sg = Y_L @ self.n_sLg
261 add_f(rgd, n_sg, f_g)
263 ft_g[:] = 0.
264 nt_sg = Y_L @ self.nt_sLg
265 add_f(rgd, nt_sg, ft_g)
267 df_ng[n, :] = f_g - ft_g
269 return df_ng
271 def expand_function(self, add_f, **kwargs):
272 f_ng = self.evaluate_function(add_f)
273 return self.expand(f_ng, **kwargs)
275 def expand_paw_correction(self, add_f, **kwargs):
276 df_ng = self.evaluate_paw_correction(add_f)
277 return self.expand(df_ng, **kwargs)
279 def expand(self, f_ng, **kwargs):
280 """Expand into real spherical harmonics."""
281 return calculate_reduced_rshe(self.rgd, f_ng, self.Y_nL, **kwargs)
284def extract_micro_setup(pawdata, D_sp) -> MicroSetup:
285 """Extract the a.e. and pseudo (spin-)densities as a MicroSetup."""
286 # Radial grid descriptor:
287 rgd = pawdata.xc_correction.rgd
288 # Spherical harmonics on the Lebedev quadrature:
289 Y_nL = pawdata.xc_correction.Y_nL
290 n_sLg, nt_sLg = calculate_atom_centered_densities(pawdata,
291 # atomic density matrix
292 D_sp)
293 return MicroSetup(rgd, Y_nL, n_sLg, nt_sLg)
296def calculate_atom_centered_densities(pawdata, D_sp):
297 """Calculate the AE and pseudo densities inside the augmentation sphere.
299 Returns
300 -------
301 n_sLg : nd.array
302 all-electron density
303 nt_sLg : nd.array
304 pseudo density
305 (s=spin, L=(l,m) spherical harmonic index, g=radial grid index)
306 """
307 n_qg = pawdata.xc_correction.n_qg
308 nt_qg = pawdata.xc_correction.nt_qg
309 nc_g = pawdata.xc_correction.nc_g
310 nct_g = pawdata.xc_correction.nct_g
312 B_pqL = pawdata.xc_correction.B_pqL
313 D_sLq = np.inner(D_sp, B_pqL.T)
314 nspins = len(D_sp)
316 n_sLg = D_sLq @ n_qg
317 nt_sLg = D_sLq @ nt_qg
319 # Add core density
320 n_sLg[:, 0] += np.sqrt(4. * np.pi) / nspins * nc_g
321 nt_sLg[:, 0] += np.sqrt(4. * np.pi) / nspins * nct_g
323 return n_sLg, nt_sLg
326class LocalPAWFTEngine:
328 def __init__(self, context, rshelmax=-1, rshewmin=None):
329 """Construct the engine."""
330 self.context = context
331 self.rshelmax = rshelmax
332 self.rshewmin = rshewmin
334 self._add_f = None
336 def calculate(self, qpd, nt_sR, R_av, micro_setups, add_f):
337 r"""Calculate the Fourier transform f(G) by splitting up the
338 calculation into a pseudo density contribution and a PAW correction
339 accounting for the difference
341 Δf[n,ñ](r) = f(n(r)) - f(ñ(r)),
343 such that:
345 f(G) = f[ñ](G) + Δf[n,ñ](G).
347 See [PRB 103, 245110 (2021)] for definitions and notation details."""
348 self._add_f = add_f
350 ft_G = self.calculate_pseudo_contribution(qpd, nt_sR)
351 fPAW_G = self.calculate_paw_corrections(qpd, R_av, micro_setups)
353 return ft_G + fPAW_G
355 def calculate_pseudo_contribution(self, qpd, nt_sR):
356 """Calculate the pseudo density contribution by performing a FFT of
357 f(ñ(r)) on the cubic real-space grid.
359 NB: This operation assumes that the function f is a slowly varrying
360 function of the pseudo density ñ(r) everywhere in space, such that
361 f(ñ(r)) is accurately described on the cubic real-space grid."""
362 # Calculate ft(r) (t=tilde=pseudo)
363 gd = qpd.gd
364 ft_R = gd.zeros()
365 self._add_f(gd, nt_sR, ft_R)
367 # FFT to reciprocal space
368 ft_G = fft_from_grid(ft_R, qpd) # G = 1D grid of |G|^2/2 < ecut
370 return ft_G
372 @timer('Calculate PAW corrections')
373 def calculate_paw_corrections(self, qpd, R_av, micro_setups):
374 r"""Calculate the PAW corrections to f(G), for each augmentation sphere
375 at a time:
376 __
377 \ /
378 Δf[n,ñ](G) = / |dr Δf_a[n_a,ñ_a](r - R_a) e^(-iG.r)
379 ‾‾ /
380 a V0
382 where Δf_a is the atom centered difference between the all electron
383 and pseudo quantities inside augmentation sphere a:
385 Δf_a[n_a,ñ_a](r) = f(n_a(r)) - f(ñ_a(r)).
386 """
387 self.context.print(' Calculating PAW corrections\n')
389 # Extract reciprocal lattice vectors
390 nG = qpd.ngmax
391 G_Gv = qpd.get_reciprocal_vectors(add_q=False)
392 assert G_Gv.shape[0] == nG
394 # Allocate output array
395 fPAW_G = np.zeros(nG, dtype=complex)
397 # Distribute plane waves
398 G_myG = self._distribute_correction(nG)
399 G_myGv = G_Gv[G_myG]
401 # Calculate and add the PAW corrections from each augmentation sphere
402 for a, (R_v, micro_setup) in enumerate(zip(R_av, micro_setups)):
403 self._add_paw_correction(a, R_v, micro_setup,
404 G_myG, G_myGv, fPAW_G)
406 self.context.comm.sum(fPAW_G)
408 return fPAW_G
410 def _distribute_correction(self, nG):
411 comm = self.context.comm
412 nGpr = (nG + comm.size - 1) // comm.size
413 Ga = min(comm.rank * nGpr, nG)
414 Gb = min(Ga + nGpr, nG)
416 return range(Ga, Gb)
418 def _add_paw_correction(self, a, R_v, micro_setup, G_myG, G_myGv, fPAW_G):
419 r"""Calculate the PAW correction of augmentation sphere a,
421 /
422 Δf_a(G) = e^(-iG.R_a) |dr Δf_a[n_a,ñ_a](r) e^(-iG.r),
423 /
424 V0
426 by expanding both the atom centered correction and the plane wave in
427 real spherical harmonics, see [PRB 103, 245110 (2021)]:
429 l a
430 __ __ R_c
431 \ \ l ^ / a
432 Δf_a(G) = e^(-iG.R_a) / / (-i) Y (G) |4πr^2 dr j(|G|r) Δf (r)
433 ‾‾ ‾‾ lm / l lm
434 l m=-l 0
436 The calculated atomic correction is then added to the output array."""
437 rgd = micro_setup.rgd
439 # Expand Δf_a[n_a,ñ_a](r) in real spherical harmonics
440 rshe, info_string = micro_setup.expand_paw_correction(
441 self._add_f, lmax=self.rshelmax, wmin=self.rshewmin)
442 self.print_rshe_info(a, info_string)
444 # Expand the plane waves in real spherical harmonics (and spherical
445 # Bessel functions)
446 (ii_MmyG,
447 j_gMmyG,
448 Y_MmyG) = self._expand_plane_waves(
449 G_myGv, rgd.r_g, rshe.L_M, rshe.l_M)
451 # Calculate the PAW correction as an integral over the radial grid
452 # and rshe coefficients
453 with self.context.timer('Integrate PAW correction'):
454 angular_coef_MmyG = ii_MmyG * Y_MmyG
455 # Radial integral, dv = 4πr^2
456 df_gM = rshe.f_gM
457 radial_coef_MmyG = np.tensordot(j_gMmyG * df_gM[..., np.newaxis],
458 rgd.dv_g, axes=([0, 0]))
459 # Angular integral (sum over l,m)
460 atomic_corr_myG = np.sum(angular_coef_MmyG * radial_coef_MmyG,
461 axis=0)
463 position_prefactor_myG = np.exp(-1j * np.inner(G_myGv, R_v))
465 # Add to output array
466 fPAW_G[G_myG] += position_prefactor_myG * atomic_corr_myG
468 def print_rshe_info(self, a, info_string):
469 """Print information about the expansion at atom a."""
470 info_string = f' RSHE of atom {a}\n' + info_string
471 self.context.print(
472 info_string.replace('\n', '\n ') + '\n')
474 @timer('Expand plane waves in real spherical harmonics')
475 def _expand_plane_waves(self, G_myGv, r_g, L_M, l_M):
476 r"""Expand plane waves in spherical Bessel functions and real spherical
477 harmonics:
478 l
479 __ __
480 -iG.r \ \ l ^ ^
481 e = 4π / / (-i) j (|G|r) Y (G) Y (r)
482 ‾‾ ‾‾ l lm lm
483 l m=-l
485 Returns
486 -------
487 ii_MmyG : nd.array
488 (-i)^l for used (l,m) coefficients M
489 j_gMmyG : nd.array
490 j_l(|G|r) for used (l,m) coefficients M
491 Y_MmyG : nd.array
492 ^
493 Y_lm(K) for used (l,m) coefficients M
494 """
495 nmyG = G_myGv.shape[0]
496 Gnorm_myG, Gdir_myGv = self._calculate_norm_and_direction(G_myGv)
498 # Setup arrays to fully vectorize computations
499 nM = len(L_M)
500 (r_gMmyG, l_gMmyG,
501 Gnorm_gMmyG) = (a.reshape(len(r_g), nM, nmyG)
502 for a in np.meshgrid(r_g, l_M, Gnorm_myG,
503 indexing='ij'))
505 with self.context.timer('Compute spherical bessel functions'):
506 # Slow step
507 j_gMmyG = spherical_jn(l_gMmyG, Gnorm_gMmyG * r_gMmyG)
509 Y_MmyG = Yarr(L_M, Gdir_myGv)
510 ii_MmyG = (-1j) ** np.repeat(l_M, nmyG).reshape((nM, nmyG))
512 return ii_MmyG, j_gMmyG, Y_MmyG
514 @staticmethod
515 def _calculate_norm_and_direction(G_myGv):
516 """Calculate the length and direction of reciprocal lattice vectors."""
517 Gnorm_myG = np.linalg.norm(G_myGv, axis=1)
518 Gdir_myGv = np.zeros_like(G_myGv)
519 mask0 = np.where(Gnorm_myG != 0.)
520 Gdir_myGv[mask0] = G_myGv[mask0] / Gnorm_myG[mask0][:, np.newaxis]
522 return Gnorm_myG, Gdir_myGv
525def fft_from_grid(f_R, qpd):
526 r"""Perform a FFT to reciprocal space:
527 __
528 / V0 \
529 f(G) = |dr f(r) e^(-iG.r) ≃ ‾‾ / f(r) e^(-iG.r)
530 / N ‾‾
531 V0 r
533 where N is the number of grid points."""
534 Q_G = qpd.Q_qG[0]
536 # Perform the FFT
537 N = np.prod(qpd.gd.N_c)
538 f_Q123 = qpd.gd.volume / N * np.fft.fftn(f_R) # Q123 = 3D grid in Q-rep
540 # Change the view of the plane-wave components from the 3D grid in the
541 # Q-representation that numpy spits out to the 1D grid in the
542 # G-representation, that GPAW relies on internally
543 f_G = f_Q123.ravel()[Q_G]
545 return f_G
548# ---------- Local functions of the (spin-)density ---------- #
551def add_total_density(gd, n_sR, n_R):
552 n_R += np.sum(n_sR, axis=0)
555def add_spin_polarization(gd, n_sR, nz_R):
556 nz_R += calculate_spin_polarization(n_sR)
559def calculate_spin_polarization(n_sR):
560 return n_sR[0] - n_sR[1]
563def add_LSDA_Wxc(gd, n_sR, Wxc_R):
564 Wxc_R += calculate_LSDA_Wxc(gd, n_sR)
567def calculate_LSDA_Wxc(gd, n_sR):
568 """Calculate W_xc^z in the local spin-density approximation.
570 For a collinear system:
572 δE_xc[n,n^z] 1
573 W_xc^z(r) = ‾‾‾‾‾‾‾‾‾‾‾‾ = ‾ [V_LSDA^↑(r) - V_LSDA^↓(r)]
574 δn^z(r) 2
575 """
576 # Allocate an array for the spin-dependent xc potential on the real
577 # space grid
578 v_sR = np.zeros(np.shape(n_sR))
580 # Calculate the spin-dependent potential
581 xc = XC('LDA')
582 xc.calculate(gd, n_sR, v_sg=v_sR)
584 return (v_sR[0] - v_sR[1]) / 2
587def add_LSDA_zeeman_energy(gd, n_sR, EZ_R):
588 """Calculate and add the LSDA Zeeman energy to the output array.
590 The Zeeman energy is defined as:
592 E_Z(r) = - B^(xc)(r) m(r) = - W_xc^z n^z(r).
593 """
594 EZ_R += - calculate_LSDA_Wxc(gd, n_sR) * calculate_spin_polarization(n_sR)
597def add_LDA_dens_fxc(gd, n_sR, fxc_R, *, fxc):
598 r"""Calculate the LDA density kernel and add it to the output array fxc_R.
600 The LDA density kernel is given by:
602 ∂^2[ϵ_xc(n,m)n] |
603 f_LDA^(00)(r) = ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾ |
604 ∂n^2 |n=n(r),m=m(r)
605 """
606 assert len(n_sR) == 1, \
607 'The density kernel is untested for spin-polarized systems'
609 if fxc == 'ALDA_x':
610 fxc_R += -1. / 3. * (3. / np.pi)**(1. / 3.) * n_sR[0]**(-2. / 3.)
611 else:
612 assert fxc in ['ALDA_X', 'ALDA']
613 kernel = LibXC(fxc[1:])
614 fxc_sR = np.zeros_like(n_sR)
615 kernel.xc.calculate_fxc_spinpaired(n_sR.ravel(), fxc_sR)
617 fxc_R += fxc_sR[0]
620def add_LSDA_trans_fxc(gd, n_sR, fxc_R, *, fxc):
621 r"""Calculate the transverse LDA kernel and add it to the output arr. fxc_R
623 The transverse LDA kernel is given by:
625 2 ∂[ϵ_xc(n,n^z)n] |
626 f_LDA^(+-)(r) = ‾‾‾ ‾‾‾‾‾‾‾‾‾‾‾‾‾‾ |
627 n^z ∂n^z |n=n(r),n^z=n^z(r)
629 V_LSDA^↑(r) - V_LSDA^↓(r)
630 = ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾
631 n^z(r)
632 """
633 assert len(n_sR) == 2 # nspins
634 nz_R = n_sR[0] - n_sR[1]
636 if fxc == 'ALDA_x':
637 fxc_R += - (6. / np.pi)**(1. / 3.) \
638 * (n_sR[0]**(1. / 3.) - n_sR[1]**(1. / 3.)) / nz_R
639 else:
640 assert fxc in ['ALDA_X', 'ALDA']
641 v_sR = np.zeros(np.shape(n_sR))
642 xc = XC(fxc[1:])
643 xc.calculate(gd, n_sR, v_sg=v_sR)
645 fxc_R += (v_sR[0] - v_sR[1]) / nz_R