Coverage for gpaw/response/fxc_kernels.py: 99%
85 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1"""Contains methods for calculating local LR-TDDFT kernels."""
3from functools import partial
4from pathlib import Path
6import numpy as np
8from gpaw.response import timer
9from gpaw.response.pw_parallelization import Blocks1D
10from gpaw.response.dyson import PWKernel
11from gpaw.response.localft import (LocalFTCalculator,
12 add_LDA_dens_fxc, add_LSDA_trans_fxc)
15class FXCKernel(PWKernel):
16 r"""Adiabatic local exchange-correlation kernel in a plane-wave basis.
18 In real-space, the adiabatic local xc-kernel matrix is given by:
20 K_xc^μν(r, r', t-t') = f_xc^μν(r) δ(r-r') δ(t-t')
22 where the local xc kernel is given by:
24 ∂^2[ϵ_xc(n,m)n] |
25 f_xc^μν(r) = ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾ |
26 ∂n^μ ∂n^ν |n=n(r),m=m(r)
28 In the plane-wave basis (and frequency domain),
30 1
31 K_xc^μν(G, G') = ‾‾ f_xc^μν(G-G'),
32 V0
34 where V0 is the cell volume.
36 Because the set of unique reciprocal lattice vector differences
37 dG = G-G' is much more compact than the full kernel matrix structure
38 Kxc_GG', we store data internally in the dG representation and only
39 unfold the data into the full kernel matrix when requested.
40 """
42 def __init__(self, fxc_dG, dG_K, GG_shape, volume):
43 """Construct the fxc kernel."""
44 assert np.prod(GG_shape) == len(dG_K), \
45 "The K index should be a flattened (G,G') composite index'"
47 self._fxc_dG = fxc_dG
48 self._dG_K = dG_K
49 self.GG_shape = GG_shape
50 self.volume = volume
52 def get_number_of_plane_waves(self):
53 assert self.GG_shape[0] == self.GG_shape[1]
54 return self.GG_shape[0]
56 def _add_to(self, x_GG):
57 """Add Kxc_GG to input array."""
58 x_GG[:] += self.get_Kxc_GG()
60 def get_Kxc_GG(self):
61 """Unfold the fxc(G-G') kernel into the Kxc_GG' kernel matrix."""
62 # Kxc(G-G') = 1 / V0 * fxc(G-G')
63 Kxc_dG = 1 / self.volume * self._fxc_dG
65 # Unfold Kxc(G-G') to the kernel matrix structure Kxc_GG'
66 Kxc_GG = Kxc_dG[self._dG_K].reshape(self.GG_shape)
68 return Kxc_GG
70 def save(self, path: Path):
71 """Save the fxc kernel in a .npz file."""
72 assert path.suffix == '.npz'
73 with open(str(path), 'wb') as fd:
74 np.savez(fd,
75 fxc_dG=self._fxc_dG,
76 dG_K=self._dG_K,
77 GG_shape=self.GG_shape,
78 volume=self.volume)
80 @staticmethod
81 def from_file(path: Path):
82 """Construct an fxc kernel from a previous calculation."""
83 assert path.suffix == '.npz'
84 npzfile = np.load(path)
86 args = [npzfile[key]
87 for key in ['fxc_dG', 'dG_K', 'GG_shape', 'volume']]
89 return FXCKernel(*args)
92class AdiabaticFXCCalculator:
93 """Calculator for adiabatic local exchange-correlation kernels."""
95 def __init__(self, localft_calc: LocalFTCalculator):
96 """Contruct the fxc calculator based on a local FT calculator."""
97 self.localft_calc = localft_calc
99 self.gs = localft_calc.gs
100 self.context = localft_calc.context
102 @staticmethod
103 def from_rshe_parameters(*args, **kwargs):
104 return AdiabaticFXCCalculator(
105 LocalFTCalculator.from_rshe_parameters(*args, **kwargs))
107 @timer('Calculate XC kernel')
108 def __call__(self, fxc, spincomponent, qpd):
109 """Calculate fxc(G-G'), which defines the kernel matrix Kxc_GG'.
111 The fxc kernel is calculated for all unique dG = G-G' reciprocal
112 lattice vectors and returned as an FXCKernel instance which can unfold
113 itself to produce the full kernel matrix Kxc_GG'.
114 """
115 # Generate a large_qpd to encompass all G-G' in qpd
116 large_ecut = 4 * qpd.ecut # G = 1D grid of |G|^2/2 < ecut
117 large_qpd = qpd.copy_with(ecut=large_ecut,
118 gammacentered=True,
119 gd=self.gs.finegd)
121 # Calculate fxc(Q) on the large plane-wave grid (Q = large grid index)
122 add_fxc = create_add_fxc(fxc, spincomponent)
123 fxc_Q = self.localft_calc(large_qpd, add_fxc)
125 # Create a mapping from the large plane-wave grid to an unfoldable mesh
126 # of all unique dG = G-G' reciprocal lattice vector differences on the
127 # qpd plane-wave representation
128 GG_shape, dG_K, Q_dG = self.create_unfoldable_Q_dG_mapping(
129 qpd, large_qpd)
131 # Map the calculated kernel fxc(Q) onto the unfoldable grid of unique
132 # reciprocal lattice vector differences fxc(dG)
133 fxc_dG = fxc_Q[Q_dG]
135 # Return the calculated kernel as an fxc kernel object
136 fxc_kernel = FXCKernel(fxc_dG, dG_K, GG_shape, qpd.gd.volume)
138 return fxc_kernel
140 @timer('Create unfoldable Q_dG mapping')
141 def create_unfoldable_Q_dG_mapping(self, qpd, large_qpd):
142 """Create mapping from Q index to the kernel matrix indeces GG'.
144 The mapping is split into two parts:
145 * Mapping from the large plane-wave representation index Q (of
146 large_qpd) to an index dG representing all unique reciprocal lattice
147 vector differences (G-G') of the original plane-wave representation
148 qpd
149 * A mapping from the dG index to a flattened K = (G, G') composite
150 kernel matrix index
152 Lastly the kernel matrix shape GG_shape is returned so that an array in
153 index K can easily be reshaped to a kernel matrix Kxc_GG'.
154 """
155 # Calculate all (G-G') reciprocal lattice vectors
156 dG_GGv = calculate_dG_GGv(qpd)
157 GG_shape = dG_GGv.shape[:2]
159 # Reshape to composite K = (G, G') index
160 dG_Kv = dG_GGv.reshape(-1, dG_GGv.shape[-1])
162 # Find unique dG-vectors
163 # We need tight control of the decimals to avoid precision artifacts
164 dG_dGv, dG_K = np.unique(dG_Kv.round(decimals=6),
165 return_inverse=True, axis=0)
167 # Create the mapping from Q-index to dG-index
168 Q_dG = self.create_Q_dG_map(large_qpd, dG_dGv)
170 return GG_shape, dG_K, Q_dG
172 @timer('Create Q_dG map')
173 def create_Q_dG_map(self, large_qpd, dG_dGv):
174 """Create mapping between (G-G') index dG and large_qpd index Q."""
175 G_Qv = large_qpd.get_reciprocal_vectors(add_q=False)
176 # Make sure to match the precision of dG_dGv
177 G_Qv = G_Qv.round(decimals=6)
179 # Distribute dG over world
180 # This is necessary because the next step is to create a K_QdGv buffer
181 # of which the norm is taken. When the number of plane-wave
182 # coefficients is large, this step becomes a memory bottleneck, hence
183 # the distribution.
184 dGblocks = Blocks1D(self.context.comm, dG_dGv.shape[0])
185 dG_mydGv = dG_dGv[dGblocks.myslice]
187 # Determine Q index for each dG index
188 diff_QmydG = np.linalg.norm(G_Qv[:, np.newaxis] - dG_mydGv[np.newaxis],
189 axis=2)
190 Q_mydG = np.argmin(diff_QmydG, axis=0)
192 # Check that all the identified Q indices produce identical reciprocal
193 # lattice vectors
194 assert np.allclose(np.diagonal(diff_QmydG[Q_mydG]), 0.), \
195 'Could not find a perfect matching reciprocal wave vector in '\
196 'large_qpd for all dG_dGv'
198 # Collect the global Q_dG map
199 Q_dG = dGblocks.all_gather(Q_mydG)
201 return Q_dG
204def create_add_fxc(fxc: str, spincomponent: str):
205 """Create an add_fxc function according to the requested functional and
206 spin component."""
207 assert fxc in ['ALDA_x', 'ALDA_X', 'ALDA']
209 if spincomponent in ['00', 'uu', 'dd']:
210 add_fxc = partial(add_LDA_dens_fxc, fxc=fxc)
211 elif spincomponent in ['+-', '-+']:
212 add_fxc = partial(add_LSDA_trans_fxc, fxc=fxc)
213 else:
214 raise ValueError(spincomponent)
216 return add_fxc
219def calculate_dG_GGv(qpd):
220 """Calculate dG_GG' = (G-G') for the plane wave basis in qpd."""
221 nG = qpd.ngmax
222 G_Gv = qpd.get_reciprocal_vectors(add_q=False)
224 dG_GGv = np.zeros((nG, nG, 3))
225 for v in range(3):
226 dG_GGv[:, :, v] = np.subtract.outer(G_Gv[:, v], G_Gv[:, v])
228 return dG_GGv