How to use the pyscf.lib.einsum function in pyscf

#

'''
Unrestricted CISD
'''

import time
from functools import reduce
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.ci import cisd
from pyscf.cc import uccsd

einsum = lib.einsum

def kernel(myci, eris, ci0=None, max_cycle=50, tol=1e-8,
           verbose=logger.INFO):
    mol = myci.mol
    diag = myci.make_diagonal(eris)
    ehf = diag[0]
    diag -= ehf

    if ci0 is None:
        ci0 = myci.get_init_guess(eris)[1]

    def op(xs):
        return [myci.contract(x, eris) for x in xs]

    def precond(x, e, *args):
        diagd = diag - (e-myci.level_shift)

t2bb += tmpbb - tmpbb.transpose(1,0,2,3)
    t2ab += lib.einsum('ie,MBea->iMaB', c1a, eris_HOVvv.conj())
    t2ab += lib.einsum('IE,maEB->mIaB', c1b, eris_HovVV.conj())
    #:tmp = einsum('ma,mbij->ijab', c1, eris.ovoo)
    #:t2 -= tmp - tmp.transpose(0,1,3,2)
    eris_Hoovo = _cp(eris.oovo)
    eris_HOOVO = _cp(eris.OOVO)
    eris_HooVO = _cp(eris.ooVO)
    eris_HOOvo = _cp(eris.OOvo)
    oovo = eris_Hoovo - eris_Hoovo.transpose(0,3,2,1)
    OOVO = eris_HOOVO - eris_HOOVO.transpose(0,3,2,1)
    tmpaa = lib.einsum('ma,mibj->ijab', c1a, oovo)
    t2aa -= tmpaa - tmpaa.transpose(0,1,3,2)
    tmpbb = lib.einsum('ma,mibj->ijab', c1b, OOVO)
    t2bb -= tmpbb - tmpbb.transpose(0,1,3,2)
    t2ab -= lib.einsum('ma,miBJ->iJaB', c1a, eris_HooVO)
    t2ab -= lib.einsum('MA,MJbi->iJbA', c1b, eris_HOOvo)

    #:t1 += fov * c0
    t1a += fova * c0
    t1b += fovb * c0
    #:t2 += eris.oovv * c0
    t2aa += eris.voov.transpose(1,2,0,3) * c0
    t2aa -= eris.voov.transpose(1,2,3,0) * c0
    t2bb += eris.VOOV.transpose(1,2,0,3) * c0
    t2bb -= eris.VOOV.transpose(1,2,3,0) * c0
    t2ab += eris.voOV.transpose(1,2,0,3) * c0
    #:t0  = numpy.einsum('ia,ia', fov, c1)
    #:t0 += numpy.einsum('ijab,ijab', eris.oovv, c2) * .25
    t0  = numpy.einsum('ia,ia', fova, c1a)
    t0 += numpy.einsum('ia,ia', fovb, c1b)
    t0 += numpy.einsum('aijb,ijab', eris.voov, c2aa) * .25

t1a += numpy.einsum('imae,me->ia', c2ab, fovb)
    t1b += numpy.einsum('imae,me->ia', c2bb, fovb)
    t1b += numpy.einsum('miea,me->ia', c2ab, fova)

    #:tmp = einsum('ijae,be->ijab', c2, fvv)
    #:t2  = tmp - tmp.transpose(0,1,3,2)
    t2aa += lib.einsum('ijae,be->ijab', c2aa, fvva*.5)
    t2bb += lib.einsum('ijae,be->ijab', c2bb, fvvb*.5)
    t2ab += lib.einsum('iJaE,BE->iJaB', c2ab, fvvb)
    t2ab += lib.einsum('iJeA,be->iJbA', c2ab, fvva)
    #:tmp = einsum('imab,mj->ijab', c2, foo)
    #:t2 -= tmp - tmp.transpose(1,0,2,3)
    t2aa -= lib.einsum('imab,mj->ijab', c2aa, fooa*.5)
    t2bb -= lib.einsum('imab,mj->ijab', c2bb, foob*.5)
    t2ab -= lib.einsum('iMaB,MJ->iJaB', c2ab, foob)
    t2ab -= lib.einsum('mIaB,mj->jIaB', c2ab, fooa)

    #:tmp = numpy.einsum('ia,bj->ijab', c1, fvo)
    #:tmp = tmp - tmp.transpose(0,1,3,2)
    #:t2 += tmp - tmp.transpose(1,0,2,3)
    t2aa += numpy.einsum('ia,bj->ijab', c1a, fvoa)
    t2bb += numpy.einsum('ia,bj->ijab', c1b, fvob)
    t2ab += numpy.einsum('ia,bj->ijab', c1a, fvob)
    t2ab += numpy.einsum('ia,bj->jiba', c1b, fvoa)

    t2aa = t2aa - t2aa.transpose(0,1,3,2)
    t2aa = t2aa - t2aa.transpose(1,0,2,3)
    t2bb = t2bb - t2bb.transpose(0,1,3,2)
    t2bb = t2bb - t2bb.transpose(1,0,2,3)

    #:t2 += 0.5*einsum('mnab,mnij->ijab', c2, eris.oooo)
    eris_oooo = _cp(eris.oooo)

