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 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 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 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 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 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 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 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 _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
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
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
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
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
def fuse(Lpq):
Lpq, chgLpq = Lpq[:naux], Lpq[naux:]
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 * 2 + 1
for i0, i1 in lib.prange(aux_loc[i], aux_loc[i+1], nd):
Lpq[i0:i1] -= chgLpq[p0:p0+nd]
return Lpq
def loop(self):
coulG = tools.get_coulG(self.cell, numpy.zeros(3), gs=mydf.gs)
ngs = len(coulG)
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], compact=True)
ao_pairs_G *= numpy.sqrt(coulG * (cell.vol / ngs ** 2)).reshape(-1, 1)
Lpq = numpy.empty((self.blockdim, ao_pairs_G.shape[1]))
for p0, p1 in lib.prange(0, ngs, self.blockdim):
Lpq[: p1 - p0] = ao_pairs_G[p0:p1].real
yield Lpq[: p1 - p0]
Lpq[: p1 - p0] = ao_pairs_G[p0:p1].imag
yield Lpq[: p1 - p0]
def prange(self, start, stop, step=None):
# affect pw_loop and ft_loop function
size = stop - start
mpi_size = mpi.pool.size
segsize = (size+mpi_size-1) // mpi_size
if step is None:
step = segsize
else:
step = min(step, segsize)
start = min(size, start + rank * segsize)
stop = min(size, start + segsize)
return lib.prange(start, stop, step)
def sr_loop(self, kpti_kptj=numpy.zeros((2,3)), max_memory=2000,
compact=True, blksize=None):
'''Short range part'''
kpti, kptj = kpti_kptj
unpack = is_zero(kpti-kptj) and not compact
is_real = is_zero(kpti_kptj)
nao = self.cell.nao_nr()
if blksize is None:
if is_real:
if unpack:
blksize = max_memory*1e6/8/(nao*(nao+1)//2+nao**2*2)
else:
blksize = max_memory*1e6/8/(nao*(nao+1)*2)
else:
blksize = max_memory*1e6/16/(nao**2*3)
blksize = max(16, min(int(blksize), self.blockdim))
logger.debug2(self, 'max_memory %d MB, blksize %d', max_memory, blksize)
if unpack:
buf = numpy.empty((blksize,nao*(nao+1)//2))
def load(Lpq, b0, b1, bufR, bufI):
Lpq = numpy.asarray(Lpq[b0:b1])
if is_real:
if unpack:
LpqR = lib.unpack_tril(Lpq, out=bufR).reshape(-1,nao**2)
else:
LpqR = Lpq
LpqI = numpy.zeros_like(LpqR)
else:
shape = Lpq.shape
if unpack:
tmp = numpy.ndarray(shape, buffer=buf)
tmp[:] = Lpq.real
LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1,nao**2)
tmp[:] = Lpq.imag
LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI, out=bufI).reshape(-1,nao**2)
else:
LpqR = numpy.ndarray(shape, buffer=bufR)
LpqR[:] = Lpq.real
LpqI = numpy.ndarray(shape, buffer=bufI)
LpqI[:] = Lpq.imag
return LpqR, LpqI
LpqR = LpqI = j3cR = j3cI = None
with self.load_Lpq(kpti_kptj) as Lpq:
naux = Lpq.shape[0]
with self.load_j3c(kpti_kptj) as j3c:
for b0, b1 in lib.prange(0, naux, blksize):
LpqR, LpqI = load(Lpq, b0, b1, LpqR, LpqI)
j3cR, j3cI = load(j3c, b0, b1, j3cR, j3cI)
yield LpqR, LpqI, j3cR, j3cI
def make_phi(pcmobj, dm, r_vdw, ui):
mol = pcmobj.mol
natm = mol.natm
coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
ngrid_1sph = coords_1sph.shape[0]
if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
dm = dm[0] + dm[1]
tril_dm = lib.pack_tril(dm + dm.T)
nao = dm.shape[0]
diagidx = numpy.arange(nao)
diagidx = diagidx * (diagidx + 1) // 2 + diagidx
tril_dm[diagidx] *= .5
atom_coords = mol.atom_coords()
atom_charges = mol.atom_charges()
extern_point_idx = ui > 0
cav_coords = (atom_coords.reshape(natm, 1, 3) +
numpy.einsum('r,gx->rgx', r_vdw, coords_1sph))
v_phi = numpy.empty((natm, ngrid_1sph))
for ia in range(natm):
# Note (-) sign is not applied to atom_charges, because (-) is explicitly
# included in rhs and L matrix
d_rs = atom_coords.reshape(-1, 1, 3) - cav_coords[ia]
v_phi[ia] = numpy.einsum('z,zp->p', atom_charges,
1. / lib.norm(d_rs, axis=2))
max_memory = pcmobj.max_memory - lib.current_memory()[0]
blksize = int(max(max_memory * 1e6 / 8 / nao**2, 400))
cav_coords = cav_coords[extern_point_idx]
v_phi_e = numpy.empty(cav_coords.shape[0])
int3c2e = mol._add_suffix('int3c2e')
cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, int3c2e)
for i0, i1 in lib.prange(0, cav_coords.shape[0], blksize):
fakemol = gto.fakemol_for_charges(cav_coords[i0:i1])
v_nj = df.incore.aux_e2(mol,
fakemol,
intor=int3c2e,
aosym='s2ij',
cintopt=cintopt)
v_phi_e[i0:i1] = numpy.einsum('x,xk->k', tril_dm, v_nj)
v_phi[extern_point_idx] -= v_phi_e
ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, pcmobj.lmax, True))
phi = -numpy.einsum('n,xn,jn,jn->jx', weights_1sph, ylm_1sph, ui, v_phi)
return phi
def pw_contract(istep, sh_range, j3cR, j3cI):
bstart, bend, ncol = sh_range
if aosym == 's2':
shls_slice = (bstart, bend, 0, bend)
else:
shls_slice = (bstart, bend, 0, cell.nbas)
for p0, p1 in lib.prange(0, ngrids, Gblksize):
dat = ft_ao._ft_aopair_kpts(cell,
Gv[p0:p1],
shls_slice,
aosym,
b,
gxyz[p0:p1],
Gvbase,
kpt,
adapted_kptjs,
out=buf)
nG = p1 - p0
for k, ji in enumerate(adapted_ji_idx):
aoao = dat[k].reshape(nG, ncol)
pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf)
pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf)
pqkR[:] = aoao.real.T
pqkI[:] = aoao.imag.T
lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k], 1)
lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k], 1)
if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k], 1)
lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k], 1)
for k, ji in enumerate(adapted_ji_idx):
if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
v = j3cR[k]
else:
v = j3cR[k] + j3cI[k] * 1j
if j2ctag == 'CD':
v = scipy.linalg.solve_triangular(j2c,
v,
lower=True,
overwrite_b=True)
feri['j3c/%d/%d' % (ji, istep)] = v
else:
feri['j3c/%d/%d' % (ji, istep)] = lib.dot(j2c, v)
# low-dimension systems
if j2c_negative is not None:
feri['j3c-/%d/%d' % (ji, istep)] = lib.dot(j2c_negative, v)
def reduce(sendbuf, op=MPI.SUM, root=0):
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.Reduce(send_seg[p0:p1], recv_seg[p0:p1], op, root)
if rank == root:
return recvbuf
else:
return sendbuf
def build_Lpq_1c_approx(mydf, auxcell):
get_Lpq = pyscf.df.mdf._make_Lpq_atomic_approx(mydf, mydf.cell, auxcell)
max_memory = mydf.max_memory - lib.current_memory()[0]
naux = auxcell.nao_nr()
nao = mydf.cell.nao_nr()
nao_pair = nao * (nao+1) // 2
blksize = max(int(max_memory*.5*1e6/8/naux/mydf.blockdim), 1) * mydf.blockdim
with h5py.File(mydf._cderi) as feri:
if 'Lpq' in feri:
del(feri['Lpq'])
chunks = (min(mydf.blockdim,naux), min(mydf.blockdim,nao_pair)) # 512K
Lpq = feri.create_dataset('Lpq/0', (naux,nao_pair), 'f8', chunks=chunks)
for p0, p1 in lib.prange(0, nao_pair, blksize):
v = get_Lpq(None, p0, p1)
Lpq[:,p0:p1] = get_Lpq(None, p0, p1)
def build_Lpq_1c_approx(mydf, auxcell):
get_Lpq = pyscf.df.mdf._make_Lpq_atomic_approx(mydf, mydf.cell, auxcell)
max_memory = mydf.max_memory - lib.current_memory()[0]
naux = auxcell.nao_nr()
nao = mydf.cell.nao_nr()
nao_pair = nao * (nao+1) // 2
blksize = max(int(max_memory*.5*1e6/8/naux/mydf.blockdim), 1) * mydf.blockdim
with h5py.File(mydf._cderi) as feri:
if 'Lpq' in feri:
del(feri['Lpq'])
chunks = (min(mydf.blockdim,naux), min(mydf.blockdim,nao_pair)) # 512K
Lpq = feri.create_dataset('Lpq/0', (naux,nao_pair), 'f8', chunks=chunks)
for p0, p1 in lib.prange(0, nao_pair, blksize):
v = get_Lpq(None, p0, p1)
Lpq[:,p0:p1] = get_Lpq(None, p0, p1)
def pw_loop(self, cell, gs=None, kpti_kptj=None, shls_slice=None,
max_memory=2000):
'''Plane wave part'''
if gs is None:
gs = self.gs
if kpti_kptj is None:
kpti = kptj = numpy.zeros(3)
else:
kpti, kptj = kpti_kptj
nao = cell.nao_nr()
gxyz = lib.cartesian_prod((numpy.append(range(gs[0]+1), range(-gs[0],0)),
numpy.append(range(gs[1]+1), range(-gs[1],0)),
numpy.append(range(gs[2]+1), range(-gs[2],0))))
invh = numpy.linalg.inv(cell._h)
Gv = 2*numpy.pi * numpy.dot(gxyz, invh)
ngs = gxyz.shape[0]
# Theoretically, hermitian symmetry can be also found for kpti == kptj:
# f_ji(G) = \int f_ji exp(-iGr) = \int f_ij^* exp(-iGr) = [f_ij(-G)]^*
# The hermi operation needs reordering the axis-0. It is inefficient
if gamma_point(kpti) and gamma_point(kptj):
aosym = 's1hermi'
else:
aosym = 's1'
blksize = min(max(16, int(max_memory*1e6*.75/16/nao**2)), 16384)
sublk = max(16, int(blksize//4))
buf = [numpy.zeros(nao*nao*blksize, dtype=numpy.complex128)]
pqkRbuf = numpy.empty(nao*nao*sublk)
pqkIbuf = numpy.empty(nao*nao*sublk)
for p0, p1 in self.prange(0, ngs, blksize):
#aoao = ft_ao.ft_aopair(cell, Gv[p0:p1], shls_slice, aosym, invh,
# gxyz[p0:p1], gs, (kpti, kptj))
aoao = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym,
invh, gxyz[p0:p1], gs,
kptj-kpti, kptj.reshape(1,3), out=buf)[0]
for i0, i1 in lib.prange(0, p1-p0, sublk):
nG = i1 - i0
pqkR = numpy.ndarray((nao,nao,nG), buffer=pqkRbuf)
pqkI = numpy.ndarray((nao,nao,nG), buffer=pqkIbuf)
pqkR[:] = aoao[i0:i1].real.transpose(1,2,0)
pqkI[:] = aoao[i0:i1].imag.transpose(1,2,0)
yield (pqkR.reshape(-1,nG),
pqkI.reshape(-1,nG), p0+i0, p0+i1)
aoao[:] = 0
def mep(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION):
"""Calculates the molecular electrostatic potential (MEP) and write out in
cube format.
Args:
mol : Mole
Molecule to calculate the electron density for.
outfile : str
Name of Cube file to be written.
dm : ndarray
Density matrix of molecule.
Kwargs:
nx : int
Number of grid point divisions in x direction.
Note this is function of the molecule's size; a larger molecule
will have a coarser representation than a smaller one for the
same value.
ny : int
Number of grid point divisions in y direction.
nz : int
Number of grid point divisions in z direction.
"""
cc = Cube(mol, nx, ny, nz, resolution)
coords = cc.get_coords()
# Nuclear potential at given points
Vnuc = 0
for i in range(mol.natm):
r = mol.atom_coord(i)
Z = mol.atom_charge(i)
rp = r - coords
Vnuc += Z / numpy.einsum('xi,xi->x', rp, rp)**.5
# Potential of electron density
Vele = numpy.empty_like(Vnuc)
for p0, p1 in lib.prange(0, Vele.size, 600):
fakemol = gto.fakemol_for_charges(coords[p0:p1])
ints = df.incore.aux_e2(mol, fakemol)
Vele[p0:p1] = numpy.einsum('ijp,ij->p', ints, dm)
MEP = Vnuc - Vele # MEP at each point
MEP = MEP.reshape(nx,ny,nz)
# Write the potential
cc.write(MEP, outfile, 'Molecular electrostatic potential in real space')
def mep(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION):
"""Calculates the molecular electrostatic potential (MEP) and write out in
cube format.
Args:
mol : Mole
Molecule to calculate the electron density for.
outfile : str
Name of Cube file to be written.
dm : ndarray
Density matrix of molecule.
Kwargs:
nx : int
Number of grid point divisions in x direction.
Note this is function of the molecule's size; a larger molecule
will have a coarser representation than a smaller one for the
same value.
ny : int
Number of grid point divisions in y direction.
nz : int
Number of grid point divisions in z direction.
"""
cc = Cube(mol, nx, ny, nz, resolution)
coords = cc.get_coords()
# Nuclear potential at given points
Vnuc = 0
for i in range(mol.natm):
r = mol.atom_coord(i)
Z = mol.atom_charge(i)
rp = r - coords
Vnuc += Z / numpy.einsum('xi,xi->x', rp, rp)**.5
# Potential of electron density
Vele = numpy.empty_like(Vnuc)
for p0, p1 in lib.prange(0, Vele.size, 600):
fakemol = gto.fakemol_for_charges(coords[p0:p1])
ints = df.incore.aux_e2(mol, fakemol)
Vele[p0:p1] = numpy.einsum('ijp,ij->p', ints, dm)
MEP = Vnuc - Vele # MEP at each point
MEP = MEP.reshape(nx, ny, nz)
# Write the potential
cc.write(MEP, outfile, 'Molecular electrostatic potential in real space')
def _sort_eri(mycc, eris, nocc, nvir, vvop, log):
cpu1 = (time.clock(), time.time())
mol = mycc.mol
nmo = nocc + nvir
if mol.symmetry:
orbsym = symm.addons.label_orb_symm(mol,
mol.irrep_id,
mol.symm_orb,
eris.mo_coeff,
check=False)
orbsym = numpy.asarray(orbsym, dtype=numpy.int32) % 10
else:
orbsym = numpy.zeros(nmo, dtype=numpy.int32)
o_sorted = _irrep_argsort(orbsym[:nocc])
v_sorted = _irrep_argsort(orbsym[nocc:])
vrank = numpy.argsort(v_sorted)
max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
max_memory = min(8000, max_memory * .9)
blksize = min(nvir, max(16,
int(max_memory * 1e6 / 8 / (nvir * nocc * nmo))))
log.debug1('_sort_eri max_memory %g blksize %d', max_memory, blksize)
dtype = vvop.dtype
with lib.call_in_background(vvop.__setitem__,
sync=not mycc.async_io) as save:
bufopv = numpy.empty((nocc, nmo, nvir), dtype=dtype)
buf1 = numpy.empty_like(bufopv)
buf = numpy.empty((nocc, nvir, nvir), dtype=dtype)
for j0, j1 in lib.prange(0, nvir, blksize):
ovov = numpy.asarray(eris.ovov[:, j0:j1])
#ovvv = numpy.asarray(eris.ovvv[:,j0:j1])
ovvv = eris.get_ovvv(slice(None), slice(j0, j1))
for j in range(j0, j1):
oov = ovov[o_sorted, j - j0]
ovv = ovvv[o_sorted, j - j0]
#if ovv.ndim == 2:
# ovv = lib.unpack_tril(ovv, out=buf)
bufopv[:, :nocc, :] = oov[:, o_sorted][:, :, v_sorted].conj()
bufopv[:, nocc:, :] = ovv[:, v_sorted][:, :, v_sorted].conj()
save(vrank[j], bufopv.transpose(2, 0, 1))
bufopv, buf1 = buf1, bufopv
cpu1 = log.timer_debug1('transpose %d:%d' % (j0, j1), *cpu1)
return orbsym
def bcast(buf, root=0):
if size == 1:
return buf
is_array = isinstance(buf, np.ndarray)
buf = np.asarray(buf, order='C')
buf = buf.astype(buf.dtype.char)
shape, mpi_dtype = comm.bcast((buf.shape, buf.dtype.char))
if rank != root:
buf = np.empty(shape, dtype=mpi_dtype)
buf_seg = np.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 if is_array else buf.ravel()[0]
def matelem_int3d_coo(self, g, v):
""" Compute matrix elements of a potential v given on the 3d grid g using blocks along the grid """
from pyscf import lib
bsize = int(min(max(160e6 / (self.norbs * 8.0), 1), g.size))
#print(__name__, bsize, g.size*self.norbs*8)
v_matelem = np.zeros((self.norbs, self.norbs))
va = v.reshape(-1)
wgts = g.weights if type(g.weights) == np.ndarray else np.repeat(
g.weights, g.size)
for s, f in lib.prange(0, g.size, bsize):
ca2o = self.comp_aos_den(
g.coords[s:f]) # compute values of atomic orbitals
v_w = (wgts[s:f] * va[s:f]).reshape((f - s, 1))
cb2vo = ca2o * v_w
v_matelem += np.dot(ca2o.T, cb2vo)
return coo_matrix(v_matelem)
def allreduce(sendbuf, root=0, op=getattr(mpi, 'SUM', None)):
if size == 1:
return sendbuf
is_array = isinstance(sendbuf, np.ndarray)
sendbuf = np.asarray(sendbuf, order='C')
sendbuf = sendbuf.astype(sendbuf.dtype.char)
shape, mpi_dtype = comm.bcast((sendbuf.shape, sendbuf.dtype.char))
assert sendbuf.shape == shape and sendbuf.dtype.char == mpi_dtype
recvbuf = np.zeros_like(sendbuf)
send_seg = np.ndarray(sendbuf.size, dtype=sendbuf.dtype, buffer=sendbuf)
recv_seg = np.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 if is_array else recvbuf.ravel()[0]
def orbital(mol, outfile, coeff, nx=80, ny=80, nz=80, resolution=RESOLUTION):
"""Calculate orbital value on real space grid and write out in cube format.
Args:
mol : Mole
Molecule to calculate the electron density for.
outfile : str
Name of Cube file to be written.
coeff : 1D array
coeff coefficient.
Kwargs:
nx : int
Number of grid point divisions in x direction.
Note this is function of the molecule's size; a larger molecule
will have a coarser representation than a smaller one for the
same value. Conflicts to keyword resolution.
ny : int
Number of grid point divisions in y direction.
nz : int
Number of grid point divisions in z direction.
resolution: float
Resolution of the mesh grid in the cube box. If resolution is
given in the input, the input nx/ny/nz have no effects. The value
of nx/ny/nz will be determined by the resolution and the cube box
size.
"""
cc = Cube(mol, nx, ny, nz, resolution)
# Compute density on the .cube grid
coords = cc.get_coords()
ngrids = cc.get_ngrids()
blksize = min(8000, ngrids)
orb_on_grid = numpy.empty(ngrids)
for ip0, ip1 in lib.prange(0, ngrids, blksize):
ao = numint.eval_ao(mol, coords[ip0:ip1])
orb_on_grid[ip0:ip1] = numpy.dot(ao, coeff)
orb_on_grid = orb_on_grid.reshape(cc.nx, cc.ny, cc.nz)
# Write out orbital to the .cube file
cc.write(orb_on_grid,
outfile,
comment='Orbital value in real space (1/Bohr^3)')
return orb_on_grid
def make_phi(pcmobj, dm, r_vdw, ui):
mol = pcmobj.mol
natm = mol.natm
coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
ngrid_1sph = coords_1sph.shape[0]
if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
dm = dm[0] + dm[1]
tril_dm = lib.pack_tril(dm+dm.T)
nao = dm.shape[0]
diagidx = numpy.arange(nao)
diagidx = diagidx*(diagidx+1)//2 + diagidx
tril_dm[diagidx] *= .5
atom_coords = mol.atom_coords()
atom_charges = mol.atom_charges()
extern_point_idx = ui > 0
cav_coords = (atom_coords.reshape(natm,1,3)
+ numpy.einsum('r,gx->rgx', r_vdw, coords_1sph))
v_phi = numpy.empty((natm,ngrid_1sph))
for ia in range(natm):
# Note (-) sign is not applied to atom_charges, because (-) is explicitly
# included in rhs and L matrix
d_rs = atom_coords.reshape(-1,1,3) - cav_coords[ia]
v_phi[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))
max_memory = pcmobj.max_memory - lib.current_memory()[0]
blksize = int(max(max_memory*1e6/8/nao**2, 400))
cav_coords = cav_coords[extern_point_idx]
v_phi_e = numpy.empty(cav_coords.shape[0])
int3c2e = mol._add_suffix('int3c2e')
for i0, i1 in lib.prange(0, cav_coords.shape[0], blksize):
fakemol = gto.fakemol_for_charges(cav_coords[i0:i1])
v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s2ij')
v_phi_e[i0:i1] = numpy.einsum('x,xk->k', tril_dm, v_nj)
v_phi[extern_point_idx] -= v_phi_e
ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, pcmobj.lmax, True))
phi = -numpy.einsum('n,xn,jn,jn->jx', weights_1sph, ylm_1sph, ui, v_phi)
return phi
def loop(self, blksize=None):
if blksize is None:
blksize = self.blockdim
kpts0 = numpy.zeros((2, 3))
coulG = tools.get_coulG(self.cell,
numpy.zeros(3),
mesh=self.mesh,
low_dim_ft_type=self.low_dim_ft_type)
ngrids = len(coulG)
ao_pairs_G = self.get_ao_pairs_G(kpts0, compact=True)
ao_pairs_G *= numpy.sqrt(coulG * (self.cell.vol / ngrids**2)).reshape(
-1, 1)
Lpq = numpy.empty((self.blockdim, ao_pairs_G.shape[1]))
for p0, p1 in lib.prange(0, ngrids, blksize):
Lpq[:p1 - p0] = ao_pairs_G[p0:p1].real
yield Lpq[:p1 - p0]
Lpq[:p1 - p0] = ao_pairs_G[p0:p1].imag
yield Lpq[:p1 - p0]
def sr_loop(self, kpti_kptj=numpy.zeros((2,3)), max_memory=2000,
compact=True):
'''Short range part'''
kpti, kptj = kpti_kptj
unpack = is_zero(kpti-kptj) and not compact
nao = self.cell.nao_nr()
if is_zero(kpti_kptj) and compact:
nao_pair = nao * (nao+1) // 2
else:
nao_pair = nao ** 2
if is_zero(kpti_kptj):
blksize = max_memory*1e6/8/(nao_pair*5+nao*(nao+1)//2)
else:
blksize = max_memory*1e6/16/(nao_pair*5+nao*(nao+1)//2)
blksize = max(16, min(int(blksize), self.blockdim))
logger.debug2(self, 'max_memory %d MB, blksize %d', max_memory, blksize)
if unpack:
buf = numpy.empty((blksize,nao*(nao+1)//2))
def load(Lpq, b0, b1, bufR, bufI):
Lpq = numpy.asarray(Lpq[b0:b1])
shape = Lpq.shape
if unpack:
tmp = numpy.ndarray(shape, buffer=buf)
tmp[:] = Lpq.real
LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1,nao**2)
tmp[:] = Lpq.imag
LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI, out=bufI).reshape(-1,nao**2)
else:
LpqR = numpy.ndarray(shape, buffer=bufR)
LpqR[:] = Lpq.real
LpqI = numpy.ndarray(shape, buffer=bufI)
LpqI[:] = Lpq.imag
return LpqR, LpqI
LpqR = LpqI = j3cR = j3cI = None
with self.load_Lpq(kpti_kptj) as Lpq:
naux = Lpq.shape[0]
with self.load_j3c(kpti_kptj) as j3c:
for b0, b1 in lib.prange(0, naux, blksize):
LpqR, LpqI = load(Lpq, b0, b1, LpqR, LpqI)
j3cR, j3cI = load(j3c, b0, b1, j3cR, j3cI)
yield LpqR, LpqI, j3cR, j3cI
def AO(mol,
C,
outfile='AO_dir/AOs',
nx=80,
ny=80,
nz=80,
resolution=None,
save=True):
counter = 1
if not os.path.exists(outfile):
os.makedirs(outfile)
else:
exists = True
while exists:
cache = outfile + '_' + str(counter)
exists = os.path.exists(cache)
counter += 1
outfile = cache
os.makedirs(cache)
cc = Cube(mol, nx, ny, nz, resolution)
# Compute density on the .cube grid
coords = cc.get_coords()
ngrids = cc.get_ngrids()
#blksize = min(8000, ngrids)
blksize = ngrids
rho = np.empty(ngrids)
for ip0, ip1 in lib.prange(0, ngrids, blksize):
ao = numint.eval_ao(mol, coords[ip0:ip1])
x = (coords[:, 0].min(), coords[:, 0].max(), nx)
y = (coords[:, 1].min(), coords[:, 1].max(), ny)
z = (coords[:, 2].min(), coords[:, 2].max(), nz)
if save:
np.savetxt(outfile + '/AOs', ao)
np.savetxt(outfile + '/C', C)
np.savetxt(outfile + '/coords', np.array([y, x, z]))
f = open(outfile + '/geom_basis', 'w')
f.write(str(mol.atom) + '_')
f.write(str(mol.basis))
f.close()
#cc.write(np.array([]),outfile + '_geom',comment = 'Geometry')
return ao, np.array([y, x, z])
def get_hcore(self, mol=None):
if mol is None: mol = self.mol
if getattr(method_class, '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')
coords = self.mm_mol.atom_coords()
charges = self.mm_mol.atom_charges()
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))
if max_memory <= 0:
blksize = 1
logger.warn(
self,
'Memory estimate for reading point charges is negative. '
'Trying to read point charges one by one.')
cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas,
mol._env, intor)
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',
cintopt=cintopt)
v += numpy.einsum('xk,k->x', j3c, -charges[i0:i1])
v = lib.unpack_tril(v)
return h1e + v
def allgather_new(sendbuf, split_recvbuf=False):
"""
gather that avoids overflow of displs.
"""
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])
dtype = mpi_dtype = numpy.result_type(*[x[1] for x in attr]).char
_assert(sendbuf.dtype.char == mpi_dtype or sendbuf.size == 0)
matched = all([x[1:] == rshape[0][1:] for x in rshape[1:]])
if matched:
elem_dtype = MPI._typedict.get(mpi_dtype)
each_count = numpy.prod(rshape[0][1:])
counts = counts // each_count
mpi_dtype = MPI.Datatype(elem_dtype).Create_contiguous(
each_count).Commit()
else:
each_count = 1
displs = numpy.append(0, numpy.cumsum(counts[:-1]))
recvbuf = numpy.empty(sum(counts * each_count), dtype=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 * each_count:p1 * each_count], mpi_dtype],
[recvbuf, counts_seg, displs + p0, mpi_dtype])
if matched:
mpi_dtype.Free()
if split_recvbuf:
return [
recvbuf[p0 * each_count:(p0 + c) * each_count].reshape(shape)
for p0, c, shape in zip(displs, counts, rshape)
]
else:
try:
return recvbuf.reshape((-1, ) + shape[1:])
except ValueError:
return recvbuf
def loop(self, blksize=None):
if self.cell.dimension < 3:
raise RuntimeError('ERIs of 1D and 2D systems are not positive '
'definite. Current API only supports postive '
'definite ERIs.')
if blksize is None:
blksize = self.blockdim
kpts0 = numpy.zeros((2,3))
coulG = tools.get_coulG(self.cell, numpy.zeros(3), mesh=self.mesh)
ngrids = len(coulG)
ao_pairs_G = self.get_ao_pairs_G(kpts0, compact=True)
ao_pairs_G *= numpy.sqrt(coulG*(self.cell.vol/ngrids**2)).reshape(-1,1)
Lpq = numpy.empty((blksize, ao_pairs_G.shape[1]))
for p0, p1 in lib.prange(0, ngrids, blksize):
Lpq[:p1-p0] = ao_pairs_G[p0:p1].real
yield Lpq[:p1-p0]
Lpq[:p1-p0] = ao_pairs_G[p0:p1].imag
yield Lpq[:p1-p0]
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 * .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 sr_loop(self, kpti_kptj=numpy.zeros((2,3)), max_memory=2000,
compact=True, transpose102=False):
'''Short range part'''
kpti, kptj = kpti_kptj
unpack = is_zero(kpti-kptj) and not compact
nao = self.cell.nao_nr()
max_memory = max(2000, self.max_memory-lib.current_memory()[0])
if is_zero(kpti_kptj):
nao_pair = nao * (nao+1) // 2
blksize = max(16, min(int(max_memory*1e6/8/nao_pair/2), self.blockdim))
else:
blksize = max(16, min(int(max_memory*1e6/16/nao**2/2), self.blockdim))
if unpack:
buf = numpy.empty((blksize,nao*(nao+1)//2))
def load(Lpq, bufR, bufI):
shape = Lpq.shape
if unpack:
tmp = numpy.ndarray(shape, buffer=buf)
tmp[:] = Lpq.real
LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1,nao**2)
tmp[:] = Lpq.imag
LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI, out=bufI).reshape(-1,nao**2)
else:
LpqR = numpy.ndarray(shape, buffer=bufR)
LpqR[:] = Lpq.real
LpqI = numpy.ndarray(shape, buffer=bufI)
LpqI[:] = Lpq.imag
if transpose102:
LpqR = numpy.asarray(LpqR.reshape(-1,nao,nao).transpose(1,0,2), order='C')
LpqI = numpy.asarray(LpqI.reshape(-1,nao,nao).transpose(1,0,2), order='C')
return LpqR, LpqI
LpqR = LpqI = j3cR = j3cI = None
with self.load_j3c(kpti_kptj) as j3c:
with self.load_Lpq(kpti_kptj) as Lpq:
naoaux = j3c.shape[0]
for b0, b1 in lib.prange(0, naoaux, blksize):
LpqR, LpqI = load(numpy.asarray(Lpq[b0:b1]), LpqR, LpqI)
j3cR, j3cI = load(numpy.asarray(j3c[b0:b1]), j3cR, j3cI)
yield LpqR, LpqI, j3cR, j3cI
def _sort_eri(mycc, eris, nocc, nvir, vvop, log):
cpu1 = (time.clock(), time.time())
mol = mycc.mol
nmo = nocc + nvir
if mol.symmetry:
orbsym = symm.addons.label_orb_symm(mol, mol.irrep_id, mol.symm_orb,
eris.mo_coeff, check=False)
orbsym = numpy.asarray(orbsym, dtype=numpy.int32) % 10
else:
orbsym = numpy.zeros(nmo, dtype=numpy.int32)
o_sorted = _irrep_argsort(orbsym[:nocc])
v_sorted = _irrep_argsort(orbsym[nocc:])
vrank = numpy.argsort(v_sorted)
max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
max_memory = min(8000, max_memory*.9)
blksize = min(nvir, max(16, int(max_memory*1e6/8/(nvir*nocc*nmo))))
log.debug1('_sort_eri max_memory %g blksize %d', max_memory, blksize)
dtype = vvop.dtype
with lib.call_in_background(vvop.__setitem__, sync=not mycc.async_io) as save:
bufopv = numpy.empty((nocc,nmo,nvir), dtype=dtype)
buf1 = numpy.empty_like(bufopv)
buf = numpy.empty((nocc,nvir,nvir), dtype=dtype)
for j0, j1 in lib.prange(0, nvir, blksize):
ovov = numpy.asarray(eris.ovov[:,j0:j1])
#ovvv = numpy.asarray(eris.ovvv[:,j0:j1])
ovvv = eris.get_ovvv(slice(None), slice(j0,j1))
for j in range(j0,j1):
oov = ovov[o_sorted,j-j0]
ovv = ovvv[o_sorted,j-j0]
#if ovv.ndim == 2:
# ovv = lib.unpack_tril(ovv, out=buf)
bufopv[:,:nocc,:] = oov[:,o_sorted][:,:,v_sorted].conj()
bufopv[:,nocc:,:] = ovv[:,v_sorted][:,:,v_sorted].conj()
save(vrank[j], bufopv.transpose(2,0,1))
bufopv, buf1 = buf1, bufopv
cpu1 = log.timer_debug1('transpose %d:%d'%(j0,j1), *cpu1)
return orbsym
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
if eris is None: eris = mycc.ao2mo()
nocc, nvir = t1.shape
fock = eris.fock
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, nvir, blksize):
eris_ovvo = eris.ovvo[:, p0:p1]
tau = t2[:, :, p0:p1]
e += 2 * numpy.einsum('ijab,iabj', tau, eris_ovvo)
e -= numpy.einsum('jiab,iabj', tau, eris_ovvo)
if abs(e.imag) > 1e-4:
logger.warn(mycc, 'Non-zero imaginary part found in QCISD energy %s',
e)
return e.real
def write_coeff(fout, mol, mo_coeff):
#fout.write('ALDET\n')
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._libcint_ctr_coeff(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.transpose(1,0,2).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('From PySCF\n')
for k in range(nmo):
mo = mo_cart[:,k]
fout.write('CANMO %d\n' % (k+1))
for i0, i1 in lib.prange(0, nprim, 5):
fout.write(' %s\n' % ' '.join('%15.8E'%x for x in mo[i0:i1]))
fout.write('END DATA\n')
def _mm_pot(mol: gto.Mole, mm_mol: gto.Mole) -> np.ndarray:
"""
this function returns the full mm potential
(adapted from: qmmm/itrf.py:get_hcore() in PySCF)
"""
# settings
coords = mm_mol.atom_coords()
charges = mm_mol.atom_charges()
blksize = BLKSIZE
# integrals
intor = 'int3c2e_cart' if mol.cart else 'int3c2e_sph'
cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas,
mol._env, intor)
# compute interaction potential
mm_pot = 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', cintopt=cintopt)
mm_pot += np.einsum('xk,k->x', j3c, -charges[i0:i1])
mm_pot = lib.unpack_tril(mm_pot)
return mm_pot
def write(self, fname, **kw):
import time
import pyscf
from pyscf import lib
""" Result: .cube file with the field in the file fname. """
if 'mol' in kw:
mol = kw['mol'] # Obligatory argument
coord = mol.atom_coords()
zz = [mol.atom_charge(ia) for ia in range(mol.natm)]
natm = mol.natm
else:
zz = kw['a2z']
natm = len(zz)
coord = kw['a2xyz']
field = kw['field'] # Obligatory argument?
comment = kw['comment'] if 'comment' in kw else 'none'
with open(fname, 'w') as f:
f.write(comment + '\n')
f.write('PySCF Version: {:s} Date: {:s}\n'.format(
pyscf.__version__, time.ctime()))
bo = self.origin - self.get_mesh_center()
f.write(('{:5d}' + '{:12.6f}' * 3 + '\n').format(natm, *bo))
for s, rr in zip(self.shape, self.rr):
f.write(('{:5d}' + '{:12.6f}' * 3 + '\n').format(s, *rr[1]))
for chg, xyz in zip(zz, coord):
f.write(
('{:5d}' + '{:12.6f}' * 4 + '\n').format(chg, chg, *xyz))
for ix in range(self.shape[0]):
for iy in range(self.shape[1]):
for iz0, iz1 in lib.prange(0, self.shape[2], 6):
f.write(('{:13.5e}' * (iz1 - iz0) +
'\n').format(*field[ix, iy, iz0:iz1]))
def get_gridss(mol, level=1, gthrd=1e-10):
Ktime = (time.clock(), time.time())
grids = dft.gen_grid.Grids(mol)
grids.level = level
grids.build()
ngrids = grids.weights.size
mask = []
for p0, p1 in lib.prange(0, ngrids, 10000):
ao_v = mol.eval_gto('GTOval', grids.coords[p0:p1])
ao_v *= grids.weights[p0:p1,None]
wao_v0 = ao_v
mask.append(numpy.any(wao_v0>gthrd, axis=1) |
numpy.any(wao_v0<-gthrd, axis=1))
mask = numpy.hstack(mask)
grids.coords = grids.coords[mask]
grids.weights = grids.weights[mask]
logger.debug(mol, 'threshold for grids screening %g', gthrd)
logger.debug(mol, 'number of grids %d', grids.weights.size)
logger.timer_debug1(mol, "Xg screening", *Ktime)
return grids
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(
f'PySCF Version: {pyscf.__version__} Date: {time.ctime()}\n')
f.write(f'{mol.natm:5d}')
f.write('%12.6f%12.6f%12.6f\n' % tuple(self.boxorig.tolist()))
dx = self.xs[-1] if len(self.xs) == 1 else self.xs[1]
dy = self.ys[-1] if len(self.ys) == 1 else self.ys[1]
dz = self.zs[-1] if len(self.zs) == 1 else self.zs[1]
delta = (self.box.T * [dx, dy, dz]).T
f.write(
f'{self.nx:5d}{delta[0,0]:12.6f}{delta[0,1]:12.6f}{delta[0,2]:12.6f}\n'
)
f.write(
f'{self.ny:5d}{delta[1,0]:12.6f}{delta[1,1]:12.6f}{delta[1,2]:12.6f}\n'
)
f.write(
f'{self.nz:5d}{delta[2,0]:12.6f}{delta[2,1]:12.6f}{delta[2,2]:12.6f}\n'
)
for ia in range(mol.natm):
atmsymb = mol.atom_symbol(ia)
f.write('%5d%12.6f' % (gto.charge(atmsymb), 0.))
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 orbital(mol, outfile, coeff, nx=80, ny=80, nz=80, resolution=RESOLUTION):
"""Calculate orbital value on real space grid and write out in cube format.
Args:
mol : Mole
Molecule to calculate the electron density for.
outfile : str
Name of Cube file to be written.
coeff : 1D array
coeff coefficient.
Kwargs:
nx : int
Number of grid point divisions in x direction.
Note this is function of the molecule's size; a larger molecule
will have a coarser representation than a smaller one for the
same value.
ny : int
Number of grid point divisions in y direction.
nz : int
Number of grid point divisions in z direction.
"""
cc = Cube(mol, nx, ny, nz, resolution)
# Compute density on the .cube grid
coords = cc.get_coords()
ngrids = cc.get_ngrids()
blksize = min(8000, ngrids)
orb_on_grid = numpy.zeros(ngrids, dtype=numpy.complex128)
for ip0, ip1 in lib.prange(0, ngrids, blksize):
aoa, aob = eval_ao(mol, coords[ip0:ip1], deriv=0)
orb_on_grid[ip0:ip1] = numpy.dot(aoa, coeff)
orb_on_grid[ip0:ip1] += numpy.dot(aob, coeff)
orb_on_grid = orb_on_grid.reshape(cc.nx,cc.ny,cc.nz)
# Write out orbital to the .cube file
cc.write(orb_on_grid, outfile, comment='Orbital value in real space (1/Bohr^3)')
def density(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION):
"""Calculates electron density and write out in cube format.
Args:
mol : Mole
Molecule to calculate the electron density for.
outfile : str
Name of Cube file to be written.
dm : ndarray
Density matrix of molecule.
Kwargs:
nx : int
Number of grid point divisions in x direction.
Note this is function of the molecule's size; a larger molecule
will have a coarser representation than a smaller one for the
same value.
ny : int
Number of grid point divisions in y direction.
nz : int
Number of grid point divisions in z direction.
"""
cc = Cube(mol, nx, ny, nz, resolution)
# Compute density on the .cube grid
coords = cc.get_coords()
ngrids = cc.get_ngrids()
blksize = min(8000, ngrids)
rho = numpy.empty(ngrids)
for ip0, ip1 in lib.prange(0, ngrids, blksize):
ao = numint.eval_ao(mol, coords[ip0:ip1])
rho[ip0:ip1] = numint.eval_rho(mol, ao, dm)
rho = rho.reshape(cc.nx, cc.ny, cc.nz)
# Write out density to the .cube file
cc.write(rho, outfile, comment='Electron density in real space (e/Bohr^3)')
def init_amps(mycc, eris=None):
eris = getattr(mycc, '_eris', None)
if eris is None:
mycc.ao2mo()
eris = mycc._eris
time0 = logger.process_clock(), logger.perf_counter()
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
t1T = np.zeros_like(eris.fock[nocc:, :nocc])
loc0, loc1 = _task_location(nvir)
t2T = np.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_vvoo = eris.xvoo[p0:p1]
t2T[p0:p1] = (
eris_vvoo /
lib.direct_sum('ia, jb -> abij', eia[:, loc0 + p0:loc0 + p1], eia))
emp2 += np.einsum('abij, abij',
t2T[p0:p1],
eris_vvoo.conj(),
optimize=True).real
eris_vvoo = None
mycc.emp2 = comm.allreduce(emp2) * 0.25
logger.info(mycc, 'Init t2, MP2 energy = %.15g', mycc.emp2)
logger.timer(mycc, 'init mp2', *time0)
mycc.t1 = t1T.T
mycc.t2 = t2T.transpose(2, 3, 0, 1)
return mycc.emp2, mycc.t1, mycc.t2
def dump_normal_mode(mol, results):
dump = mol.stdout.write
freq_wn = results['freq_wavenumber']
idx = freq_wn.real > 0
freq_wn = freq_wn.real[idx]
nfreq = freq_wn.size
r_mass = results['reduced_mass'].real[idx]
force = results['force_const_dyne'].real[idx]
vib_t = results['vib_temperature'].real[idx]
mode = results['norm_mode'].real[idx]
symbols = [mol.atom_symbol(i) for i in range(mol.natm)]
def inline(q, col0, col1):
return ''.join('%20.4f' % q[i] for i in range(col0, col1))
def mode_inline(row, col0, col1):
return ' '.join('%6.2f%6.2f%6.2f' %
(mode[i, row, 0], mode[i, row, 1], mode[i, row, 2])
for i in range(col0, col1))
for col0, col1 in lib.prange(0, nfreq, 3):
dump('Mode %s\n' % ''.join('%20d' % i
for i in range(col0, col1)))
dump('Irrep\n')
dump('Freq [cm^-1] %s\n' % inline(freq_wn, col0, col1))
dump('Reduced mass [au] %s\n' % inline(r_mass, col0, col1))
dump('Force const [Dyne/A] %s\n' % inline(force, col0, col1))
dump('Char temp [K] %s\n' % inline(vib_t, col0, col1))
#dump('IR\n')
#dump('Raman\n')
dump('Normal mode %s\n' % (' x y z' *
(col1 - col0)))
for j, at in enumerate(symbols):
dump(' %4d%4s %s\n' %
(j, at, mode_inline(j, col0, col1)))
def get_hcore(self, mol=None):
''' (QM 1e grad) + <-d/dX i|q_mm/r_mm|j>'''
if mol is None: mol = self.mol
g_qm = scf_grad.get_hcore(mol)
nao = g_qm.shape[1]
if pyscf.DEBUG:
v = 0
for i,q in enumerate(charges):
mol.set_rinv_origin(coords[i])
v += mol.intor('int1e_iprinv', comp=3) * q
else:
if mol.cart:
intor = 'int3c2e_ip1_cart'
else:
intor = 'int3c2e_ip1_sph'
nao = mol.nao
max_memory = self.max_memory - lib.current_memory()[0]
blksize = int(max(max_memory*1e6/8/nao**2, 400))
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, aosym='s1', comp=3)
v += numpy.einsum('ipqk,k->ipq', j3c, charges[i0:i1])
return g_qm + v
def _assemble(mydf, kptij_lst, j3c_jobs, gen_int3c, ft_fuse, cderi_file, fswap,
log):
t1 = (time.clock(), time.time())
cell = mydf.cell
ao_loc = cell.ao_loc_nr()
nao = ao_loc[-1]
kptis = kptij_lst[:, 0]
kptjs = kptij_lst[:, 1]
kpt_ji = kptjs - kptis
uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
aosym_s2 = numpy.einsum('ix->i', abs(kptis - kptjs)) < 1e-9
t2 = t1
j3c_workers = numpy.zeros(len(j3c_jobs), dtype=int)
#for job_id, ish0, ish1 in mpi.work_share_partition(j3c_jobs):
for job_id, ish0, ish1 in mpi.work_stealing_partition(j3c_jobs):
gen_int3c(job_id, ish0, ish1)
t2 = log.alltimer_debug2('int j3c %d' % job_id, *t2)
for k, kpt in enumerate(uniq_kpts):
ft_fuse(job_id, k, ish0, ish1)
t2 = log.alltimer_debug2('ft-fuse %d k %d' % (job_id, k), *t2)
j3c_workers[job_id] = rank
j3c_workers = mpi.allreduce(j3c_workers)
log.debug2('j3c_workers %s', j3c_workers)
t1 = log.timer_debug1('int3c and fuse', *t1)
# Pass 2
# Transpose 3-index tensor and save data in cderi_file
feri = h5py.File(cderi_file, 'w')
nauxs = [fswap['j2c/%d' % k].shape[0] for k, kpt in enumerate(uniq_kpts)]
segsize = (max(nauxs) + mpi.pool.size - 1) // mpi.pool.size
naux0 = rank * segsize
for k, kptij in enumerate(kptij_lst):
naux1 = min(nauxs[uniq_inverse[k]], naux0 + segsize)
nrow = max(0, naux1 - naux0)
if gamma_point(kptij):
dtype = 'f8'
else:
dtype = 'c16'
if aosym_s2[k]:
nao_pair = nao * (nao + 1) // 2
else:
nao_pair = nao * nao
feri.create_dataset('j3c/%d' % k, (nrow, nao_pair),
dtype,
maxshape=(None, nao_pair))
def get_segs_loc(aosym):
off0 = numpy.asarray([ao_loc[i0] for x, i0, i1 in j3c_jobs])
off1 = numpy.asarray([ao_loc[i1] for x, i0, i1 in j3c_jobs])
if aosym: # s2
dims = off1 * (off1 + 1) // 2 - off0 * (off0 + 1) // 2
else:
dims = (off1 - off0) * nao
#dims = numpy.asarray([ao_loc[i1]-ao_loc[i0] for x,i0,i1 in j3c_jobs])
dims = numpy.hstack(
[dims[j3c_workers == w] for w in range(mpi.pool.size)])
job_idx = numpy.hstack(
[numpy.where(j3c_workers == w)[0] for w in range(mpi.pool.size)])
segs_loc = numpy.append(0, numpy.cumsum(dims))
segs_loc = [(segs_loc[j], segs_loc[j + 1])
for j in numpy.argsort(job_idx)]
return segs_loc
segs_loc_s1 = get_segs_loc(False)
segs_loc_s2 = get_segs_loc(True)
job_ids = numpy.where(rank == j3c_workers)[0]
def load(k, p0, p1):
naux1 = nauxs[uniq_inverse[k]]
slices = [(min(i * segsize + p0, naux1), min(i * segsize + p1, naux1))
for i in range(mpi.pool.size)]
segs = []
for p0, p1 in slices:
val = [
fswap['j3c-chunks/%d/%d' % (job, k)][p0:p1].ravel()
for job in job_ids
]
if val:
segs.append(numpy.hstack(val))
else:
segs.append(numpy.zeros(0))
return segs
def save(k, p0, p1, segs):
segs = mpi.alltoall(segs)
naux1 = nauxs[uniq_inverse[k]]
loc0, loc1 = min(p0, naux1 - naux0), min(p1, naux1 - naux0)
nL = loc1 - loc0
if nL > 0:
if aosym_s2[k]:
segs = numpy.hstack([
segs[i0 * nL:i1 * nL].reshape(nL, -1)
for i0, i1 in segs_loc_s2
])
else:
segs = numpy.hstack([
segs[i0 * nL:i1 * nL].reshape(nL, -1)
for i0, i1 in segs_loc_s1
])
feri['j3c/%d' % k][loc0:loc1] = segs
mem_now = max(comm.allgather(lib.current_memory()[0]))
max_memory = max(2000, min(8000, mydf.max_memory - mem_now))
if numpy.all(aosym_s2):
if gamma_point(kptij_lst):
blksize = max(16, int(max_memory * .5e6 / 8 / nao**2))
else:
blksize = max(16, int(max_memory * .5e6 / 16 / nao**2))
else:
blksize = max(16, int(max_memory * .5e6 / 16 / nao**2 / 2))
log.debug1('max_momory %d MB (%d in use), blksize %d', max_memory, mem_now,
blksize)
t2 = t1
with lib.call_in_background(save) as async_write:
for k, kptji in enumerate(kptij_lst):
for p0, p1 in lib.prange(0, segsize, blksize):
segs = load(k, p0, p1)
async_write(k, p0, p1, segs)
t2 = log.timer_debug1(
'assemble k=%d %d:%d (in %d)' % (k, p0, p1, segsize), *t2)
if 'j2c-' in fswap:
j2c_kpts_lists = []
for k, kpt in enumerate(uniq_kpts):
if ('j2c-/%d' % k) in fswap:
adapted_ji_idx = numpy.where(uniq_inverse == k)[0]
j2c_kpts_lists.append(adapted_ji_idx)
for k in numpy.hstack(j2c_kpts_lists):
val = [
numpy.asarray(fswap['j3c-/%d/%d' % (job, k)]).ravel()
for job in job_ids
]
val = mpi.gather(numpy.hstack(val))
if rank == 0:
naux1 = fswap['j3c-/0/%d' % k].shape[0]
if aosym_s2[k]:
v = [
val[i0 * naux1:i1 * naux1].reshape(naux1, -1)
for i0, i1 in segs_loc_s2
]
else:
v = [
val[i0 * naux1:i1 * naux1].reshape(naux1, -1)
for i0, i1 in segs_loc_s1
]
feri['j3c-/%d' % k] = numpy.hstack(v)
if 'j3c-kptij' in feri: del (feri['j3c-kptij'])
feri['j3c-kptij'] = kptij_lst
t1 = log.alltimer_debug1('assembling j3c', *t1)
feri.close()
def ft_fuse(job_id, uniq_kptji_id, sh0, sh1):
kpt = uniq_kpts[uniq_kptji_id] # kpt = kptj - kpti
adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0]
adapted_kptjs = kptjs[adapted_ji_idx]
nkptj = len(adapted_kptjs)
j2c = numpy.asarray(fswap['j2c/%d' % uniq_kptji_id])
j2ctag = j2ctags[uniq_kptji_id]
naux0 = j2c.shape[0]
if ('j2c-/%d' % uniq_kptji_id) in fswap:
j2c_negative = numpy.asarray(fswap['j2c-/%d' % uniq_kptji_id])
else:
j2c_negative = None
if is_zero(kpt):
aosym = 's2'
else:
aosym = 's1'
if aosym == 's2' and cell.dimension == 3:
vbar = fuse(mydf.auxbar(fused_cell))
ovlp = cell.pbc_intor('int1e_ovlp', hermi=1, kpts=adapted_kptjs)
ovlp = [lib.pack_tril(s) for s in ovlp]
j3cR = [None] * nkptj
j3cI = [None] * nkptj
i0 = ao_loc[sh0]
i1 = ao_loc[sh1]
for k, idx in enumerate(adapted_ji_idx):
key = 'j3c-chunks/%d/%d' % (job_id, idx)
v = numpy.asarray(fswap[key])
if aosym == 's2' and cell.dimension == 3:
for i in numpy.where(vbar != 0)[0]:
v[i] -= vbar[i] * ovlp[k][i0 * (i0 + 1) // 2:i1 *
(i1 + 1) // 2].ravel()
j3cR[k] = numpy.asarray(v.real, order='C')
if v.dtype == numpy.complex128:
j3cI[k] = numpy.asarray(v.imag, order='C')
v = None
ncol = j3cR[0].shape[1]
Gblksize = max(16, int(max_memory * 1e6 / 16 / ncol /
(nkptj + 1))) # +1 for pqkRbuf/pqkIbuf
Gblksize = min(Gblksize, ngrids, 16384)
pqkRbuf = numpy.empty(ncol * Gblksize)
pqkIbuf = numpy.empty(ncol * Gblksize)
buf = numpy.empty(nkptj * ncol * Gblksize, dtype=numpy.complex128)
log.alldebug2('job_id %d blksize (%d,%d)', job_id, Gblksize, ncol)
wcoulG = mydf.weighted_coulG(kpt, False, mesh)
fused_cell_slice = (auxcell.nbas, fused_cell.nbas)
if aosym == 's2':
shls_slice = (sh0, sh1, 0, sh1)
else:
shls_slice = (sh0, sh1, 0, cell.nbas)
for p0, p1 in lib.prange(0, ngrids, Gblksize):
Gaux = ft_ao.ft_ao(fused_cell, Gv[p0:p1], fused_cell_slice, b,
gxyz[p0:p1], Gvbase, kpt)
Gaux *= wcoulG[p0:p1, None]
kLR = Gaux.real.copy('C')
kLI = Gaux.imag.copy('C')
Gaux = None
dat = ft_ao._ft_aopair_kpts(cell,
Gv[p0:p1],
shls_slice,
aosym,
b,
gxyz[p0:p1],
Gvbase,
kpt,
adapted_kptjs,
out=buf)
nG = p1 - p0
for k, ji in enumerate(adapted_ji_idx):
aoao = dat[k].reshape(nG, ncol)
pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf)
pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf)
pqkR[:] = aoao.real.T
pqkI[:] = aoao.imag.T
lib.dot(kLR.T, pqkR.T, -1, j3cR[k][naux:], 1)
lib.dot(kLI.T, pqkI.T, -1, j3cR[k][naux:], 1)
if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
lib.dot(kLR.T, pqkI.T, -1, j3cI[k][naux:], 1)
lib.dot(kLI.T, pqkR.T, 1, j3cI[k][naux:], 1)
kLR = kLI = None
for k, idx in enumerate(adapted_ji_idx):
if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
v = fuse(j3cR[k])
else:
v = fuse(j3cR[k] + j3cI[k] * 1j)
if j2ctag == 'CD':
v = scipy.linalg.solve_triangular(j2c,
v,
lower=True,
overwrite_b=True)
fswap['j3c-chunks/%d/%d' % (job_id, idx)][:naux0] = v
else:
fswap['j3c-chunks/%d/%d' % (job_id, idx)][:naux0] = lib.dot(
j2c, v)
# low-dimension systems
if j2c_negative is not None:
fswap['j3c-/%d/%d' % (job_id, idx)] = lib.dot(j2c_negative, v)
