Python _ao2mo.nr_e2函数代码示例

    def vind(xys):
        nz = len(xys)
        if wfnsym is not None and mol.symmetry:
            # shape(nz,2,nocc,nvir): 2 ~ X,Y
            xys = numpy.copy(xys).reshape(nz,2,nocc,nvir)
            xys[:,:,sym_forbid] = 0
        dms = numpy.empty((nz,nao,nao))
        for i in range(nz):
            x, y = xys[i].reshape(2,nocc,nvir)
            # *2 for double occupancy
            dmx = reduce(numpy.dot, (orbo, x*2, orbv.T))
            dmy = reduce(numpy.dot, (orbv, y.T*2, orbo.T))
            dms[i] = dmx + dmy  # AX + BY

        v1ao = vresp(dms)
        v1ov = _ao2mo.nr_e2(v1ao, mo_coeff, (0,nocc,nocc,nmo)).reshape(-1,nocc,nvir)
        v1vo = _ao2mo.nr_e2(v1ao, mo_coeff, (nocc,nmo,0,nocc)).reshape(-1,nvir,nocc)
        hx = numpy.empty((nz,2,nocc,nvir), dtype=v1ov.dtype)
        for i in range(nz):
            x, y = xys[i].reshape(2,nocc,nvir)
            hx[i,0] = v1ov[i]
            hx[i,0]+= numpy.einsum('sp,qs->qp', fvv, x)  # AX
            hx[i,0]-= numpy.einsum('sp,pr->sr', foo, x)  # AX
            hx[i,1] =-v1vo[i].T
            hx[i,1]-= numpy.einsum('sp,qs->qp', fvv, y)  #-AY
            hx[i,1]+= numpy.einsum('sp,pr->sr', foo, y)  #-AY

        if wfnsym is not None and mol.symmetry:
            hx[:,:,sym_forbid] = 0
        return hx.reshape(nz,-1)

 def vind(mo1):
     mo1aa, mo1ab, mo1ba, mo1bb = _split_mo1(mo1)
     dm1aa = _dm1_mo2ao(mo1aa, orbva, orboa)
     dm1ab = _dm1_mo2ao(mo1ab, orbva, orbob)
     dm1ba = _dm1_mo2ao(mo1ba, orbvb, orboa)
     dm1bb = _dm1_mo2ao(mo1bb, orbvb, orbob)
     dm1 = lib.asarray([dm1aa+dm1aa.transpose(0,2,1),
                        dm1ab+dm1ba.transpose(0,2,1),
                        dm1ba+dm1ab.transpose(0,2,1),
                        dm1bb+dm1bb.transpose(0,2,1)])
     v1 = vresp(dm1)
     v1aa = _ao2mo.nr_e2(v1[0], mo_va_oa, (0,nvira,nvira,nvira+nocca))
     v1ab = _ao2mo.nr_e2(v1[1], mo_va_ob, (0,nvira,nvira,nvira+noccb))
     v1ba = _ao2mo.nr_e2(v1[2], mo_vb_oa, (0,nvirb,nvirb,nvirb+nocca))
     v1bb = _ao2mo.nr_e2(v1[3], mo_vb_ob, (0,nvirb,nvirb,nvirb+noccb))
     v1aa = v1aa.reshape(nset,nvira,nocca)
     v1ab = v1ab.reshape(nset,nvira,noccb)
     v1ba = v1ba.reshape(nset,nvirb,nocca)
     v1bb = v1bb.reshape(nset,nvirb,noccb)
     v1aa *= eai_aa
     v1ab *= eai_ab
     v1ba *= eai_ba
     v1bb *= eai_bb
     v1mo = numpy.hstack((v1aa.reshape(nset,-1),
                          v1ab.reshape(nset,-1),
                          v1ba.reshape(nset,-1),
                          v1bb.reshape(nset,-1)))
     return v1mo.ravel()

