Coverage for gpaw/response/pair_integrator.py: 99%
114 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1import numpy as np
2from abc import ABC, abstractmethod
4from gpaw.utilities.progressbar import ProgressBar
5from gpaw.typing import Vector
7from gpaw.response import ResponseGroundStateAdapter, ResponseContext, timer
8from gpaw.response.frequencies import ComplexFrequencyDescriptor
9from gpaw.response.symmetry import QSymmetryAnalyzer
10from gpaw.response.kspair import (KohnShamKPointPair,
11 KohnShamKPointPairExtractor)
12from gpaw.response.pw_parallelization import block_partition
13from gpaw.response.pair_transitions import PairTransitions
16class PairFunction(ABC):
17 r"""Pair function data object.
19 In the GPAW response module, a pair function is understood as any function
20 which can be written as a sum over the eigenstate transitions with a given
21 crystal momentum difference q,
22 __
23 \
24 pf(q) = / pf_αα' δ_{q,q_{α',α}}
25 ‾‾
26 α,α'
28 where pf_αα' is an arbitrary matrix element.
29 """
31 def __init__(self, q_c: Vector):
32 self.q_c = np.asarray(q_c)
33 self.array = self.zeros()
35 @abstractmethod
36 def zeros(self):
37 """Generate an array of zeros, representing the pair function."""
40class DynamicPairFunction(PairFunction):
41 r"""Dynamic pair function data object.
43 A dynamic pair function is not only a function of the crystal momentum q,
44 but also of the complex frequency z = ω + iη:
45 __
46 \
47 pf(q,z) = / pf_αα'(z) δ_{q,q_{α',α}}
48 ‾‾
49 α,α'
51 Typically, this will be some generalized (linear) susceptibility, which is
52 defined by the Kubo formula,
54 i ˰ ˰
55 χ_BA(t-t') = - ‾ θ(t-t') <[B_0(t-t'), A]>_0
56 ħ
58 and can be written in its Lehmann representation as a function of frequency
59 in the upper half complex frequency plane,
61 __ ˰ ˰
62 \ <α|B|α'><α'|A|α>
63 χ_BA(z) = / ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾ (n_α - n_α')
64 ‾‾ ħz - (E_α' - E_α)
65 α,α'
67 where E_α and n_α are the eigenstate energies and occupations respectively.
69 For more information, please refer to [Skovhus T., PhD. Thesis, 2021]."""
71 def __init__(self, q_c: Vector, zd: ComplexFrequencyDescriptor):
72 self.zd = zd
73 super().__init__(q_c)
76class PairFunctionIntegrator(ABC):
77 r"""Baseclass for computing pair functions in the Kohn-Sham system.
79 The implementation is currently restricted collinear periodic crystals in
80 absence of spin-orbit coupling.
82 In the Kohn-Sham system, pair functions can be constructed
83 straight-forwardly as a sum over transitions between Kohn-Sham eigenstates
84 at k and k + q,
85 __ __ __ __
86 1 \ \ \ 1 \
87 pf(q) = ‾ / / / pf_nks,n'k+qs'(q) = ‾ / pf_T(q)
88 V ‾‾ ‾‾ ‾‾ V ‾‾
89 k n,n' s,s' T
91 where V is the crystal volume and T is a composite index encoding all
92 relevant transitions:
94 T: (n, k, s) -> (n', k + q, s')
96 The sum over transitions can be split into two steps: (1) an integral over
97 k-points k inside the 1st Brillouin Zone and (2) a sum over band and spin
98 transitions t:
100 t (composite transition index): (n, s) -> (n', s')
101 __ __ __ __
102 1 \ 1 \ \ 1 \
103 pf(q) = ‾ / pf_T(q) = ‾ / / pf_kt(q) = ‾ / (...)_k
104 V ‾‾ V ‾‾ ‾‾ V ‾‾
105 T k t k
107 In the code, the k-point integral is handled by a KPointPairIntegral
108 object, while the sum over band and spin transitions t is carried out in
109 the self.add_integrand() method, which also defines the specific pair
110 function in question.
112 KPointPairIntegral:
113 __
114 1 \
115 ‾ / (...)_k
116 V ‾‾
117 k
119 self.add_integrand():
120 __ __ __
121 \ \ \
122 (...)_k = / pf_kt(q) = / / pf_nks,n'k+qs'(q)
123 ‾‾ ‾‾ ‾‾
124 t n,n' s,s'
126 In practise, the integration is carried out by letting the
127 KPointPairIntegral extract individual KohnShamKPointPair objects, which
128 contain all relevant information about the Kohn-Sham eigenstates at k and
129 k + q for a number of specified spin and band transitions t.
130 KPointPairIntegral.weighted_kpoint_pairs() generates these kptpairs along
131 with their integral weights such that self._integrate() can construct the
132 pair functions in a flexible, yet general manner.
133 """
135 def __init__(self, gs, context, qsymmetry=True, nblocks=1):
136 """Construct the PairFunctionIntegrator
138 Parameters
139 ----------
140 gs : ResponseGroundStateAdaptable
141 context : ResponseContextInput
142 qsymmetry : QSymmetryInput
143 nblocks : int
144 Distribute the pair function into nblocks. Useful when the pair
145 function itself becomes a large array (read: memory limiting).
146 """
147 self.gs = ResponseGroundStateAdapter.from_input(gs)
148 self.context = ResponseContext.from_input(context)
149 self.qsymmetry = QSymmetryAnalyzer.from_input(qsymmetry)
151 # Communicators for distribution of memory and work
152 (self.blockcomm,
153 self.intrablockcomm) = self.create_communicators(nblocks)
154 self.nblocks = self.blockcomm.size
156 # The KohnShamKPointPairExtractor class handles extraction of k-point
157 # pairs from the ground state
158 self.kptpair_extractor = KohnShamKPointPairExtractor(
159 self.gs, self.context,
160 # Distribution of work:
161 # t-transitions are distributed through blockcomm,
162 # k-points through intrablockcomm.
163 transitions_blockcomm=self.blockcomm,
164 kpts_blockcomm=self.intrablockcomm)
166 @timer('Integrate pair function')
167 def _integrate(self, out: PairFunction, transitions: PairTransitions):
168 """In-place pair function integration
170 Parameters
171 ----------
172 out : PairFunction
173 Output data structure
174 transitions : PairTransitions
175 Band and spin transitions to integrate.
177 Returns
178 -------
179 symmetries : QSymmetries
180 """
181 symmetries, generator = self.qsymmetry.analyze(
182 out.q_c, self.gs.kpoints, self.context)
184 # Perform the actual integral as a point integral over k-point pairs
185 integral = KPointPairPointIntegral(self.kptpair_extractor, generator)
186 weighted_kptpairs = integral.weighted_kpoint_pairs(transitions)
187 pb = ProgressBar(self.context.fd) # pb with a generator is awkward
188 for _, _ in pb.enumerate([None] * integral.ni):
189 kptpair, weight = next(weighted_kptpairs)
190 if weight is not None:
191 assert kptpair is not None
192 self.add_integrand(kptpair, weight, out)
194 # Sum over the k-points, which have been distributed between processes
195 with self.context.timer('Sum over distributed k-points'):
196 self.intrablockcomm.sum(out.array)
198 return symmetries
200 @abstractmethod
201 def add_integrand(self, kptpair: KohnShamKPointPair, weight,
202 out: PairFunction):
203 """Add the relevant integrand of the outer k-point integral to the
204 output data structure 'out', weighted by 'weight' and constructed
205 from the provided KohnShamKPointPair 'kptpair'.
207 This method effectively defines the pair function in question.
208 """
210 def create_communicators(self, nblocks):
211 """Create MPI communicators to distribute the memory needed to store
212 large arrays and parallelize calculations when possible.
214 Parameters
215 ----------
216 nblocks : int
217 Separate large arrays into n different blocks. Each process
218 allocates memory for the large arrays. By allocating only a
219 fraction/block of the total arrays, the memory requirements are
220 eased.
222 Returns
223 -------
224 blockcomm : gpaw.mpi.Communicator
225 Communicate between processes belonging to different memory blocks.
226 In every communicator, there is one process for each block of
227 memory, so that all blocks are represented.
228 If nblocks < comm.size, there will be size // nblocks different
229 processes that allocate memory for the same block of the large
230 arrays. Thus, there will be also size // nblocks different block
231 communicators, grouping the processes into sets that allocate the
232 entire arrays between them.
233 intrablockcomm : gpaw.mpi.Communicator
234 Communicate between processes belonging to the same memory block.
235 There will be size // nblocks processes per memory block.
236 """
237 comm = self.context.comm
238 blockcomm, intrablockcomm = block_partition(comm, nblocks)
240 return blockcomm, intrablockcomm
242 def get_band_and_spin_transitions(self, spincomponent, nbands=None,
243 bandsummation='pairwise'):
244 """Get band and spin transitions (n, s) -> (n', s') to integrate."""
245 # Defaults to inclusion of all bands in the ground state calculator
246 if nbands is None:
247 nbands = self.gs.nbands
248 assert nbands <= self.gs.nbands
249 return PairTransitions.from_transitions_domain_arguments(
250 spincomponent, nbands, self.gs.nocc1, self.gs.nocc2,
251 self.gs.nspins, bandsummation)
253 def get_basic_info_string(self):
254 """Get basic information about the ground state and parallelization."""
255 nspins = self.gs.nspins
256 nbands, nocc1, nocc2 = self.gs.nbands, self.gs.nocc1, self.gs.nocc2
257 nk = self.gs.kd.nbzkpts
258 nik = self.gs.kd.nibzkpts
260 csize = self.context.comm.size
261 knsize = self.intrablockcomm.size
262 bsize = self.blockcomm.size
264 isl = ['',
265 'The pair function integration is based on a ground state '
266 'with:',
267 f' Number of spins: {nspins}',
268 f' Number of bands: {nbands}',
269 f' Number of completely occupied bands: {nocc1}',
270 f' Number of partially occupied bands: {nocc2}',
271 f' Number of kpoints: {nk}',
272 f' Number of irreducible kpoints: {nik}',
273 '',
274 'The pair function integration is performed in parallel with:',
275 f' comm.size: {csize}',
276 f' intrablockcomm.size: {knsize}',
277 f' blockcomm.size: {bsize}']
279 return '\n'.join(isl)
281 @staticmethod
282 def get_band_and_transitions_info_string(nbands, nt):
283 isl = [] # info string list
284 if nbands is None:
285 isl.append(' Bands included: All')
286 else:
287 isl.append(f' Number of bands included: {nbands}')
288 isl.append('Resulting in:')
289 isl.append(f' A total number of band and spin transitions of: {nt}')
290 return '\n'.join(isl)
293class KPointPairIntegral(ABC):
294 r"""Baseclass for reciprocal space integrals of the first Brillouin Zone,
295 where the integrand is a sum over transitions between any number of states
296 at the wave vectors k and k + q (referred to as a k-point pair).
298 Definition (V is the total crystal hypervolume and D is the dimension of
299 the crystal):
300 __
301 1 \ 1 /
302 ‾ / (...)_k = ‾‾‾‾‾‾ |dk (...)_k
303 V ‾‾ (2π)^D /
304 k
306 NB: In the current implementation, the dimension is fixed to 3. This is
307 sensible for pair functions which are functions of position (such as the
308 four-component Kohn-Sham susceptibility tensor), and in most circumstances
309 a change in dimensionality can be accomplished simply by adding an extra
310 prefactor to the integral elsewhere.
311 """
313 def __init__(self, kptpair_extractor, generator):
314 """Construct a KPointPairIntegral corresponding to a given q-point.
316 Parameters
317 ----------
318 kptpair_extractor : KohnShamKPointPairExtractor
319 Object responsible for extracting all relevant information about
320 the k-point pairs from the underlying ground state calculation.
321 generator : KPointDomainGenerator
322 Object responsible for generating the k-point integration domain
323 based on the symmetries of the q-point in question.
324 """
325 self.gs = kptpair_extractor.gs
326 self.kptpair_extractor = kptpair_extractor
327 self.q_c = generator.symmetries.q_c
329 # Prepare the k-point pair integral
330 bzk_kc, weight_k = self.get_kpoint_domain(generator)
331 bzk_ipc, weight_i = self.slice_kpoint_domain(bzk_kc, weight_k)
332 self._domain = (bzk_ipc, weight_i)
333 self.ni = len(weight_i)
335 def weighted_kpoint_pairs(self, transitions):
336 r"""Generate all k-point pairs in the integral along with their
337 integral weights.
339 The reciprocal space integral is estimated as the sum over a discrete
340 k-point domain. The domain will genererally depend on the integration
341 method as well as the symmetry of the crystal.
343 Definition:
344 __
345 1 / ~ 1 \ (2π)^D
346 ‾‾‾‾‾‾ |dk (...)_k = ‾‾‾‾‾‾ / ‾‾‾‾‾‾ w_kr (...)_kr
347 (2π)^D / (2π)^D ‾‾ Nk V0
348 kr
349 __
350 ~ \
351 = / iw_kr (...)_kr
352 ‾‾
353 kr
355 The sum over kr denotes the reduced k-point domain specified by the
356 integration method (a reduced selection of Nkr points from the ground
357 state k-point grid of Nk total points in the entire 1BZ). Each point
358 is weighted by its k-point volume in the ground state k-point grid
360 (2π)^D
361 kpointvol = ‾‾‾‾‾‾,
362 Nk V0
364 and an additional individual k-point weight w_kr specific to the
365 integration method (V0 denotes the cell volume). Together with the
366 integral prefactor, these make up the integral weight
368 1
369 iw_kr = ‾‾‾‾‾‾ kpointvol w_kr
370 (2π)^D
372 Parameters
373 ----------
374 transitions : PairTransitions
375 Band and spin transitions to integrate.
376 """
377 # Calculate prefactors
378 outer_prefactor = 1 / (2 * np.pi)**3
379 V = self.crystal_volume() # V = Nk * V0
380 kpointvol = (2 * np.pi)**3 / V
381 prefactor = outer_prefactor * kpointvol
383 # Generate k-point pairs
384 for k_pc, weight in zip(*self._domain):
385 if weight is None:
386 integral_weight = None
387 else:
388 integral_weight = prefactor * weight
389 kptpair = self.kptpair_extractor.get_kpoint_pairs(
390 k_pc, k_pc + self.q_c, transitions)
391 yield kptpair, integral_weight
393 @abstractmethod
394 def get_kpoint_domain(self, generator):
395 """Use the KPointDomainGenerator to define and weight the k-integral.
397 Returns
398 -------
399 bzk_kc : np.array
400 k-points to integrate in relative coordinates.
401 weight_k : np.array
402 Integral weight of each k-point in the integral.
403 """
405 def slice_kpoint_domain(self, bzk_kc, weight_k):
406 """When integrating over k-points, slice the domain in pieces with one
407 k-point per process each.
409 Returns
410 -------
411 bzk_ipc : nd.array
412 k-points (relative) coordinates for each process for each iteration
413 """
414 comm = self.kptpair_extractor.kpts_blockcomm
415 rank, size = comm.rank, comm.size
417 nk = bzk_kc.shape[0]
418 ni = (nk + size - 1) // size
419 bzk_ipc = [bzk_kc[i * size:(i + 1) * size] for i in range(ni)]
421 # Extract the weight corresponding to the process' own k-point pair
422 weight_ip = [weight_k[i * size:(i + 1) * size] for i in range(ni)]
423 weight_i = [None] * len(weight_ip)
424 for i, w_p in enumerate(weight_ip):
425 if rank in range(len(w_p)):
426 weight_i[i] = w_p[rank]
428 return bzk_ipc, weight_i
430 def crystal_volume(self):
431 """Calculate the total crystal volume, V = Nk * V0, corresponding to
432 the ground state k-point grid."""
433 return self.gs.kd.nbzkpts * self.gs.volume
436class KPointPairPointIntegral(KPointPairIntegral):
437 r"""Reciprocal space integral of k-point pairs in the first Brillouin Zone,
438 estimated as a point integral over all k-points of the ground state k-point
439 grid.
441 Definition:
442 __
443 1 / ~ 1 \ (2π)^D
444 ‾‾‾‾‾‾ |dk (...)_k = ‾‾‾‾‾‾ / ‾‾‾‾‾‾ (...)_k
445 (2π)^D / (2π)^D ‾‾ Nk V0
446 k
448 """
450 def get_kpoint_domain(self, generator):
451 # Generate k-point domain in relative coordinates
452 K_gK = generator.group_kpoints()
453 bzk_kc = np.array([self.gs.kd.bzk_kc[K_K[0]] for
454 K_K in K_gK])
455 # Generate k-point weights
456 weight_k = np.array([generator.get_kpoint_weight(k_c)
457 for k_c in bzk_kc])
458 return bzk_kc, weight_k