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# Authors: Timothy Berkelbach <tim.berkelbach@gmail.com>
17#          Qiming Sun <osirpt.sun@gmail.com>
18#
19
20'''
21Spin-restricted G0W0 approximation with exact frequency integration
22'''
23
24
25from functools import reduce
26import numpy
27import numpy as np
28from scipy.optimize import newton
29
30from pyscf import lib
31from pyscf.lib import logger
32from pyscf import ao2mo
33from pyscf import dft
34from pyscf.mp.mp2 import get_nocc, get_nmo, get_frozen_mask, _mo_without_core
35from pyscf import __config__
36
37einsum = lib.einsum
38
39BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
40MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)
41
42def kernel(gw, mo_energy, mo_coeff, td_e, td_xy, eris=None,
43           orbs=None, verbose=logger.NOTE):
44    '''GW-corrected quasiparticle orbital energies
45
46    Returns:
47        A list :  converged, mo_energy, mo_coeff
48    '''
49    # mf must be DFT; for HF use xc = 'hf'
50    mf = gw._scf
51    assert(isinstance(mf, (dft.rks.RKS      , dft.uks.UKS,
52                           dft.roks.ROKS    , dft.uks.UKS,
53                           dft.rks_symm.RKS , dft.uks_symm.UKS,
54                           dft.rks_symm.ROKS, dft.uks_symm.UKS)))
55    assert(gw.frozen == 0 or gw.frozen is None)
56
57    if eris is None:
58        eris = gw.ao2mo(mo_coeff)
59    if orbs is None:
60        orbs = range(gw.nmo)
61
62    v_mf = mf.get_veff() - mf.get_j()
63    v_mf = reduce(numpy.dot, (mo_coeff.T, v_mf, mo_coeff))
64
65    nocc = gw.nocc
66    nmo = gw.nmo
67    nvir = nmo-nocc
68
69    vk_oo = -np.einsum('piiq->pq', eris.oooo)
70    vk_ov = -np.einsum('iqpi->pq', eris.ovoo)
71    vk_vv = -np.einsum('ipqi->pq', eris.ovvo).conj()
72    vk = np.array(np.bmat([[vk_oo, vk_ov],[vk_ov.T, vk_vv]]))
73
74    nexc = len(td_e)
75    # factor of 2 for normalization, see tdscf/rhf.py
76    td_xy = 2*np.asarray(td_xy) # (nexc, 2, nocc, nvir)
77    td_z = np.sum(td_xy, axis=1).reshape(nexc,nocc,nvir)
78    tdm_oo = einsum('via,iapq->vpq', td_z, eris.ovoo)
79    tdm_ov = einsum('via,iapq->vpq', td_z, eris.ovov)
80    tdm_vv = einsum('via,iapq->vpq', td_z, eris.ovvv)
81    tdm = []
82    for oo,ov,vv in zip(tdm_oo,tdm_ov,tdm_vv):
83        tdm.append(np.array(np.bmat([[oo, ov],[ov.T, vv]])))
84    tdm = np.asarray(tdm)
85
86    conv = True
87    mo_energy = np.zeros_like(gw._scf.mo_energy)
88    for p in orbs:
89        tdm_p = tdm[:,:,p]
90        if gw.linearized:
91            ep = gw._scf.mo_energy[p]
92            sigma = get_sigma_element(gw, ep, tdm_p, tdm_p, td_e).real
93            dsigma_dw = get_sigma_deriv_element(gw, ep, tdm_p, tdm_p, td_e).real
94            zn = 1.0/(1-dsigma_dw)
95            mo_energy[p] = ep + zn*(sigma.real + vk[p,p] - v_mf[p,p])
96        else:
97            def quasiparticle(omega):
98                sigma = get_sigma_element(gw, omega, tdm_p, tdm_p, td_e)
99                return omega - gw._scf.mo_energy[p] - (sigma.real + vk[p,p] - v_mf[p,p])
100            try:
101                mo_energy[p] = newton(quasiparticle, gw._scf.mo_energy[p], tol=1e-6, maxiter=100)
102            except RuntimeError:
103                conv = False
104                mo_energy[p] = gw._scf.mo_energy[p]
105                logger.warn(gw, 'Root finding for GW eigenvalue %s did not converge. '
106                                'Setting it equal to the reference MO energy.'%(p))
107    mo_coeff = gw._scf.mo_coeff
108
109    if gw.verbose >= logger.DEBUG:
110        numpy.set_printoptions(threshold=nmo)
111        logger.debug(gw, '  GW mo_energy =\n%s', mo_energy)
112        numpy.set_printoptions(threshold=1000)
113
114    return conv, mo_energy, mo_coeff
115
116
117def get_sigma_element(gw, omega, tdm_p, tdm_q, td_e, eta=None, vir_sgn=1):
118    if eta is None:
119        eta = gw.eta
120
121    nocc = gw.nocc
122    evi = lib.direct_sum('v-i->vi', td_e, gw._scf.mo_energy[:nocc])
123    eva = lib.direct_sum('v+a->va', td_e, gw._scf.mo_energy[nocc:])
124    sigma =  np.sum( tdm_p[:,:nocc]*tdm_q[:,:nocc]/(omega + evi - 1j*eta) )
125    sigma += np.sum( tdm_p[:,nocc:]*tdm_q[:,nocc:]/(omega - eva + vir_sgn*1j*eta) )
126    return sigma
127
128
129def get_sigma_deriv_element(gw, omega, tdm_p, tdm_q, td_e, eta=None, vir_sgn=1):
130    if eta is None:
131        eta = gw.eta
132
133    nocc = gw.nocc
134    evi = lib.direct_sum('v-i->vi', td_e, gw._scf.mo_energy[:nocc])
135    eva = lib.direct_sum('v+a->va', td_e, gw._scf.mo_energy[nocc:])
136    dsigma =  -np.sum( tdm_p[:,:nocc]*tdm_q[:,:nocc]/(omega + evi - 1j*eta)**2 )
137    dsigma += -np.sum( tdm_p[:,nocc:]*tdm_q[:,nocc:]/(omega - eva + vir_sgn*1j*eta)**2 )
138    return dsigma
139
140
141def get_g(omega, mo_energy, mo_occ, eta):
142    sgn = mo_occ - 1
143    return 1.0/(omega - mo_energy + 1j*eta*sgn)
144
145
146class GWExact(lib.StreamObject):
147    '''non-relativistic restricted GW
148
149    Saved results
150
151        mo_energy :
152            Orbital energies
153        mo_coeff
154            Orbital coefficients
155    '''
156
157    eta = getattr(__config__, 'gw_gw_GW_eta', 1e-8)
158    linearized = getattr(__config__, 'gw_gw_GW_linearized', False)
159
160    def __init__(self, mf, frozen=None, tdmf=None):
161        self.mol = mf.mol
162        self._scf = mf
163        self._tdscf = tdmf
164        self.verbose = self.mol.verbose
165        self.stdout = self.mol.stdout
166        self.max_memory = mf.max_memory
167
168        self.frozen = frozen
169
170##################################################
171# don't modify the following attributes, they are not input options
172        self._nocc = None
173        self._nmo = None
174        self.mo_energy = None
175        self.mo_coeff = mf.mo_coeff
176        self.mo_occ = mf.mo_occ
177
178        keys = set(('eta', 'linearized'))
179        self._keys = set(self.__dict__.keys()).union(keys)
180
181    def dump_flags(self, verbose=None):
182        log = logger.new_logger(self, verbose)
183        log.info('')
184        log.info('******** %s ********', self.__class__)
185        log.info('method = %s', self.__class__.__name__)
186        nocc = self.nocc
187        nvir = self.nmo - nocc
188        log.info('GW nocc = %d, nvir = %d', nocc, nvir)
189        if self.frozen is not None:
190            log.info('frozen = %s', self.frozen)
191        logger.info(self, 'use perturbative linearized QP eqn = %s', self.linearized)
192        return self
193
194    @property
195    def nocc(self):
196        return self.get_nocc()
197    @nocc.setter
198    def nocc(self, n):
199        self._nocc = n
200
201    @property
202    def nmo(self):
203        return self.get_nmo()
204    @nmo.setter
205    def nmo(self, n):
206        self._nmo = n
207
208    get_nocc = get_nocc
209    get_nmo = get_nmo
210    get_frozen_mask = get_frozen_mask
211
212    def get_g0(self, omega, eta=None):
213        if eta is None:
214            eta = self.eta
215        return get_g(omega, self._scf.mo_energy, self.mo_occ, eta)
216
217    def get_g(self, omega, eta=None):
218        if eta is None:
219            eta = self.eta
220        return get_g(omega, self.mo_energy, self.mo_occ, eta)
221
222    def kernel(self, mo_energy=None, mo_coeff=None, td_e=None, td_xy=None,
223               eris=None, orbs=None):
224        if mo_coeff is None:
225            mo_coeff = self._scf.mo_coeff
226        if mo_energy is None:
227            mo_energy = self._scf.mo_energy
228        if self._tdscf is None:
229            from pyscf import tdscf
230            self._tdscf = tdscf.dRPA(self._scf)
231            nocc, nvir = self.nocc, self.nmo-self.nocc
232            self._tdscf.nstates = nocc*nvir
233            self._tdscf.verbose = 0
234            self._tdscf.kernel()
235        if td_e is None:
236            td_e = self._tdscf.e
237        if td_xy is None:
238            td_xy = self._tdscf.xy
239
240        cput0 = (logger.process_clock(), logger.perf_counter())
241        self.dump_flags()
242        self.converged, self.mo_energy, self.mo_coeff = \
243                kernel(self, mo_energy, mo_coeff, td_e, td_xy,
244                       eris=eris, orbs=orbs, verbose=self.verbose)
245
246        logger.timer(self, 'GW', *cput0)
247        return self.mo_energy
248
249    def reset(self, mol=None):
250        if mol is not None:
251            self.mol = mol
252        self._scf.reset(mol)
253        self._tdscf.reset(mol)
254        return self
255
256    def ao2mo(self, mo_coeff=None):
257        nmo = self.nmo
258        nao = self.mo_coeff.shape[0]
259        nmo_pair = nmo * (nmo+1) // 2
260        nao_pair = nao * (nao+1) // 2
261        mem_incore = (max(nao_pair**2, nmo**4) + nmo_pair**2) * 8/1e6
262        mem_now = lib.current_memory()[0]
263        if (self._scf._eri is not None and
264            (mem_incore+mem_now < self.max_memory) or self.mol.incore_anyway):
265            return _make_eris_incore(self, mo_coeff)
266
267        elif getattr(self._scf, 'with_df', None):
268            logger.warn(self, 'GW Exact detected DF being used in the HF object. '
269                        'MO integrals are computed based on the DF 3-index tensors.\n'
270                        'Developer TODO:  Write dfgw.GWExact for the '
271                        'DF-GW calculations')
272            raise NotImplementedError
273            #return _make_df_eris_outcore(self, mo_coeff)
274
275        else:
276            return _make_eris_outcore(self, mo_coeff)
277
278
279class _ChemistsERIs:
280    '''(pq|rs)
281
282    Identical to rccsd _ChemistsERIs except no vvvv.'''
283    def __init__(self, mol=None):
284        self.mol = mol
285        self.mo_coeff = None
286        self.nocc = None
287        self.fock = None
288
289        self.oooo = None
290        self.ovoo = None
291        self.oovv = None
292        self.ovvo = None
293        self.ovov = None
294        self.ovvv = None
295
296    def _common_init_(self, mycc, mo_coeff=None):
297        if mo_coeff is None:
298            mo_coeff = mycc.mo_coeff
299        self.mo_coeff = mo_coeff = _mo_without_core(mycc, mo_coeff)
300# Note: Recomputed fock matrix since SCF may not be fully converged.
301        dm = mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ)
302        fockao = mycc._scf.get_hcore() + mycc._scf.get_veff(mycc.mol, dm)
303        self.fock = reduce(numpy.dot, (mo_coeff.conj().T, fockao, mo_coeff))
304        self.nocc = mycc.nocc
305        self.mol = mycc.mol
306
307        mo_e = self.fock.diagonal()
308        try:
309            gap = abs(mo_e[:self.nocc,None] - mo_e[None,self.nocc:]).min()
310            if gap < 1e-5:
311                logger.warn(mycc, 'HOMO-LUMO gap %s too small for GW', gap)
312        except ValueError:  # gap.size == 0
313            pass
314        return self
315
316def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
317    cput0 = (logger.process_clock(), logger.perf_counter())
318    eris = _ChemistsERIs()
319    eris._common_init_(mycc, mo_coeff)
320    nocc = eris.nocc
321    nmo = eris.fock.shape[0]
322
323    if callable(ao2mofn):
324        eri1 = ao2mofn(eris.mo_coeff).reshape([nmo]*4)
325    else:
326        eri1 = ao2mo.incore.full(mycc._scf._eri, eris.mo_coeff)
327        eri1 = ao2mo.restore(1, eri1, nmo)
328    eris.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
329    eris.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
330    eris.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
331    eris.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
332    eris.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
333    eris.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
334    logger.timer(mycc, 'GW integral transformation', *cput0)
335    return eris
336
337def _make_eris_outcore(mycc, mo_coeff=None):
338    cput0 = (logger.process_clock(), logger.perf_counter())
339    log = logger.Logger(mycc.stdout, mycc.verbose)
340    eris = _ChemistsERIs()
341    eris._common_init_(mycc, mo_coeff)
342
343    mol = mycc.mol
344    mo_coeff = eris.mo_coeff
345    nocc = eris.nocc
346    nao, nmo = mo_coeff.shape
347    nvir = nmo - nocc
348    eris.feri1 = lib.H5TmpFile()
349    eris.oooo = eris.feri1.create_dataset('oooo', (nocc,nocc,nocc,nocc), 'f8')
350    eris.ovoo = eris.feri1.create_dataset('ovoo', (nocc,nvir,nocc,nocc), 'f8', chunks=(nocc,1,nocc,nocc))
351    eris.ovov = eris.feri1.create_dataset('ovov', (nocc,nvir,nocc,nvir), 'f8', chunks=(nocc,1,nocc,nvir))
352    eris.ovvo = eris.feri1.create_dataset('ovvo', (nocc,nvir,nvir,nocc), 'f8', chunks=(nocc,1,nvir,nocc))
353    eris.ovvv = eris.feri1.create_dataset('ovvv', (nocc,nvir,nvir,nvir), 'f8')
354    eris.oovv = eris.feri1.create_dataset('oovv', (nocc,nocc,nvir,nvir), 'f8', chunks=(nocc,nocc,1,nvir))
355    max_memory = max(MEMORYMIN, mycc.max_memory-lib.current_memory()[0])
356
357    ftmp = lib.H5TmpFile()
358    ao2mo.full(mol, mo_coeff, ftmp, max_memory=max_memory, verbose=log)
359    eri = ftmp['eri_mo']
360
361    nocc_pair = nocc*(nocc+1)//2
362    tril2sq = lib.square_mat_in_trilu_indices(nmo)
363    oo = eri[:nocc_pair]
364    eris.oooo[:] = ao2mo.restore(1, oo[:,:nocc_pair], nocc)
365    oovv = lib.take_2d(oo, tril2sq[:nocc,:nocc].ravel(), tril2sq[nocc:,nocc:].ravel())
366    eris.oovv[:] = oovv.reshape(nocc,nocc,nvir,nvir)
367    oo = oovv = None
368
369    tril2sq = lib.square_mat_in_trilu_indices(nmo)
370    blksize = min(nvir, max(BLKMIN, int(max_memory*1e6/8/nmo**3/2)))
371    for p0, p1 in lib.prange(0, nvir, blksize):
372        q0, q1 = p0+nocc, p1+nocc
373        off0 = q0*(q0+1)//2
374        off1 = q1*(q1+1)//2
375        buf = lib.unpack_tril(eri[off0:off1])
376
377        tmp = buf[ tril2sq[q0:q1,:nocc] - off0 ]
378        eris.ovoo[:,p0:p1] = tmp[:,:,:nocc,:nocc].transpose(1,0,2,3)
379        eris.ovvo[:,p0:p1] = tmp[:,:,nocc:,:nocc].transpose(1,0,2,3)
380        eris.ovov[:,p0:p1] = tmp[:,:,:nocc,nocc:].transpose(1,0,2,3)
381        eris.ovvv[:,p0:p1] = tmp[:,:,nocc:,nocc:].transpose(1,0,2,3)
382
383        buf = tmp = None
384    log.timer('GW integral transformation', *cput0)
385    return eris
386
387
388if __name__ == '__main__':
389    from pyscf import gto
390    mol = gto.Mole()
391    mol.verbose = 5
392    mol.atom = [
393        [8 , (0. , 0.     , 0.)],
394        [1 , (0. , -0.757 , 0.587)],
395        [1 , (0. , 0.757  , 0.587)]]
396    mol.basis = 'cc-pvdz'
397    mol.build()
398
399    mf = dft.RKS(mol)
400    mf.xc = 'hf'
401    mf.kernel()
402
403    gw = GWExact(mf)
404    gw.kernel()
405    print(gw.mo_energy)
406# [-20.10555941  -1.2264133   -0.68160937  -0.53066324  -0.44679866
407#    0.17291986   0.24457082   0.74758225   0.80045126   1.11748749
408#    1.15083528   1.19081826   1.40406946   1.43593671   1.63324755
409#    1.79839248   1.88459324   2.31461977   2.48839545   3.26047431
410#    3.32486673   3.49601314   3.77699489   4.14575936]
411
412    nocc = mol.nelectron//2
413
414    gw.linearized = True
415    gw.kernel(orbs=[nocc-1,nocc])
416    print(gw.mo_energy[nocc-1] - -0.44684106)
417    print(gw.mo_energy[nocc] - 0.17292032)
418
419