1import functools
2import itertools
3import os
4import pickle
5import warnings
6from math import pi
7
8import numpy as np
9from ase.dft.kpoints import monkhorst_pack
10from ase.parallel import paropen
11from ase.units import Ha
12from ase.utils import opencew, pickleload
13from ase.utils.timing import timer
14
15import gpaw.mpi as mpi
16from gpaw import GPAW, debug
17from gpaw.kpt_descriptor import KPointDescriptor
18from gpaw.response.chi0 import Chi0, HilbertTransform
19from gpaw.response.fxckernel_calc import calculate_kernel
20from gpaw.response.kernels import get_coulomb_kernel, get_integrated_kernel
21from gpaw.response.pair import PairDensity
22from gpaw.response.wstc import WignerSeitzTruncatedCoulomb
23from gpaw.utilities import devnull
24from gpaw.utilities.progressbar import ProgressBar
25from gpaw.wavefunctions.pw import (PWDescriptor, PWMapping,
26                                   count_reciprocal_vectors)
27from gpaw.xc.exx import EXX, select_kpts
28from gpaw.xc.fxc import set_flags
29from gpaw.xc.tools import vxc
30
31
32class G0W0(PairDensity):
33    def __init__(self, calc, filename='gw', restartfile=None,
34                 kpts=None, bands=None, relbands=None, nbands=None, ppa=False,
35                 xc='RPA', fxc_mode='GW', density_cut=1.e-6, do_GW_too=False,
36                 av_scheme=None, Eg=None,
37                 truncation=None, integrate_gamma=0,
38                 ecut=150.0, eta=0.1, E0=1.0 * Ha,
39                 domega0=0.025, omega2=10.0, q0_correction=False,
40                 anisotropy_correction=None,
41                 nblocks=1, savew=False, savepckl=True,
42                 maxiter=1, method='G0W0', mixing=0.2,
43                 world=mpi.world, ecut_extrapolation=False,
44                 nblocksmax=False, gate_voltage=None,
45                 paw_correction='brute-force'):
46
47        """G0W0 calculator.
48
49        The G0W0 calculator is used is used to calculate the quasi
50        particle energies through the G0W0 approximation for a number
51        of states.
52
53        calc: str or PAW object
54            GPAW calculator object or filename of saved calculator object.
55        filename: str
56            Base filename of output files.
57        restartfile: str
58            File that stores data necessary to restart a calculation.
59        kpts: list
60            List of indices of the IBZ k-points to calculate the quasi particle
61            energies for.
62        bands: tuple of two ints
63            Range of band indices, like (n1, n2), to calculate the quasi
64            particle energies for. Bands n where n1<=n<n2 will be
65            calculated.  Note that the second band index is not included.
66        relbands: tuple of two ints
67            Same as *bands* except that the numbers are relative to the
68            number of occupied bands.
69            E.g. (-1, 1) will use HOMO+LUMO.
70        ecut: float
71            Plane wave cut-off energy in eV.
72        ecut_extrapolation: bool or array
73            If set to True an automatic extrapolation of the selfenergy to
74            infinite cutoff will be performed based on three points
75            for the cutoff energy.
76            If an array is given, the extrapolation will be performed based on
77            the cutoff energies given in the array.
78        nbands: int
79            Number of bands to use in the calculation. If None, the number will
80            be determined from :ecut: to yield a number close to the number of
81            plane waves used.
82        ppa: bool
83            Sets whether the Godby-Needs plasmon-pole approximation for the
84            dielectric function should be used.
85        xc: str
86            Kernel to use when including vertex corrections.
87        fxc_mode: str
88            Where to include the vertex corrections; polarizability and/or
89            self-energy. 'GWP': Polarizability only, 'GWS': Self-energy only,
90            'GWG': Both.
91        density_cut: float
92            Cutoff for density when constructing kernel.
93        do_GW_too: bool
94            When carrying out a calculation including vertex corrections, it
95            is possible to get the standard GW results at the same time
96            (almost for free).
97        av_scheme: str
98            'wavevector'. Method to construct kernel. Only
99            'wavevector' has been tested and works here. The implementation
100            could be extended to include the 'density' method which has been
101            tested for total energy calculations (rALDA etc.)
102        Eg: float
103            Gap to apply in the 'JGMs' (simplified jellium-with-gap) kernel.
104            If None the DFT gap is used.
105        truncation: str
106            Coulomb truncation scheme. Can be either wigner-seitz,
107            2D, 1D, or 0D
108        integrate_gamma: int
109            Method to integrate the Coulomb interaction. 1 is a numerical
110            integration at all q-points with G=[0,0,0] - this breaks the
111            symmetry slightly. 0 is analytical integration at q=[0,0,0] only -
112            this conserves the symmetry. integrate_gamma=2 is the same as 1,
113            but the average is only carried out in the non-periodic directions.
114        E0: float
115            Energy (in eV) used for fitting in the plasmon-pole approximation.
116        domega0: float
117            Minimum frequency step (in eV) used in the generation of the non-
118            linear frequency grid.
119        omega2: float
120            Control parameter for the non-linear frequency grid, equal to the
121            frequency where the grid spacing has doubled in size.
122        gate_voltage: float
123            Shift Fermi level of ground state calculation by the
124            specified amount.
125        q0_correction: bool
126            Analytic correction to the q=0 contribution applicable to 2D
127            systems.
128        anisotropy_correction: bool
129            Old term for the q0_correction.
130        nblocks: int
131            Number of blocks chi0 should be distributed in so each core
132            does not have to store the entire matrix. This is to reduce
133            memory requirement. nblocks must be less than or equal to the
134            number of processors.
135        nblocksmax: bool
136            Cuts chi0 into as many blocks as possible to reduce memory
137            requirements as much as possible.
138        savew: bool
139            Save W to a file.
140        savepckl: bool
141            Save output to a pckl file.
142        method: str
143            G0W0 or GW0(eigenvalue selfconsistency in G) currently available.
144        maxiter: int
145            Number of iterations in a GW0 calculation.
146        mixing: float
147            Number between 0 and 1 determining how much of previous
148            iteration's eigenvalues to mix with.
149        ecut_extrapolation: bool
150            Carries out the extrapolation to infinite cutoff automatically.
151        """
152
153        if world.rank != 0:
154            txt = devnull
155        else:
156            txt = open(filename + '.txt', 'w')
157
158        p = functools.partial(print, file=txt)
159        p('  ___  _ _ _ ')
160        p(' |   || | | |')
161        p(' | | || | | |')
162        p(' |__ ||_____|')
163        p(' |___|')
164        p()
165
166        self.inputcalc = calc
167
168        if ppa and (nblocks > 1 or nblocksmax):
169            p('PPA is currently not compatible with block parallellisation. '
170              'Setting nblocks=1 and continuing.')
171            nblocks = 1
172            nblocksmax = False
173
174        if ecut_extrapolation is True:
175            pct = 0.8
176            necuts = 3
177            ecut_e = ecut * (1 + (1. / pct - 1) * np.arange(necuts)[::-1] /
178                             (necuts - 1))**(-2 / 3)
179            """It doesn't make sence to save W in this case since
180            W is calculated for different cutoffs:"""
181            assert not savew
182        elif isinstance(ecut_extrapolation, (list, np.ndarray)):
183            ecut_e = np.array(np.sort(ecut_extrapolation))
184            ecut = ecut_e[-1]
185            assert not savew
186        else:
187            ecut_e = np.array([ecut])
188
189        # Check if nblocks is compatible, adjust if not
190        if nblocksmax:
191            nblocks = world.size
192            if isinstance(calc, GPAW):
193                raise Exception('Using a calulator is not implemented at '
194                                'the moment, load from file!')
195                # nblocks_calc = calc
196            else:
197                nblocks_calc = GPAW(calc, txt=None)
198            ngmax = []
199            for q_c in nblocks_calc.wfs.kd.bzk_kc:
200                qd = KPointDescriptor([q_c])
201                pd = PWDescriptor(np.min(ecut) / Ha,
202                                  nblocks_calc.wfs.gd, complex, qd)
203                ngmax.append(pd.ngmax)
204            nG = np.min(ngmax)
205
206            while nblocks > nG**0.5 + 1 or world.size % nblocks != 0:
207                nblocks -= 1
208
209            mynG = (nG + nblocks - 1) // nblocks
210            assert mynG * (nblocks - 1) < nG
211
212        self.ecut_e = ecut_e / Ha
213
214        PairDensity.__init__(self, calc, ecut, world=world, nblocks=nblocks,
215                             gate_voltage=gate_voltage, txt=txt,
216                             paw_correction=paw_correction)
217
218        self.gate_voltage = gate_voltage
219        ecut /= Ha
220
221        self.xc = xc
222        self.density_cut = density_cut
223        self.av_scheme = av_scheme
224        self.fxc_mode = fxc_mode
225        self.do_GW_too = do_GW_too
226
227        if self.av_scheme is not None:
228            assert self.av_scheme == 'wavevector'
229
230        if not self.fxc_mode == 'GW':
231            assert self.xc != 'RPA'
232
233        if self.do_GW_too:
234            assert self.xc != 'RPA'
235            assert self.fxc_mode != 'GW'
236            if restartfile is not None:
237                raise RuntimeError('Restart function does not currently work '
238                                   'with do_GW_too=True.')
239
240        if Eg is None and self.xc == 'JGMsx':
241            from ase.dft.bandgap import get_band_gap
242            gap, k1, k2 = get_band_gap(self.calc)
243            self.Eg = gap
244        else:
245            self.Eg = Eg
246
247        set_flags(self)
248
249        self.filename = filename
250        self.restartfile = restartfile
251        self.savew = savew
252        self.savepckl = savepckl
253        self.ppa = ppa
254        self.truncation = truncation
255        self.integrate_gamma = integrate_gamma
256        self.eta = eta / Ha
257        self.E0 = E0 / Ha
258        self.domega0 = domega0 / Ha
259        self.omega2 = omega2 / Ha
260        if anisotropy_correction is not None:
261            self.ac = anisotropy_correction
262            warnings.warn('anisotropy_correction changed name to '
263                          'q0_correction. Please update your script(s).')
264        else:
265            self.ac = q0_correction
266
267        if self.ac:
268            assert self.truncation == '2D'
269            self.x0density = 0.1  # ? 0.01
270
271        self.maxiter = maxiter
272        self.method = method
273        self.mixing = mixing
274        if self.method == 'GW0':
275            assert self.maxiter > 1
276
277        self.kpts = list(select_kpts(kpts, self.calc))
278
279        if bands is not None and relbands is not None:
280            raise ValueError('Use bands or relbands!')
281
282        if relbands is not None:
283            bands = [self.calc.wfs.nvalence // 2 + b for b in relbands]
284
285        if bands is None:
286            bands = [0, self.nocc2]
287
288        self.bands = bands
289
290        self.ibzk_kc = self.calc.get_ibz_k_points()
291        nibzk = len(self.ibzk_kc)
292        self.eps0_skn = np.array([[self.calc.get_eigenvalues(kpt=k, spin=s)
293                                   for k in range(nibzk)]
294                                  for s in range(self.calc.wfs.nspins)]) / Ha
295
296        b1, b2 = bands
297        self.shape = shape = (self.calc.wfs.nspins, len(self.kpts), b2 - b1)
298        self.eps_skn = np.empty(shape)     # KS-eigenvalues
299        self.f_skn = np.empty(shape)       # occupation numbers
300        self.sigma_skn = np.zeros(shape)   # self-energies
301        self.dsigma_skn = np.zeros(shape)  # derivatives of self-energies
302        self.vxc_skn = None                # KS XC-contributions
303        self.exx_skn = None                # exact exchange contributions
304        self.Z_skn = None                  # renormalization factors
305
306        if nbands is None:
307            nbands = int(self.vol * ecut**1.5 * 2**0.5 / 3 / pi**2)
308        self.nbands = nbands
309        self.nspins = self.calc.wfs.nspins
310
311        if self.nspins != 1 and self.fxc_mode != 'GW':
312            raise RuntimeError('Including a xc kernel does currently not '
313                               'work for spinpolarized systems.')
314
315        p()
316        p('Quasi particle states:')
317        if kpts is None:
318            p('All k-points in IBZ')
319        else:
320            kptstxt = ', '.join(['{0:d}'.format(k) for k in self.kpts])
321            p('k-points (IBZ indices): [' + kptstxt + ']')
322        p('Band range: ({0:d}, {1:d})'.format(b1, b2))
323        p()
324        p('Computational parameters:')
325        if not ecut_extrapolation:
326            p('Plane wave cut-off: {0:g} eV'.format(self.ecut * Ha))
327        else:
328            p('Extrapolating to infinite plane wave cut-off using points at:')
329            p('  [%.3f, %.3f, %.3f] eV' % tuple(self.ecut_e[:] * Ha))
330        p('Number of bands: {0:d}'.format(self.nbands))
331        p('Coulomb cutoff:', self.truncation)
332        p('Broadening: {0:g} eV'.format(self.eta * Ha))
333        p()
334        p('Self-consistency method:', self.method)
335        p('fxc mode:', self.fxc_mode)
336        p('Kernel:', self.xc)
337        p('Averaging scheme:', self.av_scheme)
338        p('Do GW too:', self.do_GW_too)
339
340        if self.method == 'GW0':
341            p('Number of iterations:', self.maxiter)
342            p('Mixing:', self.mixing)
343        p()
344        kd = self.calc.wfs.kd
345
346        self.mysKn1n2 = None  # my (s, K, n1, n2) indices
347        self.distribute_k_points_and_bands(b1, b2, kd.ibz2bz_k[self.kpts])
348
349        # Find q-vectors and weights in the IBZ:
350        assert -1 not in kd.bz2bz_ks
351        offset_c = 0.5 * ((kd.N_c + 1) % 2) / kd.N_c
352        bzq_qc = monkhorst_pack(kd.N_c) + offset_c
353        self.qd = KPointDescriptor(bzq_qc)
354        self.qd.set_symmetry(self.calc.atoms, kd.symmetry)
355
356        txt.flush()
357
358    @timer('G0W0')
359    def calculate(self):
360        """Starts the G0W0 calculation.
361
362        Returns a dict with the results with the following key/value pairs:
363
364        ===========  =============================================
365        key          value
366        ===========  =============================================
367        ``f``        Occupation numbers
368        ``eps``      Kohn-Sham eigenvalues in eV
369        ``vxc``      Exchange-correlation
370                     contributions in eV
371        ``exx``      Exact exchange contributions in eV
372        ``sigma``    Self-energy contributions in eV
373        ``dsigma``   Self-energy derivatives
374        ``sigma_e``  Self-energy contributions in eV
375                     used for ecut extrapolation
376        ``Z``        Renormalization factors
377        ``qp``       Quasi particle (QP) energies in eV
378        ``iqp``      GW0/GW: QP energies for each iteration in eV
379        ===========  =============================================
380
381        All the values are ``ndarray``'s of shape
382        (spins, IBZ k-points, bands)."""
383
384        kd = self.calc.wfs.kd
385
386        self.calculate_ks_xc_contribution()
387        self.calculate_exact_exchange()
388
389        if self.restartfile is not None:
390            loaded = self.load_restart_file()
391            if not loaded:
392                self.last_q = -1
393                self.previous_sigma = 0.
394                self.previous_dsigma = 0.
395
396            else:
397                print('Reading ' + str(self.last_q + 1) +
398                      ' q-point(s) from the previous calculation: ' +
399                      self.restartfile + '.sigma.pckl', file=self.fd)
400        else:
401            self.last_q = -1
402            self.previous_sigma = 0.
403            self.previous_dsigma = 0.
404
405        self.fd.flush()
406
407        self.ite = 0
408
409        while self.ite < self.maxiter:
410            # Reset calculation
411            # self-energies
412            self.sigma_eskn = np.zeros((len(self.ecut_e), ) + self.shape)
413            # derivatives of self-energies
414            self.dsigma_eskn = np.zeros((len(self.ecut_e), ) + self.shape)
415
416            if self.do_GW_too:
417                self.sigma_GW_eskn = np.zeros((len(self.ecut_e), ) +
418                                              self.shape)
419                self.dsigma_GW_eskn = np.zeros((len(self.ecut_e), ) +
420                                               self.shape)
421
422            # Get KS eigenvalues and occupation numbers:
423            if self.ite == 0:
424                b1, b2 = self.bands
425                nibzk = self.calc.wfs.kd.nibzkpts
426                for i, k in enumerate(self.kpts):
427                    for s in range(self.nspins):
428                        u = s * nibzk + k
429                        kpt = self.calc.wfs.kpt_u[u]
430                        self.eps_skn[s, i] = kpt.eps_n[b1:b2]
431                        self.f_skn[s, i] = kpt.f_n[b1:b2] / kpt.weight
432
433                self.qp_skn = self.eps_skn.copy()
434                self.qp_iskn = np.array([self.qp_skn])
435                if self.do_GW_too:
436                    self.qp_GW_skn = self.eps_skn.copy()
437                    self.qp_GW_iskn = np.array([self.qp_GW_skn])
438
439            if self.ite > 0:
440                self.update_energies(mixing=self.mixing)
441
442            # My part of the states we want to calculate QP-energies for:
443            mykpts = [self.get_k_point(s, K, n1, n2)
444                      for s, K, n1, n2 in self.mysKn1n2]
445            nkpt = len(mykpts)
446
447            # Loop over q in the IBZ:
448            nQ = 0
449            for ie, pd0, W0, q_c, m2, W0_GW in \
450                    self.calculate_screened_potential():
451                if nQ == 0:
452                    print('Summing all q:', file=self.fd)
453                    pb = ProgressBar(self.fd)
454                for u, kpt1 in enumerate(mykpts):
455                    pb.update((nQ + 1) * u /
456                              (nkpt * self.qd.mynk * self.qd.nspins))
457                    K2 = kd.find_k_plus_q(q_c, [kpt1.K])[0]
458                    kpt2 = self.get_k_point(kpt1.s, K2, 0, m2,
459                                            block=True)
460                    k1 = kd.bz2ibz_k[kpt1.K]
461                    i = self.kpts.index(k1)
462
463                    self.calculate_q(ie, i, kpt1, kpt2, pd0, W0, W0_GW)
464                nQ += 1
465            pb.finish()
466
467            self.world.sum(self.sigma_eskn)
468            self.world.sum(self.dsigma_eskn)
469
470            if self.do_GW_too:
471                self.world.sum(self.sigma_GW_eskn)
472                self.world.sum(self.dsigma_GW_eskn)
473
474            if self.restartfile is not None and loaded:
475                self.sigma_eskn += self.previous_sigma
476                self.dsigma_eskn += self.previous_dsigma
477
478            if len(self.ecut_e) > 1:  # interpolate to infinite ecut
479                self.extrapolate_ecut()
480            else:
481                self.sigma_skn = self.sigma_eskn[0]
482                self.dsigma_skn = self.dsigma_eskn[0]
483                if self.do_GW_too:
484                    self.sigma_GW_skn = self.sigma_GW_eskn[0]
485                    self.dsigma_GW_skn = self.dsigma_GW_eskn[0]
486
487            self.Z_skn = 1 / (1 - self.dsigma_skn)
488
489            qp_skn = self.qp_skn + self.Z_skn * (
490                self.eps_skn -
491                self.vxc_skn - self.qp_skn + self.exx_skn +
492                self.sigma_skn)
493
494            self.qp_skn = qp_skn
495
496            self.qp_iskn = np.concatenate((self.qp_iskn,
497                                           np.array([self.qp_skn])))
498
499            if self.do_GW_too:
500                self.Z_GW_skn = 1 / (1 - self.dsigma_GW_skn)
501
502                qp_GW_skn = self.qp_GW_skn + self.Z_GW_skn * (
503                    self.eps_skn -
504                    self.vxc_skn - self.qp_GW_skn + self.exx_skn +
505                    self.sigma_GW_skn)
506
507                self.qp_GW_skn = qp_GW_skn
508
509            self.ite += 1
510
511        results = {'f': self.f_skn,
512                   'eps': self.eps_skn * Ha,
513                   'vxc': self.vxc_skn * Ha,
514                   'exx': self.exx_skn * Ha,
515                   'sigma': self.sigma_skn * Ha,
516                   'dsigma': self.dsigma_skn,
517                   'Z': self.Z_skn,
518                   'qp': self.qp_skn * Ha,
519                   'iqp': self.qp_iskn * Ha}
520
521        if self.do_GW_too:
522            self.results_GW = {'f': self.f_skn,
523                               'eps': self.eps_skn * Ha,
524                               'vxc': self.vxc_skn * Ha,
525                               'exx': self.exx_skn * Ha,
526                               'sigma': self.sigma_GW_skn * Ha,
527                               'dsigma': self.dsigma_GW_skn,
528                               'Z': self.Z_GW_skn,
529                               'qp': self.qp_GW_skn * Ha,
530                               'iqp': self.qp_GW_iskn * Ha}
531
532        self.print_results(results)
533
534        if len(self.ecut_e) > 1:
535            # save non-extrapolated result and R^2 value for fit quality.
536            results.update({'sigma_eskn': self.sigma_eskn * Ha,
537                            'dsigma_eskn': self.dsigma_eskn * Ha,
538                            'sigr2_skn': self.sigr2_skn,
539                            'dsigr2_skn': self.dsigr2_skn})
540
541            if self.do_GW_too:
542                self.results_GW.update(
543                    {'sigma_GW_eskn': self.sigma_GW_eskn * Ha,
544                     'dsigma_GW_eskn': self.dsigma_GW_eskn * Ha,
545                     'sigr2_GW_skn': self.sigr2_GW_skn,
546                     'dsigr2_GW_skn': self.dsigr2_GW_skn})
547
548        if self.savepckl:
549            with paropen(self.filename + '_results.pckl', 'wb') as fd:
550                pickle.dump(results, fd, 2)
551            if self.do_GW_too:
552                with paropen(self.filename + '_results_GW.pckl', 'wb') as fd:
553                    pickle.dump(self.results_GW, fd, 2)
554
555        # After we have written the results restartfile is obsolete
556        if self.restartfile is not None:
557            if self.world.rank == 0:
558                if os.path.isfile(self.restartfile + '.sigma.pckl'):
559                    os.remove(self.restartfile + '.sigma.pckl')
560
561        if self.method == 'GW0' and self.world.rank == 0:
562            for iq in range(len(self.qd.ibzk_kc)):
563                try:
564                    os.remove(self.filename + '.rank' +
565                              str(self.blockcomm.rank) + '.w.q' +
566                              str(iq) + '.pckl')
567                except OSError:
568                    pass
569
570        return results
571
572    def calculate_q(self, ie, k, kpt1, kpt2, pd0, W0, W0_GW=None):
573        """Calculates the contribution to the self-energy and its derivative
574        for a given set of k-points, kpt1 and kpt2."""
575
576        if W0_GW is None:
577            Ws = [W0]
578        else:
579            Ws = [W0, W0_GW]
580
581        wfs = self.calc.wfs
582
583        N_c = pd0.gd.N_c
584        i_cG = self.sign * np.dot(self.U_cc,
585                                  np.unravel_index(pd0.Q_qG[0], N_c))
586
587        q_c = wfs.kd.bzk_kc[kpt2.K] - wfs.kd.bzk_kc[kpt1.K]
588
589        shift0_c = q_c - self.sign * np.dot(self.U_cc, pd0.kd.bzk_kc[0])
590        assert np.allclose(shift0_c.round(), shift0_c)
591        shift0_c = shift0_c.round().astype(int)
592
593        shift_c = kpt1.shift_c - kpt2.shift_c - shift0_c
594        I_G = np.ravel_multi_index(i_cG + shift_c[:, None], N_c, 'wrap')
595
596        G_Gv = pd0.get_reciprocal_vectors()
597        pos_av = np.dot(self.spos_ac, pd0.gd.cell_cv)
598        M_vv = np.dot(pd0.gd.cell_cv.T,
599                      np.dot(self.U_cc.T,
600                             np.linalg.inv(pd0.gd.cell_cv).T))
601
602        Q_aGii = []
603        for a, Q_Gii in enumerate(self.Q_aGii):
604            x_G = np.exp(1j * np.dot(G_Gv, (pos_av[a] -
605                                            np.dot(M_vv, pos_av[a]))))
606            U_ii = self.calc.wfs.setups[a].R_sii[self.s]
607            Q_Gii = np.dot(np.dot(U_ii, Q_Gii * x_G[:, None, None]),
608                           U_ii.T).transpose(1, 0, 2)
609            if self.sign == -1:
610                Q_Gii = Q_Gii.conj()
611            Q_aGii.append(Q_Gii)
612
613        if debug:
614            self.check(ie, i_cG, shift0_c, N_c, q_c, Q_aGii)
615
616        if self.ppa:
617            calculate_sigma = self.calculate_sigma_ppa
618        else:
619            calculate_sigma = self.calculate_sigma
620
621        for n in range(kpt1.n2 - kpt1.n1):
622            ut1cc_R = kpt1.ut_nR[n].conj()
623            eps1 = kpt1.eps_n[n]
624            C1_aGi = [np.dot(Qa_Gii, P1_ni[n].conj())
625                      for Qa_Gii, P1_ni in zip(Q_aGii, kpt1.P_ani)]
626            n_mG = self.calculate_pair_densities(ut1cc_R, C1_aGi, kpt2,
627                                                 pd0, I_G)
628            if self.sign == 1:
629                n_mG = n_mG.conj()
630
631            f_m = kpt2.f_n
632            deps_m = eps1 - kpt2.eps_n
633
634            nn = kpt1.n1 + n - self.bands[0]
635
636            for jj, W in enumerate(Ws):
637                sigma, dsigma = calculate_sigma(n_mG, deps_m, f_m, W)
638                if jj == 0:
639                    self.sigma_eskn[ie, kpt1.s, k, nn] += sigma
640                    self.dsigma_eskn[ie, kpt1.s, k, nn] += dsigma
641                else:
642                    self.sigma_GW_eskn[ie, kpt1.s, k, nn] += sigma
643                    self.dsigma_GW_eskn[ie, kpt1.s, k, nn] += dsigma
644
645    def check(self, ie, i_cG, shift0_c, N_c, q_c, Q_aGii):
646        I0_G = np.ravel_multi_index(i_cG - shift0_c[:, None], N_c, 'wrap')
647        qd1 = KPointDescriptor([q_c])
648        pd1 = PWDescriptor(self.ecut_e[ie], self.calc.wfs.gd, complex, qd1)
649        G_I = np.empty(N_c.prod(), int)
650        G_I[:] = -1
651        I1_G = pd1.Q_qG[0]
652        G_I[I1_G] = np.arange(len(I0_G))
653        G_G = G_I[I0_G]
654        assert len(I0_G) == len(I1_G)
655        assert (G_G >= 0).all()
656        for a, Q_Gii in enumerate(self.initialize_paw_corrections(pd1)):
657            e = abs(Q_aGii[a] - Q_Gii[G_G]).max()
658            assert e < 1e-12
659
660    @timer('Sigma')
661    def calculate_sigma(self, n_mG, deps_m, f_m, C_swGG):
662        """Calculates a contribution to the self-energy and its derivative for
663        a given (k, k-q)-pair from its corresponding pair-density and
664        energy."""
665        o_m = abs(deps_m)
666        # Add small number to avoid zeros for degenerate states:
667        sgn_m = np.sign(deps_m + 1e-15)
668
669        # Pick +i*eta or -i*eta:
670        s_m = (1 + sgn_m * np.sign(0.5 - f_m)).astype(int) // 2
671
672        beta = (2**0.5 - 1) * self.domega0 / self.omega2
673        w_m = (o_m / (self.domega0 + beta * o_m)).astype(int)
674        m_inb = np.where(w_m < len(self.omega_w) - 1)[0]
675        o1_m = np.empty(len(o_m))
676        o2_m = np.empty(len(o_m))
677        o1_m[m_inb] = self.omega_w[w_m[m_inb]]
678        o2_m[m_inb] = self.omega_w[w_m[m_inb] + 1]
679
680        x = 1.0 / (self.qd.nbzkpts * 2 * pi * self.vol)
681        sigma = 0.0
682        dsigma = 0.0
683        # Performing frequency integration
684        for o, o1, o2, sgn, s, w, n_G in zip(o_m, o1_m, o2_m,
685                                             sgn_m, s_m, w_m, n_mG):
686
687            if w >= len(self.omega_w) - 1:
688                continue
689
690            C1_GG = C_swGG[s][w]
691            C2_GG = C_swGG[s][w + 1]
692            p = x * sgn
693            myn_G = n_G[self.Ga:self.Gb]
694
695            sigma1 = p * np.dot(np.dot(myn_G, C1_GG), n_G.conj()).imag
696            sigma2 = p * np.dot(np.dot(myn_G, C2_GG), n_G.conj()).imag
697            sigma += ((o - o1) * sigma2 + (o2 - o) * sigma1) / (o2 - o1)
698            dsigma += sgn * (sigma2 - sigma1) / (o2 - o1)
699
700        return sigma, dsigma
701
702    def calculate_screened_potential(self):
703        """Calculates the screened potential for each q-point in the 1st BZ.
704        Since many q-points are related by symmetry, the actual calculation is
705        only done for q-points in the IBZ and the rest are obtained by symmetry
706        transformations. Results are returned as a generator to that it is not
707        necessary to store a huge matrix for each q-point in the memory."""
708        # The decorator $timer('W') doesn't work for generators, do we will
709        # have to manually start and stop the timer here:
710        self.timer.start('W')
711        if self.ite == 0:
712            print('\nCalculating screened Coulomb potential', file=self.fd)
713            if self.truncation is not None:
714                print('Using %s truncated Coloumb potential' % self.truncation,
715                      file=self.fd)
716
717        if self.ppa:
718            if self.ite == 0:
719                print('Using Godby-Needs plasmon-pole approximation:',
720                      file=self.fd)
721                print('  Fitting energy: i*E0, E0 = %.3f Hartee' % self.E0,
722                      file=self.fd)
723
724            # use small imaginary frequency to avoid dividing by zero:
725            frequencies = [1e-10j, 1j * self.E0 * Ha]
726
727            parameters = {'eta': 0,
728                          'hilbert': False,
729                          'timeordered': False,
730                          'frequencies': frequencies}
731        else:
732            if self.ite == 0:
733                print('Using full frequency integration:', file=self.fd)
734                print('  domega0: {0:g}'.format(self.domega0 * Ha),
735                      file=self.fd)
736                print('  omega2: {0:g}'.format(self.omega2 * Ha),
737                      file=self.fd)
738
739            parameters = {'eta': self.eta * Ha,
740                          'hilbert': True,
741                          'timeordered': True,
742                          'domega0': self.domega0 * Ha,
743                          'omega2': self.omega2 * Ha}
744
745        self.fd.flush()
746
747        chi0 = Chi0(self.inputcalc,
748                    nbands=self.nbands,
749                    ecut=self.ecut * Ha,
750                    intraband=False,
751                    real_space_derivatives=False,
752                    txt=self.filename + '.w.txt',
753                    timer=self.timer,
754                    nblocks=self.blockcomm.size,
755                    gate_voltage=self.gate_voltage,
756                    paw_correction=self.paw_correction,
757                    **parameters)
758
759        if self.truncation == 'wigner-seitz':
760            wstc = WignerSeitzTruncatedCoulomb(
761                self.calc.wfs.gd.cell_cv,
762                self.calc.wfs.kd.N_c,
763                chi0.fd)
764        else:
765            wstc = None
766
767        self.omega_w = chi0.omega_w
768        self.omegamax = chi0.omegamax
769
770        htp = HilbertTransform(self.omega_w, self.eta, gw=True)
771        htm = HilbertTransform(self.omega_w, -self.eta, gw=True)
772
773        # Find maximum size of chi-0 matrices:
774        gd = self.calc.wfs.gd
775        nGmax = max(count_reciprocal_vectors(self.ecut, gd, q_c)
776                    for q_c in self.qd.ibzk_kc)
777        nw = len(self.omega_w)
778
779        size = self.blockcomm.size
780
781        mynGmax = (nGmax + size - 1) // size
782        mynw = (nw + size - 1) // size
783
784        # some memory sizes...
785        if self.world.rank == 0:
786            # A1_x, A2_x
787            siz = (nw * mynGmax * nGmax +
788                   max(mynw * nGmax, nw * mynGmax) * nGmax) * 16
789            sizA = (nw * nGmax * nGmax + nw * nGmax * nGmax) * 16
790            print('  memory estimate for chi0: local=%.2f MB, global=%.2f MB'
791                  % (siz / 1024**2, sizA / 1024**2), file=self.fd)
792            self.fd.flush()
793
794        # Allocate memory in the beginning and use for all q:
795        A1_x = np.empty(nw * mynGmax * nGmax, complex)
796        A2_x = np.empty(max(mynw * nGmax, nw * mynGmax) * nGmax, complex)
797
798        # Need to pause the timer in between iterations
799        self.timer.stop('W')
800        for iq, q_c in enumerate(self.qd.ibzk_kc):
801            if iq <= self.last_q:
802                continue
803
804            thisqd = KPointDescriptor([q_c])
805            pd = PWDescriptor(self.ecut, self.calc.wfs.gd, complex, thisqd)
806            nG = pd.ngmax
807            mynG = (nG + self.blockcomm.size - 1) // self.blockcomm.size
808            chi0.Ga = self.blockcomm.rank * mynG
809            chi0.Gb = min(chi0.Ga + mynG, nG)
810            if len(self.ecut_e) > 1:
811                shape = (nw, chi0.Gb - chi0.Ga, nG)
812                chi0bands_wGG = A1_x[:np.prod(shape)].reshape(shape).copy()
813                chi0bands_wGG[:] = 0.0
814
815                if np.allclose(q_c, 0.0):
816                    chi0bands_wxvG = np.zeros((nw, 2, 3, nG), complex)
817                    chi0bands_wvv = np.zeros((nw, 3, 3), complex)
818                else:
819                    chi0bands_wxvG = None
820                    chi0bands_wvv = None
821            else:
822                chi0bands_wGG = None
823                chi0bands_wxvG = None
824                chi0bands_wvv = None
825
826            m1 = chi0.nocc1
827            for ie, ecut in enumerate(self.ecut_e):
828                self.timer.start('W')
829                if self.savew:
830                    wfilename = (self.filename + '.rank' +
831                                 str(self.blockcomm.rank) +
832                                 '.w.q%d.pckl' % iq)
833                    fd = opencew(wfilename)
834                if self.savew and fd is None:
835                    # Read screened potential from file
836                    with open(wfilename, 'rb') as fd:
837                        pdi, W = pickleload(fd)
838                    # We also need to initialize the PAW corrections
839                    self.Q_aGii = self.initialize_paw_corrections(pdi)
840
841                    nG = pdi.ngmax
842                    nw = len(self.omega_w)
843                    mynG = (nG + self.blockcomm.size - 1) // \
844                        self.blockcomm.size
845                    self.Ga = self.blockcomm.rank * mynG
846                    self.Gb = min(self.Ga + mynG, nG)
847                    assert mynG * (self.blockcomm.size - 1) < nG
848                else:
849                    # First time calculation
850                    if ecut == self.ecut:
851                        # Nothing to cut away:
852                        m2 = self.nbands
853                    else:
854                        m2 = int(self.vol * ecut**1.5 * 2**0.5 / 3 / pi**2)
855
856                    pdi, W, W_GW = self.calculate_w(
857                        chi0, q_c, pd, chi0bands_wGG,
858                        chi0bands_wxvG, chi0bands_wvv,
859                        m1, m2, ecut, htp, htm, wstc,
860                        A1_x, A2_x, iq)
861                    m1 = m2
862                    if self.savew:
863                        if self.blockcomm.size > 1:
864                            thisfile = (self.filename + '.rank' +
865                                        str(self.blockcomm.rank) +
866                                        '.w.q%d.pckl' % iq)
867                            with open(thisfile, 'wb') as fd:
868                                pickle.dump((pdi, W), fd, 2)
869                        else:
870                            pickle.dump((pdi, W), fd, 2)
871
872                self.timer.stop('W')
873                # Loop over all k-points in the BZ and find those that are
874                # related to the current IBZ k-point by symmetry
875                Q1 = self.qd.ibz2bz_k[iq]
876                done = set()
877                for Q2 in self.qd.bz2bz_ks[Q1]:
878                    if Q2 >= 0 and Q2 not in done:
879                        s = self.qd.sym_k[Q2]
880                        self.s = s
881                        self.U_cc = self.qd.symmetry.op_scc[s]
882                        time_reversal = self.qd.time_reversal_k[Q2]
883                        self.sign = 1 - 2 * time_reversal
884                        Q_c = self.qd.bzk_kc[Q2]
885                        d_c = self.sign * np.dot(self.U_cc, q_c) - Q_c
886                        assert np.allclose(d_c.round(), d_c)
887                        yield ie, pdi, W, Q_c, m2, W_GW
888                        done.add(Q2)
889
890                if self.restartfile is not None:
891                    self.save_restart_file(iq)
892
893    @timer('WW')
894    def calculate_w(self, chi0, q_c, pd, chi0bands_wGG, chi0bands_wxvG,
895                    chi0bands_wvv, m1, m2, ecut, htp, htm, wstc, A1_x, A2_x,
896                    iq):
897        """Calculates the screened potential for a specified q-point."""
898
899        nw = len(self.omega_w)
900        nG = pd.ngmax
901        mynG = (nG + self.blockcomm.size - 1) // self.blockcomm.size
902        self.Ga = chi0.Ga
903        self.Gb = chi0.Gb
904        shape = (nw, chi0.Gb - chi0.Ga, nG)
905        # construct empty matrix for chi
906        chi0_wGG = A1_x[:np.prod(shape)].reshape(shape).copy()
907        chi0_wGG[:] = 0.0
908        if np.allclose(q_c, 0.0):
909            chi0_wxvG = np.zeros((nw, 2, 3, nG), complex)
910            chi0_wvv = np.zeros((nw, 3, 3), complex)
911        else:
912            chi0_wxvG = None
913            chi0_wvv = None
914
915        chi0._calculate(pd, chi0_wGG, chi0_wxvG, chi0_wvv, m1, m2,
916                        range(self.nspins), extend_head=False)
917
918        if len(self.ecut_e) > 1:
919            # Add chi from previous cutoff with remaining bands
920            chi0_wGG += chi0bands_wGG
921            chi0bands_wGG[:] = chi0_wGG.copy()
922            if np.allclose(q_c, 0.0):
923                chi0_wxvG += chi0bands_wxvG
924                chi0bands_wxvG[:] = chi0_wxvG.copy()
925                chi0_wvv += chi0bands_wvv
926                chi0bands_wvv[:] = chi0_wvv.copy()
927
928        self.Q_aGii = chi0.Q_aGii
929
930        mynw = (nw + self.blockcomm.size - 1) // self.blockcomm.size
931        if self.blockcomm.size > 1:
932            chi0_wGG = chi0.redistribute(chi0_wGG, A2_x)
933            wa = min(self.blockcomm.rank * mynw, nw)
934            wb = min(wa + mynw, nw)
935        else:
936            wa = 0
937            wb = nw
938
939        if ecut == pd.ecut:
940            nG = len(chi0_wGG[0])
941            pdi = pd
942            G2G = None
943
944        elif ecut < pd.ecut:  # construct subset chi0 matrix with lower ecut
945            pdi = PWDescriptor(ecut, pd.gd, dtype=pd.dtype,
946                               kd=pd.kd)
947            nG = pdi.ngmax
948            mynG = (nG + self.blockcomm.size - 1) // self.blockcomm.size
949            self.Ga = self.blockcomm.rank * mynG
950            self.Gb = min(self.Ga + mynG, nG)
951            nw = len(self.omega_w)
952            mynw = (nw + self.blockcomm.size - 1) // self.blockcomm.size
953
954            G2G = PWMapping(pdi, pd).G2_G1
955            chi0_wGG = chi0_wGG.take(G2G, axis=1).take(G2G, axis=2)
956
957            if chi0_wxvG is not None:
958                chi0_wxvG = chi0_wxvG.take(G2G, axis=3)
959
960            if self.Q_aGii is not None:
961                for a, Q_Gii in enumerate(self.Q_aGii):
962                    self.Q_aGii[a] = Q_Gii.take(G2G, axis=0)
963
964        if self.integrate_gamma != 0:
965            if self.integrate_gamma == 2:
966                reduced = True
967            else:
968                reduced = False
969            V0, sqrV0 = get_integrated_kernel(pdi,
970                                              self.calc.wfs.kd.N_c,
971                                              truncation=self.truncation,
972                                              reduced=reduced,
973                                              N=100)
974        elif self.integrate_gamma == 0 and np.allclose(q_c, 0):
975            bzvol = (2 * np.pi)**3 / self.vol / self.qd.nbzkpts
976            Rq0 = (3 * bzvol / (4 * np.pi))**(1. / 3.)
977            V0 = 16 * np.pi**2 * Rq0 / bzvol
978            sqrV0 = (4 * np.pi)**(1.5) * Rq0**2 / bzvol / 2
979        else:
980            pass
981
982        delta_GG = np.eye(nG)
983
984        if self.ppa:
985            einv_wGG = []
986
987        # Calculate kernel
988        fv = calculate_kernel(self, nG, self.nspins, iq, G2G)[0:nG, 0:nG]
989        # Generate fine grid in vicinity of gamma
990        if np.allclose(q_c, 0):
991            kd = self.calc.wfs.kd
992            N = 4
993            N_c = np.array([N, N, N])
994            if self.truncation is not None:
995                # Only average periodic directions if trunction is used
996                N_c[np.where(kd.N_c == 1)[0]] = 1
997            qf_qc = monkhorst_pack(N_c) / kd.N_c
998            qf_qc *= 1.0e-6
999            U_scc = kd.symmetry.op_scc
1000            qf_qc = kd.get_ibz_q_points(qf_qc, U_scc)[0]
1001            weight_q = kd.q_weights
1002            qf_qv = 2 * np.pi * np.dot(qf_qc, pd.gd.icell_cv)
1003            a_wq = np.sum([chi0_vq * qf_qv.T
1004                           for chi0_vq in
1005                           np.dot(chi0_wvv[wa:wb], qf_qv.T)], axis=1)
1006            a0_qwG = np.dot(qf_qv, chi0_wxvG[wa:wb, 0])
1007            a1_qwG = np.dot(qf_qv, chi0_wxvG[wa:wb, 1])
1008
1009        self.timer.start('Dyson eq.')
1010        # Calculate W and store it in chi0_wGG ndarray:
1011        if self.do_GW_too:
1012            chi0_GW_wGG = chi0_wGG.copy()
1013        else:
1014            chi0_GW_wGG = [0]
1015
1016        for iw, [chi0_GG, chi0_GW_GG] in enumerate(
1017            zip(chi0_wGG,
1018                itertools.cycle(chi0_GW_wGG))):
1019            if np.allclose(q_c, 0):
1020                einv_GG = np.zeros((nG, nG), complex)
1021                if self.do_GW_too:
1022                    einv_GW_GG = np.zeros((nG, nG), complex)
1023                for iqf in range(len(qf_qv)):
1024                    chi0_GG[0] = a0_qwG[iqf, iw]
1025                    chi0_GG[:, 0] = a1_qwG[iqf, iw]
1026                    chi0_GG[0, 0] = a_wq[iw, iqf]
1027                    if self.do_GW_too:
1028                        chi0_GW_GG[0] = a0_qwG[iqf, iw]
1029                        chi0_GW_GG[:, 0] = a1_qwG[iqf, iw]
1030                        chi0_GW_GG[0, 0] = a_wq[iw, iqf]
1031                    sqrV_G = get_coulomb_kernel(pdi,
1032                                                kd.N_c,
1033                                                truncation=self.truncation,
1034                                                wstc=wstc,
1035                                                q_v=qf_qv[iqf])**0.5
1036
1037#                    chi0v_GG = chi0_GG * sqrV_G * sqrV_G[:, np.newaxis]
1038#                    if self.nspins == 2:
1039#                        chi0v = np.zeros((2 * nG, 2 * nG), dtype=complex)
1040#                        for s in range(self.nspins):
1041#                            m = s * nG
1042#                            n = (s + 1) * nG
1043#                            chi0v[m:n, m:n] = chi0v_GG
1044#                    else:
1045#                    chi0v = chi0v_GG
1046
1047                    if self.fxc_mode == 'GWP':
1048                        e_GG = (np.eye(nG) -
1049                                np.dot(
1050                                    np.linalg.inv(
1051                                        np.eye(nG) -
1052                                        np.dot(chi0_GG *
1053                                               sqrV_G *
1054                                               sqrV_G[:, np.newaxis], fv) +
1055                                        chi0_GG * sqrV_G *
1056                                        sqrV_G[:, np.newaxis]),
1057                                    chi0_GG * sqrV_G *
1058                                    sqrV_G[:, np.newaxis]))
1059                    elif self.fxc_mode == 'GWS':
1060                        e_GG = np.dot(
1061                            np.linalg.inv(
1062                                np.eye(nG) +
1063                                np.dot(chi0_GG *
1064                                       sqrV_G *
1065                                       sqrV_G[:, np.newaxis], fv) -
1066                                chi0_GG * sqrV_G *
1067                                sqrV_G[:, np.newaxis]),
1068                            np.eye(nG) -
1069                            chi0_GG * sqrV_G *
1070                            sqrV_G[:, np.newaxis])
1071                    else:
1072                        e_GG = np.eye(nG) - np.dot(chi0_GG * sqrV_G *
1073                                                   sqrV_G[:, np.newaxis], fv)
1074
1075                    einv_GG += np.linalg.inv(e_GG) * weight_q[iqf]
1076
1077                    if self.do_GW_too:
1078                        e_GW_GG = (np.eye(nG) -
1079                                   chi0_GW_GG * sqrV_G *
1080                                   sqrV_G[:, np.newaxis])
1081                        einv_GW_GG += np.linalg.inv(e_GW_GG) * weight_q[iqf]
1082
1083            else:
1084                sqrV_G = get_coulomb_kernel(pdi,
1085                                            self.calc.wfs.kd.N_c,
1086                                            truncation=self.truncation,
1087                                            wstc=wstc)**0.5
1088                if self.fxc_mode == 'GWP':
1089                    e_GG = (np.eye(nG) -
1090                            np.dot(
1091                                np.linalg.inv(
1092                                    np.eye(nG) -
1093                                    np.dot(chi0_GG * sqrV_G *
1094                                           sqrV_G[:, np.newaxis], fv) +
1095                                    chi0_GG * sqrV_G *
1096                                    sqrV_G[:, np.newaxis]),
1097                                chi0_GG * sqrV_G * sqrV_G[:, np.newaxis]))
1098                elif self.fxc_mode == 'GWS':
1099                    e_GG = np.dot(
1100                        np.linalg.inv(
1101                            np.eye(nG) +
1102                            np.dot(chi0_GG * sqrV_G *
1103                                   sqrV_G[:, np.newaxis], fv) -
1104                            chi0_GG * sqrV_G *
1105                            sqrV_G[:, np.newaxis]),
1106                        np.eye(nG) - chi0_GG *
1107                        sqrV_G * sqrV_G[:, np.newaxis])
1108                else:
1109                    e_GG = (delta_GG -
1110                            np.dot(chi0_GG * sqrV_G *
1111                                   sqrV_G[:, np.newaxis], fv))
1112
1113                einv_GG = np.linalg.inv(e_GG)
1114
1115                if self.do_GW_too:
1116                    e_GW_GG = (delta_GG -
1117                               chi0_GW_GG * sqrV_G * sqrV_G[:, np.newaxis])
1118                    einv_GW_GG = np.linalg.inv(e_GW_GG)
1119
1120            if self.ppa:
1121                einv_wGG.append(einv_GG - delta_GG)
1122            else:
1123                W_GG = chi0_GG
1124                W_GG[:] = (einv_GG - delta_GG) * sqrV_G * sqrV_G[:, np.newaxis]
1125
1126                if self.do_GW_too:
1127                    W_GW_GG = chi0_GW_GG
1128                    W_GW_GG[:] = ((einv_GW_GG - delta_GG) *
1129                                  sqrV_G * sqrV_G[:, np.newaxis])
1130
1131                if self.ac and np.allclose(q_c, 0):
1132                    if iw == 0:
1133                        print_ac = True
1134                    else:
1135                        print_ac = False
1136                    self.add_q0_correction(pdi, W_GG, einv_GG,
1137                                           chi0_wxvG[wa + iw],
1138                                           chi0_wvv[wa + iw],
1139                                           sqrV_G,
1140                                           print_ac=print_ac)
1141                    if self.do_GW_too:
1142                        self.add_q0_correction(pdi, W_GW_GG, einv_GW_GG,
1143                                               chi0_wxvG[wa + iw],
1144                                               chi0_wvv[wa + iw],
1145                                               sqrV_G,
1146                                               print_ac=print_ac)
1147                elif np.allclose(q_c, 0) or self.integrate_gamma != 0:
1148                    W_GG[0, 0] = (einv_GG[0, 0] - 1.0) * V0
1149                    W_GG[0, 1:] = einv_GG[0, 1:] * sqrV_G[1:] * sqrV0
1150                    W_GG[1:, 0] = einv_GG[1:, 0] * sqrV0 * sqrV_G[1:]
1151                    if self.do_GW_too:
1152                        W_GW_GG[0, 0] = (einv_GW_GG[0, 0] - 1.0) * V0
1153                        W_GW_GG[0, 1:] = einv_GW_GG[0, 1:] * sqrV_G[1:] * sqrV0
1154                        W_GW_GG[1:, 0] = einv_GW_GG[1:, 0] * sqrV0 * sqrV_G[1:]
1155                else:
1156                    pass
1157
1158        if self.ppa:
1159            omegat_GG = self.E0 * np.sqrt(einv_wGG[1] /
1160                                          (einv_wGG[0] - einv_wGG[1]))
1161            R_GG = -0.5 * omegat_GG * einv_wGG[0]
1162            W_GG = pi * R_GG * sqrV_G * sqrV_G[:, np.newaxis]
1163            if np.allclose(q_c, 0) or self.integrate_gamma != 0:
1164                W_GG[0, 0] = pi * R_GG[0, 0] * V0
1165                W_GG[0, 1:] = pi * R_GG[0, 1:] * sqrV_G[1:] * sqrV0
1166                W_GG[1:, 0] = pi * R_GG[1:, 0] * sqrV0 * sqrV_G[1:]
1167
1168            self.timer.stop('Dyson eq.')
1169            return pdi, [W_GG, omegat_GG], None
1170
1171        if self.do_GW_too:
1172            A1_GW_x = A1_x.copy()
1173            A2_GW_x = A2_x.copy()
1174
1175        if self.blockcomm.size > 1:
1176            Wm_wGG = chi0.redistribute(chi0_wGG, A1_x)
1177        else:
1178            Wm_wGG = chi0_wGG
1179
1180        Wp_wGG = A2_x[:Wm_wGG.size].reshape(Wm_wGG.shape)
1181        Wp_wGG[:] = Wm_wGG
1182
1183        with self.timer('Hilbert transform'):
1184            htp(Wp_wGG)
1185            htm(Wm_wGG)
1186        self.timer.stop('Dyson eq.')
1187
1188        if self.do_GW_too:
1189            if self.blockcomm.size > 1:
1190                Wm_GW_wGG = chi0.redistribute(chi0_GW_wGG, A1_GW_x)
1191            else:
1192                Wm_GW_wGG = chi0_GW_wGG
1193
1194            Wp_GW_wGG = A2_GW_x[:Wm_GW_wGG.size].reshape(Wm_GW_wGG.shape)
1195            Wp_GW_wGG[:] = Wm_GW_wGG
1196
1197            htp(Wp_GW_wGG)
1198            htm(Wm_GW_wGG)
1199            GW_return = [Wp_GW_wGG, Wm_GW_wGG]
1200        else:
1201            GW_return = None
1202
1203        return pdi, [Wp_wGG, Wm_wGG], GW_return
1204
1205    @timer('Kohn-Sham XC-contribution')
1206    def calculate_ks_xc_contribution(self):
1207        name = self.filename + '.vxc.npy'
1208        fd, self.vxc_skn = self.read_contribution(name)
1209        if self.vxc_skn is None:
1210            print('Calculating Kohn-Sham XC contribution', file=self.fd)
1211            vxc_skn = vxc(self.calc, self.calc.hamiltonian.xc) / Ha
1212            n1, n2 = self.bands
1213            self.vxc_skn = vxc_skn[:, self.kpts, n1:n2]
1214            np.save(fd, self.vxc_skn)
1215
1216    @timer('EXX')
1217    def calculate_exact_exchange(self):
1218        name = self.filename + '.exx.npy'
1219        fd, self.exx_skn = self.read_contribution(name)
1220        if self.exx_skn is None:
1221            print('Calculating EXX contribution', file=self.fd)
1222            exx = EXX(self.calc, kpts=self.kpts, bands=self.bands,
1223                      txt=self.filename + '.exx.txt', timer=self.timer)
1224            exx.calculate()
1225            self.exx_skn = exx.get_eigenvalue_contributions() / Ha
1226            np.save(fd, self.exx_skn)
1227
1228    def read_contribution(self, filename):
1229        fd = opencew(filename)  # create, exclusive, write
1230        if fd is not None:
1231            # File was not there: nothing to read
1232            return fd, None
1233
1234        try:
1235            with open(filename, 'rb') as fd:
1236                x_skn = np.load(fd)
1237        except IOError:
1238            print('Removing broken file:', filename, file=self.fd)
1239        else:
1240            print('Read:', filename, file=self.fd)
1241            if x_skn.shape == self.shape:
1242                return None, x_skn
1243            print('Removing bad file (wrong shape of array):', filename,
1244                  file=self.fd)
1245
1246        if self.world.rank == 0:
1247            os.remove(filename)
1248
1249        return opencew(filename), None
1250
1251    def print_results(self, results):
1252        description = ['f:      Occupation numbers',
1253                       'eps:     KS-eigenvalues [eV]',
1254                       'vxc:     KS vxc [eV]',
1255                       'exx:     Exact exchange [eV]',
1256                       'sigma:   Self-energies [eV]',
1257                       'dsigma:  Self-energy derivatives',
1258                       'Z:       Renormalization factors',
1259                       'qp:      QP-energies [eV]']
1260
1261        print('\nResults:', file=self.fd)
1262        for line in description:
1263            print(line, file=self.fd)
1264
1265        b1, b2 = self.bands
1266        names = [line.split(':', 1)[0] for line in description]
1267        ibzk_kc = self.calc.wfs.kd.ibzk_kc
1268        for s in range(self.calc.wfs.nspins):
1269            for i, ik in enumerate(self.kpts):
1270                print('\nk-point ' +
1271                      '{0} ({1}): ({2:.3f}, {3:.3f}, {4:.3f})'.format(
1272                          i, ik, *ibzk_kc[ik]) + '                ' +
1273                      self.fxc_mode, file=self.fd)
1274                print('band' +
1275                      ''.join('{0:>8}'.format(name) for name in names),
1276                      file=self.fd)
1277                for n in range(b2 - b1):
1278                    print('{0:4}'.format(n + b1) +
1279                          ''.join('{0:8.3f}'.format(results[name][s, i, n])
1280                                  for name in names),
1281                          file=self.fd)
1282                if self.do_GW_too:
1283                    print(' ' * 67 + 'GW', file=self.fd)
1284                    for n in range(b2 - b1):
1285                        print('{0:4}'.format(n + b1) +
1286                              ''.join('{0:8.3f}'
1287                                      .format(self.results_GW[name][s, i, n])
1288                                      for name in names),
1289                              file=self.fd)
1290
1291        self.timer.write(self.fd)
1292
1293    @timer('PPA-Sigma')
1294    def calculate_sigma_ppa(self, n_mG, deps_m, f_m, W):
1295        W_GG, omegat_GG = W
1296
1297        sigma = 0.0
1298        dsigma = 0.0
1299
1300        # init variables (is this necessary?)
1301        nG = n_mG.shape[1]
1302        deps_GG = np.empty((nG, nG))
1303        sign_GG = np.empty((nG, nG))
1304        x1_GG = np.empty((nG, nG))
1305        x2_GG = np.empty((nG, nG))
1306        x3_GG = np.empty((nG, nG))
1307        x4_GG = np.empty((nG, nG))
1308        x_GG = np.empty((nG, nG))
1309        dx_GG = np.empty((nG, nG))
1310        nW_G = np.empty(nG)
1311        for m in range(np.shape(n_mG)[0]):
1312            deps_GG = deps_m[m]
1313            sign_GG = 2 * f_m[m] - 1
1314            x1_GG = 1 / (deps_GG + omegat_GG - 1j * self.eta)
1315            x2_GG = 1 / (deps_GG - omegat_GG + 1j * self.eta)
1316            x3_GG = 1 / (deps_GG + omegat_GG - 1j * self.eta * sign_GG)
1317            x4_GG = 1 / (deps_GG - omegat_GG - 1j * self.eta * sign_GG)
1318            x_GG = W_GG * (sign_GG * (x1_GG - x2_GG) + x3_GG + x4_GG)
1319            dx_GG = W_GG * (sign_GG * (x1_GG**2 - x2_GG**2) +
1320                            x3_GG**2 + x4_GG**2)
1321            nW_G = np.dot(n_mG[m], x_GG)
1322            sigma += np.vdot(n_mG[m], nW_G).real
1323            nW_G = np.dot(n_mG[m], dx_GG)
1324            dsigma -= np.vdot(n_mG[m], nW_G).real
1325
1326        x = 1 / (self.qd.nbzkpts * 2 * pi * self.vol)
1327        return x * sigma, x * dsigma
1328
1329    def save_restart_file(self, nQ):
1330        sigma_eskn_write = self.sigma_eskn.copy()
1331        dsigma_eskn_write = self.dsigma_eskn.copy()
1332        self.world.sum(sigma_eskn_write)
1333        self.world.sum(dsigma_eskn_write)
1334        data = {'last_q': nQ,
1335                'sigma_eskn': sigma_eskn_write + self.previous_sigma,
1336                'dsigma_eskn': dsigma_eskn_write + self.previous_dsigma,
1337                'kpts': self.kpts,
1338                'bands': self.bands,
1339                'nbands': self.nbands,
1340                'ecut_e': self.ecut_e,
1341                'domega0': self.domega0,
1342                'omega2': self.omega2,
1343                'integrate_gamma': self.integrate_gamma}
1344
1345        if self.world.rank == 0:
1346            with open(self.restartfile + '.sigma.pckl', 'wb') as fd:
1347                pickle.dump(data, fd, 2)
1348
1349    def load_restart_file(self):
1350        try:
1351            with open(self.restartfile + '.sigma.pckl', 'rb') as fd:
1352                data = pickleload(fd)
1353        except IOError:
1354            return False
1355        else:
1356            if (data['kpts'] == self.kpts and
1357                data['bands'] == self.bands and
1358                data['nbands'] == self.nbands and
1359                (data['ecut_e'] == self.ecut_e).all and
1360                data['domega0'] == self.domega0 and
1361                data['omega2'] == self.omega2 and
1362                data['integrate_gamma'] == self.integrate_gamma):
1363                self.last_q = data['last_q']
1364                self.previous_sigma = data['sigma_eskn']
1365                self.previous_dsigma = data['dsigma_eskn']
1366                return True
1367            else:
1368                raise ValueError(
1369                    'Restart file not compatible with parameters used in '
1370                    'current calculation. Check kpts, bands, nbands, ecut, '
1371                    'domega0, omega2, integrate_gamma.')
1372
1373    def extrapolate_ecut(self):
1374        # Do linear fit of selfenergy vs. inverse of number of plane waves
1375        # to extrapolate to infinite number of plane waves
1376        from scipy.stats import linregress
1377        print('', file=self.fd)
1378        print('Extrapolating selfenergy to infinite energy cutoff:',
1379              file=self.fd)
1380        print('  Performing linear fit to %d points' % len(self.ecut_e),
1381              file=self.fd)
1382        self.sigr2_skn = np.zeros(self.shape)
1383        self.dsigr2_skn = np.zeros(self.shape)
1384        self.sigma_skn = np.zeros(self.shape)
1385        self.dsigma_skn = np.zeros(self.shape)
1386        invN_i = self.ecut_e**(-3. / 2)
1387        for m in range(np.product(self.shape)):
1388            s, k, n = np.unravel_index(m, self.shape)
1389
1390            slope, intercept, r_value, p_value, std_err = \
1391                linregress(invN_i, self.sigma_eskn[:, s, k, n])
1392
1393            self.sigr2_skn[s, k, n] = r_value**2
1394            self.sigma_skn[s, k, n] = intercept
1395
1396            slope, intercept, r_value, p_value, std_err = \
1397                linregress(invN_i, self.dsigma_eskn[:, s, k, n])
1398
1399            self.dsigr2_skn[s, k, n] = r_value**2
1400            self.dsigma_skn[s, k, n] = intercept
1401
1402        if np.any(self.sigr2_skn < 0.9) or np.any(self.dsigr2_skn < 0.9):
1403            print('  Warning: Bad quality of linear fit for some (n,k). ',
1404                  file=self.fd)
1405            print('           Higher cutoff might be necesarry.', file=self.fd)
1406
1407        print('  Minimum R^2 = %1.4f. (R^2 Should be close to 1)' %
1408              min(np.min(self.sigr2_skn), np.min(self.dsigr2_skn)),
1409              file=self.fd)
1410
1411        if self.do_GW_too:
1412            self.sigr2_GW_skn = np.zeros(self.shape)
1413            self.dsigr2_GW_skn = np.zeros(self.shape)
1414            self.sigma_GW_skn = np.zeros(self.shape)
1415            self.dsigma_GW_skn = np.zeros(self.shape)
1416            invN_i = self.ecut_e**(-3. / 2)
1417            for m in range(np.product(self.shape)):
1418                s, k, n = np.unravel_index(m, self.shape)
1419
1420                slope, intercept, r_value, p_value, std_err = \
1421                    linregress(invN_i, self.sigma_GW_eskn[:, s, k, n])
1422
1423                self.sigr2_GW_skn[s, k, n] = r_value**2
1424                self.sigma_GW_skn[s, k, n] = intercept
1425
1426                slope, intercept, r_value, p_value, std_err = \
1427                    linregress(invN_i, self.dsigma_GW_eskn[:, s, k, n])
1428
1429                self.dsigr2_GW_skn[s, k, n] = r_value**2
1430                self.dsigma_GW_skn[s, k, n] = intercept
1431
1432            if np.any(self.sigr2_GW_skn < 0.9) or np.any(self.dsigr2_GW_skn <
1433                                                         0.9):
1434                print('  GW calculation. Warning: Bad quality of linear fit '
1435                      'for some (n,k). ',
1436                      file=self.fd)
1437                print('           Higher cutoff might be necesarry.',
1438                      file=self.fd)
1439
1440            print('  Minimum R^2 = %1.4f. (R^2 Should be close to 1)' %
1441                  min(np.min(self.sigr2_GW_skn), np.min(self.dsigr2_GW_skn)),
1442                  file=self.fd)
1443
1444    def add_q0_correction(self, pd, W_GG, einv_GG, chi0_xvG, chi0_vv,
1445                          sqrV_G, print_ac=False):
1446        from ase.dft import monkhorst_pack
1447        self.cell_cv = self.calc.wfs.gd.cell_cv
1448        self.qpts_qc = self.calc.wfs.kd.bzk_kc
1449        self.weight_q = 1.0 * np.ones(len(self.qpts_qc)) / len(self.qpts_qc)
1450        L = self.cell_cv[2, 2]
1451        vc_G0 = sqrV_G[1:]**2
1452
1453        B_GG = einv_GG[1:, 1:]
1454        u_v0G = vc_G0[np.newaxis, :]**0.5 * chi0_xvG[0, :, 1:]
1455        u_vG0 = vc_G0[np.newaxis, :]**0.5 * chi0_xvG[1, :, 1:]
1456        U_vv = -chi0_vv
1457        a_v0G = -np.dot(u_v0G, B_GG)
1458        a_vG0 = -np.dot(u_vG0, B_GG.T)
1459        A_vv = U_vv - np.dot(a_v0G, u_vG0.T)
1460        S_v0G = a_v0G
1461        S_vG0 = a_vG0
1462        L_vv = A_vv
1463
1464        # Get necessary G vectors.
1465        G_Gv = pd.get_reciprocal_vectors()[1:]
1466        G_Gv += np.array([1e-14, 1e-14, 0])
1467        G2_G = np.sum(G_Gv**2, axis=1)
1468        Gpar_G = np.sum(G_Gv[:, 0:2]**2, axis=1)**0.5
1469
1470        # Generate numerical q-point grid
1471        rcell_cv = 2 * pi * np.linalg.inv(self.cell_cv).T
1472        N_c = self.qd.N_c
1473
1474        iq = np.argmin(np.sum(self.qpts_qc**2, axis=1))
1475        assert np.allclose(self.qpts_qc[iq], 0)
1476        q0cell_cv = np.array([1, 1, 1])**0.5 * rcell_cv / N_c
1477        q0vol = abs(np.linalg.det(q0cell_cv))
1478
1479        x0density = self.x0density
1480        q0density = 2. / L * x0density
1481        npts_c = np.ceil(np.sum(q0cell_cv**2, axis=1)**0.5 /
1482                         q0density).astype(int)
1483        npts_c[2] = 1
1484        npts_c += (npts_c + 1) % 2
1485        if print_ac:
1486            print('Applying analytical 2D correction to W:',
1487                  file=self.fd)
1488            print('    Evaluating Gamma point contribution to W on a ' +
1489                  '%dx%dx%d grid' % tuple(npts_c), file=self.fd)
1490
1491        qpts_qc = monkhorst_pack(npts_c)
1492        qgamma = np.argmin(np.sum(qpts_qc**2, axis=1))
1493
1494        qpts_qv = np.dot(qpts_qc, q0cell_cv)
1495        qpts_q = np.sum(qpts_qv**2, axis=1)**0.5
1496        qpts_q[qgamma] = 1e-14
1497        qdir_qv = qpts_qv / qpts_q[:, np.newaxis]
1498        qdir_qvv = qdir_qv[:, :, np.newaxis] * qdir_qv[:, np.newaxis, :]
1499        nq = len(qpts_qc)
1500        q0area = q0vol / q0cell_cv[2, 2]
1501        dq0 = q0area / nq
1502        dq0rad = (dq0 / pi)**0.5
1503        R = L / 2.
1504        x0area = q0area * R**2
1505        dx0rad = dq0rad * R
1506
1507        exp_q = 4 * pi * (1 - np.exp(-qpts_q * R))
1508        dv_G = ((pi * L * G2_G * np.exp(-Gpar_G * R) * np.cos(G_Gv[:, 2] * R) -
1509                 4 * pi * Gpar_G * (1 - np.exp(-Gpar_G * R) *
1510                                    np.cos(G_Gv[:, 2] * R))) /
1511                (G2_G**1.5 * Gpar_G *
1512                 (4 * pi * (1 - np.exp(-Gpar_G * R) *
1513                            np.cos(G_Gv[:, 2] * R)))**0.5))
1514
1515        dv_Gv = dv_G[:, np.newaxis] * G_Gv
1516
1517        # Add corrections
1518        W_GG[:, 0] = 0.0
1519        W_GG[0, :] = 0.0
1520
1521        A_q = np.sum(qdir_qv * np.dot(qdir_qv, L_vv), axis=1)
1522        frac_q = 1. / (1 + exp_q * A_q)
1523
1524        # HEAD:
1525        w00_q = -(exp_q / qpts_q)**2 * A_q * frac_q
1526        w00_q[qgamma] = 0.0
1527        W_GG[0, 0] = w00_q.sum() / nq
1528        Axy = 0.5 * (L_vv[0, 0] + L_vv[1, 1])  # in-plane average
1529        a0 = 4 * pi * Axy + 1
1530
1531        W_GG[0, 0] += -((a0 * dx0rad - np.log(a0 * dx0rad + 1)) /
1532                        a0**2 / x0area * 2 * np.pi * (2 * pi * L)**2 * Axy)
1533
1534        # WINGS:
1535        u_q = -exp_q / qpts_q * frac_q
1536        W_GG[1:, 0] = 1. / nq * np.dot(
1537            np.sum(qdir_qv * u_q[:, np.newaxis], axis=0),
1538            S_vG0 * sqrV_G[np.newaxis, 1:])
1539
1540        W_GG[0, 1:] = 1. / nq * np.dot(
1541            np.sum(qdir_qv * u_q[:, np.newaxis], axis=0),
1542            S_v0G * sqrV_G[np.newaxis, 1:])
1543
1544        # BODY:
1545        # Constant corrections:
1546        W_GG[1:, 1:] += 1. / nq * sqrV_G[1:, None] * sqrV_G[None, 1:] * \
1547            np.tensordot(S_v0G, np.dot(S_vG0.T,
1548                                       np.sum(-qdir_qvv *
1549                                              exp_q[:, None, None] *
1550                                              frac_q[:, None, None],
1551                                              axis=0)), axes=(0, 1))
1552        u_vvv = np.tensordot(u_q[:, None] * qpts_qv, qdir_qvv, axes=(0, 0))
1553        # Gradient corrections:
1554        W_GG[1:, 1:] += 1. / nq * np.sum(
1555            dv_Gv[:, :, None] * np.tensordot(
1556                S_v0G, np.tensordot(u_vvv, S_vG0 * sqrV_G[None, 1:],
1557                                    axes=(2, 0)), axes=(0, 1)), axis=1)
1558
1559    def update_energies(self, mixing):
1560        """Updates the energies of the calculator with the quasi-particle
1561        energies."""
1562        shifts_skn = np.zeros(self.shape)
1563        na, nb = self.bands
1564        i1 = len(self.qp_iskn) - 2
1565        i2 = i1 + 1
1566        if i1 < 0:
1567            i1 = 0
1568
1569        for kpt in self.calc.wfs.kpt_u:
1570            s = kpt.s
1571            if kpt.k in self.kpts:
1572                ik = self.kpts.index(kpt.k)
1573                eps1_n = self.mixer(self.qp_iskn[i1, s, ik],
1574                                    self.qp_iskn[i2, s, ik],
1575                                    mixing)
1576                kpt.eps_n[na:nb] = eps1_n
1577                """
1578                Should we do something smart with the bands outside the
1579                interval?
1580                Here we shift the unoccupied bands not included by the average
1581                change of the top-most band and the occupied by the
1582                bottom-most band included
1583                """
1584                shifts_skn[s, ik] = (eps1_n - self.eps0_skn[s, kpt.k, na:nb])
1585
1586        for kpt in self.calc.wfs.kpt_u:
1587            s = kpt.s
1588            if kpt.k in self.kpts:
1589                ik = self.kpts.index(kpt.k)
1590                kpt.eps_n[:na] = (self.eps0_skn[s, kpt.k, :na] +
1591                                  np.mean(shifts_skn[s, :, 0]))
1592                kpt.eps_n[nb:] = (self.eps0_skn[s, kpt.k, nb:] +
1593                                  np.mean(shifts_skn[s, :, -1]))
1594            else:
1595                """
1596                kpt.eps_n[:na] = (self.eps0_skn[s, kpt.k, :na] +
1597                                  np.mean(shifts_skn[s, :, 0]))
1598                kpt.eps_n[na:nb] = (self.eps0_skn[s, kpt.k, na:nb] +
1599                                    np.mean(shifts_skn[s, :], axis=0))
1600                kpt.eps_n[nb:] = (self.eps0_skn[s, kpt.k, nb:] +
1601                                  np.mean(shifts_skn[s, :, -1]))
1602                """
1603                pass
1604
1605    def mixer(self, e0_skn, e1_skn, mixing=1.0):
1606        """Mix energies."""
1607        return e0_skn + mixing * (e1_skn - e0_skn)
1608