Python lib.direct_sum函数代码示例

def solve_mo1(sscobj, mo_energy=None, mo_coeff=None, mo_occ=None,
              h1=None, s1=None, with_cphf=None):
    cput1 = (time.clock(), time.time())
    log = logger.Logger(sscobj.stdout, sscobj.verbose)
    if mo_energy is None: mo_energy = sscobj._scf.mo_energy
    if mo_coeff  is None: mo_coeff = sscobj._scf.mo_coeff
    if mo_occ    is None: mo_occ = sscobj._scf.mo_occ
    if with_cphf is None: with_cphf = sscobj.cphf

    mol = sscobj.mol
    if h1 is None:
        atmlst = sorted(set([j for i,j in sscobj.nuc_pair]))
        h1a, h1b = make_h1_pso(mol, sscobj._scf.mo_coeff, mo_occ, atmlst)
    else:
        h1a, h1b = h1
    h1a = numpy.asarray(h1a)
    h1b = numpy.asarray(h1b)

    if with_cphf:
        if callable(with_cphf):
            vind = with_cphf
        else:
            vind = gen_vind(sscobj._scf, mo_coeff, mo_occ)
        mo1, mo_e1 = ucphf.solve(vind, mo_energy, mo_occ, (h1a,h1b), None,
                                 sscobj.max_cycle_cphf, sscobj.conv_tol,
                                 verbose=log)
    else:
        eai_aa = lib.direct_sum('i-a->ai', mo_energy[0][mo_occ[0]>0], mo_energy[0][mo_occ[0]==0])
        eai_bb = lib.direct_sum('i-a->ai', mo_energy[1][mo_occ[1]>0], mo_energy[1][mo_occ[1]==0])
        mo1 = (h1a * (1/eai_aa), h1b * (1/eai_bb))
        mo_e1 = None

    logger.timer(sscobj, 'solving mo1 eqn', *cput1)
    return mo1, mo_e1

    def init_amps(self, eris=None):
        time0 = time.clock(), time.time()
        if eris is None:
            eris = self.ao2mo(self.mo_coeff)
        nocca, noccb = self.nocc

        fova = eris.focka[:nocca,nocca:]
        fovb = eris.fockb[:noccb,noccb:]
        mo_ea_o = eris.mo_energy[0][:nocca]
        mo_ea_v = eris.mo_energy[0][nocca:]
        mo_eb_o = eris.mo_energy[1][:noccb]
        mo_eb_v = eris.mo_energy[1][noccb:]
        eia_a = lib.direct_sum('i-a->ia', mo_ea_o, mo_ea_v)
        eia_b = lib.direct_sum('i-a->ia', mo_eb_o, mo_eb_v)

        t1a = fova.conj() / eia_a
        t1b = fovb.conj() / eia_b

        eris_ovov = np.asarray(eris.ovov)
        eris_OVOV = np.asarray(eris.OVOV)
        eris_ovOV = np.asarray(eris.ovOV)
        t2aa = eris_ovov.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_a, eia_a)
        t2ab = eris_ovOV.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_a, eia_b)
        t2bb = eris_OVOV.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_b, eia_b)
        t2aa = t2aa - t2aa.transpose(0,1,3,2)
        t2bb = t2bb - t2bb.transpose(0,1,3,2)
        e  =      np.einsum('iJaB,iaJB', t2ab, eris_ovOV)
        e += 0.25*np.einsum('ijab,iajb', t2aa, eris_ovov)
        e -= 0.25*np.einsum('ijab,ibja', t2aa, eris_ovov)
        e += 0.25*np.einsum('ijab,iajb', t2bb, eris_OVOV)
        e -= 0.25*np.einsum('ijab,ibja', t2bb, eris_OVOV)
        self.emp2 = e.real
        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, (t1a,t1b), (t2aa,t2ab,t2bb)

