Python lib.current_memory函数代码示例

def build_Lpq_pbc(mydf, auxcell, kptij_lst):
    """Fitting coefficients for auxiliary functions"""
    kpts_ji = kptij_lst[:, 1] - kptij_lst[:, 0]
    uniq_kpts, uniq_index, uniq_inverse = unique(kpts_ji)
    max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]))
    if mydf.metric.upper() == "S":
        outcore.aux_e2(
            mydf.cell, auxcell, mydf._cderi, "cint3c1e_sph", kptij_lst=kptij_lst, dataname="Lpq", max_memory=max_memory
        )
        s_aux = auxcell.pbc_intor("cint1e_ovlp_sph", hermi=1, kpts=uniq_kpts)
    else:  # mydf.metric.upper() == 'T'
        outcore.aux_e2(
            mydf.cell,
            auxcell,
            mydf._cderi,
            "cint3c1e_p2_sph",
            kptij_lst=kptij_lst,
            dataname="Lpq",
            max_memory=max_memory,
        )
        s_aux = [x * 2 for x in auxcell.pbc_intor("cint1e_kin_sph", hermi=1, kpts=uniq_kpts)]

    s_aux = [scipy.linalg.cho_factor(x) for x in s_aux]

    max_memory = mydf.max_memory - lib.current_memory()[0]
    naux = auxcell.nao_nr()
    blksize = max(int(max_memory * 0.5 * 1e6 / 16 / naux / mydf.blockdim), 1) * mydf.blockdim
    with h5py.File(mydf._cderi) as feri:
        for k, where in enumerate(uniq_inverse):
            s_k = s_aux[where]
            key = "Lpq/%d" % k
            Lpq = feri[key]
            nao_pair = Lpq.shape[1]
            for p0, p1 in lib.prange(0, nao_pair, blksize):
                Lpq[:, p0:p1] = scipy.linalg.cho_solve(s_k, Lpq[:, p0:p1])

        def get_hcore(self, mol=None):
            if mol is None: mol = self.mol
            if getattr(scf_method, 'get_hcore', None):
                h1e = method_class.get_hcore(self, mol)
            else:  # DO NOT modify post-HF objects to avoid the MM charges applied twice
                raise RuntimeError('mm_charge function cannot be applied on post-HF methods')

            if pyscf.DEBUG:
                v = 0
                for i,q in enumerate(charges):
                    mol.set_rinv_origin(coords[i])
                    v += mol.intor('int1e_rinv') * -q
            else:
                if mol.cart:
                    intor = 'int3c2e_cart'
                else:
                    intor = 'int3c2e_sph'
                nao = mol.nao
                max_memory = self.max_memory - lib.current_memory()[0]
                blksize = int(min(max_memory*1e6/8/nao**2, 200))
                v = 0
                for i0, i1 in lib.prange(0, charges.size, blksize):
                    fakemol = gto.fakemol_for_charges(coords[i0:i1])
                    j3c = df.incore.aux_e2(mol, fakemol, intor=intor, aosym='s2ij')
                    v += numpy.einsum('xk,k->x', j3c, -charges[i0:i1])
                v = lib.unpack_tril(v)
            return h1e + v

