Python lib.prange函数代码示例

def _contract_plain(mydf, mos, coulG, phase, max_memory):
    cell = mydf.cell
    moiT, mojT, mokT, molT = mos
    nmoi, nmoj, nmok, nmol = [x.shape[0] for x in mos]
    ngrids = moiT.shape[1]
    wcoulG = coulG * (cell.vol/ngrids)
    dtype = numpy.result_type(phase, *mos)
    eri = numpy.empty((nmoi*nmoj,nmok*nmol), dtype=dtype)

    blksize = int(min(max(nmoi,nmok), (max_memory*1e6/16 - eri.size)/2/ngrids/max(nmoj,nmol)+1))
    assert blksize > 0
    buf0 = numpy.empty((blksize,max(nmoj,nmol),ngrids), dtype=dtype)
    buf1 = numpy.ndarray((blksize,nmoj,ngrids), dtype=dtype, buffer=buf0)
    buf2 = numpy.ndarray((blksize,nmol,ngrids), dtype=dtype, buffer=buf0)
    for p0, p1 in lib.prange(0, nmoi, blksize):
        mo_pairs = numpy.einsum('ig,jg->ijg', moiT[p0:p1].conj()*phase,
                                mojT, out=buf1[:p1-p0])
        mo_pairs_G = tools.fft(mo_pairs.reshape(-1,ngrids), mydf.mesh)
        mo_pairs = None
        mo_pairs_G*= wcoulG
        v = tools.ifft(mo_pairs_G, mydf.mesh)
        mo_pairs_G = None
        v *= phase.conj()
        if dtype == numpy.double:
            v = numpy.asarray(v.real, order='C')
        for q0, q1 in lib.prange(0, nmok, blksize):
            mo_pairs = numpy.einsum('ig,jg->ijg', mokT[q0:q1].conj(),
                                    molT, out=buf2[:q1-q0])
            eri[p0*nmoj:p1*nmoj,q0*nmol:q1*nmol] = lib.dot(v, mo_pairs.reshape(-1,ngrids).T)
        v = None
    return eri

def write_mo(fout, mol, mo_coeff, mo_energy=None, mo_occ=None):
    nmo = mo_coeff.shape[1]
    mo_cart = []
    centers = []
    types = []
    exps = []
    p0 = 0
    for ib in range(mol.nbas):
        ia = mol.bas_atom(ib)
        l = mol.bas_angular(ib)
        es = mol.bas_exp(ib)
        c = mol.bas_ctr_coeff(ib)
        np, nc = c.shape
        nd = nc*(2*l+1)
        mosub = mo_coeff[p0:p0+nd].reshape(-1,nc,nmo)
        c2s = gto.cart2sph(l)
        mosub = numpy.einsum('yki,cy,pk->pci', mosub, c2s, c)
        mo_cart.append(mosub.reshape(-1,nmo))

        for t in TYPE_MAP[l]:
            types.append([t]*np)
        ncart = gto.len_cart(l)
        exps.extend([es]*ncart)
        centers.extend([ia+1]*(np*ncart))
        p0 += nd
    mo_cart = numpy.vstack(mo_cart)
    centers = numpy.hstack(centers)
    types = numpy.hstack(types)
    exps = numpy.hstack(exps)
    nprim, nmo = mo_cart.shape

    fout.write('title line\n')
    fout.write('GAUSSIAN            %3d MOL ORBITALS    %3d PRIMITIVES      %3d NUCLEI\n'
               % (mo_cart.shape[1], mo_cart.shape[0], mol.natm))
    for ia in range(mol.natm):
        x, y, z = mol.atom_coord(ia)
        fout.write('%-4s %-4d (CENTRE%3d) %12.8f %12.8f %12.8f  CHARGE = %.1f\n'
                   % (mol.atom_pure_symbol(ia), ia, ia, x, y, z,
                      mol.atom_charge(ia)))
    for i0, i1 in lib.prange(0, nprim, 20):
        fout.write('CENTRE ASSIGNMENTS  %s\n' % ''.join('%3d'%x for x in centers[i0:i1]))
    for i0, i1 in lib.prange(0, nprim, 20):
        fout.write('TYPE ASSIGNMENTS    %s\n' % ''.join('%3d'%x for x in types[i0:i1]))
    for i0, i1 in lib.prange(0, nprim, 5):
        fout.write('EXPONENTS %s\n' % ' '.join('%14.7E'%x for x in exps[i0:i1]))

    for k in range(nmo):
        mo = mo_cart[:,k]
        if mo_energy is None or mo_occ is None:
            fout.write('CANMO %d\n'
                       (k, mo_occ[k], mo_energy[k]))
        else:
            fout.write('MO  %-4d                  OCC NO = %12.8f ORB. ENERGY = %12.8f\n' %
                       (k, mo_occ[k], mo_energy[k]))
        for i0, i1 in lib.prange(0, nprim, 5):
            fout.write('%s\n' % ' '.join('%15.8E'%x for x in mo[i0:i1]))

    def pw_loop(self, mol, auxmol, gs, shls_slice=None, max_memory=2000):
        '''Plane wave part'''
        if isinstance(gs, int):
            gs = [gs]*3
        naux = auxmol.nao_nr()
        Gv, Gvbase, kws = non_uniform_kgrids(gs)
        nxyz = [i*2 for i in gs]
        gxyz = lib.cartesian_prod([range(i) for i in nxyz])
        kk = numpy.einsum('ki,ki->k', Gv, Gv)
        idx = numpy.argsort(kk)[::-1]