def make_intermediates(mycc, t1, t2, eris):
    saved = ccsd_lambda.make_intermediates(mycc, t1, t2, eris)

    nocc, nvir = t1.shape
    eris_ovvv = _cp(eris.ovvv)
    eris_ovvv = lib.unpack_tril(eris_ovvv.reshape(nocc*nvir,-1))
    eris_ovvv = eris_ovvv.reshape(nocc,nvir,nvir,nvir)
    mo_e = mycc._scf.mo_energy
    eia = lib.direct_sum('i-a->ia',mo_e[:nocc], mo_e[nocc:])
    d3 = lib.direct_sum('ia,jb,kc->ijkabc', eia, eia, eia)
    eris_ovoo = eris.ovoo
    w =(numpy.einsum('iabf,kjcf->ijkabc', eris_ovvv, t2)
      - numpy.einsum('iajm,mkbc->ijkabc', eris_ovoo, t2)) / d3
    v = numpy.einsum('iajb,kc->ijkabc', eris.ovov, t1) / d3 * .5
    w = p6_(w)
    v = p6_(v)
    rwv = r6_(w*2+v)
    jov = numpy.einsum('jbkc,ijkabc->ia', eris.ovov, r6_(w))
    joovv = numpy.einsum('iabf,ijkabc->kjcf', eris_ovvv, rwv)
    joovv-= numpy.einsum('iajm,ijkabc->mkbc', eris.ovoo, rwv)
    joovv = joovv + joovv.transpose(1,0,3,2)

    saved.jov = jov
    saved.joovv = joovv
    return saved

def kernel(mp, mo_energy, mo_coeff, nocc, ioblk=256, verbose=None):
    nmo = mo_coeff.shape[1]
    nvir = nmo - nocc
    auxmol = df.incore.format_aux_basis(mp.mol, mp.auxbasis)
    naoaux = auxmol.nao_nr()

    iolen = max(int(ioblk*1e6/8/(nvir*nocc)), 160)

    eia = lib.direct_sum('i-a->ia', mo_energy[:nocc], mo_energy[nocc:])
    t2 = None
    emp2 = 0
    with mp.ao2mo(mo_coeff, nocc) as fov:
        for p0, p1 in prange(0, naoaux, iolen):
            logger.debug(mp, 'Load cderi block %d:%d', p0, p1)
            qov = numpy.array(fov[p0:p1], copy=False)
            for i in range(nocc):
                buf = numpy.dot(qov[:,i*nvir:(i+1)*nvir].T,
                                qov).reshape(nvir,nocc,nvir)
                gi = numpy.array(buf, copy=False)
                gi = gi.reshape(nvir,nocc,nvir).transpose(1,2,0)
                t2i = gi/lib.direct_sum('jb+a->jba', eia, eia[i])
                # 2*ijab-ijba
                theta = gi*2 - gi.transpose(0,2,1)
                emp2 += numpy.einsum('jab,jab', t2i, theta)

    return emp2, t2