def get_pp_loc_part1(mydf, cell, kpts):
    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = t0 = (time.clock(), time.time())
    nkpts = len(kpts)

    gs = mydf.gs
    nao = cell.nao_nr()
    Gv = cell.get_Gv(gs)
    SI = cell.get_SI(Gv)
    vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
    vpplocG = -1./cell.vol * numpy.einsum('ij,ij->j', SI, vpplocG)
    kpt_allow = numpy.zeros(3)
    real = gamma_point(kpts)

    if real:
        vloc = numpy.zeros((nkpts,nao**2))
    else:
        vloc = numpy.zeros((nkpts,nao**2), dtype=numpy.complex128)
    max_memory = mydf.max_memory - lib.current_memory()[0]
    for k, pqkR, pqkI, p0, p1 \
            in mydf.ft_loop(cell, mydf.gs, kpt_allow, kpts, max_memory=max_memory):
        vG = vpplocG[p0:p1]
        if not real:
            vloc[k] += numpy.einsum('k,xk->x', vG.real, pqkI) * 1j
            vloc[k] += numpy.einsum('k,xk->x', vG.imag, pqkR) *-1j
        vloc[k] += numpy.einsum('k,xk->x', vG.real, pqkR)
        vloc[k] += numpy.einsum('k,xk->x', vG.imag, pqkI)
        pqkR = pqkI = None
    t1 = log.timer_debug1('contracting vloc part1', *t1)
    return vloc.reshape(-1,nao,nao)

    def build(self):
        t0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)
        mol = self.mol
        auxmol = self.auxmol = incore.format_aux_basis(self.mol, self.auxbasis)
        nao = mol.nao_nr()
        naux = auxmol.nao_nr()
        nao_pair = nao*(nao+1)//2

        max_memory = (self.max_memory - lib.current_memory()[0]) * .8
        if nao_pair*naux*3*8/1e6 < max_memory:
            self._cderi = incore.cholesky_eri(mol, auxmol=auxmol, verbose=log)
        else:
            if not isinstance(self._cderi, str):
                if isinstance(self._cderi_file, str):
                    self._cderi = self._cderi_file
                else:
                    self._cderi = self._cderi_file.name
            outcore.cholesky_eri(mol, self._cderi, auxmol=auxmol, verbose=log)
            if nao_pair*nao*8/1e6 < max_memory:
                with addons.load(self._cderi) as feri:
                    cderi = numpy.asarray(feri)
                self._cderi = cderi
            log.timer_debug1('Generate density fitting integrals', *t0)

        return self

def kernel(casci, mo_coeff=None, ci0=None, verbose=logger.NOTE):
    '''UHF-CASCI solver
    '''
    if mo_coeff is None: mo_coeff = casci.mo_coeff
    log = logger.new_logger(casci, verbose)
    t0 = (time.clock(), time.time())
    log.debug('Start uhf-based CASCI')

    ncas = casci.ncas
    nelecas = casci.nelecas
    ncore = casci.ncore
    mo_core, mo_cas, mo_vir = extract_orbs(mo_coeff, ncas, nelecas, ncore)

    # 1e
    h1eff, energy_core = casci.h1e_for_cas(mo_coeff)
    log.debug('core energy = %.15g', energy_core)
    t1 = log.timer('effective h1e in CAS space', *t0)

    # 2e
    eri_cas = casci.get_h2eff(mo_cas)
    t1 = log.timer('integral transformation to CAS space', *t1)

    # FCI
    max_memory = max(400, casci.max_memory-lib.current_memory()[0])
    e_tot, fcivec = casci.fcisolver.kernel(h1eff, eri_cas, ncas, nelecas,
                                           ci0=ci0, verbose=log,
                                           max_memory=max_memory,
                                           ecore=energy_core)

    t1 = log.timer('FCI solver', *t1)
    e_cas = e_tot - energy_core
    return e_tot, e_cas, fcivec

def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
    if kpts_band is not None:
        return get_j_for_bands(mydf, dm_kpts, hermi, kpts, kpts_band)

    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]
    vj_kpts = numpy.zeros((nset,nkpts,nao,nao), dtype=numpy.complex128)
    kpt_allow = numpy.zeros(3)
    mesh = mydf.mesh
    coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
    max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
    weight = 1./len(kpts)
    dmsC = dms.conj()
    for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
        vG = [0] * nset
        #:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
        for k, aoao in enumerate(aoaoks):
            for i in range(nset):
                rho = numpy.einsum('ij,Lij->L', dmsC[i,k],
                                   aoao.reshape(-1,nao,nao)).conj()
                vG[i] += rho * coulG[p0:p1]
        for i in range(nset):
            vG[i] *= weight
        for k, aoao in enumerate(aoaoks):
            for i in range(nset):
                vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i],
                                             aoao.reshape(-1,nao,nao))
    aoao = aoaoks = p0 = p1 = None

    if gamma_point(kpts):
        vj_kpts = vj_kpts.real.copy()
    return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)