def _trans_aapp_(mo, ncore, ncas, fload, ao_loc=None):
    nmo = mo[0].shape[1]
    nocc = (ncore[0] + ncas, ncore[1] + ncas)
    c_nmo = ctypes.c_int(nmo)

    klshape = (0, nmo, 0, nmo)

    japcv = numpy.empty((ncas,nmo,ncore[0],nmo-ncore[0]))
    aapp = numpy.empty((ncas,ncas,nmo,nmo))
    aaPP = numpy.empty((ncas,ncas,nmo,nmo))
    appa = numpy.empty((ncas,nmo,nmo,ncas))
    apPA = numpy.empty((ncas,nmo,nmo,ncas))
    apCV = numpy.empty((ncas,nmo,ncore[1],nmo-ncore[1]))
    ppp = numpy.empty((nmo,nmo,nmo))
    for i in range(ncas):
        buf = _ao2mo.nr_e2(fload(ncore[0]+i), mo[0], klshape,
                            aosym='s4', mosym='s2', ao_loc=ao_loc)
        lib.unpack_tril(buf, out=ppp)
        aapp[i] = ppp[ncore[0]:nocc[0]]
        appa[i] = ppp[:,:,ncore[0]:nocc[0]]
        #japcp = avcp * 2 - acpv.transpose(0,2,1,3) - avcp.transpose(0,3,2,1)
        japcv[i] = ppp[:,:ncore[0],ncore[0]:] * 2 \
                 - ppp[:ncore[0],:,ncore[0]:].transpose(1,0,2) \
                 - ppp[ncore[0]:,:ncore[0],:].transpose(2,1,0)

        buf = _ao2mo.nr_e2(fload(ncore[0]+i), mo[1], klshape,
                            aosym='s4', mosym='s2', ao_loc=ao_loc)
        lib.unpack_tril(buf, out=ppp)
        aaPP[i] = ppp[ncore[0]:nocc[0]]
        apPA[i] = ppp[:,:,ncore[1]:nocc[1]]
        apCV[i] = ppp[:,:ncore[1],ncore[1]:]

    return aapp, aaPP, appa, apPA, japcv, apCV

 def vind(mo1):
     mo1a = mo1.reshape(-1,nova+novb)[:,:nova].reshape(-1,nvira,nocca)
     mo1b = mo1.reshape(-1,nova+novb)[:,nova:].reshape(-1,nvirb,noccb)
     nset = mo1a.shape[0]
     dm1a = _dm1_mo2ao(mo1a, orbva, orboa)
     dm1b = _dm1_mo2ao(mo1b, orbvb, orbob)
     dm1 = numpy.vstack([dm1a-dm1a.transpose(0,2,1),
                         dm1b-dm1b.transpose(0,2,1)])
     v1 = vresp(dm1)
     v1a = _ao2mo.nr_e2(v1[    :nset], mo_va_oa, (0,nvira,nvira,nvira+nocca))
     v1b = _ao2mo.nr_e2(v1[nset:    ], mo_vb_ob, (0,nvirb,nvirb,nvirb+noccb))
     v1mo = numpy.hstack((v1a.reshape(nset,-1), v1b.reshape(nset,-1)))
     return v1mo.ravel()