def kernel(mp, mo_energy=None, mo_coeff=None, eris=None, with_t2=WITH_T2,
           verbose=logger.NOTE):
    if mo_energy is None or mo_coeff is None:
        moidx = mp.get_frozen_mask()
        mo_coeff = None
        mo_energy = (mp.mo_energy[0][moidx[0]], mp.mo_energy[1][moidx[1]])
    else:
        # For backward compatibility.  In pyscf-1.4 or earlier, mp.frozen is
        # not supported when mo_energy or mo_coeff is given.
        assert(mp.frozen is 0 or mp.frozen is None)

    if eris is None: eris = mp.ao2mo(mo_coeff)

    nocca, noccb = mp.get_nocc()
    nmoa, nmob = mp.get_nmo()
    nvira, nvirb = nmoa-nocca, nmob-noccb
    mo_ea, mo_eb = mo_energy
    eia_a = mo_ea[:nocca,None] - mo_ea[None,nocca:]
    eia_b = mo_eb[:noccb,None] - mo_eb[None,noccb:]

    if with_t2:
        dtype = eris.ovov.dtype
        t2aa = numpy.empty((nocca,nocca,nvira,nvira), dtype=dtype)
        t2ab = numpy.empty((nocca,noccb,nvira,nvirb), dtype=dtype)
        t2bb = numpy.empty((noccb,noccb,nvirb,nvirb), dtype=dtype)
        t2 = (t2aa,t2ab,t2bb)
    else:
        t2 = None

    emp2 = 0.0
    for i in range(nocca):
        eris_ovov = numpy.asarray(eris.ovov[i*nvira:(i+1)*nvira])
        eris_ovov = eris_ovov.reshape(nvira,nocca,nvira).transpose(1,0,2)
        t2i = eris_ovov.conj()/lib.direct_sum('a+jb->jab', eia_a[i], eia_a)
        emp2 += numpy.einsum('jab,jab', t2i, eris_ovov) * .5
        emp2 -= numpy.einsum('jab,jba', t2i, eris_ovov) * .5
        if with_t2:
            t2aa[i] = t2i - t2i.transpose(0,2,1)

        eris_ovov = numpy.asarray(eris.ovOV[i*nvira:(i+1)*nvira])
        eris_ovov = eris_ovov.reshape(nvira,noccb,nvirb).transpose(1,0,2)
        t2i = eris_ovov.conj()/lib.direct_sum('a+jb->jab', eia_a[i], eia_b)
        emp2 += numpy.einsum('JaB,JaB', t2i, eris_ovov)
        if with_t2:
            t2ab[i] = t2i

    for i in range(noccb):
        eris_ovov = numpy.asarray(eris.OVOV[i*nvirb:(i+1)*nvirb])
        eris_ovov = eris_ovov.reshape(nvirb,noccb,nvirb).transpose(1,0,2)
        t2i = eris_ovov.conj()/lib.direct_sum('a+jb->jab', eia_b[i], eia_b)
        emp2 += numpy.einsum('jab,jab', t2i, eris_ovov) * .5
        emp2 -= numpy.einsum('jab,jba', t2i, eris_ovov) * .5
        if with_t2:
            t2bb[i] = t2i - t2i.transpose(0,2,1)

    return emp2.real, t2

def get_sigma_element(gw, omega, tdm_p, tdm_q, td_e, eta=None, vir_sgn=1):
    if eta is None:
        eta = gw.eta

    nocc = gw.nocc
    evi = lib.direct_sum('v-i->vi', td_e, gw._scf.mo_energy[:nocc])
    eva = lib.direct_sum('v+a->va', td_e, gw._scf.mo_energy[nocc:])
    sigma =  np.sum( tdm_p[:,:nocc]*tdm_q[:,:nocc]/(omega + evi - 1j*eta) )
    sigma += np.sum( tdm_p[:,nocc:]*tdm_q[:,nocc:]/(omega - eva + vir_sgn*1j*eta) )
    return sigma