def _ERIS(mc, mo, method='incore'):
    nmo = mo.shape[1]
    ncore = mc.ncore
    ncas = mc.ncas

    mem_incore, mem_outcore, mem_basic = mc_ao2mo._mem_usage(ncore, ncas, nmo)
    mem_now = lib.current_memory()[0]
    if (method == 'incore' and mc._scf._eri is not None and
        (mem_incore+mem_now < mc.max_memory*.9) or
        mc.mol.incore_anyway):
        ppaa, papa, pacv, cvcv = trans_e1_incore(mc, mo)
    else:
        max_memory = max(2000, mc.max_memory-mem_now)
        ppaa, papa, pacv, cvcv = \
                trans_e1_outcore(mc, mo, max_memory=max_memory,
                                 verbose=mc.verbose)

    dmcore = numpy.dot(mo[:,:ncore], mo[:,:ncore].T)
    vj, vk = mc._scf.get_jk(mc.mol, dmcore)
    vhfcore = reduce(numpy.dot, (mo.T, vj*2-vk, mo))

    eris = {}
    eris['vhf_c'] = vhfcore
    eris['ppaa'] = ppaa
    eris['papa'] = papa
    eris['pacv'] = pacv
    eris['cvcv'] = cvcv
    eris['h1eff'] = reduce(numpy.dot, (mo.T, mc.get_hcore(), mo)) + vhfcore
    return eris

def gen_hop(hobj, mo_energy=None, mo_coeff=None, mo_occ=None, verbose=None):
    log = logger.new_logger(hobj, verbose)
    mol = hobj.mol
    mf = hobj.base

    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None:    mo_occ = mf.mo_occ
    if mo_coeff is None:  mo_coeff = mf.mo_coeff

    natm = mol.natm
    nao, nmo = mo_coeff.shape
    mocc = mo_coeff[:,mo_occ>0]
    nocc = mocc.shape[1]

    atmlst = range(natm)
    max_memory = max(2000, hobj.max_memory - lib.current_memory()[0])
    de2 = hobj.partial_hess_elec(mo_energy, mo_coeff, mo_occ, atmlst,
                                 max_memory, log)
    de2 += hobj.hess_nuc()

    # Compute H1 integrals and store in hobj.chkfile
    hobj.make_h1(mo_coeff, mo_occ, hobj.chkfile, atmlst, log)

    aoslices = mol.aoslice_by_atom()
    s1a = -mol.intor('int1e_ipovlp', comp=3)

    fvind = gen_vind(mf, mo_coeff, mo_occ)
    def h_op(x):
        x = x.reshape(natm,3)
        hx = numpy.einsum('abxy,ax->by', de2, x)
        h1ao = 0
        s1ao = 0
        for ia in range(natm):
            shl0, shl1, p0, p1 = aoslices[ia]
            h1ao_i = lib.chkfile.load(hobj.chkfile, 'scf_f1ao/%d' % ia)
            h1ao += numpy.einsum('x,xij->ij', x[ia], h1ao_i)
            s1ao_i = numpy.zeros((3,nao,nao))
            s1ao_i[:,p0:p1] += s1a[:,p0:p1]
            s1ao_i[:,:,p0:p1] += s1a[:,p0:p1].transpose(0,2,1)
            s1ao += numpy.einsum('x,xij->ij', x[ia], s1ao_i)

        s1vo = reduce(numpy.dot, (mo_coeff.T, s1ao, mocc))
        h1vo = reduce(numpy.dot, (mo_coeff.T, h1ao, mocc))
        mo1, mo_e1 = cphf.solve(fvind, mo_energy, mo_occ, h1vo, s1vo)
        mo1 = numpy.dot(mo_coeff, mo1)
        mo_e1 = mo_e1.reshape(nocc,nocc)
        dm1 = numpy.einsum('pi,qi->pq', mo1, mocc)
        dme1 = numpy.einsum('pi,qi,i->pq', mo1, mocc, mo_energy[mo_occ>0])
        dme1 = dme1 + dme1.T + reduce(numpy.dot, (mocc, mo_e1.T, mocc.T))

        for ja in range(natm):
            q0, q1 = aoslices[ja][2:]
            h1ao = lib.chkfile.load(hobj.chkfile, 'scf_f1ao/%s'%ja)
            hx[ja] += numpy.einsum('xpq,pq->x', h1ao, dm1) * 4
            hx[ja] -= numpy.einsum('xpq,pq->x', s1a[:,q0:q1], dme1[q0:q1]) * 2
            hx[ja] -= numpy.einsum('xpq,qp->x', s1a[:,q0:q1], dme1[:,q0:q1]) * 2
        return hx.ravel()

    hdiag = numpy.einsum('aaxx->ax', de2).ravel()
    return h_op, hdiag

    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 = 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 = 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 kernel(casci, mo_coeff=None, ci0=None, verbose=logger.NOTE):
    '''CASCI solver
    '''
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(casci.stdout, verbose)
    if mo_coeff is None: mo_coeff = casci.mo_coeff
    t0 = (time.clock(), time.time())
    log.debug('Start CASCI')

    ncas = casci.ncas
    nelecas = casci.nelecas

    # 2e
    eri_cas = casci.get_h2eff(mo_coeff)
    t1 = log.timer('integral transformation to CAS space', *t0)

    # 1e
    h1eff, energy_core = casci.get_h1eff(mo_coeff)
    log.debug('core energy = %.15g', energy_core)
    t1 = log.timer('effective h1e in CAS space', *t1)

    # FCI
    max_memory = max(400, casci.max_memory-lib.current_memory()[0])
    e_tot, fcivec = casci.fcisolver.kernel(h1eff, eri_cas, ncas, nelecas,
                                           ci0=ci0, verbose=log,
                                           max_memory=max_memory,
                                           ecore=energy_core)

    t1 = log.timer('FCI solver', *t1)
    e_cas = e_tot - energy_core
    return e_tot, e_cas, fcivec