if (method == "adc(2)-x" or method == "adc(3)"):
        
               r2_a = r2 - r2.transpose(0,2,1).copy()
               t2_2 = adc.t2[1]

               eris_oooo = eris.oooo
               eris_oovv = eris.oovv
               eris_ovvo = eris.ovvo
               
               s[s2:f2] -= 0.5*lib.einsum('kijl,ali->ajk',eris_oooo, r2, optimize = True).reshape(-1)
               s[s2:f2] -= 0.5*lib.einsum('klji,ail->ajk',eris_oooo ,r2, optimize = True).reshape(-1)
               
               s[s2:f2] += 0.5*lib.einsum('klba,bjl->ajk',eris_oovv,r2,optimize = True).reshape(-1)
               
               s[s2:f2] +=  0.5*lib.einsum('jabl,bkl->ajk',eris_ovvo,r2_a,optimize = True).reshape(-1)
               s[s2:f2] +=  0.5*lib.einsum('jlba,blk->ajk',eris_oovv,r2,optimize = True).reshape(-1)
               s[s2:f2] -=  0.5*lib.einsum('jabl,blk->ajk',eris_ovvo,r2,optimize = True).reshape(-1)
               
               s[s2:f2] += 0.5*lib.einsum('kiba,bji->ajk',eris_oovv,r2,optimize = True).reshape(-1)
               
               s[s2:f2] += 0.5*lib.einsum('jiba,bik->ajk',eris_oovv,r2,optimize = True).reshape(-1)
               s[s2:f2] -= 0.5*lib.einsum('jabi,bik->ajk',eris_ovvo,r2,optimize = True).reshape(-1)
               s[s2:f2] -= 0.5*lib.einsum('jabi,bik->ajk',eris_ovvo,r2_a,optimize = True).reshape(-1)
               
        if (method == "adc(3)"):

               eris_ovoo = eris.ovoo

################ ADC(3) i - kja block and ajk - i ############################

               eris_ovvv = radc_ao2mo.unpack_eri_1(eris.ovvv, nvir)
               t2_1_a = t2_1 - t2_1.transpose(1,0,2,3).copy()

Hr1 -= lib.einsum('mnie,mnae->ia', imds.woOoV, r2aaba)
    Hr1 -= lib.einsum('mnie,mnae->ia', imds.woOoV, r2baaa)
    tmp = lib.einsum('mnie,me->ni', imds.woOoV, r1)
    tmp = lib.einsum('ni,njab->ijab', tmp, t2)
    Hr2baaa += tmp
    Hr2aaba += tmp
    for p0,p1 in lib.prange(0, nvir, nocc):
        tmp = lib.einsum('ejab,ie->ijab', np.asarray(imds.wvOvV[p0:p1]), r1[:,p0:p1])
        Hr2baaa += tmp
        Hr2aaba += tmp
    tmp = None

    oVVo = np.asarray(imds.woVVo)
    Hr1 += np.einsum('maei,me->ia', oVVo, r1)
    Hr2baaa += lib.einsum('mbej,miea->jiba', oVVo, r2baaa)
    Hr2aaba += lib.einsum('mbej,miea->jiba', oVVo, r2aaba)
    oVvO = np.asarray(imds.woVvO)
    Hr2baaa += lib.einsum('mbej,imae->ijab', oVvO, r2aaba)
    Hr2aaba += lib.einsum('mbej,imae->ijab', oVvO, r2baaa)
    oVvO += oVVo
    Hr2baaa += lib.einsum('mbej,imae->ijab', oVvO, r2baaa)
    Hr2aaba += lib.einsum('mbej,imae->ijab', oVvO, r2aaba)
    oVvO = oVVo = None

    eris_ovov = np.asarray(eris.ovov)
    tau = _make_tau(t2, t1, t1)
    tmp1baaa = lib.einsum('menf,ijef->mnij', eris_ovov, tau2aaba)
    tmp1aaba = lib.einsum('menf,ijef->mnij', eris_ovov, tau2baaa)
    Hr2baaa += .5*lib.einsum('mnij,mnab->ijab', tmp1aaba, tau)
    Hr2aaba += .5*lib.einsum('mnij,mnab->ijab', tmp1baaa, tau)
    tau2 = tmp1baaa = tmp1aaba = None

