Coverage for gpaw/xc/rpa.py: 98%
267 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1from __future__ import annotations
2from dataclasses import dataclass
3import functools
4from time import ctime
6import numpy as np
7from ase.units import Hartree
8from scipy.special import p_roots
10import gpaw.mpi as mpi
11from gpaw.response import timer
12from gpaw.response.chi0 import Chi0Calculator
13from gpaw.response.coulomb_kernels import CoulombKernel
14from gpaw.response.frequencies import FrequencyDescriptor
15from gpaw.response.pair import get_gs_and_context
18def default_ecut_extrapolation(ecut, extrapolate):
19 return ecut * (1 + 0.5 * np.arange(extrapolate))**(-2 / 3)
22class GCut:
23 def __init__(self, cut_G):
24 self._cut_G = cut_G
26 @property
27 def nG(self):
28 return len(self._cut_G)
30 def spin_cut(self, array_GG, ns):
31 # Strange special case for spin-repeated arrays.
32 # Maybe we can get rid of this.
33 if self._cut_G is None:
34 return array_GG
36 cut_sG = np.tile(self._cut_G, ns)
37 cut_sG[self.nG:] += len(array_GG) // ns
38 array_GG = array_GG.take(cut_sG, 0).take(cut_sG, 1)
39 return array_GG
41 def cut(self, array, axes=(0,)):
42 if self._cut_G is None:
43 return array
45 for axis in axes:
46 array = array.take(self._cut_G, axis)
47 return array
50def initialize_q_points(kd, qsym):
51 bzq_qc = kd.get_bz_q_points(first=True)
53 if not qsym:
54 ibzq_qc = bzq_qc
55 weight_q = np.ones(len(bzq_qc)) / len(bzq_qc)
56 else:
57 U_scc = kd.symmetry.op_scc
58 ibzq_qc = kd.get_ibz_q_points(bzq_qc, U_scc)[0]
59 weight_q = kd.q_weights
60 return ibzq_qc, weight_q
63@dataclass
64class RPAIntegral:
65 omega_w: np.ndarray
66 weight_w: np.ndarray
67 ibzq_qc: np.ndarray
68 weight_q: np.ndarray
69 ecut_i: np.ndarray
71 @property
72 def nq(self):
73 """Number of q-points."""
74 return len(self.weight_q)
76 @property
77 def nw(self):
78 """Number of frequencies."""
79 nw = len(self.omega_w)
80 assert len(self.weight_w) == nw
81 return nw
83 @property
84 def ni(self):
85 """Number of plane-wave cutoffs to extrapolate over."""
86 return len(self.ecut_i)
88 def integrate_frequencies(self, in_xwi):
89 out_xi = self.weight_w @ in_xwi
90 return out_xi
92 def integrate_qpoints(self, in_xqi):
93 out_xi = self.weight_q @ in_xqi
94 return out_xi
97@dataclass
98class RPAData:
99 integral: RPAIntegral
101 def __post_init__(self):
102 self.energy_qwi = np.zeros(self.shape)
104 @property
105 def shape(self):
106 return (self.integral.nq, self.integral.nw, self.integral.ni)
108 def contribution_from_qpoint(self, q, i=-1):
109 return self.integral.integrate_frequencies(self.energy_qwi[q, :, i])
111 @property
112 def energy_wqi(self):
113 return np.swapaxes(self.energy_qwi, 0, 1)
115 @property
116 def energy_qi(self):
117 """Correlation energy contribution, E_c(q)."""
118 return self.integral.integrate_frequencies(self.energy_qwi)
120 @property
121 def energy_wi(self):
122 """Correlation energy contribution, E_c(ω)."""
123 return self.integral.integrate_qpoints(self.energy_wqi)
125 @property
126 def energy_i(self):
127 """Correlation energy E_c as a function of the plane-wave cutoff."""
128 return self.integral.integrate_qpoints(self.energy_qi)
131class RPACalculator:
132 def __init__(
133 self,
134 gs,
135 *,
136 context,
137 ecut,
138 frequencies,
139 weights,
140 qsym=True,
141 skip_gamma=False,
142 truncation=None,
143 nblocks=1,
144 calculate_q=None
145 ):
146 self.gs = gs
147 self.context = context
149 # Normalize RPA integral inputs
150 if isinstance(ecut, (float, int)):
151 ecut = default_ecut_extrapolation(ecut, extrapolate=6)
152 ecut_i = np.asarray(np.sort(ecut)) / Hartree
153 # TODO: We should avoid this requirement.
154 # thosk notes: it might work now (after some extensive clean-up of the
155 # frequency integration), but I have not tested it
156 assert len(frequencies) % nblocks == 0
157 # We should actually have a kpoint descriptor for the qpoints.
158 # We are badly failing at making use of the existing tools by reducing
159 # the qpoints to dumb arrays.
160 ibzq_qc, weight_q = initialize_q_points(gs.kd, qsym)
161 # Collect information about the RPA integral on a single object (a.u.)
162 self.integral = RPAIntegral(
163 omega_w=frequencies / Hartree,
164 weight_w=weights / Hartree / (2 * np.pi),
165 ibzq_qc=ibzq_qc,
166 weight_q=weight_q,
167 ecut_i=ecut_i,
168 )
170 self.chi0calc = Chi0Calculator(
171 self.gs, self.context.with_txt('chi0.txt'),
172 nblocks=nblocks,
173 wd=FrequencyDescriptor(1j * self.integral.omega_w),
174 eta=0.0,
175 intraband=False,
176 hilbert=False,
177 ecut=max(self.integral.ecut_i) * Hartree)
178 self.coulomb = CoulombKernel.from_gs(gs, truncation=truncation)
180 self.skip_gamma = skip_gamma
181 # This is a super weird way of achieving inheritance...
182 if calculate_q is None:
183 calculate_q = self.calculate_q_rpa
184 self.calculate_q = calculate_q
186 def calculate(self, *, nbands=None) -> np.ndarray:
187 """Calculate the RPA correlation energy as a function of cutoff."""
188 data = self.calculate_all_contributions(nbands=nbands)
189 return data.energy_i * Hartree # energies in eV
191 def calculate_all_contributions(
192 self, *, nbands=None, spin=False) -> RPAData:
193 """Calculate RPA correlation energy contributions.
195 nbands: int
196 Number of bands (defaults to number of plane-waves).
197 spin: bool
198 Separate spin in response function.
199 (Only needed for beyond RPA methods that inherit this function).
200 """
202 p = functools.partial(self.context.print, flush=False)
204 if nbands is None:
205 p('Response function bands : Equal to number of plane waves')
206 else:
207 p('Response function bands : %s' % nbands)
208 p('Plane wave cutoffs (eV) :', end='')
209 for e in self.integral.ecut_i:
210 p(f' {e * Hartree:.3f}', end='')
211 p()
212 p(self.coulomb.description())
213 p('', flush=True)
215 self.context.timer.start('RPA')
217 data = RPAData(self.integral)
218 ecutmax = max(self.integral.ecut_i)
219 for q, q_c in enumerate(self.integral.ibzq_qc):
220 if np.allclose(q_c, 0.0) and self.skip_gamma:
221 p('Not calculating E_c(q) at Gamma', end='\n')
222 continue
224 chi0_s = [self.chi0calc.create_chi0(q_c)]
225 if spin:
226 chi0_s.append(self.chi0calc.create_chi0(q_c))
228 qpd = chi0_s[0].qpd
229 nG = qpd.ngmax
231 # First not completely filled band:
232 m1 = self.gs.nocc1
233 p(f'# {q} - {ctime().split()[-2]}')
234 p('q = [%1.3f %1.3f %1.3f]' % tuple(q_c))
236 for i, ecut in enumerate(self.integral.ecut_i):
237 if ecut == ecutmax:
238 # Nothing to cut away:
239 gcut = GCut(None)
240 m2 = nbands or nG
241 else:
242 gcut = GCut(np.arange(nG)[qpd.G2_qG[0] <= 2 * ecut])
243 m2 = gcut.nG
245 p('E_cut = %d eV / Bands = %d:' % (ecut * Hartree, m2),
246 end='\n', flush=True)
247 p('E_c(q) = ', end='', flush=False)
248 data.energy_qwi[q, :, i] = self.calculate_q(
249 chi0_s, m1, m2, gcut
250 )
251 energy = data.contribution_from_qpoint(q, i=i)
252 p('%.3f eV' % (energy * Hartree), flush=True)
253 m1 = m2
255 p()
257 e_i = data.energy_i
258 p('==========================================================')
259 p()
260 p('Total correlation energy:')
261 for e_cut, e in zip(self.integral.ecut_i, e_i):
262 p(f'{e_cut * Hartree:6.0f}: {e * Hartree:6.4f} eV')
263 p()
265 if len(e_i) > 1:
266 self.extrapolate(e_i, self.integral.ecut_i)
268 p('Calculation completed at: ', ctime())
269 p()
271 self.context.timer.stop('RPA')
272 self.context.write_timer()
274 return data
276 @timer('chi0(q)')
277 def calculate_q_rpa(self, chi0_s, m1, m2, gcut):
278 chi0 = chi0_s[0]
279 self.chi0calc.update_chi0(
280 chi0, m1=m1, m2=m2, spins=range(self.chi0calc.gs.nspins))
281 qpd = chi0.qpd
282 chi0_wGG = chi0.body.get_distributed_frequencies_array()
283 wblocks = chi0.body.get_distributed_frequencies_blocks1d()
284 # Calculate RPA energy contributions (as a function of w)
285 if chi0.qpd.optical_limit:
286 chi0_wvv = chi0.chi0_Wvv[wblocks.myslice]
287 chi0_wxvG = chi0.chi0_WxvG[wblocks.myslice]
288 energy_w = self.calculate_optical_limit_rpa_energies(
289 qpd, chi0_wGG, chi0_wvv, chi0_wxvG, gcut
290 )
291 else:
292 energy_w = self.calculate_rpa_energies(qpd, chi0_wGG, gcut)
293 return wblocks.all_gather(energy_w)
295 def calculate_optical_limit_rpa_energies(
296 self, qpd, chi0_wGG, chi0_wvv, chi0_wxvG, gcut):
297 """Calculate correlation energy from chi0 in the optical limit."""
298 from gpaw.response.gamma_int import GammaIntegral
300 gamma_int = GammaIntegral(self.coulomb, qpd=qpd)
302 energy_w = []
303 for chi0_GG, chi0_vv, chi0_xvG in zip(chi0_wGG, chi0_wvv, chi0_wxvG):
304 # Integrate over the optical wave vector volume
305 energy = 0.
306 for qweight, sqrtV_G, chi0_mapping in gamma_int:
307 chi0p_GG = chi0_mapping(chi0_GG, chi0_vv, chi0_xvG)
308 energy += qweight * single_rpa_energy(
309 chi0p_GG, gcut.cut(sqrtV_G), gcut)
310 energy_w.append(energy)
311 return np.array(energy_w)
313 def calculate_rpa_energies(self, qpd, chi0_wGG, gcut):
314 """Evaluate correlation energy from chi0 at finite q."""
315 sqrtV_G = gcut.cut(self.coulomb.sqrtV(qpd, q_v=None))
316 return np.array([
317 single_rpa_energy(chi0_GG, sqrtV_G, gcut) for chi0_GG in chi0_wGG
318 ])
320 def extrapolate(self, e_i, ecut_i):
321 self.context.print('Extrapolated energies:', flush=False)
322 ex_i = []
323 for i in range(len(e_i) - 1):
324 e1, e2 = e_i[i:i + 2]
325 x1, x2 = ecut_i[i:i + 2]**-1.5
326 ex = (e1 * x2 - e2 * x1) / (x2 - x1)
327 ex_i.append(ex)
329 self.context.print(' %4.0f -%4.0f: %5.3f eV' %
330 (ecut_i[i] * Hartree, ecut_i[i + 1]
331 * Hartree, ex * Hartree), flush=False)
332 self.context.print('')
334 return e_i * Hartree
337def single_rpa_energy(chi0_GG, sqrtV_G, gcut):
338 nG = len(sqrtV_G)
339 chi0_GG = gcut.cut(chi0_GG, [0, 1])
340 e_GG = np.eye(nG) - chi0_GG * sqrtV_G * sqrtV_G[:, np.newaxis]
341 e = np.log(np.linalg.det(e_GG)) + nG - np.trace(e_GG)
342 return e.real
345def get_gauss_legendre_points(nw=16, frequency_max=800.0, frequency_scale=2.0):
346 y_w, weights_w = p_roots(nw)
347 y_w = y_w.real
348 ys = 0.5 - 0.5 * y_w
349 ys = ys[::-1]
350 w = (-np.log(1 - ys))**frequency_scale
351 w *= frequency_max / w[-1]
352 alpha = (-np.log(1 - ys[-1]))**frequency_scale / frequency_max
353 transform = (-np.log(1 - ys))**(frequency_scale - 1) \
354 / (1 - ys) * frequency_scale / alpha
355 return w, weights_w * transform / 2
358class RPACorrelation(RPACalculator):
359 def __init__(self, calc, xc='RPA',
360 nlambda=None,
361 nfrequencies=16, frequency_max=800.0, frequency_scale=2.0,
362 frequencies=None, weights=None,
363 world=mpi.world,
364 txt='-',
365 truncation: str | None = None,
366 **kwargs):
367 """Creates the RPACorrelation object
369 calc: str or calculator object
370 The string should refer to the .gpw file contaning KS orbitals
371 xc: str
372 Exchange-correlation kernel. This is only different from RPA when
373 this object is constructed from a different module - e.g. fxc.py
374 skip_gamma: bool
375 If True, skip q = [0,0,0] from the calculation
376 qsym: bool
377 Use symmetry to reduce q-points
378 nlambda: int
379 Number of lambda points. Only used for numerical coupling
380 constant integration involved when called from fxc.py
381 nfrequencies: int
382 Number of frequency points used in the Gauss-Legendre integration
383 frequency_max: float
384 Largest frequency point in Gauss-Legendre integration
385 frequency_scale: float
386 Determines density of frequency points at low frequencies. A slight
387 increase to e.g. 2.5 or 3.0 improves convergence wth respect to
388 frequency points for metals
389 frequencies: list
390 List of frequencies for user-specified frequency integration
391 weights: list
392 list of weights (integration measure) for a user specified
393 frequency grid. Must be specified and have the same length as
394 frequencies if frequencies is not None
395 truncation: str or None
396 Coulomb truncation scheme. Can be None, '0D' or '2D'. If None
397 and the system is a molecule then '0D' will be used.
398 world: communicator
399 nblocks: int
400 Number of parallelization blocks. Frequency parallelization
401 can be specified by setting nblocks=nfrequencies and is useful
402 for memory consuming calculations
403 ecut: float or list of floats
404 Plane-wave cutoff(s) in eV.
405 txt: str
406 txt file for saving and loading contributions to the correlation
407 energy from different q-points
408 """
409 gs, context = get_gs_and_context(calc=calc, txt=txt, world=world,
410 timer=None)
412 if frequencies is None:
413 frequencies, weights = get_gauss_legendre_points(nfrequencies,
414 frequency_max,
415 frequency_scale)
416 user_spec = False
417 else:
418 assert weights is not None
419 user_spec = True
421 if truncation is None and not gs.pbc.any():
422 truncation = '0D'
424 super().__init__(gs=gs, context=context,
425 frequencies=frequencies, weights=weights,
426 truncation=truncation,
427 **kwargs)
429 self.print_initialization(xc, frequency_scale, nlambda, user_spec)
431 def print_initialization(self, xc, frequency_scale, nlambda, user_spec):
432 p = functools.partial(self.context.print, flush=False)
433 p('----------------------------------------------------------')
434 p('Non-self-consistent %s correlation energy' % xc)
435 p('----------------------------------------------------------')
436 p('Started at: ', ctime())
437 p()
438 p('Atoms :',
439 self.gs.atoms.get_chemical_formula(mode='hill'))
440 p('Ground state XC functional :', self.gs.xcname)
441 p('Valence electrons :', self.gs.nvalence)
442 p('Number of bands :', self.gs.bd.nbands)
443 p('Number of spins :', self.gs.nspins)
444 p('Number of k-points :', len(self.gs.kd.bzk_kc))
445 p('Number of irreducible k-points :', len(self.gs.kd.ibzk_kc))
446 p('Number of irreducible q-points :', len(self.integral.ibzq_qc))
447 p()
448 for q, weight in zip(self.integral.ibzq_qc, self.integral.weight_q):
449 p(' q: [%1.4f %1.4f %1.4f] - weight: %1.3f' %
450 (q[0], q[1], q[2], weight))
451 p()
452 p('----------------------------------------------------------')
453 p('----------------------------------------------------------')
454 p()
455 if nlambda is None:
456 p('Analytical coupling constant integration')
457 else:
458 p('Numerical coupling constant integration using', nlambda,
459 'Gauss-Legendre points')
460 p()
461 p('Frequencies')
462 if not user_spec:
463 p(' Gauss-Legendre integration with %s frequency points' %
464 len(self.integral.omega_w))
465 p(' Transformed from [0,oo] to [0,1] using e^[-aw^(1/B)]')
466 p(' Highest frequency point at %5.1f eV and B=%1.1f' %
467 (self.integral.omega_w[-1] * Hartree, frequency_scale))
468 else:
469 p(' User specified frequency integration with',
470 len(self.integral.omega_w), 'frequency points')
471 p()
472 p('Parallelization')
473 p(' Total number of CPUs : % s' % self.context.comm.size)
474 blockcomm = self.chi0calc.chi0_body_calc.blockcomm
475 p(' G-vector decomposition : % s' % blockcomm.size)
476 kncomm = self.chi0calc.chi0_body_calc.kncomm
477 p(' K-point/band decomposition : % s' % kncomm.size)
478 self.context.print('')