def gamma1_intermediates(mycc, t1, t2, l1, l2):
    nocc, nvir = t1.shape
    doo =-numpy.einsum('ja,ia->ij', l1, t1)
    dvv = numpy.einsum('ia,ib->ab', l1, t1)
    dvo = l1.T
    xtv = numpy.einsum('ie,me->im', t1, l1)
    dov = t1 - numpy.einsum('im,ma->ia', xtv, t1)
    #:doo -= numpy.einsum('jkab,ikab->ij', l2, theta)
    #:dvv += numpy.einsum('jica,jicb->ab', l2, theta)
    #:xt1  = numpy.einsum('mnef,inef->mi', l2, make_theta(t2))
    #:xt2  = numpy.einsum('mnaf,mnef->ea', l2, make_theta(t2))
    #:dov += numpy.einsum('imae,me->ia', make_theta(t2), l1)
    #:dov -= numpy.einsum('ma,ie,me->ia', t1, t1, l1)
    #:dov -= numpy.einsum('mi,ma->ia', xt1, t1)
    #:dov -= numpy.einsum('ie,ae->ia', t1, xt2)
    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = nocc*nvir**2
    blksize = max(ccsd.BLKMIN, int(max_memory*.95e6/8/unit))
    for p0, p1 in prange(0, nocc, blksize):
        theta = make_theta(t2[p0:p1])
        doo[p0:p1] -= lib.dot(theta.reshape(p1-p0,-1), l2.reshape(nocc,-1).T)
        dov[p0:p1] += numpy.einsum('imae,me->ia', theta, l1)
        xt1 = lib.dot(l2.reshape(nocc,-1), theta.reshape(p1-p0,-1).T)
        dov[p0:p1] -= numpy.einsum('mi,ma->ia', xt1, t1)
        xt2 = lib.dot(theta.reshape(-1,nvir).T, l2[p0:p1].reshape(-1,nvir))
        dov -= numpy.einsum('ie,ae->ia', t1, xt2)
        dvv += lib.dot(l2[p0:p1].reshape(-1,nvir).T, theta.reshape(-1,nvir))
    return doo, dov, dvo, dvv

 def _is_mem_enough(self):
     nao = self.cell.nao_nr()
     if abs(self.kpt).sum() < 1e-9:
         mem_need = nao**4*8/4/1e6
     else:
         mem_need = nao**4*16/1e6
     return mem_need + lib.current_memory()[0] < self.max_memory*.95