def trans_e1_incore(eri_ao, mo, ncore, ncas):
    nmo = mo.shape[1]
    nocc = ncore + ncas
    eri1 = pyscf.ao2mo.incore.half_e1(eri_ao, (mo,mo[:,:nocc]), compact=False)
    eri1 = eri1.reshape(nmo,nocc,-1)

    klppshape = (0, nmo, 0, nmo)
    klpashape = (0, nmo, ncore, nocc)
    aapp = numpy.empty((ncas,ncas,nmo,nmo))
    for i in range(ncas):
        _ao2mo.nr_e2(eri1[ncore+i,ncore:nocc], mo, klppshape,
                      aosym='s4', mosym='s1', out=aapp[i])
    ppaa = pyscf.lib.transpose(aapp.reshape(ncas*ncas,-1)).reshape(nmo,nmo,ncas,ncas)
    aapp = None

    papa = numpy.empty((nmo,ncas,nmo,ncas))
    for i in range(nmo):
        _ao2mo.nr_e2(eri1[i,ncore:nocc], mo, klpashape,
                      aosym='s4', mosym='s1', out=papa[i])

    pp = numpy.empty((nmo,nmo))
    j_cp = numpy.zeros((ncore,nmo))
    k_pc = numpy.zeros((nmo,ncore))
    for i in range(ncore):
        _ao2mo.nr_e2(eri1[i,i:i+1], mo, klppshape, aosym='s4', mosym='s1', out=pp)
        j_cp[i] = pp.diagonal()
    j_pc = j_cp.T.copy()

    pp = numpy.empty((ncore,ncore))
    for i in range(nmo):
        klshape = (i, i+1, 0, ncore)
        _ao2mo.nr_e2(eri1[i,:ncore], mo, klshape, aosym='s4', mosym='s1', out=pp)
        k_pc[i] = pp.diagonal()
    return j_pc, k_pc, ppaa, papa

 def vind(mo1):
     dm1 = _dm1_mo2ao(mo1.reshape(nset,nvir,nocc), orbv, orbo*2)  # *2 for double occupancy
     dm1 = dm1 + dm1.transpose(0,2,1)
     v1 = vresp(dm1)
     v1 = _ao2mo.nr_e2(v1, mo_v_o, (0,nvir,nvir,nmo)).reshape(nset,nvir,nocc)
     v1 *= eai
     return v1.ravel()

    def get_vind(self, zs):
        mol = self.mol
        mo_coeff = self._scf.mo_coeff
        mo_energy = self._scf.mo_energy
        nao, nmo = mo_coeff.shape
        nocc = (self._scf.mo_occ>0).sum()
        nvir = nmo - nocc
        orbv = mo_coeff[:,nocc:]
        orbo = mo_coeff[:,:nocc]
        eai = pyscf.lib.direct_sum('a-i->ai', mo_energy[nocc:], mo_energy[:nocc])
        dai = numpy.sqrt(eai).ravel()

        nz = len(zs)
        dmvo = numpy.empty((nz,nao,nao))
        for i, z in enumerate(zs):
            dm = reduce(numpy.dot, (orbv, (dai*z).reshape(nvir,nocc), orbo.T))
            dmvo[i] = dm + dm.T # +cc for A+B and K_{ai,jb} in A == K_{ai,bj} in B

        mem_now = pyscf.lib.current_memory()[0]
        max_memory = max(2000, self.max_memory*.9-mem_now)
        v1ao = _contract_xc_kernel(self, self._scf.xc, dmvo,
                                   singlet=self.singlet, max_memory=max_memory)

        if self.singlet:
            vj = self._scf.get_j(mol, dmvo, hermi=1)
            v1ao += vj * 2

        v1vo = _ao2mo.nr_e2(v1ao, mo_coeff, (nocc,nmo,0,nocc)).reshape(-1,nvir*nocc)
        edai = eai.ravel() * dai
        for i, z in enumerate(zs):
            # numpy.sqrt(eai) * (eai*dai*z + v1vo)
            v1vo[i] += edai*z
            v1vo[i] *= dai
        return v1vo.reshape(nz,-1)