def solve_withs1(fvind, mo_energy, mo_occ, h1, s1,
                 max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
    '''For field dependent basis. First order overlap matrix is non-zero.
    The first order orbitals are set to
    C^1_{ij} = -1/2 S1
    e1 = h1 - s1*e0 + (e0_j-e0_i)*c1 + vhf[c1]

    Kwargs:
        hermi : boolean
            Whether the matrix defined by fvind is Hermitian or not.

    Returns:
        First order orbital coefficients (in MO basis) and first order orbital
        energy matrix
    '''
    log = logger.new_logger(verbose=verbose)
    t0 = (time.clock(), time.time())

    occidx = mo_occ > 0
    viridx = mo_occ == 0
    e_a = mo_energy[viridx]
    e_i = mo_energy[occidx]
    e_ai = 1 / lib.direct_sum('a-i->ai', e_a, e_i)
    nvir, nocc = e_ai.shape
    nmo = nocc + nvir

    s1 = s1.reshape(-1,nmo,nocc)
    hs = mo1base = h1.reshape(-1,nmo,nocc) - s1*e_i
    mo_e1 = hs[:,occidx,:].copy()

    mo1base[:,viridx] *= -e_ai
    mo1base[:,occidx] = -s1[:,occidx] * .5

    def vind_vo(mo1):
        v = fvind(mo1.reshape(h1.shape)).reshape(-1,nmo,nocc)
        v[:,viridx,:] *= e_ai
        v[:,occidx,:] = 0
        return v.ravel()
    mo1 = lib.krylov(vind_vo, mo1base.ravel(),
                     tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log)
    mo1 = mo1.reshape(mo1base.shape)
    log.timer('krylov solver in CPHF', *t0)

    v1mo = fvind(mo1.reshape(h1.shape)).reshape(-1,nmo,nocc)
    mo1[:,viridx] = mo1base[:,viridx] - v1mo[:,viridx]*e_ai

    # mo_e1 has the same symmetry as the first order Fock matrix (hermitian or
    # anti-hermitian). mo_e1 = v1mo - s1*lib.direct_sum('i+j->ij',e_i,e_i)
    mo_e1 += mo1[:,occidx] * lib.direct_sum('i-j->ij', e_i, e_i)
    mo_e1 += v1mo[:,occidx,:]

    if h1.ndim == 3:
        return mo1, mo_e1
    else:
        return mo1.reshape(h1.shape), mo_e1.reshape(nocc,nocc)

    def get_init_guess(self, eris=None, nroots=1, diag=None):
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        nocca, noccb = self.nocc
        mo_ea, mo_eb = eris.mo_energy
        eia_a = mo_ea[:nocca,None] - mo_ea[None,nocca:]
        eia_b = mo_eb[:noccb,None] - mo_eb[None,noccb:]
        t1a = eris.focka[:nocca,nocca:].conj() / eia_a
        t1b = eris.fockb[:noccb,noccb:].conj() / eia_b

        eris_ovov = _cp(eris.ovov)
        eris_ovOV = _cp(eris.ovOV)
        eris_OVOV = _cp(eris.OVOV)
        t2aa = eris_ovov.transpose(0,2,1,3) - eris_ovov.transpose(0,2,3,1)
        t2bb = eris_OVOV.transpose(0,2,1,3) - eris_OVOV.transpose(0,2,3,1)
        t2ab = eris_ovOV.transpose(0,2,1,3).copy()
        t2aa = t2aa.conj()
        t2ab = t2ab.conj()
        t2bb = t2bb.conj()
        t2aa /= lib.direct_sum('ia+jb->ijab', eia_a, eia_a)
        t2ab /= lib.direct_sum('ia+jb->ijab', eia_a, eia_b)
        t2bb /= lib.direct_sum('ia+jb->ijab', eia_b, eia_b)

        emp2  = numpy.einsum('iajb,ijab', eris_ovov, t2aa) * .25
        emp2 -= numpy.einsum('jaib,ijab', eris_ovov, t2aa) * .25
        emp2 += numpy.einsum('iajb,ijab', eris_OVOV, t2bb) * .25
        emp2 -= numpy.einsum('jaib,ijab', eris_OVOV, t2bb) * .25
        emp2 += numpy.einsum('iajb,ijab', eris_ovOV, t2ab)
        self.emp2 = emp2.real
        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)

        if abs(emp2) < 1e-3 and (abs(t1a).sum()+abs(t1b).sum()) < 1e-3:
            t1a = 1e-1 / eia_a
            t1b = 1e-1 / eia_b

        ci_guess = amplitudes_to_cisdvec(1, (t1a,t1b), (t2aa,t2ab,t2bb))

        if nroots > 1:
            civec_size = ci_guess.size
            ci1_size = t1a.size + t1b.size
            dtype = ci_guess.dtype
            nroots = min(ci1_size+1, nroots)

            if diag is None:
                idx = range(1, nroots)
            else:
                idx = diag[:ci1_size+1].argsort()[1:nroots]  # exclude HF determinant

            ci_guess = [ci_guess]
            for i in idx:
                g = numpy.zeros(civec_size, dtype)
                g[i] = 1.0
                ci_guess.append(g)

        return self.emp2, ci_guess

