1#!/usr/bin/env python
2# Copyright 2014-2021 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'''
20PBC spin-unrestricted G0W0-AC QP eigenvalues with k-point sampling
21GW-AC is recommended for valence states only, and is inaccuarate for core states.
22
23Method:
24    See T. Zhu and G.K.-L. Chan, arxiv:2007.03148 (2020) for details
25    Compute Sigma on imaginary frequency with density fitting,
26    then analytically continued to real frequency.
27    Gaussian density fitting must be used (FFTDF and MDF are not supported).
28'''
29
30from functools import reduce
31import numpy
32import numpy as np
33import h5py
34from scipy.optimize import newton, least_squares
35
36from pyscf import lib
37from pyscf.lib import logger
38from pyscf.ao2mo import _ao2mo
39from pyscf.ao2mo.incore import _conc_mos
40from pyscf.pbc import df, dft, scf
41from pyscf.pbc.cc.kccsd_uhf import get_nocc, get_nmo, get_frozen_mask
42from pyscf import __config__
43
44einsum = lib.einsum
45
46def kernel(gw, mo_energy, mo_coeff, orbs=None,
47           kptlist=None, nw=None, verbose=logger.NOTE):
48    '''
49    GW-corrected quasiparticle orbital energies
50    Returns:
51        A list :  converged, mo_energy, mo_coeff
52    '''
53    mf = gw._scf
54    if gw.frozen is None:
55        frozen = 0
56    else:
57        frozen = gw.frozen
58    assert (frozen == 0)
59
60    nmoa, nmob = gw.nmo
61    nocca, noccb = gw.nocc
62
63    if orbs is None:
64        orbs = range(nmoa)
65    if kptlist is None:
66        kptlist = range(gw.nkpts)
67
68    nkpts = gw.nkpts
69    nklist = len(kptlist)
70
71    # v_xc
72    dm = np.array(mf.make_rdm1())
73    v_mf = np.array(mf.get_veff())
74    vj = np.array(mf.get_j(dm_kpts=dm))
75    v_mf[0] = v_mf[0] - (vj[0] + vj[1])
76    v_mf[1] = v_mf[1] - (vj[0] + vj[1])
77    for s in range(2):
78        for k in range(nkpts):
79            v_mf[s,k] = reduce(numpy.dot, (mo_coeff[s,k].T.conj(), v_mf[s,k], mo_coeff[s,k]))
80
81    # v_hf from DFT/HF density
82    if gw.fc:
83        exxdiv = 'ewald'
84    else:
85        exxdiv = None
86    uhf = scf.KUHF(gw.mol, gw.kpts, exxdiv=exxdiv)
87    uhf.with_df = gw.with_df
88    uhf.with_df._cderi = gw.with_df._cderi
89    vk = uhf.get_veff(gw.mol,dm_kpts=dm)
90    vj = uhf.get_j(gw.mol,dm_kpts=dm)
91    vk[0] = vk[0] - (vj[0] + vj[1])
92    vk[1] = vk[1] - (vj[0] + vj[1])
93    for s in range(2):
94        for k in range(nkpts):
95            vk[s,k] = reduce(numpy.dot, (mo_coeff[s,k].T.conj(), vk[s,k], mo_coeff[s,k]))
96
97    # Grids for integration on imaginary axis
98    freqs,wts = _get_scaled_legendre_roots(nw)
99
100    # Compute self-energy on imaginary axis i*[0,iw_cutoff]
101    sigmaI, omega = get_sigma_diag(gw, orbs, kptlist, freqs, wts, iw_cutoff=5.)
102
103    # Analytic continuation
104    coeff_a = []
105    coeff_b = []
106    if gw.ac == 'twopole':
107        for k in range(nklist):
108            coeff_a.append(AC_twopole_diag(sigmaI[0,k], omega[0], orbs, nocca))
109            coeff_b.append(AC_twopole_diag(sigmaI[1,k], omega[1], orbs, noccb))
110    elif gw.ac == 'pade':
111        for k in range(nklist):
112            coeff_a_tmp, omega_fit_a = AC_pade_thiele_diag(sigmaI[0,k], omega[0])
113            coeff_b_tmp, omega_fit_b = AC_pade_thiele_diag(sigmaI[1,k], omega[1])
114            coeff_a.append(coeff_a_tmp)
115            coeff_b.append(coeff_b_tmp)
116        omega_fit = np.asarray((omega_fit_a, omega_fit_b))
117    coeff = np.asarray((coeff_a, coeff_b))
118
119    conv = True
120    # This code does not support metals
121    homo = -99.
122    lumo = 99.
123    mo_energy = np.asarray(mf.mo_energy)
124    for k in range(nkpts):
125        if homo < max(mo_energy[0,k][nocca-1],mo_energy[1,k][noccb-1]):
126            homo = max(mo_energy[0,k][nocca-1],mo_energy[1,k][noccb-1])
127        if lumo > min(mo_energy[0,k][nocca],mo_energy[1,k][noccb]):
128            lumo = min(mo_energy[0,k][nocca],mo_energy[1,k][noccb])
129    ef = (homo+lumo)/2.
130
131    mo_energy = np.zeros_like(np.array(mf.mo_energy))
132    for s in range(2):
133        for k in range(nklist):
134            kn = kptlist[k]
135            for p in orbs:
136                if gw.linearized:
137                    # linearized G0W0
138                    de = 1e-6
139                    ep = mf.mo_energy[s][kn][p]
140                    #TODO: analytic sigma derivative
141                    if gw.ac == 'twopole':
142                        sigmaR = two_pole(ep-ef, coeff[s,k,:,p-orbs[0]]).real
143                        dsigma = two_pole(ep-ef+de, coeff[s,k,:,p-orbs[0]]).real - sigmaR.real
144                    elif gw.ac == 'pade':
145                        sigmaR = pade_thiele(ep-ef, omega_fit[s,p-orbs[0]], coeff[s,k,:,p-orbs[0]]).real
146                        dsigma = pade_thiele(ep-ef+de, omega_fit[s,p-orbs[0]],
147                                             coeff[s,k,:,p-orbs[0]]).real - sigmaR.real
148                    zn = 1.0/(1.0-dsigma/de)
149                    e = ep + zn*(sigmaR.real + vk[s,kn,p,p].real - v_mf[s,kn,p,p].real)
150                    mo_energy[s,kn,p] = e
151                else:
152                    # self-consistently solve QP equation
153                    def quasiparticle(omega):
154                        if gw.ac == 'twopole':
155                            sigmaR = two_pole(omega-ef, coeff[s,k,:,p-orbs[0]]).real
156                        elif gw.ac == 'pade':
157                            sigmaR = pade_thiele(omega-ef, omega_fit[s,p-orbs[0]], coeff[s,k,:,p-orbs[0]]).real
158                        return omega - mf.mo_energy[s][kn][p] - (sigmaR.real + vk[s,kn,p,p].real - v_mf[s,kn,p,p].real)
159                    try:
160                        e = newton(quasiparticle, mf.mo_energy[s][kn][p], tol=1e-6, maxiter=100)
161                        mo_energy[s,kn,p] = e
162                    except RuntimeError:
163                        conv = False
164    mo_coeff = mf.mo_coeff
165
166    if gw.verbose >= logger.DEBUG:
167        numpy.set_printoptions(threshold=nmoa)
168        for k in range(nkpts):
169            logger.debug(gw, '  GW mo_energy spin-up @ k%d =\n%s', k,mo_energy[0,k])
170        for k in range(nkpts):
171            logger.debug(gw, '  GW mo_energy spin-down @ k%d =\n%s', k,mo_energy[1,k])
172        numpy.set_printoptions(threshold=1000)
173
174    return conv, mo_energy, mo_coeff
175
176def get_rho_response(gw, omega, mo_energy, Lpq, kL, kidx):
177    '''
178    Compute density response function in auxiliary basis at freq iw
179    '''
180    spin, nkpts, naux, nmo, nmo = Lpq.shape
181    nocca, noccb = gw.nocc
182    kpts = gw.kpts
183    kscaled = gw.mol.get_scaled_kpts(kpts)
184    kscaled -= kscaled[0]
185
186    # Compute Pi for kL
187    Pi = np.zeros((naux,naux),dtype=np.complex128)
188    for i, kpti in enumerate(kpts):
189        # Find ka that conserves with ki and kL (-ki+ka+kL=G)
190        a = kidx[i]
191        eia_a = mo_energy[0,i,:nocca,None] - mo_energy[0,a,None,nocca:]
192        eia_b = mo_energy[1,i,:noccb,None] - mo_energy[1,a,None,noccb:]
193        eia_a = eia_a/(omega**2+eia_a*eia_a)
194        eia_b = eia_b/(omega**2+eia_b*eia_b)
195        Pia_a = einsum('Pia,ia->Pia',Lpq[0,i][:,:nocca,nocca:],eia_a)
196        Pia_b = einsum('Pia,ia->Pia',Lpq[1,i][:,:noccb,noccb:],eia_b)
197        # Response from both spin-up and spin-down density
198        Pi += 2./nkpts * (einsum('Pia,Qia->PQ',Pia_a,Lpq[0,i][:,:nocca,nocca:].conj()) +
199                          einsum('Pia,Qia->PQ',Pia_b,Lpq[1,i][:,:noccb,noccb:].conj()))
200    return Pi
201
202def get_sigma_diag(gw, orbs, kptlist, freqs, wts, iw_cutoff=None, max_memory=8000):
203    '''
204    Compute GW correlation self-energy (diagonal elements) in MO basis
205    on imaginary axis
206    '''
207    mo_energy = np.array(gw._scf.mo_energy)
208    mo_coeff = np.array(gw._scf.mo_coeff)
209    nocca, noccb = gw.nocc
210    nmoa, nmob = gw.nmo
211    nkpts = gw.nkpts
212    kpts = gw.kpts
213    nklist = len(kptlist)
214    nw = len(freqs)
215    norbs = len(orbs)
216    mydf = gw.with_df
217
218    # possible kpts shift
219    kscaled = gw.mol.get_scaled_kpts(kpts)
220    kscaled -= kscaled[0]
221
222    # This code does not support metals
223    homo = -99.
224    lumo = 99.
225    for k in range(nkpts):
226        if homo < max(mo_energy[0,k][nocca-1],mo_energy[1,k][noccb-1]):
227            homo = max(mo_energy[0,k][nocca-1],mo_energy[1,k][noccb-1])
228        if lumo > min(mo_energy[0,k][nocca],mo_energy[1,k][noccb]):
229            lumo = min(mo_energy[0,k][nocca],mo_energy[1,k][noccb])
230    if (lumo-homo)<1e-3:
231        logger.warn(gw, 'Current KUGW is not supporting metals!')
232    ef = (homo+lumo)/2.
233
234    # Integration on numerical grids
235    if iw_cutoff is not None:
236        nw_sigma = sum(iw < iw_cutoff for iw in freqs) + 1
237    else:
238        nw_sigma = nw + 1
239
240    # Compute occ for -iw and vir for iw separately
241    # to avoid branch cuts in analytic continuation
242    omega_occ = np.zeros((nw_sigma),dtype=np.complex128)
243    omega_vir = np.zeros((nw_sigma),dtype=np.complex128)
244    omega_occ[1:] = -1j*freqs[:(nw_sigma-1)]
245    omega_vir[1:] = 1j*freqs[:(nw_sigma-1)]
246    orbs_occ_a = [i for i in orbs if i < nocca]
247    orbs_occ_b = [i for i in orbs if i < noccb]
248    norbs_occ_a = len(orbs_occ_a)
249    norbs_occ_b = len(orbs_occ_b)
250
251    emo_occ_a = np.zeros((nkpts,nmoa,nw_sigma),dtype=np.complex128)
252    emo_occ_b = np.zeros((nkpts,nmob,nw_sigma),dtype=np.complex128)
253    emo_vir_a = np.zeros((nkpts,nmoa,nw_sigma),dtype=np.complex128)
254    emo_vir_b = np.zeros((nkpts,nmob,nw_sigma),dtype=np.complex128)
255    for k in range(nkpts):
256        emo_occ_a[k] = omega_occ[None,:] + ef - mo_energy[0,k][:,None]
257        emo_occ_b[k] = omega_occ[None,:] + ef - mo_energy[1,k][:,None]
258        emo_vir_a[k] = omega_vir[None,:] + ef - mo_energy[0,k][:,None]
259        emo_vir_b[k] = omega_vir[None,:] + ef - mo_energy[1,k][:,None]
260
261    sigma = np.zeros((2,nklist,norbs,nw_sigma),dtype=np.complex128)
262    omega = np.zeros((2,norbs,nw_sigma),dtype=np.complex128)
263    for s in range(2):
264        for p in range(norbs):
265            orbp = orbs[p]
266            if orbp < gw.nocc[s]:
267                omega[s,p] = omega_occ.copy()
268            else:
269                omega[s,p] = omega_vir.copy()
270
271    if gw.fc:
272        # Set up q mesh for q->0 finite size correction
273        q_pts = np.array([1e-3,0,0]).reshape(1,3)
274        q_abs = gw.mol.get_abs_kpts(q_pts)
275
276        # Get qij = 1/sqrt(Omega) * < psi_{ik} | e^{iqr} | psi_{ak-q} > at q: (nkpts, nocc, nvir)
277        qij = get_qij(gw, q_abs[0], mo_coeff)
278
279    for kL in range(nkpts):
280        # Lij: (2, ki, L, i, j) for looping every kL
281        #Lij = np.zeros((2,nkpts,naux,nmoa,nmoa),dtype=np.complex128)
282        Lij = []
283        # kidx: save kj that conserves with kL and ki (-ki+kj+kL=G)
284        # kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G)
285        kidx = np.zeros((nkpts),dtype=np.int64)
286        kidx_r = np.zeros((nkpts),dtype=np.int64)
287        for i, kpti in enumerate(kpts):
288            for j, kptj in enumerate(kpts):
289                # Find (ki,kj) that satisfies momentum conservation with kL
290                kconserv = -kscaled[i] + kscaled[j] + kscaled[kL]
291                is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12
292                if is_kconserv:
293                    kidx[i] = j
294                    kidx_r[j] = i
295                    logger.debug(gw, "Read Lpq (kL: %s / %s, ki: %s, kj: %s)"%(kL+1, nkpts, i, j))
296                    Lij_out_a = None
297                    Lij_out_b = None
298                    # Read (L|pq) and ao2mo transform to (L|ij)
299                    Lpq = []
300                    for LpqR, LpqI, sign \
301                            in mydf.sr_loop([kpti, kptj], max_memory=0.1*gw._scf.max_memory, compact=False):
302                        Lpq.append(LpqR+LpqI*1.0j)
303                    Lpq = np.vstack(Lpq).reshape(-1,nmoa**2)
304                    moija, ijslicea = _conc_mos(mo_coeff[0,i], mo_coeff[0,j])[2:]
305                    moijb, ijsliceb = _conc_mos(mo_coeff[1,i], mo_coeff[1,j])[2:]
306                    tao = []
307                    ao_loc = None
308                    Lij_out_a = _ao2mo.r_e2(Lpq, moija, ijslicea, tao, ao_loc, out=Lij_out_a)
309                    tao = []
310                    ao_loc = None
311                    Lij_out_b = _ao2mo.r_e2(Lpq, moijb, ijsliceb, tao, ao_loc, out=Lij_out_b)
312                    Lij.append(np.asarray((Lij_out_a.reshape(-1,nmoa,nmoa),Lij_out_b.reshape(-1,nmob,nmob))))
313
314        Lij = np.asarray(Lij)
315        Lij = Lij.transpose(1,0,2,3,4)
316        naux = Lij.shape[2]
317
318        if kL == 0:
319            for w in range(nw):
320                # body dielectric matrix eps_body
321                Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx)
322                eps_body_inv = np.linalg.inv(np.eye(naux)-Pi)
323
324                if gw.fc:
325                    # head dielectric matrix eps_00
326                    Pi_00 = get_rho_response_head(gw, freqs[w], mo_energy, qij)
327                    eps_00 = 1. - 4. * np.pi/np.linalg.norm(q_abs[0])**2 * Pi_00
328
329                    # wings dielectric matrix eps_P0
330                    Pi_P0 = get_rho_response_wing(gw, freqs[w], mo_energy, Lij, qij)
331                    eps_P0 = -np.sqrt(4.*np.pi) / np.linalg.norm(q_abs[0]) * Pi_P0
332
333                    # inverse dielectric matrix
334                    eps_inv_00 = 1./(eps_00 - np.dot(np.dot(eps_P0.conj(),eps_body_inv),eps_P0))
335                    eps_inv_P0 = -eps_inv_00 * np.dot(eps_body_inv, eps_P0)
336
337                    # head correction
338                    Del_00 = 2./np.pi * (6.*np.pi**2/gw.mol.vol/nkpts)**(1./3.) * (eps_inv_00 - 1.)
339
340                eps_inv_PQ = eps_body_inv
341                g0_occ_a = wts[w] * emo_occ_a / (emo_occ_a**2+freqs[w]**2)
342                g0_occ_b = wts[w] * emo_occ_b / (emo_occ_b**2+freqs[w]**2)
343                g0_vir_a = wts[w] * emo_vir_a / (emo_vir_a**2+freqs[w]**2)
344                g0_vir_b = wts[w] * emo_vir_b / (emo_vir_b**2+freqs[w]**2)
345                for k in range(nklist):
346                    kn = kptlist[k]
347                    # Find km that conserves with kn and kL (-km+kn+kL=G)
348                    km = kidx_r[kn]
349                    Qmn_a = einsum('Pmn,PQ->Qmn',Lij[0,km][:,:,orbs].conj(),eps_inv_PQ-np.eye(naux))
350                    Qmn_b = einsum('Pmn,PQ->Qmn',Lij[1,km][:,:,orbs].conj(),eps_inv_PQ-np.eye(naux))
351                    Wmn_a = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn_a,Lij[0,km][:,:,orbs])
352                    Wmn_b = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn_b,Lij[1,km][:,:,orbs])
353
354                    sigma[0,k][:norbs_occ_a] += -einsum('mn,mw->nw',Wmn_a[:,:norbs_occ_a],g0_occ_a[km])/np.pi
355                    sigma[1,k][:norbs_occ_b] += -einsum('mn,mw->nw',Wmn_b[:,:norbs_occ_b],g0_occ_b[km])/np.pi
356                    sigma[0,k][norbs_occ_a:] += -einsum('mn,mw->nw',Wmn_a[:,norbs_occ_a:],g0_vir_a[km])/np.pi
357                    sigma[1,k][norbs_occ_b:] += -einsum('mn,mw->nw',Wmn_b[:,norbs_occ_b:],g0_vir_b[km])/np.pi
358
359                    if gw.fc:
360                        # apply head correction
361                        assert(kn == km)
362                        sigma[0,k][:norbs_occ_a] += -Del_00 * g0_occ_a[kn][orbs][:norbs_occ_a] /np.pi
363                        sigma[0,k][norbs_occ_a:] += -Del_00 * g0_vir_a[kn][orbs][norbs_occ_a:] /np.pi
364                        sigma[1,k][:norbs_occ_b] += -Del_00 * g0_occ_b[kn][orbs][:norbs_occ_b] /np.pi
365                        sigma[1,k][norbs_occ_b:] += -Del_00 * g0_vir_b[kn][orbs][norbs_occ_b:] /np.pi
366
367                        # apply wing correction
368                        Wn_P0_a = einsum('Pnm,P->nm',Lij[0,kn],eps_inv_P0).diagonal()
369                        Wn_P0_b = einsum('Pnm,P->nm',Lij[1,kn],eps_inv_P0).diagonal()
370                        Wn_P0_a = Wn_P0_a.real * 2.
371                        Wn_P0_b = Wn_P0_b.real * 2.
372                        Del_P0_a = np.sqrt(gw.mol.vol/4./np.pi**3) * (6.*np.pi**2/gw.mol.vol/nkpts)**(2./3.) * Wn_P0_a[orbs]  # noqa: E501
373                        Del_P0_b = np.sqrt(gw.mol.vol/4./np.pi**3) * (6.*np.pi**2/gw.mol.vol/nkpts)**(2./3.) * Wn_P0_b[orbs]  # noqa: E501
374                        sigma[0,k][:norbs_occ_a] += -einsum('n,nw->nw',Del_P0_a[:norbs_occ_a],g0_occ_a[kn][orbs][:norbs_occ_a]) /np.pi  # noqa: E501
375                        sigma[0,k][norbs_occ_a:] += -einsum('n,nw->nw',Del_P0_a[norbs_occ_a:],g0_vir_a[kn][orbs][norbs_occ_a:]) /np.pi  # noqa: E501
376                        sigma[1,k][:norbs_occ_b] += -einsum('n,nw->nw',Del_P0_b[:norbs_occ_b],g0_occ_b[kn][orbs][:norbs_occ_b]) /np.pi  # noqa: E501
377                        sigma[1,k][norbs_occ_b:] += -einsum('n,nw->nw',Del_P0_b[norbs_occ_b:],g0_vir_b[kn][orbs][norbs_occ_b:]) /np.pi  # noqa: E501
378        else:
379            for w in range(nw):
380                Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx)
381                Pi_inv = np.linalg.inv(np.eye(naux)-Pi)-np.eye(naux)
382                g0_occ_a = wts[w] * emo_occ_a / (emo_occ_a**2+freqs[w]**2)
383                g0_occ_b = wts[w] * emo_occ_b / (emo_occ_b**2+freqs[w]**2)
384                g0_vir_a = wts[w] * emo_vir_a / (emo_vir_a**2+freqs[w]**2)
385                g0_vir_b = wts[w] * emo_vir_b / (emo_vir_b**2+freqs[w]**2)
386                for k in range(nklist):
387                    kn = kptlist[k]
388                    # Find km that conserves with kn and kL (-km+kn+kL=G)
389                    km = kidx_r[kn]
390                    Qmn_a = einsum('Pmn,PQ->Qmn',Lij[0,km][:,:,orbs].conj(),Pi_inv)
391                    Qmn_b = einsum('Pmn,PQ->Qmn',Lij[1,km][:,:,orbs].conj(),Pi_inv)
392                    Wmn_a = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn_a,Lij[0,km][:,:,orbs])
393                    Wmn_b = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn_b,Lij[1,km][:,:,orbs])
394
395                    sigma[0,k][:norbs_occ_a] += -einsum('mn,mw->nw',Wmn_a[:,:norbs_occ_a],g0_occ_a[km])/np.pi
396                    sigma[1,k][:norbs_occ_b] += -einsum('mn,mw->nw',Wmn_b[:,:norbs_occ_b],g0_occ_b[km])/np.pi
397                    sigma[0,k][norbs_occ_a:] += -einsum('mn,mw->nw',Wmn_a[:,norbs_occ_a:],g0_vir_a[km])/np.pi
398                    sigma[1,k][norbs_occ_b:] += -einsum('mn,mw->nw',Wmn_b[:,norbs_occ_b:],g0_vir_b[km])/np.pi
399
400    return sigma, omega
401
402def get_rho_response_head(gw, omega, mo_energy, qij):
403    '''
404    Compute head (G=0, G'=0) density response function in auxiliary basis at freq iw
405    '''
406    qij_a, qij_b = qij
407    nocca, noccb = gw.nocc
408    kpts = gw.kpts
409    nkpts = len(kpts)
410
411    # Compute Pi head
412    Pi_00 = 0j
413    for i, kpti in enumerate(kpts):
414        eia_a = mo_energy[0,i,:nocca,None] - mo_energy[0,i,None,nocca:]
415        eia_b = mo_energy[1,i,:noccb,None] - mo_energy[1,i,None,noccb:]
416        eia_a = eia_a/(omega**2+eia_a*eia_a)
417        eia_b = eia_b/(omega**2+eia_b*eia_b)
418        Pi_00 += 2./nkpts * (einsum('ia,ia->',eia_a,qij_a[i].conj()*qij_a[i]) +
419                             einsum('ia,ia->',eia_b,qij_b[i].conj()*qij_b[i]))
420    return Pi_00
421
422def get_rho_response_wing(gw, omega, mo_energy, Lpq, qij):
423    '''
424    Compute wing (G=P, G'=0) density response function in auxiliary basis at freq iw
425    '''
426    qij_a, qij_b = qij
427    spin, nkpts, naux, nmo, nmo = Lpq.shape
428    nocca, noccb = gw.nocc
429    kpts = gw.kpts
430    nkpts = len(kpts)
431
432    # Compute Pi wing
433    Pi = np.zeros(naux,dtype=np.complex128)
434    for i, kpti in enumerate(kpts):
435        eia_a = mo_energy[0,i,:nocca,None] - mo_energy[0,i,None,nocca:]
436        eia_b = mo_energy[1,i,:noccb,None] - mo_energy[1,i,None,noccb:]
437        eia_a = eia_a/(omega**2+eia_a*eia_a)
438        eia_b = eia_b/(omega**2+eia_b*eia_b)
439        eia_q_a = eia_a * qij_a[i].conj()
440        eia_q_b = eia_b * qij_b[i].conj()
441        Pi += 2./nkpts * (einsum('Pia,ia->P',Lpq[0,i][:,:nocca,nocca:],eia_q_a) +
442                          einsum('Pia,ia->P',Lpq[1,i][:,:noccb,noccb:],eia_q_b))
443    return Pi
444
445def get_qij(gw, q, mo_coeff, uniform_grids=False):
446    '''
447    Compute qij = 1/Omega * |< psi_{ik} | e^{iqr} | psi_{ak-q} >|^2 at q: (nkpts, nocc, nvir)
448    through kp perturbtation theory
449    Ref: Phys. Rev. B 83, 245122 (2011)
450    '''
451    nocca, noccb = gw.nocc
452    nmoa, nmob = gw.nmo
453    nvira = nmoa - nocca
454    nvirb = nmob - noccb
455    kpts = gw.kpts
456    nkpts = len(kpts)
457    cell = gw.mol
458    mo_energy = np.asarray(gw._scf.mo_energy)
459
460    if uniform_grids:
461        mydf = df.FFTDF(cell, kpts=kpts)
462        coords = cell.gen_uniform_grids(mydf.mesh)
463    else:
464        coords, weights = dft.gen_grid.get_becke_grids(cell,level=4)
465    ngrid = len(coords)
466
467    qij_a = np.zeros((nkpts,nocca,nvira),dtype=np.complex128)
468    qij_b = np.zeros((nkpts,noccb,nvirb),dtype=np.complex128)
469    for i, kpti in enumerate(kpts):
470        ao_p = dft.numint.eval_ao(cell, coords, kpt=kpti, deriv=1)
471        ao = ao_p[0]
472        ao_grad = ao_p[1:4]
473        if uniform_grids:
474            ao_ao_grad = einsum('mg,xgn->xmn',ao.T.conj(),ao_grad) * cell.vol / ngrid
475        else:
476            ao_ao_grad = einsum('g,mg,xgn->xmn',weights,ao.T.conj(),ao_grad)
477        q_ao_ao_grad = -1j * einsum('x,xmn->mn',q,ao_ao_grad)
478        q_mo_mo_grad_a = np.dot(np.dot(mo_coeff[0,i][:,:nocca].T.conj(), q_ao_ao_grad), mo_coeff[0,i][:,nocca:])
479        q_mo_mo_grad_b = np.dot(np.dot(mo_coeff[1,i][:,:noccb].T.conj(), q_ao_ao_grad), mo_coeff[1,i][:,noccb:])
480        enm_a = 1./(mo_energy[0,i][nocca:,None] - mo_energy[0,i][None,:nocca])
481        enm_b = 1./(mo_energy[1,i][noccb:,None] - mo_energy[1,i][None,:noccb])
482        dens_a = enm_a.T * q_mo_mo_grad_a
483        dens_b = enm_b.T * q_mo_mo_grad_b
484        qij_a[i] = dens_a / np.sqrt(cell.vol)
485        qij_b[i] = dens_b / np.sqrt(cell.vol)
486
487    return (qij_a, qij_b)
488
489def _get_scaled_legendre_roots(nw):
490    """
491    Scale nw Legendre roots, which lie in the
492    interval [-1, 1], so that they lie in [0, inf)
493    Ref: www.cond-mat.de/events/correl19/manuscripts/ren.pdf
494
495    Returns:
496        freqs : 1D ndarray
497        wts : 1D ndarray
498    """
499    freqs, wts = np.polynomial.legendre.leggauss(nw)
500    x0 = 0.5
501    freqs_new = x0*(1.+freqs)/(1.-freqs)
502    wts = wts*2.*x0/(1.-freqs)**2
503    return freqs_new, wts
504
505def _get_clenshaw_curtis_roots(nw):
506    """
507    Clenshaw-Curtis qaudrature on [0,inf)
508    Ref: J. Chem. Phys. 132, 234114 (2010)
509    Returns:
510        freqs : 1D ndarray
511        wts : 1D ndarray
512    """
513    freqs = np.zeros(nw)
514    wts = np.zeros(nw)
515    a = 0.2
516    for w in range(nw):
517        t = (w+1.0)/nw * np.pi/2.
518        freqs[w] = a / np.tan(t)
519        if w != nw-1:
520            wts[w] = a*np.pi/2./nw/(np.sin(t)**2)
521        else:
522            wts[w] = a*np.pi/4./nw/(np.sin(t)**2)
523    return freqs[::-1], wts[::-1]
524
525def two_pole_fit(coeff, omega, sigma):
526    cf = coeff[:5] + 1j*coeff[5:]
527    f = cf[0] + cf[1]/(omega+cf[3]) + cf[2]/(omega+cf[4]) - sigma
528    f[0] = f[0]/0.01
529    return np.array([f.real,f.imag]).reshape(-1)
530
531def two_pole(freqs, coeff):
532    cf = coeff[:5] + 1j*coeff[5:]
533    return cf[0] + cf[1]/(freqs+cf[3]) + cf[2]/(freqs+cf[4])
534
535def AC_twopole_diag(sigma, omega, orbs, nocc):
536    """
537    Analytic continuation to real axis using a two-pole model
538    Returns:
539        coeff: 2D array (ncoeff, norbs)
540    """
541    norbs, nw = sigma.shape
542    coeff = np.zeros((10,norbs))
543    for p in range(norbs):
544        if orbs[p] < nocc:
545            x0 = np.array([0, 1, 1, 1, -1, 0, 0, 0, -1.0, -0.5])
546        else:
547            x0 = np.array([0, 1, 1, 1, -1, 0, 0, 0, 1.0, 0.5])
548        #TODO: analytic gradient
549        xopt = least_squares(two_pole_fit, x0, jac='3-point', method='trf', xtol=1e-10,
550                             gtol = 1e-10, max_nfev=2000, verbose=0, args=(omega[p], sigma[p]))
551        if xopt.success is False:
552            print('WARN: 2P-Fit Orb %d not converged, cost function %e'%(p,xopt.cost))
553        coeff[:,p] = xopt.x.copy()
554    return coeff
555
556def thiele(fn,zn):
557    nfit = len(zn)
558    g = np.zeros((nfit,nfit),dtype=np.complex128)
559    g[:,0] = fn.copy()
560    for i in range(1,nfit):
561        g[i:,i] = (g[i-1,i-1]-g[i:,i-1])/((zn[i:]-zn[i-1])*g[i:,i-1])
562    a = g.diagonal()
563    return a
564
565def pade_thiele(freqs,zn,coeff):
566    nfit = len(coeff)
567    X = coeff[-1]*(freqs-zn[-2])
568    for i in range(nfit-1):
569        idx = nfit-i-1
570        X = coeff[idx]*(freqs-zn[idx-1])/(1.+X)
571    X = coeff[0]/(1.+X)
572    return X
573
574def AC_pade_thiele_diag(sigma, omega):
575    """
576    Analytic continuation to real axis using a Pade approximation
577    from Thiele's reciprocal difference method
578    Reference: J. Low Temp. Phys. 29, 179 (1977)
579    Returns:
580        coeff: 2D array (ncoeff, norbs)
581        omega: 2D array (norbs, npade)
582    """
583    idx = range(1,40,6)
584    sigma1 = sigma[:,idx].copy()
585    sigma2 = sigma[:,(idx[-1]+4)::4].copy()
586    sigma = np.hstack((sigma1,sigma2))
587    omega1 = omega[:,idx].copy()
588    omega2 = omega[:,(idx[-1]+4)::4].copy()
589    omega = np.hstack((omega1,omega2))
590    norbs, nw = sigma.shape
591    npade = nw // 2
592    coeff = np.zeros((npade*2,norbs),dtype=np.complex128)
593    for p in range(norbs):
594        coeff[:,p] = thiele(sigma[p,:npade*2], omega[p,:npade*2])
595
596    return coeff, omega[:,:npade*2]
597
598
599class KUGWAC(lib.StreamObject):
600
601    linearized = getattr(__config__, 'gw_gw_GW_linearized', False)
602    # Analytic continuation: pade or twopole
603    ac = getattr(__config__, 'gw_gw_GW_ac', 'pade')
604    # Whether applying finite size corrections
605    fc = getattr(__config__, 'gw_gw_GW_fc', True)
606
607    def __init__(self, mf, frozen=0):
608        self.mol = mf.mol
609        self._scf = mf
610        self.verbose = self.mol.verbose
611        self.stdout = self.mol.stdout
612        self.max_memory = mf.max_memory
613
614        #TODO: implement frozen orbs
615        if frozen > 0:
616            raise NotImplementedError
617        self.frozen = frozen
618
619        # DF-KGW must use GDF integrals
620        if getattr(mf, 'with_df', None):
621            self.with_df = mf.with_df
622        else:
623            raise NotImplementedError
624        self._keys.update(['with_df'])
625
626##################################################
627# don't modify the following attributes, they are not input options
628        self._nocc = None
629        self._nmo = None
630        self.kpts = mf.kpts
631        self.nkpts = len(self.kpts)
632        # self.mo_energy: GW quasiparticle energy, not scf mo_energy
633        self.mo_energy = None
634        self.mo_coeff = mf.mo_coeff
635        self.mo_occ = mf.mo_occ
636        self.sigma = None
637
638        keys = set(('linearized','ac','fc'))
639        self._keys = set(self.__dict__.keys()).union(keys)
640
641    def dump_flags(self):
642        log = logger.Logger(self.stdout, self.verbose)
643        log.info('')
644        log.info('******** %s ********', self.__class__)
645        log.info('method = %s', self.__class__.__name__)
646        nocca, noccb = self.nocc
647        nmoa, nmob = self.nmo
648        nvira = nmoa - nocca
649        nvirb = nmob - noccb
650        nkpts = self.nkpts
651        log.info('GW (nocca, noccb) = (%d, %d), (nvira, nvirb) = (%d, %d), nkpts = %d',
652                 nocca, noccb, nvira, nvirb, nkpts)
653        if self.frozen is not None:
654            log.info('frozen orbitals %s', str(self.frozen))
655        logger.info(self, 'use perturbative linearized QP eqn = %s', self.linearized)
656        logger.info(self, 'analytic continuation method = %s', self.ac)
657        logger.info(self, 'GW finite size corrections = %s', self.fc)
658        return self
659
660    @property
661    def nocc(self):
662        return self.get_nocc()
663    @nocc.setter
664    def nocc(self, n):
665        self._nocc = n
666
667    @property
668    def nmo(self):
669        return self.get_nmo()
670    @nmo.setter
671    def nmo(self, n):
672        self._nmo = n
673
674    get_nocc = get_nocc
675    get_nmo = get_nmo
676    get_frozen_mask = get_frozen_mask
677
678    def kernel(self, mo_energy=None, mo_coeff=None, orbs=None, kptlist=None, nw=100):
679        """
680        Input:
681            kptlist: self-energy k-points
682            orbs: self-energy orbs
683            nw: grid number
684        Output:
685            mo_energy: GW quasiparticle energy
686        """
687        if mo_coeff is None:
688            mo_coeff = np.array(self._scf.mo_coeff)
689        if mo_energy is None:
690            mo_energy = np.array(self._scf.mo_energy)
691
692        nmoa, nmob = self.nmo
693        naux = self.with_df.get_naoaux()
694        nkpts = self.nkpts
695        mem_incore = (3*nkpts*nmoa**2*naux) * 16/1e6
696        mem_now = lib.current_memory()[0]
697        if (mem_incore + mem_now > 0.99*self.max_memory):
698            logger.warn(self, 'Memory may not be enough!')
699            raise NotImplementedError
700
701        cput0 = (logger.process_clock(), logger.perf_counter())
702        self.dump_flags()
703        self.converged, self.mo_energy, self.mo_coeff = \
704                kernel(self, mo_energy, mo_coeff, orbs=orbs,
705                       kptlist=kptlist, nw=nw, verbose=self.verbose)
706
707        logger.warn(self, 'GW QP energies may not be sorted from min to max')
708        logger.timer(self, 'GW', *cput0)
709        return self.mo_energy
710
711if __name__ == '__main__':
712    from pyscf.pbc import gto
713    from pyscf.pbc.lib import chkfile
714    import os
715    cell = gto.Cell()
716    cell.build(
717        unit = 'B',
718        a = [[ 0.,          6.74027466,  6.74027466],
719             [ 6.74027466,  0.,          6.74027466],
720             [ 6.74027466,  6.74027466,  0.        ]],
721        atom = '''H 0 0 0
722                  H 1.68506866 1.68506866 1.68506866
723                  H 3.37013733 3.37013733 3.37013733''',
724        basis = 'gth-dzvp',
725        pseudo = 'gth-pade',
726        verbose = 5,
727        charge = 0,
728        spin = 1)
729
730    cell.spin = cell.spin * 3
731    kpts = cell.make_kpts([3,1,1],scaled_center=[0,0,0])
732    gdf = df.GDF(cell, kpts)
733    gdf_fname = 'h3_ints_311.h5'
734    gdf._cderi_to_save = gdf_fname
735    if not os.path.isfile(gdf_fname):
736        gdf.build()
737
738    chkfname = 'h_311.chk'
739    if os.path.isfile(chkfname):
740        kmf = scf.KUHF(cell, kpts, exxdiv=None)
741        kmf.with_df = gdf
742        kmf.with_df._cderi = gdf_fname
743        data = chkfile.load(chkfname, 'scf')
744        kmf.__dict__.update(data)
745    else:
746        kmf = scf.KUHF(cell, kpts, exxdiv=None)
747        kmf.with_df = gdf
748        kmf.with_df._cderi = gdf_fname
749        kmf.conv_tol = 1e-12
750        kmf.chkfile = chkfname
751        kmf.kernel()
752
753    gw = KUGWAC(kmf)
754    gw.linearized = False
755    gw.ac = 'pade'
756    gw.fc = False
757    nocca, noccb = gw.nocc
758    gw.kernel(kptlist=[0,1,2],orbs=range(0,nocca+3))
759    print(gw.mo_energy)
760    assert((abs(gw.mo_energy[0][0][nocca-1]--0.28012813))<1e-5)
761    assert((abs(gw.mo_energy[0][0][nocca]-0.13748876))<1e-5)
762    assert((abs(gw.mo_energy[0][1][nocca-1]--0.29515851))<1e-5)
763    assert((abs(gw.mo_energy[0][1][nocca]-0.14128011))<1e-5)
764    assert((abs(gw.mo_energy[1][0][noccb-1]--0.33991721))<1e-5)
765    assert((abs(gw.mo_energy[1][0][noccb]-0.10578847))<1e-5)
766    assert((abs(gw.mo_energy[1][1][noccb-1]--0.33547973))<1e-5)
767    assert((abs(gw.mo_energy[1][1][noccb]-0.08053408))<1e-5)
768
769    gw.fc = True
770    nocca, noccb = gw.nocc
771    gw.kernel(kptlist=[0,1,2],orbs=range(0,nocca+3))
772    print(gw.mo_energy)
773    assert((abs(gw.mo_energy[0][0][nocca-1]--0.40244058))<1e-5)
774    assert((abs(gw.mo_energy[0][0][nocca]-0.13618348))<1e-5)
775    assert((abs(gw.mo_energy[0][1][nocca-1]--0.41743063))<1e-5)
776    assert((abs(gw.mo_energy[0][1][nocca]-0.13997427))<1e-5)
777    assert((abs(gw.mo_energy[1][0][noccb-1]--0.46133481))<1e-5)
778    assert((abs(gw.mo_energy[1][0][noccb]-0.1044926))<1e-5)
779    assert((abs(gw.mo_energy[1][1][noccb-1]--0.4568894))<1e-5)
780    assert((abs(gw.mo_energy[1][1][noccb]-0.07922511))<1e-5)
781