def _trans(mo, ncore, ncas, fload, cvcv=None, ao_loc=None):
    nao, nmo = mo.shape
    nocc = ncore + ncas
    nvir = nmo - nocc
    nav = nmo - ncore

    if cvcv is None:
        cvcv = numpy.zeros((ncore*nvir,ncore*nvir))
    pacv = numpy.empty((nmo,ncas,ncore*nvir))
    aapp = numpy.empty((ncas,ncas,nmo*nmo))
    papa = numpy.empty((nmo,ncas,nmo*ncas))
    vcv = numpy.empty((nav,ncore*nvir))
    apa = numpy.empty((ncas,nmo*ncas))
    vpa = numpy.empty((nav,nmo*ncas))
    app = numpy.empty((ncas,nmo*nmo))
    for i in range(ncore):
        buf = fload(i, i+1)
        klshape = (0, ncore, nocc, nmo)
        _ao2mo.nr_e2(buf, mo, klshape,
                      aosym='s4', mosym='s1', out=vcv, ao_loc=ao_loc)
        cvcv[i*nvir:(i+1)*nvir] = vcv[ncas:]
        pacv[i] = vcv[:ncas]

        klshape = (0, nmo, ncore, nocc)
        _ao2mo.nr_e2(buf[:ncas], mo, klshape,
                      aosym='s4', mosym='s1', out=apa, ao_loc=ao_loc)
        papa[i] = apa
    for i in range(ncas):
        buf = fload(ncore+i, ncore+i+1)
        klshape = (0, ncore, nocc, nmo)
        _ao2mo.nr_e2(buf, mo, klshape,
                      aosym='s4', mosym='s1', out=vcv, ao_loc=ao_loc)
        pacv[ncore:,i] = vcv

        klshape = (0, nmo, ncore, nocc)
        _ao2mo.nr_e2(buf, mo, klshape,
                      aosym='s4', mosym='s1', out=vpa, ao_loc=ao_loc)
        papa[ncore:,i] = vpa

        klshape = (0, nmo, 0, nmo)
        _ao2mo.nr_e2(buf[:ncas], mo, klshape,
                      aosym='s4', mosym='s1', out=app, ao_loc=ao_loc)
        aapp[i] = app
    #pyscf.lib.transpose(aapp.reshape(ncas**2, -1), inplace=True)
    ppaa = pyscf.lib.transpose(aapp.reshape(ncas**2,-1))
    return (ppaa.reshape(nmo,nmo,ncas,ncas), papa.reshape(nmo,ncas,nmo,ncas),
            pacv.reshape(nmo,ncas,ncore,nvir), cvcv)

 def loop_ao2mo(self, mo_coeff, nocc):
     mo = numpy.asarray(mo_coeff, order='F')
     nmo = mo.shape[1]
     ijslice = (0, nocc, nocc, nmo)
     Lov = None
     for eri1 in self._scf.with_df.loop():
         Lov = _ao2mo.nr_e2(eri1, mo, ijslice, aosym='s2', out=Lov)
         yield Lov

 def ao2mo(self, mo_coeffs, compact=True):
     if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
         mo_coeffs = (mo_coeffs,) * 4
     ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact)
     klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact)
     mo_eri = numpy.zeros((nij_pair,nkl_pair))
     sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
            iden_coeffs(mo_coeffs[1], mo_coeffs[3]))
     Lij = Lkl = None
     for eri1 in self.loop():
         Lij = _ao2mo.nr_e2(eri1, moij, ijslice, aosym='s2', mosym=ijmosym, out=Lij)
         if sym:
             Lkl = Lij
         else:
             Lkl = _ao2mo.nr_e2(eri1, mokl, klslice, aosym='s2', mosym=klmosym, out=Lkl)
         lib.dot(Lij.T, Lkl, 1, mo_eri, 1)
     return mo_eri