def update_amps(mycc, t1, t2, l1, l2, eris=None, saved=None):
    if eris is None: eris = ccsd._ERIS(mycc)
    if saved is None: saved = make_intermediates(mycc, t1, t2, eris)

    nocc, nvir = t1.shape
    l1, l2 = ccsd_lambda.update_amps(mycc, t1, t2, l1, l2, eris, saved)

    mo_e = eris.fock.diagonal()
    eia = lib.direct_sum('i-a->ia', mo_e[:nocc], mo_e[nocc:])
    l1 += saved.jov/eia * .5

    eijab = lib.direct_sum('ia+jb->ijab', eia, eia)
    l2 += (saved.joovv*(2./3)+saved.joovv.transpose(1,0,2,3)*(1./3)) / eijab
    return l1, l2

def make_intermediates(mycc, t1, t2, eris):
    imds = ccsd_lambda.make_intermediates(mycc, t1, t2, eris)

    nocc, nvir = t1.shape
    eris_ovvv = numpy.asarray(eris.get_ovvv())
    eris_ovoo = numpy.asarray(eris.ovoo)
    eris_ovov = numpy.asarray(eris.ovov)

    mo_e = eris.mo_energy
    eia = lib.direct_sum('i-a->ia', mo_e[:nocc], mo_e[nocc:])
    d3 = lib.direct_sum('ia,jb,kc->ijkabc', eia, eia, eia)

    def p6(t):
        t1 = t + t.transpose(0,2,1,3,5,4)
        return t1 + t1.transpose(1,0,2,4,3,5) + t1.transpose(1,2,0,4,5,3)

    def r6(w):
        return (4 * w + w.transpose(0,1,2,4,5,3) + w.transpose(0,1,2,5,3,4)
                - 2 * w.transpose(0,1,2,5,4,3) - 2 * w.transpose(0,1,2,3,5,4)
                - 2 * w.transpose(0,1,2,4,3,5))

    w =(numpy.einsum('iafb,kjcf->ijkabc', eris_ovvv.conj(), t2)
      - numpy.einsum('iajm,mkbc->ijkabc', eris_ovoo.conj(), t2)) / d3
    v =(numpy.einsum('iajb,kc->ijkabc', eris_ovov.conj(), t1)
      + numpy.einsum('ck,ijab->ijkabc', eris.fock[nocc:,:nocc], t2)) / d3
    w = p6(w)
    v = p6(v)

    imds.l1_t = numpy.einsum('jbkc,ijkabc->ia', eris_ovov, r6(w)).conj() / eia * .5

    def as_r6(m):
        # When making derivative over t2, r6 should be called on the 6-index
        # tensor. It gives the equation for lambda2, but not corresponding to
        # the lambda equation used by RCCSD-lambda code.  A transformation was
        # applied in RCCSD-lambda equation  F(lambda)_{ijab} = 0:
        #       2/3 * # F(lambda)_{ijab} + 1/3 * F(lambda)_{jiab} = 0
        # Combining this transformation with r6 operation, leads to the
        # transformation code below
        return m * 2 - m.transpose(0,1,2,5,4,3) - m.transpose(0,1,2,3,5,4)

    m = as_r6(w * 2 + v * .5)
    joovv = numpy.einsum('kfbe,ijkaef->ijab', eris_ovvv, m.conj())
    joovv-= numpy.einsum('ncmj,imnabc->ijab', eris_ovoo, m.conj())
    joovv = joovv + joovv.transpose(1,0,3,2)
    rw = as_r6(w)
    joovv+= numpy.einsum('kc,ijkabc->ijab', eris.fock[:nocc,nocc:], rw.conj())
    imds.l2_t = joovv / lib.direct_sum('ia+jb->ijab', eia, eia)

    return imds

