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-unrestricted G0W0 approximation with analytic continuation 21 22This implementation has N^4 scaling, and is faster than GW-CD (N^4) 23and analytic GW (N^6) methods. 24GW-AC is recommended for valence states only, and is inaccuarate for core states. 25 26Method: 27 See T. Zhu and G.K.-L. Chan, arxiv:2007.03148 (2020) for details 28 Compute Sigma on imaginary frequency with density fitting, 29 then analytically continued to real frequency 30 31Useful References: 32 J. Chem. Theory Comput. 12, 3623-3635 (2016) 33 New J. Phys. 14, 053020 (2012) 34''' 35 36from functools import reduce 37import numpy 38import numpy as np 39import h5py 40from scipy.optimize import newton, least_squares 41 42from pyscf import lib 43from pyscf.lib import logger 44from pyscf.ao2mo import _ao2mo 45from pyscf import df, scf 46from pyscf.mp.ump2 import get_nocc, get_nmo, get_frozen_mask 47from pyscf import __config__ 48 49einsum = lib.einsum 50 51def kernel(gw, mo_energy, mo_coeff, Lpq=None, orbs=None, 52 nw=None, vhf_df=False, verbose=logger.NOTE): 53 ''' 54 GW-corrected quasiparticle orbital energies 55 Returns: 56 A list : converged, mo_energy, mo_coeff 57 ''' 58 mf = gw._scf 59 mol = gw.mol 60 if gw.frozen is None: 61 frozen = 0 62 else: 63 frozen = gw.frozen 64 65 assert (isinstance(frozen, int)) 66 67 nocca, noccb = gw.nocc 68 nmoa, nmob = gw.nmo 69 # only support frozen core 70 assert (frozen < nocca and frozen < noccb) 71 72 if Lpq is None: 73 Lpq = gw.ao2mo(mo_coeff) 74 if orbs is None: 75 orbs = range(nmoa) 76 else: 77 orbs = [x - frozen for x in orbs] 78 if orbs[0] < 0: 79 logger.warn(gw, 'GW orbs must be larger than frozen core!') 80 raise RuntimeError 81 82 # v_xc 83 v_mf = mf.get_veff() 84 vj = mf.get_j() 85 v_mf[0] = v_mf[0] - (vj[0] + vj[1]) 86 v_mf[1] = v_mf[1] - (vj[0] + vj[1]) 87 v_mf_frz = np.zeros((2, nmoa-frozen, nmob-frozen)) 88 for s in range(2): 89 v_mf_frz[s] = reduce(numpy.dot, (mo_coeff[s].T, v_mf[s], mo_coeff[s])) 90 v_mf = v_mf_frz 91 92 # v_hf from DFT/HF density 93 if vhf_df and frozen == 0: 94 # density fitting vk 95 vk = np.zeros_like(v_mf) 96 vk[0] = -einsum('Lni,Lim->nm',Lpq[0,:,:,:nocca],Lpq[0,:,:nocca,:]) 97 vk[1] = -einsum('Lni,Lim->nm',Lpq[1,:,:,:noccb],Lpq[1,:,:noccb,:]) 98 else: 99 # exact vk without density fitting 100 dm = mf.make_rdm1() 101 uhf = scf.UHF(mol) 102 vk = uhf.get_veff(mol, dm) 103 vj = uhf.get_j(mol, dm) 104 vk[0] = vk[0] - (vj[0] + vj[1]) 105 vk[1] = vk[1] - (vj[0] + vj[1]) 106 vk_frz = np.zeros((2, nmoa-frozen, nmob-frozen)) 107 for s in range(2): 108 vk_frz[s] = reduce(numpy.dot, (mo_coeff[s].T, vk[s], mo_coeff[s])) 109 vk = vk_frz 110 111 # Grids for integration on imaginary axis 112 freqs,wts = _get_scaled_legendre_roots(nw) 113 114 # Compute self-energy on imaginary axis i*[0,iw_cutoff] 115 sigmaI,omega = get_sigma_diag(gw, orbs, Lpq, freqs, wts, iw_cutoff=5.) 116 117 # Analytic continuation 118 if gw.ac == 'twopole': 119 coeff_a = AC_twopole_diag(sigmaI[0], omega[0], orbs, nocca) 120 coeff_b = AC_twopole_diag(sigmaI[1], omega[1], orbs, noccb) 121 elif gw.ac == 'pade': 122 coeff_a, omega_fit_a = AC_pade_thiele_diag(sigmaI[0], omega[0]) 123 coeff_b, omega_fit_b = AC_pade_thiele_diag(sigmaI[1], omega[1]) 124 omega_fit = np.asarray((omega_fit_a,omega_fit_b)) 125 coeff = np.asarray((coeff_a,coeff_b)) 126 127 conv = True 128 homo = max(mo_energy[0][nocca-1], mo_energy[1][noccb-1]) 129 lumo = min(mo_energy[0][nocca], mo_energy[1][noccb]) 130 ef = (homo+lumo)/2. 131 mf_mo_energy = mo_energy.copy() 132 mo_energy = np.zeros_like(np.asarray(gw._scf.mo_energy)) 133 for s in range(2): 134 for p in orbs: 135 if gw.linearized: 136 # linearized G0W0 137 de = 1e-6 138 ep = mf_mo_energy[s][p] 139 #TODO: analytic sigma derivative 140 if gw.ac == 'twopole': 141 sigmaR = two_pole(ep-ef, coeff[s,:,p-orbs[0]]).real 142 dsigma = two_pole(ep-ef+de, coeff[s,:,p-orbs[0]]).real - sigmaR.real 143 elif gw.ac == 'pade': 144 sigmaR = pade_thiele(ep-ef, omega_fit[s,p-orbs[0]], coeff[s,:,p-orbs[0]]).real 145 dsigma = pade_thiele(ep-ef+de, omega_fit[s,p-orbs[0]], coeff[s,:,p-orbs[0]]).real - sigmaR.real 146 zn = 1.0/(1.0-dsigma/de) 147 e = ep + zn*(sigmaR.real + vk[s,p,p] - v_mf[s,p,p]) 148 mo_energy[s,p+frozen] = e 149 else: 150 # self-consistently solve QP equation 151 def quasiparticle(omega): 152 if gw.ac == 'twopole': 153 sigmaR = two_pole(omega-ef, coeff[s,:,p-orbs[0]]).real 154 elif gw.ac == 'pade': 155 sigmaR = pade_thiele(omega-ef, omega_fit[s,p-orbs[0]], coeff[s,:,p-orbs[0]]).real 156 return omega - mf_mo_energy[s][p] - (sigmaR.real + vk[s,p,p] - v_mf[s,p,p]) 157 try: 158 e = newton(quasiparticle, mf_mo_energy[s][p], tol=1e-6, maxiter=100) 159 mo_energy[s,p+frozen] = e 160 except RuntimeError: 161 conv = False 162 mo_coeff = gw._scf.mo_coeff 163 164 if gw.verbose >= logger.DEBUG: 165 numpy.set_printoptions(threshold=nmoa) 166 logger.debug(gw, ' GW mo_energy spin-up =\n%s', mo_energy[0]) 167 logger.debug(gw, ' GW mo_energy spin-down =\n%s', mo_energy[1]) 168 numpy.set_printoptions(threshold=1000) 169 170 return conv, mo_energy, mo_coeff 171 172def get_rho_response(omega, mo_energy, Lpqa, Lpqb): 173 ''' 174 Compute density response function in auxiliary basis at freq iw 175 ''' 176 naux, nocca, nvira = Lpqa.shape 177 naux, noccb, nvirb = Lpqb.shape 178 eia_a = mo_energy[0,:nocca,None] - mo_energy[0,None,nocca:] 179 eia_b = mo_energy[1,:noccb,None] - mo_energy[1,None,noccb:] 180 eia_a = eia_a/(omega**2+eia_a*eia_a) 181 eia_b = eia_b/(omega**2+eia_b*eia_b) 182 Pia_a = einsum('Pia,ia->Pia',Lpqa,eia_a) 183 Pia_b = einsum('Pia,ia->Pia',Lpqb,eia_b) 184 # Response from both spin-up and spin-down density 185 Pi = 2.* (einsum('Pia,Qia->PQ',Pia_a,Lpqa) + einsum('Pia,Qia->PQ',Pia_b,Lpqb)) 186 187 return Pi 188 189def get_sigma_diag(gw, orbs, Lpq, freqs, wts, iw_cutoff=None): 190 ''' 191 Compute GW correlation self-energy (diagonal elements) 192 in MO basis on imaginary axis 193 ''' 194 mo_energy = _mo_energy_without_core(gw, gw._scf.mo_energy) 195 nocca, noccb = gw.nocc 196 nmoa, nmob = gw.nmo 197 nw = len(freqs) 198 naux = Lpq[0].shape[0] 199 norbs = len(orbs) 200 201 # TODO: Treatment of degeneracy 202 homo = max(mo_energy[0][nocca-1], mo_energy[1][noccb-1]) 203 lumo = min(mo_energy[0][nocca], mo_energy[1][noccb]) 204 if (lumo-homo) < 1e-3: 205 logger.warn(gw, 'GW not well-defined for degeneracy!') 206 ef = (homo+lumo)/2. 207 208 # Integration on numerical grids 209 if iw_cutoff is not None: 210 nw_sigma = sum(iw < iw_cutoff for iw in freqs) + 1 211 else: 212 nw_sigma = nw + 1 213 214 # Compute occ for -iw and vir for iw separately 215 # to avoid branch cuts in analytic continuation 216 omega_occ = np.zeros((nw_sigma), dtype=np.complex128) 217 omega_vir = np.zeros((nw_sigma), dtype=np.complex128) 218 omega_occ[0] = 0.j 219 omega_vir[0] = 0.j 220 omega_occ[1:] = -1j*freqs[:(nw_sigma-1)] 221 omega_vir[1:] = 1j*freqs[:(nw_sigma-1)] 222 orbs_occ_a = [i for i in orbs if i < nocca] 223 orbs_occ_b = [i for i in orbs if i < noccb] 224 norbs_occ_a = len(orbs_occ_a) 225 norbs_occ_b = len(orbs_occ_b) 226 227 emo_occ_a = omega_occ[None,:] + ef - mo_energy[0,:,None] 228 emo_occ_b = omega_occ[None,:] + ef - mo_energy[1,:,None] 229 emo_vir_a = omega_vir[None,:] + ef - mo_energy[0,:,None] 230 emo_vir_b = omega_vir[None,:] + ef - mo_energy[1,:,None] 231 232 sigma = np.zeros((2,norbs,nw_sigma),dtype=np.complex128) 233 omega = np.zeros((2,norbs,nw_sigma),dtype=np.complex128) 234 for s in range(2): 235 for p in range(norbs): 236 orbp = orbs[p] 237 if orbp < gw.nocc[s]: 238 omega[s,p] = omega_occ.copy() 239 else: 240 omega[s,p] = omega_vir.copy() 241 242 for w in range(nw): 243 Pi = get_rho_response(freqs[w], mo_energy, Lpq[0,:,:nocca,nocca:], Lpq[1,:,:noccb,noccb:]) 244 Pi_inv = np.linalg.inv(np.eye(naux)-Pi)-np.eye(naux) 245 g0_occ_a = wts[w] * emo_occ_a / (emo_occ_a**2+freqs[w]**2) 246 g0_vir_a = wts[w] * emo_vir_a / (emo_vir_a**2+freqs[w]**2) 247 g0_occ_b = wts[w] * emo_occ_b / (emo_occ_b**2+freqs[w]**2) 248 g0_vir_b = wts[w] * emo_vir_b / (emo_vir_b**2+freqs[w]**2) 249 250 Qnm_a = einsum('Pnm,PQ->Qnm',Lpq[0][:,orbs,:],Pi_inv) 251 Qnm_b = einsum('Pnm,PQ->Qnm',Lpq[1][:,orbs,:],Pi_inv) 252 Wmn_a = einsum('Qnm,Qmn->mn',Qnm_a,Lpq[0][:,:,orbs]) 253 Wmn_b = einsum('Qnm,Qmn->mn',Qnm_b,Lpq[1][:,:,orbs]) 254 255 sigma[0,:norbs_occ_a] += -einsum('mn,mw->nw',Wmn_a[:,:norbs_occ_a],g0_occ_a)/np.pi 256 sigma[0,norbs_occ_a:] += -einsum('mn,mw->nw',Wmn_a[:,norbs_occ_a:],g0_vir_a)/np.pi 257 sigma[1,:norbs_occ_b] += -einsum('mn,mw->nw',Wmn_b[:,:norbs_occ_b],g0_occ_b)/np.pi 258 sigma[1,norbs_occ_b:] += -einsum('mn,mw->nw',Wmn_b[:,norbs_occ_b:],g0_vir_b)/np.pi 259 260 return sigma, omega 261 262def _get_scaled_legendre_roots(nw): 263 """ 264 Scale nw Legendre roots, which lie in the 265 interval [-1, 1], so that they lie in [0, inf) 266 Ref: www.cond-mat.de/events/correl19/manuscripts/ren.pdf 267 268 Returns: 269 freqs : 1D ndarray 270 wts : 1D ndarray 271 """ 272 freqs, wts = np.polynomial.legendre.leggauss(nw) 273 x0 = 0.5 274 freqs_new = x0*(1.+freqs)/(1.-freqs) 275 wts = wts*2.*x0/(1.-freqs)**2 276 return freqs_new, wts 277 278def _get_clenshaw_curtis_roots(nw): 279 """ 280 Clenshaw-Curtis qaudrature on [0,inf) 281 Ref: J. Chem. Phys. 132, 234114 (2010) 282 Returns: 283 freqs : 1D ndarray 284 wts : 1D ndarray 285 """ 286 freqs = np.zeros(nw) 287 wts = np.zeros(nw) 288 a = 0.2 289 for w in range(nw): 290 t = (w+1.0)/nw * np.pi/2. 291 freqs[w] = a / np.tan(t) 292 if w != nw-1: 293 wts[w] = a*np.pi/2./nw/(np.sin(t)**2) 294 else: 295 wts[w] = a*np.pi/4./nw/(np.sin(t)**2) 296 return freqs[::-1], wts[::-1] 297 298def two_pole_fit(coeff, omega, sigma): 299 cf = coeff[:5] + 1j*coeff[5:] 300 f = cf[0] + cf[1]/(omega+cf[3]) + cf[2]/(omega+cf[4]) - sigma 301 f[0] = f[0]/0.01 302 return np.array([f.real,f.imag]).reshape(-1) 303 304def two_pole(freqs, coeff): 305 cf = coeff[:5] + 1j*coeff[5:] 306 return cf[0] + cf[1]/(freqs+cf[3]) + cf[2]/(freqs+cf[4]) 307 308def AC_twopole_diag(sigma, omega, orbs, nocc): 309 """ 310 Analytic continuation to real axis using a two-pole model 311 Returns: 312 coeff: 2D array (ncoeff, norbs) 313 """ 314 norbs, nw = sigma.shape 315 coeff = np.zeros((10,norbs)) 316 for p in range(norbs): 317 # target = np.array([sigma[p].real,sigma[p].imag]).reshape(-1) 318 if orbs[p] < nocc: 319 x0 = np.array([0, 1, 1, 1, -1, 0, 0, 0, -1.0, -0.5]) 320 else: 321 x0 = np.array([0, 1, 1, 1, -1, 0, 0, 0, 1.0, 0.5]) 322 #TODO: analytic gradient 323 xopt = least_squares(two_pole_fit, x0, jac='3-point', method='trf', xtol=1e-10, 324 gtol = 1e-10, max_nfev=1000, verbose=0, args=(omega[p], sigma[p])) 325 if xopt.success is False: 326 print('WARN: 2P-Fit Orb %d not converged, cost function %e'%(p,xopt.cost)) 327 coeff[:,p] = xopt.x.copy() 328 return coeff 329 330def thiele(fn,zn): 331 nfit = len(zn) 332 g = np.zeros((nfit,nfit),dtype=np.complex128) 333 g[:,0] = fn.copy() 334 for i in range(1,nfit): 335 g[i:,i] = (g[i-1,i-1]-g[i:,i-1])/((zn[i:]-zn[i-1])*g[i:,i-1]) 336 a = g.diagonal() 337 return a 338 339def pade_thiele(freqs,zn,coeff): 340 nfit = len(coeff) 341 X = coeff[-1]*(freqs-zn[-2]) 342 for i in range(nfit-1): 343 idx = nfit-i-1 344 X = coeff[idx]*(freqs-zn[idx-1])/(1.+X) 345 X = coeff[0]/(1.+X) 346 return X 347 348def AC_pade_thiele_diag(sigma, omega): 349 """ 350 Analytic continuation to real axis using a Pade approximation 351 from Thiele's reciprocal difference method 352 Reference: J. Low Temp. Phys. 29, 179 (1977) 353 Returns: 354 coeff: 2D array (ncoeff, norbs) 355 omega: 2D array (norbs, npade) 356 """ 357 idx = range(1,40,6) 358 sigma1 = sigma[:,idx].copy() 359 sigma2 = sigma[:,(idx[-1]+4)::4].copy() 360 sigma = np.hstack((sigma1,sigma2)) 361 omega1 = omega[:,idx].copy() 362 omega2 = omega[:,(idx[-1]+4)::4].copy() 363 omega = np.hstack((omega1,omega2)) 364 norbs, nw = sigma.shape 365 npade = nw // 2 366 coeff = np.zeros((npade*2,norbs),dtype=np.complex128) 367 for p in range(norbs): 368 coeff[:,p] = thiele(sigma[p,:npade*2], omega[p,:npade*2]) 369 370 return coeff, omega[:,:npade*2] 371 372def _mo_energy_without_core(gw, mo_energy): 373 moidx = get_frozen_mask(gw) 374 mo_energy = (mo_energy[0][moidx[0]], mo_energy[1][moidx[1]]) 375 return np.asarray(mo_energy) 376 377def _mo_without_core(gw, mo): 378 moidx = get_frozen_mask(gw) 379 mo = (mo[0][:,moidx[0]], mo[1][:,moidx[1]]) 380 return np.asarray(mo) 381 382class UGWAC(lib.StreamObject): 383 384 linearized = getattr(__config__, 'gw_ugw_UGW_linearized', False) 385 # Analytic continuation: pade or twopole 386 ac = getattr(__config__, 'gw_ugw_UGW_ac', 'pade') 387 388 def __init__(self, mf, frozen=None): 389 self.mol = mf.mol 390 self._scf = mf 391 self.verbose = self.mol.verbose 392 self.stdout = self.mol.stdout 393 self.max_memory = mf.max_memory 394 395 self.frozen = frozen 396 397 # DF-GW must use density fitting integrals 398 if getattr(mf, 'with_df', None): 399 self.with_df = mf.with_df 400 else: 401 self.with_df = df.DF(mf.mol) 402 self.with_df.auxbasis = df.make_auxbasis(mf.mol, mp2fit=True) 403 self._keys.update(['with_df']) 404 405################################################## 406# don't modify the following attributes, they are not input options 407 self._nocc = None 408 self._nmo = None 409 # self.mo_energy: GW quasiparticle energy, not scf mo_energy 410 self.mo_energy = None 411 self.mo_coeff = mf.mo_coeff 412 self.mo_occ = mf.mo_occ 413 self.sigma = None 414 415 keys = set(('linearized','ac')) 416 self._keys = set(self.__dict__.keys()).union(keys) 417 418 def dump_flags(self): 419 log = logger.Logger(self.stdout, self.verbose) 420 log.info('') 421 log.info('******** %s ********', self.__class__) 422 log.info('method = %s', self.__class__.__name__) 423 nocca, noccb = self.nocc 424 nmoa, nmob = self.nmo 425 nvira = nmoa - nocca 426 nvirb = nmob - noccb 427 log.info('GW (nocca, noccb) = (%d, %d), (nvira, nvirb) = (%d, %d)', 428 nocca, noccb, nvira, nvirb) 429 if self.frozen is not None: 430 log.info('frozen orbitals %s', str(self.frozen)) 431 logger.info(self, 'use perturbative linearized QP eqn = %s', self.linearized) 432 logger.info(self, 'analytic continuation method = %s', self.ac) 433 return self 434 435 @property 436 def nocc(self): 437 return self.get_nocc() 438 @nocc.setter 439 def nocc(self, n): 440 self._nocc = n 441 442 @property 443 def nmo(self): 444 return self.get_nmo() 445 @nmo.setter 446 def nmo(self, n): 447 self._nmo = n 448 449 get_nocc = get_nocc 450 get_nmo = get_nmo 451 get_frozen_mask = get_frozen_mask 452 453 def kernel(self, mo_energy=None, mo_coeff=None, Lpq=None, orbs=None, nw=100, vhf_df=False): 454 """ 455 Input: 456 orbs: self-energy orbs 457 nw: grid number 458 vhf_df: whether using density fitting for HF exchange 459 Output: 460 mo_energy: GW quasiparticle energy 461 """ 462 if mo_coeff is None: 463 mo_coeff = _mo_without_core(self, self._scf.mo_coeff) 464 if mo_energy is None: 465 mo_energy = _mo_energy_without_core(self, self._scf.mo_energy) 466 467 cput0 = (logger.process_clock(), logger.perf_counter()) 468 self.dump_flags() 469 self.converged, self.mo_energy, self.mo_coeff = \ 470 kernel(self, mo_energy, mo_coeff, 471 Lpq=Lpq, orbs=orbs, nw=nw, vhf_df=vhf_df, verbose=self.verbose) 472 473 logger.warn(self, 'GW QP energies may not be sorted from min to max') 474 logger.timer(self, 'GW', *cput0) 475 return self.mo_energy 476 477 def ao2mo(self, mo_coeff=None): 478 nmoa, nmob = self.nmo 479 nao = self.mo_coeff[0].shape[0] 480 naux = self.with_df.get_naoaux() 481 mem_incore = (nmoa**2*naux + nmob**2*naux + nao**2*naux) * 8/1e6 482 mem_now = lib.current_memory()[0] 483 484 moa = numpy.asarray(mo_coeff[0], order='F') 485 mob = numpy.asarray(mo_coeff[1], order='F') 486 ijslicea = (0, nmoa, 0, nmoa) 487 ijsliceb = (0, nmob, 0, nmob) 488 Lpqa = None 489 Lpqb = None 490 if (mem_incore + mem_now < 0.99*self.max_memory) or self.mol.incore_anyway: 491 Lpqa = _ao2mo.nr_e2(self.with_df._cderi, moa, ijslicea, aosym='s2', out=Lpqa) 492 Lpqb = _ao2mo.nr_e2(self.with_df._cderi, mob, ijsliceb, aosym='s2', out=Lpqb) 493 return np.asarray((Lpqa.reshape(naux,nmoa,nmoa),Lpqb.reshape(naux,nmob,nmob))) 494 else: 495 logger.warn(self, 'Memory may not be enough!') 496 raise NotImplementedError 497 498 499if __name__ == '__main__': 500 from pyscf import gto, dft 501 mol = gto.Mole() 502 mol.verbose = 4 503 mol.atom = 'O 0 0 0' 504 mol.basis = 'aug-cc-pvdz' 505 mol.spin = 2 506 mol.build() 507 508 mf = dft.UKS(mol) 509 mf.xc = 'pbe0' 510 mf.kernel() 511 512 nocca = (mol.nelectron + mol.spin)//2 513 noccb = mol.nelectron - nocca 514 nmo = len(mf.mo_energy[0]) 515 nvira = nmo - nocca 516 nvirb = nmo - noccb 517 518 gw = UGWAC(mf) 519 gw.frozen = 0 520 gw.linearized = False 521 gw.ac = 'pade' 522 gw.kernel(orbs=range(nocca-3,nocca+3)) 523 assert (abs(gw.mo_energy[0][nocca-1]- -0.521932084529) < 1e-5) 524 assert (abs(gw.mo_energy[0][nocca] -0.167547592784) < 1e-5) 525 assert (abs(gw.mo_energy[1][noccb-1]- -0.464605523684) < 1e-5) 526 assert (abs(gw.mo_energy[1][noccb]- -0.0133557793765) < 1e-5) 527