1#!/usr/bin/env python
2# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16# Author: Tianyu Zhu <zhutianyu1991@gmail.com>
17#
18
19'''
20Spin-restricted G0W0 approximation with contour deformation
21
22This implementation has the same scaling (N^4) as GW-AC, more robust but slower.
23GW-CD is particularly recommended for accurate core and high-energy states.
24
25Method:
26    See T. Zhu and G.K.-L. Chan, arxiv:2007.03148 (2020) for details
27    Compute Sigma directly on real axis with density fitting
28    through a contour deformation method
29
30Useful References:
31    J. Chem. Theory Comput. 14, 4856-4869 (2018)
32'''
33
34from functools import reduce
35import numpy
36import numpy as np
37import h5py
38from scipy.optimize import newton, least_squares
39
40from pyscf import lib
41from pyscf.lib import logger
42from pyscf.ao2mo import _ao2mo
43from pyscf import df, scf
44from pyscf.mp.mp2 import get_nocc, get_nmo, get_frozen_mask
45from pyscf import __config__
46
47einsum = lib.einsum
48
49def kernel(gw, mo_energy, mo_coeff, Lpq=None, orbs=None,
50           nw=None, vhf_df=False, verbose=logger.NOTE):
51    '''GW-corrected quasiparticle orbital energies
52
53    Returns:
54        A list :  converged, mo_energy, mo_coeff
55    '''
56    mf = gw._scf
57    if gw.frozen is None:
58        frozen = 0
59    else:
60        frozen = gw.frozen
61    assert frozen == 0
62
63    if Lpq is None:
64        Lpq = gw.ao2mo(mo_coeff)
65    if orbs is None:
66        orbs = range(gw.nmo)
67
68    # v_xc
69    v_mf = mf.get_veff() - mf.get_j()
70    v_mf = reduce(numpy.dot, (mo_coeff.T, v_mf, mo_coeff))
71
72    nocc = gw.nocc
73    nmo = gw.nmo
74
75    # v_hf from DFT/HF density
76    if vhf_df: # and frozen == 0:
77        # density fitting for vk
78        vk = -einsum('Lni,Lim->nm',Lpq[:,:,:nocc],Lpq[:,:nocc,:])
79    else:
80        # exact vk without density fitting
81        dm = mf.make_rdm1()
82        rhf = scf.RHF(gw.mol)
83        vk = rhf.get_veff(gw.mol,dm) - rhf.get_j(gw.mol,dm)
84        vk = reduce(numpy.dot, (mo_coeff.T, vk, mo_coeff))
85
86    # Grids for integration on imaginary axis
87    freqs,wts = _get_scaled_legendre_roots(nw)
88
89    # Compute Wmn(iw) on imaginary axis
90    logger.debug(gw, "Computing the imaginary part")
91    Wmn = get_WmnI_diag(gw, orbs, Lpq, freqs)
92
93    conv = True
94    mo_energy = np.zeros_like(gw._scf.mo_energy)
95    for p in orbs:
96        if gw.linearized:
97            # FIXME
98            logger.warn(gw,'linearization with CD leads to wrong quasiparticle energy')
99            raise NotImplementedError
100        else:
101            def quasiparticle(omega):
102                sigma = get_sigma_diag(gw, omega, p, Lpq, Wmn[:,p-orbs[0],:], freqs, wts).real
103                return omega - gw._scf.mo_energy[p] - (sigma.real + vk[p,p] - v_mf[p,p])
104            try:
105                if p < nocc:
106                    delta = -1e-2
107                else:
108                    delta = 1e-2
109                e = newton(quasiparticle, gw._scf.mo_energy[p]+delta, tol=1e-6, maxiter=50)
110                logger.debug(gw, "Computing poles for QP (orb: %s)"%(p))
111                mo_energy[p] = e
112            except RuntimeError:
113                conv = False
114    mo_coeff = gw._scf.mo_coeff
115
116    if gw.verbose >= logger.DEBUG:
117        numpy.set_printoptions(threshold=nmo)
118        logger.debug(gw, '  GW mo_energy =\n%s', mo_energy)
119        numpy.set_printoptions(threshold=1000)
120
121    return conv, mo_energy, mo_coeff
122
123def get_sigma_diag(gw, ep, p, Lpq, Wmn, freqs, wts):
124    '''
125    Compute self-energy on real axis using contour deformation
126    '''
127    nocc = gw.nocc
128    ef = (gw._scf.mo_energy[nocc-1] + gw._scf.mo_energy[nocc])/2.
129    sign = np.sign(ef-gw._scf.mo_energy)
130    sigmaI = get_sigmaI_diag(gw, ep, Wmn, sign, freqs, wts)
131    sigmaR = get_sigmaR_diag(gw, ep, p, ef, Lpq)
132    return sigmaI + sigmaR
133
134
135def get_rho_response(omega, mo_energy, Lpq):
136    '''
137    Compute density response function in auxiliary basis at freq iw
138    '''
139    naux, nocc, nvir = Lpq.shape
140    eia = mo_energy[:nocc,None] - mo_energy[None,nocc:]
141    eia = eia/(omega**2+eia*eia)
142    Pia = einsum('Pia,ia->Pia',Lpq,eia)
143    # Response from both spin-up and spin-down density
144    Pi = 4. * einsum('Pia,Qia->PQ',Pia,Lpq)
145    return Pi
146
147def get_WmnI_diag(gw, orbs, Lpq, freqs):
148    '''
149    Compute W_mn(iw) on imarginary axis grids
150    Return:
151        Wmn: (Nmo, Norbs, Nw)
152    '''
153    mo_energy = gw._scf.mo_energy
154    nocc = gw.nocc
155    nmo = gw.nmo
156    nw = len(freqs)
157    naux = Lpq.shape[0]
158
159    norbs = len(orbs)
160    Wmn = np.zeros((nmo,norbs,nw))
161    for w in range(nw):
162        Pi = get_rho_response(freqs[w], mo_energy, Lpq[:,:nocc,nocc:])
163        Pi_inv = np.linalg.inv(np.eye(naux)-Pi)-np.eye(naux)
164        Qnm = einsum('Pnm,PQ->Qnm',Lpq[:,orbs,:],Pi_inv)
165        Wmn[:,:,w] = einsum('Qnm,Qmn->mn',Qnm,Lpq[:,:,orbs])
166
167    return Wmn
168
169def get_sigmaI_diag(gw, omega, Wmn, sign, freqs, wts):
170    '''
171    Compute self-energy by integrating on imaginary axis
172    '''
173    mo_energy = gw._scf.mo_energy
174    emo = omega - 1j*gw.eta*sign - mo_energy
175    g0 = wts[None,:]*emo[:,None] / ((emo**2)[:,None]+(freqs**2)[None,:])
176    sigma = -einsum('mw,mw',g0,Wmn)/np.pi
177
178    return sigma
179
180def get_rho_response_R(gw, omega, Lpq):
181    '''
182    Compute density response function in auxiliary basis at poles
183    '''
184    naux, nocc, nvir = Lpq.shape
185    mo_energy = gw._scf.mo_energy
186    eia = mo_energy[:nocc,None] - mo_energy[None,nocc:]
187    eia = 1./(omega+eia+2j*gw.eta) + 1./(-omega+eia)
188    Pia = einsum('Pia,ia->Pia',Lpq,eia)
189    # Response from both spin-up and spin-down density
190    Pi = 2. * einsum('Pia,Qia->PQ',Pia,Lpq)
191    return Pi
192
193def get_sigmaR_diag(gw, omega, orbp, ef, Lpq):
194    '''
195    Compute self-energy for poles inside coutour
196    (more and more expensive away from Fermi surface)
197    '''
198    mo_energy = gw._scf.mo_energy
199    nocc = gw.nocc
200    naux = Lpq.shape[0]
201
202    if omega > ef:
203        fm = 1.0
204        idx = np.where((mo_energy<omega) & (mo_energy>ef))[0]
205    else:
206        fm = -1.0
207        idx = np.where((mo_energy>omega) & (mo_energy<ef))[0]
208
209    sigmaR = 0j
210    if len(idx) > 0:
211        for m in idx:
212            em = mo_energy[m] - omega
213            Pi = get_rho_response_R(gw, abs(em), Lpq[:,:nocc,nocc:])
214            Pi_inv = np.linalg.inv(np.eye(naux)-Pi)-np.eye(naux)
215            sigmaR += fm * np.dot(np.dot(Lpq[:,orbp,m],Pi_inv), Lpq[:,m,orbp])
216
217    return sigmaR
218
219def _get_scaled_legendre_roots(nw):
220    """
221    Scale nw Legendre roots, which lie in the
222    interval [-1, 1], so that they lie in [0, inf)
223    Ref: www.cond-mat.de/events/correl19/manuscripts/ren.pdf
224
225    Returns:
226        freqs : 1D ndarray
227        wts : 1D ndarray
228    """
229    freqs, wts = np.polynomial.legendre.leggauss(nw)
230    x0 = 0.5
231    freqs_new = x0*(1.+freqs)/(1.-freqs)
232    wts = wts*2.*x0/(1.-freqs)**2
233    return freqs_new, wts
234
235def _get_clenshaw_curtis_roots(nw):
236    """
237    Clenshaw-Curtis qaudrature on [0,inf)
238    Ref: J. Chem. Phys. 132, 234114 (2010)
239    Returns:
240        freqs : 1D ndarray
241        wts : 1D ndarray
242    """
243    freqs = np.zeros(nw)
244    wts = np.zeros(nw)
245    a = 0.2
246    for w in range(nw):
247        t = (w+1.0)/nw * np.pi/2.
248        freqs[w] = a / np.tan(t)
249        if w != nw-1:
250            wts[w] = a*np.pi/2./nw/(np.sin(t)**2)
251        else:
252            wts[w] = a*np.pi/4./nw/(np.sin(t)**2)
253    return freqs[::-1], wts[::-1]
254
255
256class GWCD(lib.StreamObject):
257
258    eta = getattr(__config__, 'gw_gw_GW_eta', 1e-3)
259    linearized = getattr(__config__, 'gw_gw_GW_linearized', False)
260
261    def __init__(self, mf, frozen=None):
262        self.mol = mf.mol
263        self._scf = mf
264        self.verbose = self.mol.verbose
265        self.stdout = self.mol.stdout
266        self.max_memory = mf.max_memory
267
268        self.frozen = frozen
269        #TODO: implement frozen orbs
270        if not (self.frozen is None or self.frozen == 0):
271            raise NotImplementedError
272
273        # DF-GW must use density fitting integrals
274        if getattr(mf, 'with_df', None):
275            self.with_df = mf.with_df
276        else:
277            self.with_df = df.DF(mf.mol)
278            self.with_df.auxbasis = df.make_auxbasis(mf.mol, mp2fit=True)
279        self._keys.update(['with_df'])
280
281##################################################
282# don't modify the following attributes, they are not input options
283        self._nocc = None
284        self._nmo = None
285        # self.mo_energy: GW quasiparticle energy, not scf mo_energy
286        self.mo_energy = None
287        self.mo_coeff = mf.mo_coeff
288        self.mo_occ = mf.mo_occ
289        self.sigma = None
290
291        keys = set(('eta', 'linearized'))
292        self._keys = set(self.__dict__.keys()).union(keys)
293
294    def dump_flags(self):
295        log = logger.Logger(self.stdout, self.verbose)
296        log.info('')
297        log.info('******** %s ********', self.__class__)
298        log.info('method = %s', self.__class__.__name__)
299        nocc = self.nocc
300        nvir = self.nmo - nocc
301        log.info('GW nocc = %d, nvir = %d', nocc, nvir)
302        if self.frozen is not None:
303            log.info('frozen = %s', self.frozen)
304        logger.info(self, 'use perturbative linearized QP eqn = %s', self.linearized)
305        return self
306
307    @property
308    def nocc(self):
309        return self.get_nocc()
310    @nocc.setter
311    def nocc(self, n):
312        self._nocc = n
313
314    @property
315    def nmo(self):
316        return self.get_nmo()
317    @nmo.setter
318    def nmo(self, n):
319        self._nmo = n
320
321    get_nocc = get_nocc
322    get_nmo = get_nmo
323    get_frozen_mask = get_frozen_mask
324
325    def kernel(self, mo_energy=None, mo_coeff=None, Lpq=None, orbs=None, nw=100, vhf_df=False):
326        """
327        Input:
328            orbs: self-energy orbs
329            nw: grid number
330        Output:
331            mo_energy: GW quasiparticle energy
332        """
333        if mo_coeff is None:
334            mo_coeff = self._scf.mo_coeff
335        if mo_energy is None:
336            mo_energy = self._scf.mo_energy
337
338        cput0 = (logger.process_clock(), logger.perf_counter())
339        self.dump_flags()
340        self.converged, self.mo_energy, self.mo_coeff = \
341                kernel(self, mo_energy, mo_coeff,
342                       Lpq=Lpq, orbs=orbs, nw=nw, vhf_df=vhf_df, verbose=self.verbose)
343
344        logger.timer(self, 'GW', *cput0)
345        return self.mo_energy
346
347    def ao2mo(self, mo_coeff=None):
348        if mo_coeff is None:
349            mo_coeff = self.mo_coeff
350        nmo = self.nmo
351        naux = self.with_df.get_naoaux()
352        mem_incore = (2*nmo**2*naux) * 8/1e6
353        mem_now = lib.current_memory()[0]
354
355        mo = numpy.asarray(mo_coeff, order='F')
356        ijslice = (0, nmo, 0, nmo)
357        Lpq = None
358        if (mem_incore + mem_now < 0.99*self.max_memory) or self.mol.incore_anyway:
359            Lpq = _ao2mo.nr_e2(self.with_df._cderi, mo, ijslice, aosym='s2', out=Lpq)
360            return Lpq.reshape(naux,nmo,nmo)
361        else:
362            logger.warn(self, 'Memory may not be enough!')
363            raise NotImplementedError
364
365
366if __name__ == '__main__':
367    from pyscf import gto, dft, tddft
368    mol = gto.Mole()
369    mol.verbose = 4
370    mol.atom = [
371        [8 , (0. , 0.     , 0.)],
372        [1 , (0. , -0.7571 , 0.5861)],
373        [1 , (0. , 0.7571 , 0.5861)]]
374    mol.basis = 'def2-svp'
375    mol.build()
376
377    mf = dft.RKS(mol)
378    mf.xc = 'pbe'
379    mf.kernel()
380
381    nocc = mol.nelectron//2
382    nmo = mf.mo_energy.size
383    nvir = nmo-nocc
384
385    gw = GWCD(mf)
386    gw.kernel(orbs=range(0,nocc+3))
387    print(gw.mo_energy)
388    assert(abs(gw.mo_energy[nocc-1]--0.41284735)<1e-5)
389    assert(abs(gw.mo_energy[nocc]-0.16574524)<1e-5)
390    assert(abs(gw.mo_energy[0]--19.53387986)<1e-5)
391