def solve_withs1(fvind, mo_energy, mo_occ, h1, s1,
                 max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
    ''' C^1_{ij} = -1/2 S1
    e1 = h1 - s1*e0 + (e0_j-e0_i)*c1 + vhf[c1]
    '''
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(sys.stdout, verbose)
    t0 = (time.clock(), time.time())

    occidx = mo_occ > 0
    viridx = mo_occ == 0
    e_a = mo_energy[viridx]
    e_i = mo_energy[occidx]
    e_ai = 1 / lib.direct_sum('a-i->ai', e_a, e_i)
    nvir, nocc = e_ai.shape
    nmo = nocc + nvir

# Do not reshape h1 because h1 may be a 2D array (nvir,nocc)
    s1 = s1.reshape(-1,nmo,nocc)

    hs = mo1base = h1.reshape(-1,nmo,nocc) - s1*e_i
    mo_e1 = hs[:,occidx,:].copy()

    mo1base[:,viridx] *= -e_ai
    mo1base[:,occidx] = -s1[:,occidx] * .5

    def vind_vo(mo1):
        v = fvind(mo1.reshape(h1.shape)).reshape(-1,nmo,nocc)
        v[:,viridx,:] *= e_ai
        v[:,occidx,:] = 0
        return v.ravel()
    mo1 = lib.krylov(vind_vo, mo1base.ravel(),
                     tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log)
    mo1 = mo1.reshape(mo1base.shape)
    log.timer('krylov solver in CPHF', *t0)

    v_mo = fvind(mo1.reshape(h1.shape)).reshape(-1,nmo,nocc)
    mo1[:,viridx] = mo1base[:,viridx] - v_mo[:,viridx]*e_ai

    mo_e1 += mo1[:,occidx] * lib.direct_sum('i-j->ij', e_i, e_i)
    mo_e1 += v_mo[:,occidx,:]

    if h1.ndim == 3:
        return mo1, mo_e1
    else:
        return mo1.reshape(h1.shape), mo_e1.reshape(nocc,nocc)

    def fupdate(t1, t2, istep, normt, de, adiis):
        nocc, nvir = t1.shape
        nov = nocc*nvir
        moidx = numpy.ones(mycc.mo_energy.size, dtype=numpy.bool)
        if isinstance(mycc.frozen, (int, numpy.integer)):
            moidx[:mycc.frozen] = False
        else:
            moidx[mycc.frozen] = False
        mo_e = mycc.mo_energy[moidx]
        eia = mo_e[:nocc,None] - mo_e[None,nocc:]
        if (istep > mycc.diis_start_cycle and
            abs(de) < mycc.diis_start_energy_diff):
            if mycc.t1 is None:
                mycc.t1 = t1
                mycc.t2 = t2
            else:
                tbuf = numpy.empty(nov*(nov+1))
                tbuf[:nov] = ((t1-mycc.t1)*eia).ravel()
                pbuf = tbuf[nov:].reshape(nocc,nocc,nvir,nvir)
                for i in range(nocc):
                    pbuf[i] = (t2[i]-mycc.t2[i]) * lib.direct_sum('jb,a->jba', eia, eia[i])
                adiis.push_err_vec(tbuf)
                tbuf = numpy.empty(nov*(nov+1))
                tbuf[:nov] = t1.ravel()
                tbuf[nov:] = t2.ravel()
                t1.data = tbuf.data # release memory
                t2.data = tbuf.data

                tbuf = adiis.update(tbuf)
                mycc.t1 = t1 = tbuf[:nov].reshape(nocc,nvir)
                mycc.t2 = t2 = tbuf[nov:].reshape(nocc,nocc,nvir,nvir)
            logger.debug(mycc, 'DIIS for step %d', istep)
        return t1, t2

def gen_hop_rhf_external(mf):
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nmo = mo_coeff.shape[1]
    nocc = numpy.count_nonzero(mo_occ)
    nvir = nmo - nocc
    nov = nocc * nvir

    eri_mo = ao2mo.full(mol, mo_coeff)
    eri_mo = ao2mo.restore(1, eri_mo, nmo)
    eai = lib.direct_sum('a-i->ai', mo_energy[nocc:], mo_energy[:nocc])
    # A
    h = numpy.einsum('ckld->kcld', eri_mo[nocc:,:nocc,:nocc,nocc:]) * 2
    h-= numpy.einsum('cdlk->kcld', eri_mo[nocc:,nocc:,:nocc,:nocc])
    for a in range(nvir):
        for i in range(nocc):
            h[i,a,i,a] += eai[a,i]
    # B
    h-= numpy.einsum('ckdl->kcld', eri_mo[nocc:,:nocc,nocc:,:nocc]) * 2
    h+= numpy.einsum('cldk->kcld', eri_mo[nocc:,:nocc,nocc:,:nocc])
    h1 = h.transpose(1,0,3,2).reshape(nov,nov)
    def hop1(x):
        return h1.dot(x)

    h =-numpy.einsum('cdlk->kcld', eri_mo[nocc:,nocc:,:nocc,:nocc])
    for a in range(nvir):
        for i in range(nocc):
            h[i,a,i,a] += eai[a,i]
    h-= numpy.einsum('cldk->kcld', eri_mo[nocc:,:nocc,nocc:,:nocc])
    h2 = h.transpose(1,0,3,2).reshape(nov,nov)
    def hop2(x):
        return h2.dot(x)
    return hop1, hop2

def solve_nos1(fvind, mo_energy, mo_occ, h1,
               max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(sys.stdout, verbose)
    t0 = (time.clock(), time.time())

    e_a = mo_energy[mo_occ==0]
    e_i = mo_energy[mo_occ>0]
    e_ai = 1 / lib.direct_sum('a-i->ai', e_a, e_i)
    nocc = e_i.size
    nvir = e_a.size
    nmo = nocc + nvir

    mo1base = h1 * -e_ai

    def vind_vo(mo1):
        v = fvind(mo1.reshape(h1.shape)).reshape(h1.shape)
        v *= e_ai
        return v.ravel()
    mo1 = lib.krylov(vind_vo, mo1base.ravel(),
                     tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log)
    log.timer('krylov solver in CPHF', *t0)
    return mo1.reshape(h1.shape), None

def kernel(mp, mo_energy=None, mo_coeff=None, eris=None, with_t2=WITH_T2,
           verbose=logger.NOTE):
    if mo_energy is None or mo_coeff is None:
        mo_coeff = mp2._mo_without_core(mp, mp.mo_coeff)
        mo_energy = mp2._mo_energy_without_core(mp, mp.mo_energy)
    else:
        # For backward compatibility.  In pyscf-1.4 or earlier, mp.frozen is
        # not supported when mo_energy or mo_coeff is given.
        assert(mp.frozen is 0 or mp.frozen is None)

    nocc = mp.nocc
    nvir = mp.nmo - nocc
    eia = mo_energy[:nocc,None] - mo_energy[None,nocc:]

    t2 = None
    emp2 = 0
    for istep, qov in enumerate(mp.loop_ao2mo(mo_coeff, nocc)):
        logger.debug(mp, 'Load cderi step %d', istep)
        for i in range(nocc):
            buf = numpy.dot(qov[:,i*nvir:(i+1)*nvir].T,
                            qov).reshape(nvir,nocc,nvir)
            gi = numpy.array(buf, copy=False)
            gi = gi.reshape(nvir,nocc,nvir).transpose(1,0,2)
            t2i = gi/lib.direct_sum('jb+a->jba', eia, eia[i])
            emp2 += numpy.einsum('jab,jab', t2i, gi) * 2
            emp2 -= numpy.einsum('jab,jba', t2i, gi)

    return emp2, t2