#        idx = idx[(kk[idx] < 300.) & (kk[idx] > 1e-4)]  # ~ Cut high energy plain waves
#        log.debug('Cut grids %d to %d', Gv.shape[0], len(idx))
        kk = kk[idx]
        Gv = Gv[idx]
        kws = kws[idx]
        gxyz = gxyz[idx]
        coulG = .5/numpy.pi**2 * kws / kk

        if shls_slice is None:
            ni = nj = mol.nao_nr()
        else:
            ao_loc = mol.ao_loc_nr()
            ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
            nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
        nij = ni * nj
        blksize = min(max(16, int(max_memory*1e6*.7/16/nij)), 16384)
        sublk = max(16,int(blksize//4))
        pqkRbuf = numpy.empty(nij*sublk)
        pqkIbuf = numpy.empty(nij*sublk)
        LkRbuf = numpy.empty(naux*sublk)
        LkIbuf = numpy.empty(naux*sublk)

        for p0, p1 in lib.prange(0, coulG.size, blksize):
            aoao = ft_ao.ft_aopair(mol, Gv[p0:p1], shls_slice, 's1',
                                   Gvbase, gxyz[p0:p1], nxyz)
            aoaux = ft_ao.ft_ao(auxmol, Gv[p0:p1], None, Gvbase, gxyz[p0:p1], nxyz)

            for i0, i1 in lib.prange(0, p1-p0, sublk):
                nG = i1 - i0
                pqkR = numpy.ndarray((ni,nj,nG), buffer=pqkRbuf)
                pqkI = numpy.ndarray((ni,nj,nG), buffer=pqkIbuf)
                LkR = numpy.ndarray((naux,nG), buffer=LkRbuf)
                LkI = numpy.ndarray((naux,nG), buffer=LkIbuf)
                pqkR[:] = aoao[i0:i1].real.transpose(1,2,0)
                pqkI[:] = aoao[i0:i1].imag.transpose(1,2,0)
                LkR [:] = aoaux[i0:i1].real.T
                LkI [:] = aoaux[i0:i1].imag.T
                yield (pqkR.reshape(-1,nG), LkR, pqkI.reshape(-1,nG), LkI,
                       coulG[p0+i0:p0+i1])
            aoao = aoaux = None

        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 recv(source=0, tag=0):
    shape, dtype = comm.recv(source=source, tag=tag)
    recvbuf = numpy.empty(shape, dtype=dtype)
    recv_seg = numpy.ndarray(recvbuf.size, dtype=recvbuf.dtype, buffer=recvbuf)
    for p0, p1 in lib.prange(0, recvbuf.size, BLKSIZE):
        comm.Recv(recv_seg[p0:p1], source=source, tag=tag)
    return recvbuf

def send(sendbuf, dest=0, tag=0):
    sendbuf = numpy.asarray(sendbuf, order='C')
    comm.send((sendbuf.shape, sendbuf.dtype), dest=dest, tag=tag)
    send_seg = numpy.ndarray(sendbuf.size, dtype=sendbuf.dtype, buffer=sendbuf)
    for p0, p1 in lib.prange(0, sendbuf.size, BLKSIZE):
        comm.Send(send_seg[p0:p1], dest=dest, tag=tag)
    return sendbuf

def allgather(sendbuf, split_recvbuf=False):
    sendbuf = numpy.asarray(sendbuf, order='C')
    shape = sendbuf.shape
    attr = comm.allgather((shape, sendbuf.dtype.char))
    rshape = [x[0] for x in attr]
    counts = numpy.array([numpy.prod(x) for x in rshape])
    mpi_dtype = numpy.result_type(*[x[1] for x in attr]).char
    _assert(sendbuf.dtype.char == mpi_dtype or sendbuf.size == 0)

    displs = numpy.append(0, numpy.cumsum(counts[:-1]))
    recvbuf = numpy.empty(sum(counts), dtype=mpi_dtype)

    sendbuf = sendbuf.ravel()
    for p0, p1 in lib.prange(0, numpy.max(counts), BLKSIZE):
        counts_seg = _segment_counts(counts, p0, p1)
        comm.Allgatherv([sendbuf[p0:p1], mpi_dtype],
                        [recvbuf, counts_seg, displs+p0, mpi_dtype])
    if split_recvbuf:
        return [recvbuf[p0:p0+c].reshape(shape)
                for p0,c,shape in zip(displs,counts,rshape)]
    else:
        try:
            return recvbuf.reshape((-1,) + shape[1:])
        except ValueError:
            return recvbuf

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 write(self, field, fname, comment=None):
        """  Result: .cube file with the field in the file fname.  """
        assert(field.ndim == 3)
        assert(field.shape == (self.nx, self.ny, self.nz))
        if comment is None:
            comment = 'Generic field? Supply the optional argument "comment" to define this line'

        mol = self.mol
        coord = mol.atom_coords()
        with open(fname, 'w') as f:
            f.write(comment+'\n')
            f.write('PySCF Version: %s  Date: %s\n' % (pyscf.__version__, time.ctime()))
            f.write('%5d' % mol.natm)
            f.write('%12.6f%12.6f%12.6f\n' % tuple(self.boxorig.tolist()))
            f.write('%5d%12.6f%12.6f%12.6f\n' % (self.nx, self.xs[1], 0, 0))
            f.write('%5d%12.6f%12.6f%12.6f\n' % (self.ny, 0, self.ys[1], 0))
            f.write('%5d%12.6f%12.6f%12.6f\n' % (self.nz, 0, 0, self.zs[1]))
            for ia in range(mol.natm):
                chg = mol.atom_charge(ia)
                f.write('%5d%12.6f'% (chg, chg))
                f.write('%12.6f%12.6f%12.6f\n' % tuple(coord[ia]))

            for ix in range(self.nx):
                for iy in range(self.ny):
                    for iz0, iz1 in lib.prange(0, self.nz, 6):
                        fmt = '%13.5E' * (iz1-iz0) + '\n'
                        f.write(fmt % tuple(field[ix,iy,iz0:iz1].tolist()))

    def extrapolate(self, nd=None):
        if nd is None:
            nd = self.get_num_vec()
        if nd == 0:
            raise RuntimeError('No vector found in DIIS object.')

        h = self._H[:nd+1,:nd+1].copy()
        h[1:,1:] = mpi.comm.allreduce(self._H[1:nd+1,1:nd+1])
        g = numpy.zeros(nd+1, h.dtype)
        g[0] = 1

        w, v = scipy.linalg.eigh(h)
        if numpy.any(abs(w)<1e-14):
            logger.debug(self, 'Singularity found in DIIS error vector space.')
            idx = abs(w)>1e-14
            c = numpy.dot(v[:,idx]*(1./w[idx]), numpy.dot(v[:,idx].T.conj(), g))
        else:
            try:
                c = numpy.linalg.solve(h, g)
            except numpy.linalg.linalg.LinAlgError as e:
                logger.warn(self, ' diis singular, eigh(h) %s', w)
                raise e
        logger.debug1(self, 'diis-c %s', c)

        xnew = None
        for i, ci in enumerate(c[1:]):
            xi = self.get_vec(i)
            if xnew is None:
                xnew = numpy.zeros(xi.size, c.dtype)
            for p0, p1 in lib.prange(0, xi.size, lib.diis.BLOCK_SIZE):
                xnew[p0:p1] += xi[p0:p1] * ci
        return xnew

        def fuse(Lpq):
            Lpq, chgLpq = Lpq[:naux], Lpq[naux:]
            if Lpq.ndim == 1:
                npq = 1
                Lpq_sph = numpy.empty(naux_sph, dtype=Lpq.dtype)
            else:
                npq = Lpq.shape[1]
                Lpq_sph = numpy.empty((naux_sph,npq), dtype=Lpq.dtype)
            if Lpq.dtype == numpy.complex:
                npq *= 2  # c2s_fn supports double only, *2 to handle complex
            for i in range(auxcell.nbas):
                l  = auxcell.bas_angular(i)
                ia = auxcell.bas_atom(i)
                p0 = modchg_offset[ia,l]
                if p0 >= 0:
                    nd = (l+1) * (l+2) // 2
                    c0, c1 = aux_loc[i], aux_loc[i+1]
                    s0, s1 = aux_loc_sph[i], aux_loc_sph[i+1]
                    for i0, i1 in lib.prange(c0, c1, nd):
                        Lpq[i0:i1] -= chgLpq[p0:p0+nd]

                    if l < 2:
                        Lpq_sph[s0:s1] = Lpq[c0:c1]
                    else:
                        Lpq_cart = numpy.asarray(Lpq[c0:c1], order='C')
                        c2s_fn(Lpq_sph[s0:s1].ctypes.data_as(ctypes.c_void_p),
                               ctypes.c_int(npq * auxcell.bas_nctr(i)),
                               Lpq_cart.ctypes.data_as(ctypes.c_void_p),
                               ctypes.c_int(l))
            return Lpq_sph

    def __init__(self, mycc):
        self._cc = mycc
        self.daemon = None

        nocc, nvir = mycc.t1.shape
        nmo = nocc + nvir
        nvir_seg = (nvir + mpi.pool.size - 1) // mpi.pool.size
        max_memory = mycc.max_memory - lib.current_memory()[0]
        max_memory = max(0, max_memory - nocc**3*13*lib.num_threads()*8/1e6)
        blksize = max(BLKMIN, (max_memory*.5e6/8/(6*nmo*nocc))**.5 - nocc/4)
        blksize = int(min(comm.allgather(min(nvir/6+2, nvir_seg/2+1, blksize))))
        logger.debug1(mycc, 'GlobalDataHandler blksize %s', blksize)

        self.vranges = []
        self.data_partition = []
        self.segment_location = {}
        for task in range(mpi.pool.size):
            p0 = nvir_seg * task
            p1 = min(nvir, p0 + nvir_seg)
            self.vranges.append((p0, p1))

            for j0, j1 in lib.prange(p0, p1, blksize):
                self.data_partition.append((j0, j1))
                self.segment_location[j0] = task
        logger.debug1(mycc, 'data_partition %s', self.data_partition)
        logger.debug1(mycc, 'segment_location %s', self.segment_location)

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_nuc_less_accurate(mydf, kpts=None):
    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = t0 = (time.clock(), time.time())
    if kpts is None:
        kpts_lst = numpy.zeros((1, 3))
    else:
        kpts_lst = numpy.reshape(kpts, (-1, 3))
    nkpts = len(kpts_lst)
    if mydf._cderi is None:
        mydf.build()
    cell = mydf.cell
    fused_cell, fuse = fuse_auxcell_(mydf, mydf.auxcell)

    nao = cell.nao_nr()
    charge = -cell.atom_charges()
    j2c = pgto.intor_cross("cint2c2e_sph", fused_cell, _fake_nuc(cell))
    jaux = j2c.dot(charge)
    jaux -= charge.sum() * mydf.auxbar(fused_cell)
    Gv = cell.get_Gv(mydf.gs)
    SI = cell.get_SI(Gv)
    # The normal nuclues have been considered in function get_gth_vlocG_part1
    # The result vG is the potential in G-space for erf part of the pp nuclues and
    # "numpy.dot(charge, SI) * coulG" for normal nuclues.
    vpplocG = pgto.pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
    vG = -1.0 / cell.vol * numpy.einsum("ij,ij->j", SI, vpplocG)
    kpt_allow = numpy.zeros(3)

    if is_zero(kpts_lst):
        vj = numpy.zeros((nkpts, nao ** 2))
    else:
        vj = numpy.zeros((nkpts, nao ** 2), dtype=numpy.complex128)
    max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
    for k, pqkR, pqkI, p0, p1 in mydf.ft_loop(cell, mydf.gs, kpt_allow, kpts_lst, max_memory=max_memory):
        if not gamma_point(kpts_lst[k]):
            vj[k] += numpy.einsum("k,xk->x", vG.real, pqkI) * 1j
            vj[k] += numpy.einsum("k,xk->x", vG.imag, pqkR) * -1j
        vj[k] += numpy.einsum("k,xk->x", vG.real, pqkR)
        vj[k] += numpy.einsum("k,xk->x", vG.imag, pqkI)
        pqkR = pqkI = None

    Gv = cell.get_Gv(mydf.gs)
    aoaux = ft_ao.ft_ao(fused_cell, Gv)
    jaux -= numpy.einsum("x,xj->j", vG.real, aoaux.real)
    jaux -= numpy.einsum("x,xj->j", vG.imag, aoaux.imag)
    jaux = fuse(jaux)

    vj = vj.reshape(-1, nao, nao)
    for k, kpt in enumerate(kpts_lst):
        with mydf.load_Lpq((kpt, kpt)) as Lpq:
            v = 0
            for p0, p1 in lib.prange(0, jaux.size, mydf.blockdim):
                v += numpy.dot(jaux[p0:p1], numpy.asarray(Lpq[p0:p1]))
            if gamma_point(kpt):
                vj[k] += lib.unpack_tril(numpy.asarray(v.real, order="C"))
            else:
                vj[k] += lib.unpack_tril(v)

    if kpts is None or numpy.shape(kpts) == (3,):
        vj = vj[0]
    return vj

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)