本文整理汇总了Python中pyscf.lib.prange函数的典型用法代码示例。如果您正苦于以下问题:Python prange函数的具体用法?Python prange怎么用?Python prange使用的例子?那么恭喜您, 这里精选的函数代码示例或许可以为您提供帮助。
在下文中一共展示了prange函数的20个代码示例,这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞,您的评价将有助于我们的系统推荐出更棒的Python代码示例。
示例1: _contract_plain
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
开发者ID:chrinide,项目名称:pyscf,代码行数:31,代码来源:fft_ao2mo.py
示例2: write_mo
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]))
开发者ID:berquist,项目名称:pyscf,代码行数:56,代码来源:wfn_format.py
示例3: pw_loop
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
开发者ID:berquist,项目名称:pyscf,代码行数:50,代码来源:mdf.py
示例4: get_hcore
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
开发者ID:chrinide,项目名称:pyscf,代码行数:27,代码来源:itrf.py
示例5: recv
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
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:7,代码来源:mpi.py
示例6: send
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
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:7,代码来源:mpi.py
示例7: allgather
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
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:25,代码来源:mpi.py
示例8: energy
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
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:26,代码来源:ccsd.py
示例9: write
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()))
开发者ID:chrinide,项目名称:pyscf,代码行数:27,代码来源:cubegen.py
示例10: extrapolate
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
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:32,代码来源:diis.py
示例11: fuse
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
开发者ID:chrinide,项目名称:pyscf,代码行数:30,代码来源:df.py
示例12: __init__
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)
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:26,代码来源:ccsd_t.py
示例13: build_Lpq_pbc
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])
开发者ID:ushnishray,项目名称:pyscf,代码行数:35,代码来源:mdf.py
示例14: get_nuc_less_accurate
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
开发者ID:ushnishray,项目名称:pyscf,代码行数:60,代码来源:mdf.py
示例15: init_amps
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)
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:29,代码来源:ccsd.py
示例16: _make_eris_outcore
def _make_eris_outcore(mycc, mo_coeff=None):
cput0 = (time.clock(), time.time())
log = logger.Logger(mycc.stdout, mycc.verbose)
eris = _ChemistsERIs()
eris._common_init_(mycc, mo_coeff)
mol = mycc.mol
mo_coeff = eris.mo_coeff
nocc = eris.nocc
nao, nmo = mo_coeff.shape
nvir = nmo - nocc
orbo = mo_coeff[:,:nocc]
orbv = mo_coeff[:,nocc:]
nvpair = nvir * (nvir+1) // 2
eris.feri1 = lib.H5TmpFile()
eris.oooo = eris.feri1.create_dataset('oooo', (nocc,nocc,nocc,nocc), 'f8')
eris.ovoo = eris.feri1.create_dataset('ovoo', (nocc,nvir,nocc,nocc), 'f8', chunks=(nocc,1,nocc,nocc))
eris.ovov = eris.feri1.create_dataset('ovov', (nocc,nvir,nocc,nvir), 'f8', chunks=(nocc,1,nocc,nvir))
eris.ovvo = eris.feri1.create_dataset('ovvo', (nocc,nvir,nvir,nocc), 'f8', chunks=(nocc,1,nvir,nocc))
eris.ovvv = eris.feri1.create_dataset('ovvv', (nocc,nvir,nvir,nvir), 'f8')
eris.oovv = eris.feri1.create_dataset('oovv', (nocc,nocc,nvir,nvir), 'f8', chunks=(nocc,nocc,1,nvir))
eris.vvvv = eris.feri1.create_dataset('vvvv', (nvir,nvir,nvir,nvir), 'f8')
max_memory = max(MEMORYMIN, mycc.max_memory-lib.current_memory()[0])
ftmp = lib.H5TmpFile()
ao2mo.full(mol, mo_coeff, ftmp, max_memory=max_memory, verbose=log)
eri = ftmp['eri_mo']
nocc_pair = nocc*(nocc+1)//2
tril2sq = lib.square_mat_in_trilu_indices(nmo)
oo = eri[:nocc_pair]
eris.oooo[:] = ao2mo.restore(1, oo[:,:nocc_pair], nocc)
oovv = lib.take_2d(oo, tril2sq[:nocc,:nocc].ravel(), tril2sq[nocc:,nocc:].ravel())
eris.oovv[:] = oovv.reshape(nocc,nocc,nvir,nvir)
oo = oovv = None
tril2sq = lib.square_mat_in_trilu_indices(nmo)
blksize = min(nvir, max(BLKMIN, int(max_memory*1e6/8/nmo**3/2)))
for p0, p1 in lib.prange(0, nvir, blksize):
q0, q1 = p0+nocc, p1+nocc
off0 = q0*(q0+1)//2
off1 = q1*(q1+1)//2
buf = lib.unpack_tril(eri[off0:off1])
tmp = buf[ tril2sq[q0:q1,:nocc] - off0 ]
eris.ovoo[:,p0:p1] = tmp[:,:,:nocc,:nocc].transpose(1,0,2,3)
eris.ovvo[:,p0:p1] = tmp[:,:,nocc:,:nocc].transpose(1,0,2,3)
eris.ovov[:,p0:p1] = tmp[:,:,:nocc,nocc:].transpose(1,0,2,3)
eris.ovvv[:,p0:p1] = tmp[:,:,nocc:,nocc:].transpose(1,0,2,3)
tmp = buf[ tril2sq[q0:q1,nocc:q1] - off0 ]
eris.vvvv[p0:p1,:p1] = tmp[:,:,nocc:,nocc:]
if p0 > 0:
eris.vvvv[:p0,p0:p1] = tmp[:,:p0,nocc:,nocc:].transpose(1,0,2,3)
buf = tmp = None
log.timer('CCSD integral transformation', *cput0)
return eris
开发者ID:chrinide,项目名称:pyscf,代码行数:57,代码来源:rccsd.py
示例17: get_vk_s4
def get_vk_s4(mol, dm):
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
vk = numpy.zeros((nao,nao))
bas_groups = list(lib.prange(0, mol.nbas, 3))
for ip, (ish0, ish1) in enumerate(bas_groups):
for jp, (jsh0, jsh1) in enumerate(bas_groups[:ip]):
for kp, (ksh0, ksh1) in enumerate(bas_groups):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[i0:i1,k0:k1],
dm[j0:j1,l0:l1],
dm[i0:i1,l0:l1]]
scripts = ['ijkl,jk->il',
'ijkl,ik->jl',
'ijkl,jl->ik',
'ijkl,il->jk']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
vk[i0:i1,k0:k1] += kparts[2]
vk[j0:j1,k0:k1] += kparts[3]
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1], dm[i0:i1,k0:k1]],
scripts=['ijkl,jk->il', 'ijkl,ik->jl'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
jsh0, jsh1 = ish0, ish1
for kp, (ksh0, ksh1) in enumerate(bas_groups):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1], dm[j0:j1,l0:l1]],
scripts=['ijkl,jk->il', 'ijkl,jl->ik'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,k0:k1] += kparts[1]
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1]],
scripts=['ijkl,jk->il'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
return vk
开发者ID:chrinide,项目名称:pyscf,代码行数:56,代码来源:test_jk.py
示例18: bcast
def bcast(buf, root=0):
buf = numpy.asarray(buf, order='C')
shape, dtype = comm.bcast((buf.shape, buf.dtype.char))
if rank != root:
buf = numpy.empty(shape, dtype=dtype)
buf_seg = numpy.ndarray(buf.size, dtype=buf.dtype, buffer=buf)
for p0, p1 in lib.prange(0, buf.size, BLKSIZE):
comm.Bcast(buf_seg[p0:p1], root)
return buf
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:10,代码来源:mpi.py
示例19: gather
def gather(sendbuf, root=0, split_recvbuf=False):
# if pool.debug:
# if rank == 0:
# res = [sendbuf]
# for k in range(1, pool.size):
# dat = comm.recv(source=k)
# res.append(dat)
# return numpy.vstack([x for x in res if len(x) > 0])
# else:
# comm.send(sendbuf, dest=0)
# return sendbuf
sendbuf = numpy.asarray(sendbuf, order='C')
shape = sendbuf.shape
size_dtype = comm.allgather((shape, sendbuf.dtype.char))
rshape = [x[0] for x in size_dtype]
counts = numpy.array([numpy.prod(x) for x in rshape])
mpi_dtype = numpy.result_type(*[x[1] for x in size_dtype]).char
_assert(sendbuf.dtype == mpi_dtype or sendbuf.size == 0)
if rank == root:
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.Gatherv([sendbuf[p0:p1], mpi_dtype],
[recvbuf, counts_seg, displs+p0, mpi_dtype], root)
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
else:
send_seg = sendbuf.ravel()
for p0, p1 in lib.prange(0, numpy.max(counts), BLKSIZE):
comm.Gatherv([send_seg[p0:p1], mpi_dtype], None, root)
return sendbuf
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:43,代码来源:mpi.py
示例20: allreduce
def allreduce(sendbuf, op=MPI.SUM):
sendbuf = numpy.asarray(sendbuf, order='C')
shape, mpi_dtype = comm.bcast((sendbuf.shape, sendbuf.dtype.char))
_assert(sendbuf.shape == shape and sendbuf.dtype.char == mpi_dtype)
recvbuf = numpy.zeros_like(sendbuf)
send_seg = numpy.ndarray(sendbuf.size, dtype=sendbuf.dtype, buffer=sendbuf)
recv_seg = numpy.ndarray(recvbuf.size, dtype=recvbuf.dtype, buffer=recvbuf)
for p0, p1 in lib.prange(0, sendbuf.size, BLKSIZE):
comm.Allreduce(send_seg[p0:p1], recv_seg[p0:p1], op)
return recvbuf
开发者ID:sunqm,项目名称:mpi4pyscf,代码行数:11,代码来源:mpi.py
注:本文中的pyscf.lib.prange函数示例由纯净天空整理自Github/MSDocs等源码及文档管理平台,相关代码片段筛选自各路编程大神贡献的开源项目,源码版权归原作者所有,传播和使用请参考对应项目的License;未经允许,请勿转载。 |
请发表评论