M_ij -= 0.25 *  lib.einsum('i,jlde,ilde->ij',e_occ,t2_1, t2_2,optimize=True)
        M_ij -= 0.25 *  lib.einsum('i,jlde,ilde->ij',e_occ,t2_1, t2_2,optimize=True)

        M_ij -= 0.25 *  lib.einsum('j,jlde,ilde->ij',e_occ,t2_1_a, t2_2_a,optimize=True)
        M_ij -= 0.25 *  lib.einsum('j,jlde,ilde->ij',e_occ,t2_1, t2_2,optimize=True)
        M_ij -= 0.25 *  lib.einsum('j,jlde,ilde->ij',e_occ,t2_1, t2_2,optimize=True)

        M_ij -= 0.25 *  lib.einsum('j,ilde,jlde->ij',e_occ,t2_1_a, t2_2_a,optimize=True)
        M_ij -= 0.25 *  lib.einsum('j,ilde,jlde->ij',e_occ,t2_1, t2_2,optimize=True)
        M_ij -= 0.25 *  lib.einsum('j,ilde,jlde->ij',e_occ,t2_1, t2_2,optimize=True)

        M_ij -= lib.einsum('lmde,jldf,mefi->ij',t2_1_a, t2_1_a, eris_ovvo,optimize = True)
        M_ij += lib.einsum('lmde,jldf,mife->ij',t2_1_a, t2_1_a, eris_oovv,optimize = True)
        M_ij += lib.einsum('mled,jlfd,mefi->ij',t2_1, t2_1, eris_ovvo ,optimize = True)
        M_ij -= lib.einsum('mled,jlfd,mife->ij',t2_1, t2_1, eris_oovv ,optimize = True)
        M_ij -= lib.einsum('lmde,jldf,mefi->ij',t2_1, t2_1_a, eris_ovvo,optimize = True)
        M_ij -= lib.einsum('mlde,jldf,mife->ij',t2_1, t2_1, eris_oovv ,optimize = True)
        M_ij += lib.einsum('lmde,jlfd,mefi->ij',t2_1_a, t2_1, eris_ovvo ,optimize = True)

        M_ij -= lib.einsum('lmde,ildf,mefj->ij',t2_1_a, t2_1_a, eris_ovvo ,optimize = True)
        M_ij += lib.einsum('lmde,ildf,mjfe->ij',t2_1_a, t2_1_a, eris_oovv ,optimize = True)
        M_ij += lib.einsum('mled,ilfd,mefj->ij',t2_1, t2_1, eris_ovvo ,optimize = True)
        M_ij -= lib.einsum('mled,ilfd,mjfe->ij',t2_1, t2_1, eris_oovv ,optimize = True)
        M_ij -= lib.einsum('lmde,ildf,mefj->ij',t2_1, t2_1_a, eris_ovvo,optimize = True)
        M_ij -= lib.einsum('mlde,ildf,mjfe->ij',t2_1, t2_1, eris_oovv ,optimize = True)
        M_ij += lib.einsum('lmde,ilfd,mefj->ij',t2_1_a, t2_1, eris_ovvo ,optimize = True)

        M_ij += 0.25*lib.einsum('lmde,jnde,limn->ij',t2_1_a, t2_1_a,eris_oooo, optimize = True)
        M_ij -= 0.25*lib.einsum('lmde,jnde,lnmi->ij',t2_1_a, t2_1_a,eris_oooo, optimize = True)
        M_ij += lib.einsum('lmde,jnde,limn->ij',t2_1 ,t2_1, eris_oooo, optimize = True)

        eris_vvvv = eris.vvvv

