def eaccsd_diag(eom, imds=None):
if imds is None: imds = eom.make_imds()
t1, t2 = imds.t1, imds.t2
nocc, nvir = t1.shape
foo = imds.eris.foo
fvv = imds.eris.fvv
Hr1 = imds.Lvv.diagonal()
Hr1 = tensor(Hr1)
if eom.partition == 'mp':
jab = fvv.diagonal()[None, :, None]
jab = jab + fvv.diagonal()[None, None, :]
jab = jab - foo.diagonal()[:, None, None]
Hr2 = tensor(jab)
else:
jab = imds.Lvv.diagonal().reshape(1, nvir, 1)
jab = jab + imds.Lvv.diagonal().reshape(1, 1, nvir)
jab = jab - imds.Loo.diagonal().reshape(nocc, 1, 1)
wab = ctf.einsum("abab->ab", imds.Wvvvv.array)
wjb = ctf.einsum('jbjb->jb', imds.Wovov.array)
wjb2 = ctf.einsum('jbbj->jb', imds.Wovvo.array)
wja = ctf.einsum('jaja->ja', imds.Wovov.array)
jab = jab + wab.reshape(1, nvir, nvir)
jab = jab - wjb.reshape(nocc, 1, nvir)
jab = jab + 2 * wjb2.reshape(nocc, 1, nvir)
jab -= ctf.einsum('jb,ab->jab', wjb2, ctf.eye(nvir))
jab = jab - wja.reshape(nocc, nvir, 1)
Hr2 = tensor(jab)
Hr2 -= 2 * lib.einsum('ijab,ijab->jab', t2, imds.Woovv)
Hr2 += lib.einsum('ijab,ijba->jab', t2, imds.Woovv)
vector = eom.amplitudes_to_vector(Hr1, Hr2)
return vector
def energy(cc, t1, t2, eris):
e = 2 * lib.einsum('ia,ia', eris.fov, t1)
tau = lib.einsum('ia,jb->ijab', t1, t1)
tau += t2
e += 2 * lib.einsum('ijab,iajb', tau, eris.ovov)
e += -lib.einsum('ijab,ibja', tau, eris.ovov)
return e.real
def ipccsd_diag(eom, imds=None):
if imds is None: imds = eom.make_imds()
t1, t2 = imds.t1, imds.t2
nocc, nvir = t1.shape
foo = imds.eris.foo
fvv = imds.eris.fvv
Hr1 = -imds.Loo.diagonal()
Hr1 = tensor(Hr1)
if eom.partition == 'mp':
ijb = -foo.diagonal().reshape(nocc, 1, 1)
ijb = ijb - foo.diagonal().reshape(1, nocc, 1)
ijb = ijb + fvv.diagonal().reshape(1, 1, nvir)
Hr2 = tensor(ijb)
else:
wij = ctf.einsum('ijij->ij', imds.Woooo.array)
wjb = ctf.einsum('jbjb->jb', imds.Wovov.array)
wjb2 = ctf.einsum('jbbj->jb', imds.Wovvo.array)
wib = ctf.einsum('ibib->ib', imds.Wovov.array)
ijb = imds.Lvv.diagonal().reshape(1, 1, nvir)
ijb = ijb - imds.Loo.diagonal().reshape(nocc, 1, 1)
ijb = ijb - imds.Loo.diagonal().reshape(1, nocc, 1)
#print(ijb.shape, wjb.shape)
ijb = ijb + wij.reshape(nocc, nocc, 1)
ijb = ijb - wjb.reshape(1, nocc, nvir)
ijb = ijb + 2 * wjb2.reshape(1, nocc, nvir)
ijb = ijb - ctf.einsum('ij,jb->ijb', ctf.eye(nocc), wjb2)
ijb = ijb - wib.reshape(nocc, 1, nvir)
Hr2 = tensor(ijb)
Hr2 -= 2. * lib.einsum('ijcb,jibc->ijb', t2, imds.Woovv)
Hr2 += lib.einsum('ijcb,ijbc->ijb', t2, imds.Woovv)
vector = eom.amplitudes_to_vector(Hr1, Hr2)
return vector
def init_amps(self, eris=None):
if eris is None: eris = self.ao2mo()
log = Logger(self.stdout, self.verbose)
t1 = eris.fov.conj() / eris.eia
t2 = eris.ovov.transpose(0, 2, 1, 3).conj() / eris.eijab
self.emp2 = 2 * lib.einsum('ijab,iajb', t2, eris.ovov)
self.emp2 -= lib.einsum('ijab,ibja', t2, eris.ovov)
log.info('Init t2, MP2 energy = %.15g', self.emp2.real)
return self.emp2, t1, t2
def energy(cc, t1, t2, eris):
log = Logger(cc.stdout, cc.verbose)
nkpts = cc.nkpts
e = 2*lib.einsum('ia,ia', eris.fov, t1)
tau = lib.einsum('ia,jb->ijab',t1,t1)
tau += t2
e += 2*lib.einsum('ijab,iajb', tau, eris.ovov)
e += -lib.einsum('ijab,ibja', tau, eris.ovov)
if abs(e.imag)>1e-4:
log.warn('Non-zero imaginary part found in KRCCSD energy %s', e)
return e.real / nkpts
def init_amps(self, eris):
time0 = time.clock(), time.time()
log = Logger(self.stdout, self.verbose)
nocc = self.nocc
nvir = self.nmo - nocc
kpts = self.kpts
nkpts = self.nkpts
gvec = self._scf.cell.reciprocal_vectors()
sym1 = ['+-', [kpts,]*2, None, gvec]
t1 = lib.zeros([nocc,nvir], eris.dtype, sym1, symlib=self.symlib, verbose=self.SYMVERBOSE)
t2 = eris.ovov.transpose(0,2,1,3).conj() / eris.eijab
self.emp2 = 2*lib.einsum('ijab,iajb', t2, eris.ovov)
self.emp2 -= lib.einsum('ijab,ibja', t2, eris.ovov)
self.emp2 = self.emp2.real/nkpts
log.info('Init t2, MP2 energy (with fock eigenvalue shift) = %.15g', self.emp2)
log.timer('init mp2', *time0)
return self.emp2, t1, t2
def update_amps(cc, t1, t2, eris):
# Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004) Eqs.(35)-(36)
cput1 = cput0 = (time.clock(), time.time())
log = Logger(cc.stdout, cc.verbose)
fov = eris.fov.copy()
foo = eris.foo.copy()
fvv = eris.fvv.copy()
Foo = imd.cc_Foo(t1, t2, eris)
Fvv = imd.cc_Fvv(t1, t2, eris)
Fov = imd.cc_Fov(t1, t2, eris)
# Move energy terms to the other side
Foo -= eris._foo
Fvv -= eris._fvv
# T1 equation
t1new = fov.conj().copy()
tmp = lib.einsum('kc,ka->ac', fov, t1)
t1new += -2 * lib.einsum('ac,ic->ia', tmp, t1)
t1new += lib.einsum('ac,ic->ia', Fvv, t1)
t1new += -lib.einsum('ki,ka->ia', Foo, t1)
t1new += 2 * lib.einsum('kc,kica->ia', Fov, t2)
t1new += -lib.einsum('kc,ikca->ia', Fov, t2)
tmp = lib.einsum('kc,ic->ki', Fov, t1)
t1new += lib.einsum('ki,ka->ia', tmp, t1)
t1new += 2 * lib.einsum('kcai,kc->ia', eris.ovvo, t1)
t1new += -lib.einsum('kiac,kc->ia', eris.oovv, t1)
t1new += 2 * lib.einsum('kdac,ikcd->ia', eris.ovvv, t2)
t1new += -lib.einsum('kcad,ikcd->ia', eris.ovvv, t2)
tmp = lib.einsum('kdac,kd->ac', eris.ovvv, t1)
t1new += 2 * lib.einsum('ac,ic->ia', tmp, t1)
tmp = lib.einsum('kcad,kd->ca', eris.ovvv, t1)
t1new += -lib.einsum('ca,ic->ia', tmp, t1)
t1new += -2 * lib.einsum('kilc,klac->ia', eris.ooov, t2)
t1new += lib.einsum('likc,klac->ia', eris.ooov, t2)
tmp = lib.einsum('kilc,lc->ki', eris.ooov, t1)
t1new += -2 * lib.einsum('ki,ka->ia', tmp, t1)
tmp = lib.einsum('likc,lc->ik', eris.ooov, t1)
t1new += lib.einsum('ik,ka->ia', tmp, t1)
cput1 = log.timer("updating t1", *cput1)
# T2 equation
t2new = eris.ovov.conj().transpose(0, 2, 1, 3).copy()
if cc.cc2:
Woooo2 = eris.oooo.transpose(0, 2, 1, 3).copy()
Woooo2 += lib.einsum('kilc,jc->klij', eris.ooov, t1)
Woooo2 += lib.einsum('ljkc,ic->klij', eris.ooov, t1)
tmp = lib.einsum('kcld,ic->kild', eris.ovov, t1)
Woooo2 += lib.einsum('kild,jd->klij', tmp, t1)
tmp = lib.einsum('klij,ka->alij', Woooo2, t1)
t2new += lib.einsum('alij,lb->ijab', tmp, t1)
Wvvvv = lib.einsum('kcbd,ka->abcd', eris.ovvv, -t1)
Wvvvv = Wvvvv + Wvvvv.transpose(1, 0, 3, 2)
Wvvvv += eris.vvvv.transpose(0, 2, 1, 3)
tmp = lib.einsum('abcd,ic->abid', Wvvvv, t1)
t2new += lib.einsum('abid,jd->ijab', tmp, t1)
Lvv2 = fvv - lib.einsum('kc,ka->ac', fov, t1)
Lvv2 -= fvv.diagonal(preserve_shape=True)
tmp = lib.einsum('ac,ijcb->ijab', Lvv2, t2)
t2new += (tmp + tmp.transpose(1, 0, 3, 2))
Loo2 = foo + lib.einsum('kc,ic->ki', fov, t1)
Loo2 -= foo.diagonal(preserve_shape=True)
tmp = lib.einsum('ki,kjab->ijab', Loo2, t2)
t2new -= (tmp + tmp.transpose(1, 0, 3, 2))
else:
Loo = imd.Loo(t1, t2, eris)
Lvv = imd.Lvv(t1, t2, eris)
Loo -= eris._foo
Lvv -= eris._fvv
Woooo = imd.cc_Woooo(t1, t2, eris)
Wvoov = imd.cc_Wvoov(t1, t2, eris)
Wvovo = imd.cc_Wvovo(t1, t2, eris)
Wvvvv = imd.cc_Wvvvv(t1, t2, eris)
tau = t2 + lib.einsum('ia,jb->ijab', t1, t1)
t2new += lib.einsum('klij,klab->ijab', Woooo, tau)
t2new += lib.einsum('abcd,ijcd->ijab', Wvvvv, tau)
tmp = lib.einsum('ac,ijcb->ijab', Lvv, t2)
t2new += (tmp + tmp.transpose(1, 0, 3, 2))
tmp = lib.einsum('ki,kjab->ijab', Loo, t2)
t2new -= (tmp + tmp.transpose(1, 0, 3, 2))
tmp = 2 * lib.einsum('akic,kjcb->ijab', Wvoov, t2)
tmp -= lib.einsum('akci,kjcb->ijab', Wvovo, t2)
t2new += (tmp + tmp.transpose(1, 0, 3, 2))
tmp = lib.einsum('akic,kjbc->ijab', Wvoov, t2)
t2new -= (tmp + tmp.transpose(1, 0, 3, 2))
tmp = lib.einsum('bkci,kjac->ijab', Wvovo, t2)
t2new -= (tmp + tmp.transpose(1, 0, 3, 2))
tmp2 = lib.einsum('kibc,ka->abic', eris.oovv, -t1)
tmp2 += eris.ovvv.conj().transpose(1, 3, 0, 2)
tmp = lib.einsum('abic,jc->ijab', tmp2, t1)
t2new += (tmp + tmp.transpose(1, 0, 3, 2))
tmp2 = lib.einsum('kcai,jc->akij', eris.ovvo, t1)
tmp2 += eris.ooov.transpose(3, 1, 2, 0).conj()
tmp = lib.einsum('akij,kb->ijab', tmp2, t1)
t2new -= (tmp + tmp.transpose(1, 0, 3, 2))
t1new /= eris.eia
t2new /= eris.eijab
cput1 = log.timer("updating t2", *cput1)
return t1new, t2new
def eaccsd_matvec(eom, vector, kshift, imds=None, diag=None):
# Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1994) Eqs.(30)-(31)
if imds is None: imds = eom.make_imds()
r1, r2 = eom.vector_to_amplitudes(vector, kshift)
# Eq. (30)
# 1p-1p block
Hr1 = lib.einsum('ac,c->a', imds.Lvv, r1)
# 1p-2p1h block
Hr1 += lib.einsum('ld,lad->a', 2. * imds.Fov, r2)
Hr1 += lib.einsum('ld,lda->a', -imds.Fov, r2)
Hr1 += 2 * lib.einsum('alcd,lcd->a', imds.Wvovv, r2)
Hr1 += -lib.einsum('aldc,lcd->a', imds.Wvovv, r2)
# Eq. (31)
# 2p1h-1p block
Hr2 = lib.einsum('abcj,c->jab', imds.Wvvvo, r1)
# 2p1h-2p1h block
if eom.partition == 'mp':
foo = imds.eris.foo
fvv = imds.eris.fvv
Hr2 += lib.einsum('ac,jcb->jab', fvv, r2)
Hr2 += lib.einsum('bd,jad->jab', fvv, r2)
Hr2 += -lib.einsum('lj,lab->jab', foo, r2)
elif eom.partition == 'full':
Hr2 += eom._eaccsd_diag_matrix2 * r2
else:
Hr2 += lib.einsum('ac,jcb->jab', imds.Lvv, r2)
Hr2 += lib.einsum('bd,jad->jab', imds.Lvv, r2)
Hr2 += -lib.einsum('lj,lab->jab', imds.Loo, r2)
Hr2 += 2 * lib.einsum('lbdj,lad->jab', imds.Wovvo, r2)
Hr2 += -lib.einsum('lbjd,lad->jab', imds.Wovov, r2)
Hr2 += -lib.einsum('lajc,lcb->jab', imds.Wovov, r2)
Hr2 += -lib.einsum('lbcj,lca->jab', imds.Wovvo, r2)
Hr2 += lib.einsum('abcd,jcd->jab', imds.Wvvvv, r2)
tmp = (2 * lib.einsum('klcd,lcd->k', imds.Woovv, r2) -
lib.einsum('kldc,lcd->k', imds.Woovv, r2))
Hr2 += -lib.einsum('k,kjab->jab', tmp, imds.t2)
vector = eom.amplitudes_to_vector(Hr1, Hr2)
return vector
def ipccsd_matvec(eom, vector, kshift, imds=None, diag=None):
# Ref: Nooijen and Snijders, J. Chem. Phys. 102, 1681 (1995) Eqs.(8)-(9)
if imds is None: imds = eom.make_imds()
r1, r2 = eom.vector_to_amplitudes(vector, kshift)
# 1h-1h block
Hr1 = -lib.einsum('ki,k->i', imds.Loo, r1)
#1h-2h1p block
Hr1 += 2 * lib.einsum('ld,ild->i', imds.Fov, r2)
Hr1 += -lib.einsum('kd,kid->i', imds.Fov, r2)
Hr1 += -2 * lib.einsum('klid,kld->i', imds.Wooov, r2)
Hr1 += lib.einsum('lkid,kld->i', imds.Wooov, r2)
# 2h1p-1h block
Hr2 = -lib.einsum('kbij,k->ijb', imds.Wovoo, r1)
# 2h1p-2h1p block
if eom.partition == 'mp':
foo = self.eris.foo
fvv = self.eris.fvv
Hr2 += lib.einsum('bd,ijd->ijb', fvv, r2)
Hr2 += -lib.einsum('ki,kjb->ijb', foo, r2)
Hr2 += -lib.einsum('lj,ilb->ijb', foo, r2)
elif eom.partition == 'full':
Hr2 += self._ipccsd_diag_matrix2 * r2
else:
Hr2 += lib.einsum('bd,ijd->ijb', imds.Lvv, r2)
Hr2 += -lib.einsum('ki,kjb->ijb', imds.Loo, r2)
Hr2 += -lib.einsum('lj,ilb->ijb', imds.Loo, r2)
Hr2 += lib.einsum('klij,klb->ijb', imds.Woooo, r2)
Hr2 += 2 * lib.einsum('lbdj,ild->ijb', imds.Wovvo, r2)
Hr2 += -lib.einsum('kbdj,kid->ijb', imds.Wovvo, r2)
Hr2 += -lib.einsum('lbjd,ild->ijb', imds.Wovov, r2) #typo in Ref
Hr2 += -lib.einsum('kbid,kjd->ijb', imds.Wovov, r2)
tmp = 2 * lib.einsum('lkdc,kld->c', imds.Woovv, r2)
tmp += -lib.einsum('kldc,kld->c', imds.Woovv, r2)
Hr2 += -lib.einsum('c,ijcb->ijb', tmp, imds.t2)
vector = eom.amplitudes_to_vector(Hr1, Hr2)
return vector
def _make_df_eris(mycc, eris):
from cc_sym import mpigdf
mydf = mycc._scf.with_df
mo_coeff = eris.mo_coeff
if mydf.j3c is None: mydf.build()
gvec = mydf.cell.reciprocal_vectors()
nocc, nmo = mycc.nocc, mycc.nmo
nvir = nmo - nocc
kpts = mydf.kpts
nkpts = len(kpts)
nao, naux = mydf.j3c.shape[2:]
ijL = ctf.zeros([nkpts,nkpts,nocc,nocc,naux], dtype=mydf.j3c.dtype)
iaL = ctf.zeros([nkpts,nkpts,nocc,nvir,naux], dtype=mydf.j3c.dtype)
aiL = ctf.zeros([nkpts,nkpts,nvir,nocc,naux], dtype=mydf.j3c.dtype)
abL = ctf.zeros([nkpts,nkpts,nvir,nvir,naux], dtype=mydf.j3c.dtype)
jobs = []
for ki in range(nkpts):
for kj in range(ki,nkpts):
jobs.append([ki,kj])
tasks = static_partition(jobs)
ntasks = max(comm.allgather((len(tasks))))
idx_j3c = np.arange(nao**2*naux)
idx_ooL = np.arange(nocc**2*naux)
idx_ovL = np.arange(nocc*nvir*naux)
idx_vvL = np.arange(nvir**2*naux)
log = Logger(mydf.stdout, mydf.verbose)
cput1 = cput0 = (time.clock(), time.time())
for itask in range(ntasks):
if itask >= len(tasks):
mydf.j3c.read([])
ijL.write([], [])
iaL.write([], [])
aiL.write([], [])
abL.write([], [])
ijL.write([], [])
iaL.write([], [])
aiL.write([], [])
abL.write([], [])
continue
ki, kj = tasks[itask]
ijid, ijdagger = mpigdf.get_member(kpts[ki], kpts[kj], mydf.kptij_lst)
uvL = mydf.j3c.read(ijid*idx_j3c.size+idx_j3c).reshape(nao,nao,naux)
if ijdagger: uvL = uvL.transpose(1,0,2).conj()
pvL = np.einsum("up,uvL->pvL", mo_coeff[ki].conj(), uvL, optimize=True)
uvL = None
pqL = np.einsum('vq,pvL->pqL', mo_coeff[kj], pvL, optimize=True)
off = ki * nkpts + kj
ijL.write(off*idx_ooL.size+idx_ooL, pqL[:nocc,:nocc].ravel())
iaL.write(off*idx_ovL.size+idx_ovL, pqL[:nocc,nocc:].ravel())
aiL.write(off*idx_ovL.size+idx_ovL, pqL[nocc:,:nocc].ravel())
abL.write(off*idx_vvL.size+idx_vvL, pqL[nocc:,nocc:].ravel())
off = kj * nkpts + ki
pqL = pqL.transpose(1,0,2).conj()
ijL.write(off*idx_ooL.size+idx_ooL, pqL[:nocc,:nocc].ravel())
iaL.write(off*idx_ovL.size+idx_ovL, pqL[:nocc,nocc:].ravel())
aiL.write(off*idx_ovL.size+idx_ovL, pqL[nocc:,:nocc].ravel())
abL.write(off*idx_vvL.size+idx_vvL, pqL[nocc:,nocc:].ravel())
cput1 = log.timer("j3c transformation", *cput1)
sym1 = ["+-+", [kpts,]*3, None, gvec]
sym2 = ["+--", [kpts,]*3, None, gvec]
ooL = tensor(ijL, sym1, verbose=mycc.SYMVERBOSE)
ovL = tensor(iaL, sym1, verbose=mycc.SYMVERBOSE)
voL = tensor(aiL, sym1, verbose=mycc.SYMVERBOSE)
vvL = tensor(abL, sym1, verbose=mycc.SYMVERBOSE)
ooL2 = tensor(ijL, sym2, verbose=mycc.SYMVERBOSE)
ovL2 = tensor(iaL, sym2, verbose=mycc.SYMVERBOSE)
voL2 = tensor(aiL, sym2, verbose=mycc.SYMVERBOSE)
vvL2 = tensor(abL, sym2, verbose=mycc.SYMVERBOSE)
eris.oooo = lib.einsum('ijg,klg->ijkl', ooL, ooL2) / nkpts
eris.ooov = lib.einsum('ijg,kag->ijka', ooL, ovL2) / nkpts
eris.oovv = lib.einsum('ijg,abg->ijab', ooL, vvL2) / nkpts
eris.ovvo = lib.einsum('iag,bjg->iabj', ovL, voL2) / nkpts
eris.ovov = lib.einsum('iag,jbg->iajb', ovL, ovL2) / nkpts
eris.ovvv = lib.einsum('iag,bcg->iabc', ovL, vvL2) / nkpts
eris.vvvv = lib.einsum('abg,cdg->abcd', vvL, vvL2) / nkpts
cput1 = log.timer("integral transformation", *cput1)
def _make_fftdf_eris(mycc, eris):
mydf = mycc._scf.with_df
mo_coeff = eris.mo_coeff
kpts = mycc.kpts
logger = Logger(mycc.stdout, mycc.verbose)
cell = mydf.cell
gvec = cell.reciprocal_vectors()
nao = cell.nao_nr()
coords = cell.gen_uniform_grids(mydf.mesh)
ngrids = len(coords)
nkpts = len(kpts)
nocc, nmo = mycc.nocc, mycc.nmo
nvir = nmo - nocc
cput1 = cput0 = (time.clock(), time.time())
ijG = ctf.zeros([nkpts,nkpts,nocc,nocc,ngrids], dtype=np.complex128)
iaG = ctf.zeros([nkpts,nkpts,nocc,nvir,ngrids], dtype=np.complex128)
abG = ctf.zeros([nkpts,nkpts,nvir,nvir,ngrids], dtype=np.complex128)
ijR = ctf.zeros([nkpts,nkpts,nocc,nocc,ngrids], dtype=np.complex128)
iaR = ctf.zeros([nkpts,nkpts,nocc,nvir,ngrids], dtype=np.complex128)
aiR = ctf.zeros([nkpts,nkpts,nvir,nocc,ngrids], dtype=np.complex128)
abR = ctf.zeros([nkpts,nkpts,nvir,nvir,ngrids], dtype=np.complex128)
jobs = []
for ki in range(nkpts):
for kj in range(ki,nkpts):
jobs.append([ki,kj])
tasks = list(static_partition(jobs))
ntasks = max(comm.allgather(len(tasks)))
idx_ooG = np.arange(nocc*nocc*ngrids)
idx_ovG = np.arange(nocc*nvir*ngrids)
idx_vvG = np.arange(nvir*nvir*ngrids)
for itask in range(ntasks):
if itask >= len(tasks):
ijR.write([], [])
iaR.write([], [])
aiR.write([], [])
abR.write([], [])
ijR.write([], [])
iaR.write([], [])
aiR.write([], [])
abR.write([], [])
ijG.write([], [])
iaG.write([], [])
abG.write([], [])
ijG.write([], [])
iaG.write([], [])
abG.write([], [])
continue
ki, kj = tasks[itask]
kpti, kptj = kpts[ki], kpts[kj]
ao_kpti = mydf._numint.eval_ao(cell, coords, kpti)[0]
ao_kptj = mydf._numint.eval_ao(cell, coords, kptj)[0]
q = kptj - kpti
coulG = tools.get_coulG(cell, q, mesh=mydf.mesh)
wcoulG = coulG * (cell.vol/ngrids)
fac = np.exp(-1j * np.dot(coords, q))
mo_kpti = np.dot(ao_kpti, mo_coeff[ki]).T
mo_kptj = np.dot(ao_kptj, mo_coeff[kj]).T
mo_pairs = np.einsum('ig,jg->ijg', mo_kpti.conj(), mo_kptj)
mo_pairs_G = tools.fft(mo_pairs.reshape(-1,ngrids)*fac, mydf.mesh)
off = ki * nkpts + kj
ijR.write(off*idx_ooG.size+idx_ooG, mo_pairs[:nocc,:nocc].ravel())
iaR.write(off*idx_ovG.size+idx_ovG, mo_pairs[:nocc,nocc:].ravel())
aiR.write(off*idx_ovG.size+idx_ovG, mo_pairs[nocc:,:nocc].ravel())
abR.write(off*idx_vvG.size+idx_vvG, mo_pairs[nocc:,nocc:].ravel())
off = kj * nkpts + ki
mo_pairs = mo_pairs.transpose(1,0,2).conj()
ijR.write(off*idx_ooG.size+idx_ooG, mo_pairs[:nocc,:nocc].ravel())
iaR.write(off*idx_ovG.size+idx_ovG, mo_pairs[:nocc,nocc:].ravel())
aiR.write(off*idx_ovG.size+idx_ovG, mo_pairs[nocc:,:nocc].ravel())
abR.write(off*idx_vvG.size+idx_vvG, mo_pairs[nocc:,nocc:].ravel())
mo_pairs = None
mo_pairs_G*= wcoulG
v = tools.ifft(mo_pairs_G, mydf.mesh)
v *= fac.conj()
v = v.reshape(nmo,nmo,ngrids)
off = ki * nkpts + kj
ijG.write(off*idx_ooG.size+idx_ooG, v[:nocc,:nocc].ravel())
iaG.write(off*idx_ovG.size+idx_ovG, v[:nocc,nocc:].ravel())
abG.write(off*idx_vvG.size+idx_vvG, v[nocc:,nocc:].ravel())
off = kj * nkpts + ki
v = v.transpose(1,0,2).conj()
ijG.write(off*idx_ooG.size+idx_ooG, v[:nocc,:nocc].ravel())
iaG.write(off*idx_ovG.size+idx_ovG, v[:nocc,nocc:].ravel())
abG.write(off*idx_vvG.size+idx_vvG, v[nocc:,nocc:].ravel())
cput1 = logger.timer("Generating ijG", *cput1)
sym1 = ["+-+", [kpts,]*3, None, gvec]
sym2 = ["+--", [kpts,]*3, None, gvec]
ooG = tensor(ijG, sym1, verbose=mycc.SYMVERBOSE)
ovG = tensor(iaG, sym1, verbose=mycc.SYMVERBOSE)
vvG = tensor(abG, sym1, verbose=mycc.SYMVERBOSE)
ooR = tensor(ijR, sym2, verbose=mycc.SYMVERBOSE)
ovR = tensor(iaR, sym2, verbose=mycc.SYMVERBOSE)
voR = tensor(aiR, sym2, verbose=mycc.SYMVERBOSE)
vvR = tensor(abR, sym2, verbose=mycc.SYMVERBOSE)
eris.oooo = lib.einsum('ijg,klg->ijkl', ooG, ooR)/ nkpts
eris.ooov = lib.einsum('ijg,kag->ijka', ooG, ovR)/ nkpts
eris.oovv = lib.einsum('ijg,abg->ijab', ooG, vvR)/ nkpts
ooG = ooR = ijG = ijR = None
eris.ovvo = lib.einsum('iag,bjg->iabj', ovG, voR)/ nkpts
eris.ovov = lib.einsum('iag,jbg->iajb', ovG, ovR)/ nkpts
ovR = iaR = voR = aiR = None
eris.ovvv = lib.einsum('iag,bcg->iabc', ovG, vvR)/ nkpts
ovG = iaG = None
eris.vvvv = lib.einsum('abg,cdg->abcd', vvG, vvR)/ nkpts
cput1 = logger.timer("ijG to eri", *cput1)