Coverage for gpaw/response/susceptibility.py: 99%
227 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1from __future__ import annotations
3from pathlib import Path
5import numpy as np
7from ase.units import Hartree
9from gpaw.response.frequencies import ComplexFrequencyDescriptor
10from gpaw.response.pw_parallelization import Blocks1D
11from gpaw.response.pair_functions import Chi, get_pw_coordinates
12from gpaw.response.qpd import SingleQPWDescriptor
13from gpaw.response.chiks import ChiKSCalculator
14from gpaw.response.coulomb_kernels import NewCoulombKernel
15from gpaw.response.fxc_kernels import FXCKernel, AdiabaticFXCCalculator
16from gpaw.response.dyson import (DysonSolver, HXCKernel, HXCScaling, PWKernel,
17 NoKernel)
20class ChiFactory:
21 r"""User interface to calculate individual elements of the four-component
22 susceptibility tensor χ^μν, see [PRB 103, 245110 (2021)]."""
24 def __init__(self,
25 chiks_calc: ChiKSCalculator,
26 fxc_calculator: AdiabaticFXCCalculator | None = None):
27 """Contruct a many-body susceptibility factory."""
28 self.chiks_calc = chiks_calc
29 self.gs = chiks_calc.gs
30 self.context = chiks_calc.context
31 self.dyson_solver = DysonSolver(self.context)
33 # If no fxc_calculator is supplied, fall back to default
34 if fxc_calculator is None:
35 fxc_calculator = AdiabaticFXCCalculator.from_rshe_parameters(
36 self.gs, self.context)
37 else:
38 assert fxc_calculator.gs is chiks_calc.gs
39 assert fxc_calculator.context is chiks_calc.context
40 self.fxc_calculator = fxc_calculator
42 # Prepare a buffer for the fxc kernels
43 self.fxc_kernel_cache: dict[str, FXCKernel] = {}
45 def __call__(self, spincomponent, q_c, complex_frequencies,
46 fxc: str | None = None, hxc_scaling: HXCScaling | None = None,
47 txt=None) -> tuple[Chi, Chi]:
48 r"""Calculate a given element (spincomponent) of the four-component
49 Kohn-Sham susceptibility tensor and construct a corresponding many-body
50 susceptibility object within a given approximation to the
51 exchange-correlation kernel.
53 Parameters
54 ----------
55 spincomponent : str
56 Spin component (μν) of the susceptibility.
57 Currently, '00', 'uu', 'dd', '+-' and '-+' are implemented.
58 q_c : list or ndarray
59 Wave vector
60 complex_frequencies : np.array or ComplexFrequencyDescriptor
61 Array of complex frequencies to evaluate the response function at
62 or a descriptor of those frequencies.
63 fxc : str or None
64 Approximation to the (local) xc kernel. If left as None, xc-effects
65 are neglected from the Dyson equation (RPA).
66 Other choices: ALDA, ALDA_X, ALDA_x
67 hxc_scaling : HXCScaling (or None, if irrelevant)
68 Supply an HXCScaling object to scale the hxc kernel.
69 txt : str
70 Save output of the calculation of this specific component into
71 a file with the filename of the given input.
72 """
73 # Initiate new output file, if supplied
74 if txt is not None:
75 self.context.new_txt_and_timer(txt)
77 # Print to output file
78 self.context.print('---------------', flush=False)
79 self.context.print('Calculating susceptibility spincomponent='
80 f'{spincomponent} with q_c={q_c}', flush=False)
81 self.context.print('---------------')
83 # Calculate chiks
84 chiks = self.calculate_chiks(spincomponent, q_c, complex_frequencies)
85 # Construct the hxc kernel
86 hxc_kernel = self.get_hxc_kernel(fxc, spincomponent, chiks.qpd)
87 # Solve dyson equation
88 chi = self.dyson_solver(chiks, hxc_kernel, hxc_scaling=hxc_scaling)
90 return chiks, chi
92 def get_hxc_kernel(self, fxc: str | None, spincomponent: str,
93 qpd: SingleQPWDescriptor) -> HXCKernel:
94 return HXCKernel(
95 hartree_kernel=self.get_hartree_kernel(spincomponent, qpd),
96 xc_kernel=self.get_xc_kernel(fxc, spincomponent, qpd))
98 def get_hartree_kernel(self, spincomponent: str,
99 qpd: SingleQPWDescriptor) -> PWKernel:
100 if spincomponent in ['+-', '-+']:
101 # No Hartree term in Dyson equation
102 return NoKernel.from_qpd(qpd)
103 else:
104 return NewCoulombKernel.from_qpd(
105 qpd, N_c=self.gs.kd.N_c, pbc_c=self.gs.atoms.get_pbc())
107 def get_xc_kernel(self, fxc: str | None, spincomponent: str,
108 qpd: SingleQPWDescriptor) -> PWKernel:
109 """Get the requested xc-kernel object."""
110 if fxc is None:
111 # No xc-kernel
112 return NoKernel.from_qpd(qpd)
114 if qpd.gammacentered:
115 # When using a gamma-centered plane-wave basis, we can reuse the
116 # fxc kernel for all q-vectors. Thus, we keep a cache of calculated
117 # kernels
118 key = f'{fxc},{spincomponent}'
119 if key not in self.fxc_kernel_cache:
120 self.fxc_kernel_cache[key] = self.fxc_calculator(
121 fxc, spincomponent, qpd)
122 fxc_kernel = self.fxc_kernel_cache[key]
123 else:
124 # Always compute the kernel
125 fxc_kernel = self.fxc_calculator(fxc, spincomponent, qpd)
127 return fxc_kernel
129 def calculate_chiks(self, spincomponent, q_c, complex_frequencies):
130 """Calculate the Kohn-Sham susceptibility."""
131 q_c = np.asarray(q_c)
132 if isinstance(complex_frequencies, ComplexFrequencyDescriptor):
133 zd = complex_frequencies
134 else:
135 zd = ComplexFrequencyDescriptor.from_array(complex_frequencies)
137 # Perform actual calculation
138 chiks = self.chiks_calc.calculate(spincomponent, q_c, zd)
139 # Distribute frequencies over world
140 chiks = chiks.copy_with_global_frequency_distribution()
142 return chiks
145def spectral_decomposition(chi, pos_eigs=1, neg_eigs=0):
146 """Decompose the susceptibility in terms of spectral functions.
148 The full spectrum of induced excitations is extracted and separated into
149 contributions corresponding to the pos_eigs and neg_eigs largest positive
150 and negative eigenvalues respectively.
151 """
152 # Initiate an EigendecomposedSpectrum object with the full spectrum
153 full_spectrum = EigendecomposedSpectrum.from_chi(chi)
155 # Separate the positive and negative eigenvalues for each frequency
156 Apos = full_spectrum.get_positive_eigenvalue_spectrum()
157 Aneg = full_spectrum.get_negative_eigenvalue_spectrum()
159 # Keep only a fixed number of eigenvalues
160 Apos = Apos.reduce_number_of_eigenvalues(pos_eigs)
161 Aneg = Aneg.reduce_number_of_eigenvalues(neg_eigs)
163 return Apos, Aneg
166class EigendecomposedSpectrum:
167 """Data object for eigendecomposed susceptibility spectra."""
169 def __init__(self, omega_w, G_Gc, s_we, v_wGe, A_w=None,
170 wblocks: Blocks1D | None = None):
171 """Construct the EigendecomposedSpectrum.
173 Parameters
174 ----------
175 omega_w : np.array
176 Global array of frequencies in eV.
177 G_Gc : np.array
178 Reciprocal lattice vectors in relative coordinates.
179 s_we : np.array
180 Sorted eigenvalues (in decreasing order) at all frequencies.
181 Here, e is the eigenvalue index.
182 v_wGe : np.array
183 Eigenvectors for corresponding to the (sorted) eigenvalues. With
184 all eigenvalues present in the representation, v_Ge should
185 constitute the unitary transformation matrix between the eigenbasis
186 and the plane-wave representation.
187 A_w : np.array or None
188 Full spectral weight as a function of frequency. If given as None,
189 A_w will be calculated as the sum of all eigenvalues (equal to the
190 trace of the spectrum, if no eigenvalues have been discarded).
191 wblocks : Blocks1D
192 Frequency block parallelization, if any.
193 """
194 self.omega_w = omega_w
195 self.G_Gc = G_Gc
197 self.s_we = s_we
198 self.v_wGe = v_wGe
200 self._A_w = A_w
201 if wblocks is None:
202 # Create a serial Blocks1D instance
203 from gpaw.mpi import serial_comm
204 wblocks = Blocks1D(serial_comm, len(omega_w))
205 self.wblocks = wblocks
207 @classmethod
208 def from_chi(cls, chi):
209 """Construct the eigendecomposed spectrum of a given susceptibility.
211 The spectrum of induced excitations, S_GG'^(μν)(q,ω), which are encoded
212 in a given susceptibility, can be extracted directly from its the
213 dissipative part:
215 1
216 S_GG'^(μν)(q,ω) = - ‾ χ_GG'^(μν")(q,ω)
217 π
218 """
219 assert chi.distribution == 'zGG'
221 # Extract the spectrum of induced excitations
222 chid = chi.copy_dissipative_part()
223 S_wGG = - chid.array / np.pi
225 # Extract frequencies (in eV) and reciprocal lattice vectors
226 omega_w = chid.zd.omega_w * Hartree
227 G_Gc = get_pw_coordinates(chid.qpd)
229 return cls.from_spectrum(omega_w, G_Gc, S_wGG, wblocks=chid.blocks1d)
231 @classmethod
232 def from_spectrum(cls, omega_w, G_Gc, S_wGG, wblocks=None):
233 """Perform an eigenvalue decomposition of a given spectrum."""
234 # Find eigenvalues and eigenvectors of the spectrum
235 s_wK, v_wGK = np.linalg.eigh(S_wGG)
237 # Sort by spectral intensity (eigenvalues in descending order)
238 sorted_indices_wK = np.argsort(-s_wK)
239 s_we = np.take_along_axis(s_wK, sorted_indices_wK, axis=1)
240 v_wGe = np.take_along_axis(
241 v_wGK, sorted_indices_wK[:, np.newaxis, :], axis=2)
243 return cls(omega_w, G_Gc, s_we, v_wGe, wblocks=wblocks)
245 @classmethod
246 def from_file(cls, filename):
247 """Construct the eigendecomposed spectrum from a .npz file."""
248 assert Path(filename).suffix == '.npz', filename
249 npz = np.load(filename)
250 return cls(npz['omega_w'], npz['G_Gc'],
251 npz['s_we'], npz['v_wGe'], A_w=npz['A_w'])
253 def write(self, filename):
254 """Write the eigendecomposed spectrum as a .npz file."""
255 assert Path(filename).suffix == '.npz', filename
257 # Gather data from the different blocks of frequencies to root
258 s_we = self.wblocks.gather(self.s_we)
259 v_wGe = self.wblocks.gather(self.v_wGe)
260 A_w = self.wblocks.gather(self.A_w)
262 # Let root write the spectrum to a pickle file
263 if self.wblocks.blockcomm.rank == 0:
264 np.savez(filename, omega_w=self.omega_w, G_Gc=self.G_Gc,
265 s_we=s_we, v_wGe=v_wGe, A_w=A_w)
267 @property
268 def nG(self):
269 return self.G_Gc.shape[0]
271 @property
272 def neigs(self):
273 return self.s_we.shape[1]
275 @property
276 def A_w(self):
277 if self._A_w is None:
278 self._A_w = np.nansum(self.s_we, axis=1)
279 return self._A_w
281 @property
282 def A_wGG(self):
283 """Generate the spectrum from the eigenvalues and eigenvectors."""
284 A_wGG = np.empty((self.wblocks.nlocal, self.nG, self.nG),
285 dtype=complex)
286 for w, (s_e, v_Ge) in enumerate(zip(self.s_we, self.v_wGe)):
287 emask = ~np.isnan(s_e)
288 svinv_eG = s_e[emask][:, np.newaxis] * np.conj(v_Ge.T[emask])
289 A_wGG[w] = v_Ge[:, emask] @ svinv_eG
290 return A_wGG
292 def new_nan_arrays(self, neigs):
293 """Allocate new eigenvalue and eigenvector arrays filled with np.nan.
294 """
295 s_we = np.empty((self.wblocks.nlocal, neigs), dtype=self.s_we.dtype)
296 v_wGe = np.empty((self.wblocks.nlocal, self.nG, neigs),
297 dtype=self.v_wGe.dtype)
298 s_we[:] = np.nan
299 v_wGe[:] = np.nan
300 return s_we, v_wGe
302 def get_positive_eigenvalue_spectrum(self):
303 """Create a new EigendecomposedSpectrum from the positive eigenvalues.
305 This is especially useful in order to separate the full spectrum of
306 induced excitations, see [PRB 103, 245110 (2021)],
308 S_GG'^μν(q,ω) = A_GG'^μν(q,ω) - A_(-G'-G)^νμ(-q,-ω)
310 into the ν and μ components of the spectrum. Since the spectral
311 function A_GG'^μν(q,ω) is positive definite or zero (in regions without
312 excitations), A_GG'^μν(q,ω) simply corresponds to the positive
313 eigenvalue contribution to the full spectrum S_GG'^μν(q,ω).
314 """
315 # Find the maximum number of positive eigenvalues across the entire
316 # frequency range
317 if self.wblocks.nlocal > 0:
318 pos_we = self.s_we > 0.
319 npos_max = int(np.max(np.sum(pos_we, axis=1)))
320 else:
321 npos_max = 0
322 npos_max = self.wblocks.blockcomm.max_scalar(npos_max)
324 # Allocate new arrays, using np.nan for padding (the number of positive
325 # eigenvalues might vary with frequency)
326 s_we, v_wGe = self.new_nan_arrays(npos_max)
328 # Fill arrays with the positive eigenvalue data
329 for w, (s_e, v_Ge) in enumerate(zip(self.s_we, self.v_wGe)):
330 pos_e = s_e > 0.
331 npos = np.sum(pos_e)
332 s_we[w, :npos] = s_e[pos_e]
333 v_wGe[w, :, :npos] = v_Ge[:, pos_e]
335 return EigendecomposedSpectrum(self.omega_w, self.G_Gc, s_we, v_wGe,
336 wblocks=self.wblocks)
338 def get_negative_eigenvalue_spectrum(self):
339 """Create a new EigendecomposedSpectrum from the negative eigenvalues.
341 The spectrum is created by reversing and negating the spectrum,
343 -S_GG'^μν(q,-ω) = -A_GG'^μν(q,-ω) + A_(-G'-G)^νμ(-q,ω),
345 from which the spectral function A_GG'^νμ(q,ω) can be extracted as the
346 positive eigenvalue contribution, thanks to the reciprocity relation
348 ˍˍ
349 χ_GG'^μν(q,ω) = χ_(-G'-G)^νμ(-q,ω),
350 ˍ
351 in which n^μ(r) denotes the hermitian conjugate [n^μ(r)]^†, and which
352 is valid for μν ∊ {00,0z,zz,+-} in collinear systems without spin-orbit
353 coupling.
354 """
355 # Negate the spectral function, its frequencies and reverse the order
356 # of eigenvalues
357 omega_w = - self.omega_w
358 s_we = - self.s_we[:, ::-1]
359 v_wGe = self.v_wGe[..., ::-1]
360 inverted_spectrum = EigendecomposedSpectrum(omega_w, self.G_Gc,
361 s_we, v_wGe,
362 wblocks=self.wblocks)
364 return inverted_spectrum.get_positive_eigenvalue_spectrum()
366 def reduce_number_of_eigenvalues(self, neigs):
367 """Create a new spectrum with only the neigs largest eigenvalues.
369 The returned EigendecomposedSpectrum is constructed to retain knowledge
370 of the full spectral weight of the unreduced spectrum through the A_w
371 attribute.
372 """
373 assert self.nG >= neigs
374 # Check that the available eigenvalues are in descending order
375 assert all([np.all(np.logical_not(s_e[1:] - s_e[:-1] > 0.))
376 for s_e in self.s_we]), \
377 'Eigenvalues needs to be sorted in descending order!'
379 # Create new output arrays with the requested number of eigenvalues,
380 # using np.nan for padding
381 s_we, v_wGe = self.new_nan_arrays(neigs)
382 # In reality, there may be less actual eigenvalues than requested,
383 # since the user usually does not know how many e.g. negative
384 # eigenvalues persist on the positive frequency axis (or vice-versa).
385 # Fill in available eigenvalues up to the requested number.
386 neigs = min(neigs, self.neigs)
387 s_we[:, :neigs] = self.s_we[:, :neigs]
388 v_wGe[..., :neigs] = self.v_wGe[..., :neigs]
390 return EigendecomposedSpectrum(self.omega_w, self.G_Gc, s_we, v_wGe,
391 # Keep the full spectral weight
392 A_w=self.A_w,
393 wblocks=self.wblocks)
395 def get_eigenmode_lineshapes(self, nmodes=1, wmask=None):
396 """Extract the spectral lineshapes of the eigenmodes.
398 The spectral lineshape is calculated as the inner product
400 a^μν_n(q,ω) = <v^μν_n(q)| A^μν(q,ω) |v^μν_n(q)>
402 where the eigenvectors |v^μν_n> diagonalize the full spectral function
403 at an appropriately chosen frequency ω_m:
405 S^μν(q,ω_m) |v^μν_n(q)> = s^μν_n(q,ω_m) |v^μν_n(q)>
406 """
407 wm = self.get_eigenmode_frequency(nmodes=nmodes, wmask=wmask)
408 return self.get_eigenmodes_at_frequency(wm, nmodes=nmodes)
410 def get_eigenmode_frequency(self, nmodes=1, wmask=None):
411 """Get the frequency at which to extract the eigenmodes.
413 Generally, we chose the frequency ω_m to maximize the minimum
414 eigenvalue difference
416 ω_m(q) = maxmin_ωn[s^μν_n(q,ω) - s^μν_n+1(q,ω)]
418 where n only runs over the desired number of modes (and the eigenvalues
419 are sorted in descending order).
421 However, in the case where only a single mode is extracted, we use the
422 frequency at which the eigenvalue is maximal.
423 """
424 assert nmodes <= self.neigs
425 # Use wmask to specify valid eigenmode frequencies
426 wblocks = self.wblocks
427 if wmask is None:
428 wmask = np.ones(self.wblocks.N, dtype=bool)
429 if nmodes == 1:
430 # Find frequency where the eigenvalue is maximal
431 s_w = wblocks.all_gather(self.s_we[:, 0])
432 wm = np.nanargmax(s_w * wmask) # skip np.nan padding
433 else:
434 # Find frequency with maximum minimal difference between size of
435 # eigenvalues
436 ds_we = np.array([self.s_we[:, e] - self.s_we[:, e + 1]
437 for e in range(nmodes - 1)]).T
438 dsmin_w = np.min(ds_we, axis=1)
439 dsmin_w = wblocks.all_gather(dsmin_w)
440 wm = np.nanargmax(dsmin_w * wmask) # skip np.nan padding
442 return wm
444 def get_eigenmodes_at_frequency(self, wm, nmodes=1):
445 """Extract the eigenmodes at a specific frequency index wm."""
446 v_Gm = self.get_eigenvectors_at_frequency(wm, nmodes=nmodes)
447 return self._get_eigenmode_lineshapes(v_Gm)
449 def get_eigenvectors_at_frequency(self, wm, nmodes=1):
450 """Extract the eigenvectors at a specific frequency index wm."""
451 wblocks = self.wblocks
452 root, wmlocal = wblocks.find_global_index(wm)
453 if wblocks.blockcomm.rank == root:
454 v_Ge = self.v_wGe[wmlocal]
455 v_Gm = np.ascontiguousarray(v_Ge[:, :nmodes])
456 else:
457 v_Gm = np.empty((self.nG, nmodes), dtype=complex)
458 wblocks.blockcomm.broadcast(v_Gm, root)
460 return v_Gm
462 def _get_eigenmode_lineshapes(self, v_Gm):
463 """Extract the eigenmode lineshape based on the mode eigenvectors."""
464 wblocks = self.wblocks
465 A_wGG = self.A_wGG
466 a_wm = np.empty((wblocks.nlocal, v_Gm.shape[1]), dtype=float)
467 for m, v_G in enumerate(v_Gm.T):
468 a_w = np.conj(v_G) @ A_wGG @ v_G
469 assert np.allclose(a_w.imag, 0.)
470 a_wm[:, m] = a_w.real
471 a_wm = wblocks.all_gather(a_wm)
473 return a_wm
475 def write_full_spectral_weight(self, filename):
476 A_w = self.wblocks.gather(self.A_w)
477 if self.wblocks.blockcomm.rank == 0:
478 write_full_spectral_weight(filename, self.omega_w, A_w)
480 def write_eigenmode_lineshapes(self, filename, **kwargs):
481 a_wm = self.get_eigenmode_lineshapes(**kwargs)
482 if self.wblocks.blockcomm.rank == 0:
483 write_eigenmode_lineshapes(filename, self.omega_w, a_wm)
486def write_full_spectral_weight(filename, omega_w, A_w):
487 """Write the sum of spectral weights A(ω) to a file."""
488 with open(filename, 'w') as fd:
489 print('# {:>11}, {:>11}'.format('omega [eV]', 'A(w)'), file=fd)
490 for omega, A in zip(omega_w, A_w):
491 print(f' {omega:11.6f}, {A:11.6f}', file=fd)
494def read_full_spectral_weight(filename):
495 """Read a stored full spectral weight file."""
496 data = np.loadtxt(filename, delimiter=',')
497 omega_w = np.array(data[:, 0], float)
498 A_w = np.array(data[:, 1], float)
499 return omega_w, A_w
502def write_eigenmode_lineshapes(filename, omega_w, a_wm):
503 """Write the eigenmode lineshapes a^μν_n(ω) to a file."""
504 with open(filename, 'w') as fd:
505 # Print header
506 header = '# {:>11}'.format('omega [eV]')
507 for m in range(a_wm.shape[1]):
508 header += ', {:>11}'.format(f'a_{m}(w)')
509 print(header, file=fd)
510 # Print data
511 for omega, a_m in zip(omega_w, a_wm):
512 data = f' {omega:11.6f}'
513 for a in a_m:
514 data += f', {a:11.6f}'
515 print(data, file=fd)
518def read_eigenmode_lineshapes(filename):
519 """Read a stored eigenmode lineshapes file."""
520 data = np.loadtxt(filename, delimiter=',')
521 omega_w = np.array(data[:, 0], float)
522 a_wm = np.array(data[:, 1:], float)
523 return omega_w, a_wm