t2_1_a_t = t2_1_a.reshape(nocc,nocc,-1)
               r2_a = r2_a.reshape(nocc,-1)
               temp =  0.25 * lib.einsum('lmp,jp->lmj',t2_1_a_t,r2_a)
               s[s1:f1] += lib.einsum('lmj,lamj->a',temp, eris_ovoo, optimize=True)
               s[s1:f1] -= lib.einsum('lmj,malj->a',temp, eris_ovoo, optimize=True)

               temp_1 = -lib.einsum('lmzw,jzw->jlm',t2_1,r2)
               s[s1:f1] -= lib.einsum('jlm,lamj->a',temp_1, eris_ovoo, optimize=True)

               r2_a = r2_a.reshape(nocc,nvir,nvir)
               temp_s_a = np.zeros_like(r2)
               temp_s_a = lib.einsum('jlwd,jzw->lzd',t2_1_a,r2_a,optimize=True)
               temp_s_a += lib.einsum('ljdw,jzw->lzd',t2_1,r2,optimize=True)

               temp_s_a_1 = np.zeros_like(r2)
               temp_s_a_1 = -lib.einsum('jlzd,jwz->lwd',t2_1_a,r2_a,optimize=True)
               temp_s_a_1 += -lib.einsum('ljdz,jwz->lwd',t2_1,r2,optimize=True)

               eris_ovvv = radc_ao2mo.unpack_eri_1(eris.ovvv, nvir)

               temp_1_1 = lib.einsum('ldxb,b->lxd', eris_ovvv,r1,optimize=True)
               temp_1_1 -= lib.einsum('lbxd,b->lxd', eris_ovvv,r1,optimize=True)
               temp_2_1 = lib.einsum('ldxb,b->lxd', eris_ovvv,r1,optimize=True)

               s[s1:f1] += 0.5*lib.einsum('lzd,ldza->a',temp_s_a,eris_ovvv,optimize=True)
               s[s1:f1] -= 0.5*lib.einsum('lzd,lazd->a',temp_s_a,eris_ovvv,optimize=True)
               s[s1:f1] -= 0.5*lib.einsum('lwd,ldwa->a',temp_s_a_1,eris_ovvv,optimize=True)
               s[s1:f1] += 0.5*lib.einsum('lwd,lawd->a',temp_s_a_1,eris_ovvv,optimize=True)

               temp_1 = np.zeros_like(r2)
               temp_1 = lib.einsum('jlwd,jzw->lzd',t2_1,r2_a,optimize=True)
               temp_1 += lib.einsum('jlwd,jzw->lzd',t2_1_a,r2,optimize=True)

# Compute first-order doubles t2 (tijab)

    v2e_oovv = eris_ovov.transpose(0,2,1,3).copy()

    t2_1 = v2e_oovv/D2

    # Compute second-order singles t1 (tij)

    t2_1_a = t2_1 - t2_1.transpose(1,0,2,3).copy()

    eris_ovvv = radc_ao2mo.unpack_eri_1(eris.ovvv, nvir)
    t1_2 = 0.5*lib.einsum('kdac,ikcd->ia',eris_ovvv,t2_1_a,optimize=True)
    t1_2 -= 0.5*lib.einsum('kcad,ikcd->ia',eris_ovvv,t2_1_a,optimize=True)
    t1_2 += lib.einsum('kdac,ikcd->ia',eris_ovvv,t2_1,optimize=True)
    del eris_ovvv
    t1_2 -= 0.5*lib.einsum('lcki,klac->ia',eris_ovoo,t2_1_a,optimize=True)
    t1_2 -= 0.5*lib.einsum('kcli,lkac->ia',eris_ovoo,t2_1_a,optimize=True)
    t1_2 -= lib.einsum('lcki,klac->ia',eris_ovoo,t2_1,optimize=True)

    t1_2 = t1_2/D1

    t2_2 = None
    t1_3 = None

    if (myadc.method == "adc(2)-x" or myadc.method == "adc(3)"):

    # Compute second-order doubles t2 (tijab)

        eris_oooo = eris.oooo
        eris_ovvo = eris.ovvo
 
        temp = t2_1.reshape(nocc*nocc,nvir*nvir)