def init_amps(mycc, eris=None):
    eris = getattr(mycc, '_eris', None)
    if eris is None:
        mycc.ao2mo()
        eris = mycc._eris

    time0 = time.clock(), time.time()
    mo_e = eris.mo_energy
    nocc = mycc.nocc
    nvir = mo_e.size - nocc
    eia = mo_e[:nocc,None] - mo_e[None,nocc:]
    t1T = eris.fock[nocc:,:nocc] / eia.T
    loc0, loc1 = _task_location(nvir)

    t2T = numpy.empty((loc1-loc0,nvir,nocc,nocc))
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1))))
    emp2 = 0
    for p0, p1 in lib.prange(0, loc1-loc0, blksize):
        eris_ovov = eris.ovov[:,p0:p1]
        t2T[p0:p1] = (eris_ovov.transpose(1,3,0,2) /
                      lib.direct_sum('ia,jb->abij', eia[:,p0+loc0:p1+loc0], eia))
        emp2 += 2 * numpy.einsum('abij,iajb', t2T[p0:p1], eris_ovov)
        emp2 -=     numpy.einsum('abji,iajb', t2T[p0:p1], eris_ovov)

    mycc.emp2 = comm.allreduce(emp2)
    logger.info(mycc, 'Init t2, MP2 energy = %.15g', mycc.emp2)
    logger.timer(mycc, 'init mp2', *time0)
    return mycc.emp2, t1T.T, t2T.transpose(2,3,0,1)

def energy(mycc, t1=None, t2=None, eris=None):
    '''CCSD correlation energy'''
    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    eris = getattr(mycc, '_eris', None)
    if eris is None:
        mycc.ao2mo()
        eris = mycc._eris

    nocc, nvir = t1.shape
    t2T = t2.transpose(2,3,0,1)
    fock = eris.fock
    loc0, loc1 = _task_location(nvir)
    e = numpy.einsum('ia,ia', fock[:nocc,nocc:], t1) * 2
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1))))
    for p0, p1 in lib.prange(0, loc1-loc0, blksize):
        eris_ovov = eris.ovov[:,p0:p1]
        tau = t2T[p0:p1] + numpy.einsum('ia,jb->abij', t1[:,p0+loc0:p1+loc0], t1)
        e += 2 * numpy.einsum('abij,iajb', tau, eris_ovov)
        e -=     numpy.einsum('abji,iajb', tau, eris_ovov)
    e = comm.allreduce(e)

    if rank == 0 and abs(e.imag) > 1e-4:
        logger.warn(mycc, 'Non-zero imaginary part found in CCSD energy %s', e)
    return e.real

def kernel(casci, mo_coeff=None, ci0=None, verbose=logger.NOTE):
    '''CASCI solver
    '''
    if mo_coeff is None: mo_coeff = casci.mo_coeff
    log = logger.new_logger(casci, verbose)
    t0 = (time.clock(), time.time())
    log.debug('Start CASCI')

    ncas = casci.ncas
    nelecas = casci.nelecas

    # 2e
    eri_cas = casci.get_h2eff(mo_coeff)
    t1 = log.timer('integral transformation to CAS space', *t0)

    # 1e
    h1eff, energy_core = casci.get_h1eff(mo_coeff)
    log.debug('core energy = %.15g', energy_core)
    t1 = log.timer('effective h1e in CAS space', *t1)

    if h1eff.shape[0] != ncas:
        raise RuntimeError('Active space size error. nmo=%d ncore=%d ncas=%d' %
                           (mo_coeff.shape[1], casci.ncore, ncas))

    # FCI
    max_memory = max(400, casci.max_memory-lib.current_memory()[0])
    e_tot, fcivec = casci.fcisolver.kernel(h1eff, eri_cas, ncas, nelecas,
                                           ci0=ci0, verbose=log,
                                           max_memory=max_memory,
                                           ecore=energy_core)

    t1 = log.timer('FCI solver', *t1)
    e_cas = e_tot - energy_core
    return e_tot, e_cas, fcivec