def _trans_aapp_(mo, ncore, ncas, fload, ao_loc=None):
    nmo = mo[0].shape[1]
    nocc = (ncore[0] + ncas, ncore[1] + ncas)
    c_nmo = ctypes.c_int(nmo)
    funpack = pyscf.lib.numpy_helper._np_helper.NPdunpack_tril

    klshape = (0, nmo, 0, nmo)

    japcv = numpy.empty((ncas,nmo,ncore[0],nmo-ncore[0]))
    aapp = numpy.empty((ncas,ncas,nmo,nmo))
    aaPP = numpy.empty((ncas,ncas,nmo,nmo))
    appa = numpy.empty((ncas,nmo,nmo,ncas))
    apPA = numpy.empty((ncas,nmo,nmo,ncas))
    apCV = numpy.empty((ncas,nmo,ncore[1],nmo-ncore[1]))
    ppp = numpy.empty((nmo,nmo,nmo))
    for i in range(ncas):
        buf = _ao2mo.nr_e2(fload(ncore[0]+i), mo[0], klshape,
                            aosym='s4', mosym='s2', ao_loc=ao_loc)
        for j in range(nmo):
            funpack(c_nmo, buf[j].ctypes.data_as(ctypes.c_void_p),
                    ppp[j].ctypes.data_as(ctypes.c_void_p), ctypes.c_int(1))
        aapp[i] = ppp[ncore[0]:nocc[0]]
        appa[i] = ppp[:,:,ncore[0]:nocc[0]]
        #japcp = avcp * 2 - acpv.transpose(0,2,1,3) - avcp.transpose(0,3,2,1)
        japcv[i] = ppp[:,:ncore[0],ncore[0]:] * 2 \
                 - ppp[:ncore[0],:,ncore[0]:].transpose(1,0,2) \
                 - ppp[ncore[0]:,:ncore[0],:].transpose(2,1,0)

        buf = _ao2mo.nr_e2(fload(ncore[0]+i), mo[1], klshape,
                            aosym='s4', mosym='s2', ao_loc=ao_loc)
        for j in range(nmo):
            funpack(c_nmo, buf[j].ctypes.data_as(ctypes.c_void_p),
                    ppp[j].ctypes.data_as(ctypes.c_void_p), ctypes.c_int(1))
        aaPP[i] = ppp[ncore[0]:nocc[0]]
        apPA[i] = ppp[:,:,ncore[1]:nocc[1]]
        apCV[i] = ppp[:,:ncore[1],ncore[1]:]

    return aapp, aaPP, appa, apPA, japcv, apCV

 def vind(zs):
     nz = len(zs)
     dmov = numpy.empty((nz,nao,nao))
     for i, z in enumerate(zs):
         # *2 for double occupancy
         dm = reduce(numpy.dot, (orbo, (d_ia*z).reshape(nocc,nvir)*2, orbv.T))
         dmov[i] = dm + dm.T # +cc for A+B and K_{ai,jb} in A == K_{ai,bj} in B
     v1ao = vresp(dmov)
     v1ov = _ao2mo.nr_e2(v1ao, mo_coeff, (0,nocc,nocc,nmo)).reshape(-1,nocc*nvir)
     for i, z in enumerate(zs):
         # numpy.sqrt(e_ia) * (e_ia*d_ia*z + v1ov)
         v1ov[i] += ed_ia*z
         v1ov[i] *= d_ia
     return v1ov.reshape(nz,-1)

        def vind(zs):
            nz = len(zs)
            if wfnsym is not None and mol.symmetry:
                zs = numpy.copy(zs)
                zs[:,sym_forbid] = 0
            dmov = numpy.empty((2,nz,nao,nao))
            for i in range(nz):
                z = d_ia * zs[i]
                za = z[:nocca*nvira].reshape(nocca,nvira)
                zb = z[nocca*nvira:].reshape(noccb,nvirb)
                dm = reduce(numpy.dot, (orboa, za, orbva.T))
                dmov[0,i] = dm + dm.T
                dm = reduce(numpy.dot, (orbob, zb, orbvb.T))
                dmov[1,i] = dm + dm.T

            v1ao = vresp(dmov)
            v1a = _ao2mo.nr_e2(v1ao[0], mo_coeff[0], (0,nocca,nocca,nmo))
            v1b = _ao2mo.nr_e2(v1ao[1], mo_coeff[1], (0,noccb,noccb,nmo))
            hx = numpy.hstack((v1a.reshape(nz,-1), v1b.reshape(nz,-1)))
            for i, z in enumerate(zs):
                hx[i] += ed_ia * z
                hx[i] *= d_ia
            return hx

    def loop_ao2mo(self, mo_coeff, nocc):
        mo = numpy.asarray(mo_coeff, order='F')
        nmo = mo.shape[1]
        ijslice = (0, nocc, nocc, nmo)
        Lov = None
        with_df = self.with_df

        nvir = nmo - nocc
        naux = with_df.get_naoaux()
        mem_now = lib.current_memory()[0]
        max_memory = max(2000, self.max_memory*.9-mem_now)
        blksize = int(min(naux, max(with_df.blockdim,
                                    (max_memory*1e6/8-nocc*nvir**2*2)/(nocc*nvir))))
        for eri1 in with_df.loop(blksize=blksize):
            Lov = _ao2mo.nr_e2(eri1, mo, ijslice, aosym='s2', out=Lov)
            yield Lov

        def _exact_paaa(self, mo, u, out=None):
            if self.with_df:
                nmo = mo.shape[1]
                ncore = self.ncore
                ncas = self.ncas
                nocc = ncore + ncas
                mo1 = numpy.dot(mo, u)
                mo1_cas = mo1[:,ncore:nocc]

                paaa = numpy.zeros((nmo*ncas,ncas*ncas))
                moij = numpy.asarray(numpy.hstack((mo1, mo1_cas)), order='F')
                ijshape = (0, nmo, nmo, nmo+ncas)
                for eri1 in self.with_df.loop():
                    bufpa = _ao2mo.nr_e2(eri1, moij, ijshape, 's2', 's1')
                    bufaa = numpy.asarray(buf1[ncore:nocc,:], order='C')
                    pyscf.lib.dot(bufpa.T, bufaa, 1, paaa, 1)
                return paaa.reshape(nmo,ncas,ncas,ncas)
            else:
                return casscf_class._exact_paaa(self, mol, u, out)