#:tmpab = lib.einsum('mebf,miaf->eiab', eris_ovvv, t2)
        #:tmpab = tmpab + tmpab.transpose(0,1,3,2) * .5
        #:tmpab-= lib.einsum('mfbe,mifa->eiba', ovvv, theta) * .5
        #:self.wvOvV += eris_ovvv.transpose(2,0,3,1).conj()
        #:self.wvOvV -= tmpab
        tmp1ab -= eris_ovvo.transpose(0,2,1,3)
        tmp1abba -= eris_oovv.transpose(0,2,3,1)
        eris_ovvo = eris_oovv = None
        tau = _make_tau(t2, t1, t1)
        mem_now = lib.current_memory()[0]
        max_memory = max(0, lib.param.MAX_MEMORY - mem_now)
        blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvir**3*6))))
        for i0,i1 in lib.prange(0, nocc, blksize):
            wvOvV = lib.einsum('meni,mnab->eiab', eris_ovoo[:,:,:,i0:i1], tau)

            wvOvV -= lib.einsum('me,miab->eiab', self.Fov, t2[:,i0:i1])
            tmpab = lib.einsum('ma,mbei->eiab', t1, tmp1ab[:,:,:,i0:i1])
            tmpab+= lib.einsum('ma,mbei->eiba', t1, tmp1abba[:,:,:,i0:i1])
            wvOvV += tmpab

            for p0,p1 in lib.prange(0, nocc, blksize):
                ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
                if p0 == i0:
                    wvOvV += ovvv.transpose(2,0,3,1).conj()
                tmpab = lib.einsum('mebf,miaf->eiab', ovvv, t2[p0:p1,i0:i1])
                tmpab = tmpab + tmpab.transpose(0,1,3,2) * .5
                ovvv = ovvv*2 - ovvv.transpose(0,3,2,1)
                tmpab-= lib.einsum('mfbe,mifa->eiba', ovvv, theta[p0:p1,i0:i1]) * .5
                wvOvV -= tmpab
                ovvv = tmpab = None
            self.wvOvV[:,i0:i1] = wvOvV

Hr2abb[ki,kj] -= imds.FOO[kj].diagonal()[None,:,None]
                Hr2baa[ki,kj] -= imds.FOO[ki].diagonal()[:,None,None]
                Hr2baa[ki,kj] -= imds.Foo[kj].diagonal()[None,:,None]
                Hr2bbb[ki,kj] -= imds.FOO[ki].diagonal()[:,None,None]
                Hr2bbb[ki,kj] -= imds.FOO[kj].diagonal()[None,:,None]

        for ki, kj in itertools.product(range(nkpts), repeat=2):
        #for ki in range(nkpts):
        #    for kj in range(nkpts):
            Hr2aaa[ki, kj] += lib.einsum('iijj->ij', imds.Woooo[ki, ki, kj])[:,:,None]
            Hr2abb[ki, kj] += lib.einsum('iiJJ->iJ', imds.WooOO[ki, ki, kj])[:,:,None]
            Hr2bbb[ki, kj] += lib.einsum('IIJJ->IJ', imds.WOOOO[ki, ki, kj])[:,:,None]
            Hr2baa[ki, kj] += lib.einsum('jjII->Ij', imds.WooOO[kj, kj, ki])[:,:,None]

            kb = kconserv[ki, kshift, kj]
            Hr2aaa[ki,kj] -= lib.einsum('iejb,jibe->ijb', imds.Wovov[ki,kshift,kj], t2aa[kj,ki,kb])
            Hr2abb[ki,kj] -= lib.einsum('ieJB,iJeB->iJB', imds.WovOV[ki,kshift,kj], t2ab[ki,kj,kshift])
            Hr2baa[ki,kj] -= lib.einsum('jbIE,jIbE->Ijb', imds.WovOV[kj,kb,ki], t2ab[kj,ki,kb])
            Hr2bbb[ki,kj] -= lib.einsum('IEJB,JIBE->IJB', imds.WOVOV[ki,kshift,kj], t2bb[kj,ki,kb])

            Hr2aaa[ki, kj] += lib.einsum('ibbi->ib', imds.Wovvo[ki, kb, kb])[:,None,:]
            Hr2aaa[ki, kj] += lib.einsum('jbbj->jb', imds.Wovvo[kj, kb, kb])[None,:,:]

            Hr2baa[ki, kj] += lib.einsum('jbbj->jb', imds.Wovvo[kj, kb, kb])[None,:,:]
            Hr2baa[ki, kj] -= lib.einsum('IIbb->Ib', imds.WOOvv[ki, ki, kb])[:,None,:]

            Hr2abb[ki, kj] += lib.einsum('JBBJ->JB', imds.WOVVO[kj, kb, kb])[None,:,:]
            Hr2abb[ki, kj] -= lib.einsum('iiBB->iB', imds.WooVV[ki, ki, kb])[:,None,:]

            Hr2bbb[ki, kj] += lib.einsum('IBBI->IB', imds.WOVVO[ki, kb, kb])[:,None,:]
            Hr2bbb[ki, kj] += lib.einsum('JBBJ->JB', imds.WOVVO[kj, kb, kb])[None,:,:]