def grids2d_int3c_jobs(cell, auxcell, kptij_lst, chunks, aosym_s2):
ao_loc = cell.ao_loc_nr()
if aosym_s2:
segs = ao_loc[1:]*(ao_loc[1:]+1)//2 - ao_loc[:-1]*(ao_loc[:-1]+1)//2
ij_ranges = balance_segs(segs, chunks[0]*chunks[1])
else:
segs = ao_loc[1:]-ao_loc[:-1]
ij_ranges = balance_segs(segs, chunks[0])
jobs = [(job_id, i0, i1) for job_id, (i0, i1, x) in enumerate(ij_ranges)]
return jobs
def grids2d_int3c_jobs(cell, auxcell, kptij_lst, chunks, aosym_s2):
ao_loc = cell.ao_loc_nr()
if aosym_s2:
segs = ao_loc[1:] * (ao_loc[1:] +
1) // 2 - ao_loc[:-1] * (ao_loc[:-1] + 1) // 2
ij_ranges = balance_segs(segs, chunks[0] * chunks[1])
else:
segs = ao_loc[1:] - ao_loc[:-1]
ij_ranges = balance_segs(segs, chunks[0])
jobs = [(job_id, i0, i1) for job_id, (i0, i1, x) in enumerate(ij_ranges)]
return jobs
def gen_int3c(auxcell, job_id, ish0, ish1):
dataname = 'j3c-chunks/%d' % job_id
if dataname in feri:
del (feri[dataname])
i0 = ao_loc[ish0]
i1 = ao_loc[ish1]
dii = i1 * (i1 + 1) // 2 - i0 * (i0 + 1) // 2
dij = (i1 - i0) * nao
if j_only:
buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * dii + dii)))
else:
buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * dij + dij)))
auxranges = balance_segs(aux_loc[1:] - aux_loc[:-1], buflen)
buflen = max([x[2] for x in auxranges])
buf = numpy.empty(nkptij * dij * buflen, dtype=dtype)
buf1 = numpy.empty(dij * buflen, dtype=dtype)
naux = aux_loc[-1]
for kpt_id, kptij in enumerate(kptij_lst):
key = '%s/%d' % (dataname, kpt_id)
if aosym_s2[kpt_id]:
shape = (naux, dii)
else:
shape = (naux, dij)
if gamma_point(kptij):
feri.create_dataset(key, shape, 'f8')
else:
feri.create_dataset(key, shape, 'c16')
naux0 = 0
for istep, auxrange in enumerate(auxranges):
log.alldebug2("aux_e2 job_id %d step %d", job_id, istep)
sh0, sh1, nrow = auxrange
sub_slice = (ish0, ish1, 0, cell.nbas, sh0, sh1)
if j_only:
mat = numpy.ndarray((nkptij, dii, nrow),
dtype=dtype,
buffer=buf)
else:
mat = numpy.ndarray((nkptij, dij, nrow),
dtype=dtype,
buffer=buf)
mat = int3c(sub_slice, mat)
for k, kptij in enumerate(kptij_lst):
h5dat = feri['%s/%d' % (dataname, k)]
v = lib.transpose(mat[k], out=buf1)
if not j_only and aosym_s2[k]:
idy = idxb[i0 * (i0 + 1) // 2:i1 *
(i1 + 1) // 2] - i0 * nao
out = numpy.ndarray((nrow, dii),
dtype=v.dtype,
buffer=mat[k])
v = numpy.take(v, idy, axis=1, out=out)
if gamma_point(kptij):
h5dat[naux0:naux0 + nrow] = v.real
else:
h5dat[naux0:naux0 + nrow] = v
naux0 += nrow
def _guess_shell_ranges(mol, buflen, aosym):
from pyscf.ao2mo.outcore import balance_segs
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
if 's2' in aosym:
segs = ao_loc[1:]*(ao_loc[1:]+1)//2 - ao_loc[:-1]*(ao_loc[:-1]+1)//2
else:
segs = (ao_loc[1:]-ao_loc[:-1])*nao
return balance_segs(segs, buflen)
def _guess_shell_ranges(mol, buflen, aosym):
from pyscf.ao2mo.outcore import balance_segs
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
if 's2' in aosym:
segs = ao_loc[1:] * (ao_loc[1:] +
1) // 2 - ao_loc[:-1] * (ao_loc[:-1] + 1) // 2
else:
segs = (ao_loc[1:] - ao_loc[:-1]) * nao
return balance_segs(segs, buflen)
def gen_int3c(job_id, ish0, ish1):
dataname = 'j3c-chunks/%d' % job_id
i0 = ao_loc[ish0]
i1 = ao_loc[ish1]
dii = i1*(i1+1)//2 - i0*(i0+1)//2
dij = (i1 - i0) * nao
if j_only:
buflen = max(8, int(max_memory*1e6/16/(nkptij*dii+dii)))
else:
buflen = max(8, int(max_memory*1e6/16/(nkptij*dij+dij)))
auxranges = balance_segs(aux_loc[1:]-aux_loc[:-1], buflen)
buflen = max([x[2] for x in auxranges])
buf = numpy.empty(nkptij*dij*buflen, dtype=dtype)
buf1 = numpy.empty(dij*buflen, dtype=dtype)
naux = aux_loc[-1]
for kpt_id, kptij in enumerate(kptij_lst):
key = '%s/%d' % (dataname, kpt_id)
if aosym_s2[kpt_id]:
shape = (naux, dii)
else:
shape = (naux, dij)
if gamma_point(kptij):
fswap.create_dataset(key, shape, 'f8')
else:
fswap.create_dataset(key, shape, 'c16')
naux0 = 0
for istep, auxrange in enumerate(auxranges):
log.alldebug2("aux_e1 job_id %d step %d", job_id, istep)
sh0, sh1, nrow = auxrange
sub_slice = (ish0, ish1, 0, cell.nbas, sh0, sh1)
if j_only:
mat = numpy.ndarray((nkptij,dii,nrow), dtype=dtype, buffer=buf)
else:
mat = numpy.ndarray((nkptij,dij,nrow), dtype=dtype, buffer=buf)
mat = int3c(sub_slice, mat)
for k, kptij in enumerate(kptij_lst):
h5dat = fswap['%s/%d'%(dataname,k)]
v = lib.transpose(mat[k], out=buf1)
if not j_only and aosym_s2[k]:
idy = idxb[i0*(i0+1)//2:i1*(i1+1)//2] - i0 * nao
out = numpy.ndarray((nrow,dii), dtype=v.dtype, buffer=mat[k])
v = numpy.take(v, idy, axis=1, out=out)
if gamma_point(kptij):
h5dat[naux0:naux0+nrow] = v.real
else:
h5dat[naux0:naux0+nrow] = v
naux0 += nrow
def aux_e1(cell,
auxcell,
erifile,
intor='int3c2e',
aosym='s2ij',
comp=None,
kptij_lst=None,
dataname='eri_mo',
shls_slice=None,
max_memory=2000,
verbose=0):
r'''3-center AO integrals (L|ij) with double lattice sum:
\sum_{lm} (L[0]|i[l]j[m]), where L is the auxiliary basis.
Three-index integral tensor (kptij_idx, naux, nao_pair) or four-index
integral tensor (kptij_idx, comp, naux, nao_pair) are stored on disk.
Args:
kptij_lst : (*,2,3) array
A list of (kpti, kptj)
'''
intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor),
comp)
if isinstance(erifile, h5py.Group):
feri = erifile
elif h5py.is_hdf5(erifile):
feri = h5py.File(erifile, 'a')
else:
feri = h5py.File(erifile, 'w')
if dataname in feri:
del (feri[dataname])
if dataname + '-kptij' in feri:
del (feri[dataname + '-kptij'])
if kptij_lst is None:
kptij_lst = numpy.zeros((1, 2, 3))
feri[dataname + '-kptij'] = kptij_lst
if shls_slice is None:
shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)
ao_loc = cell.ao_loc_nr()
aux_loc = auxcell.ao_loc_nr(auxcell.cart
or 'ssc' in intor)[:shls_slice[5] + 1]
ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
nkptij = len(kptij_lst)
nii = (ao_loc[shls_slice[1]] * (ao_loc[shls_slice[1]] + 1) // 2 -
ao_loc[shls_slice[0]] * (ao_loc[shls_slice[0]] + 1) // 2)
nij = ni * nj
kpti = kptij_lst[:, 0]
kptj = kptij_lst[:, 1]
aosym_ks2 = abs(kpti - kptj).sum(axis=1) < KPT_DIFF_TOL
j_only = numpy.all(aosym_ks2)
#aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
aosym_ks2 &= aosym[:2] == 's2'
for k, kptij in enumerate(kptij_lst):
key = '%s/%d' % (dataname, k)
if gamma_point(kptij):
dtype = 'f8'
else:
dtype = 'c16'
if aosym_ks2[k]:
nao_pair = nii
else:
nao_pair = nij
if comp == 1:
shape = (naux, nao_pair)
else:
shape = (comp, naux, nao_pair)
feri.create_dataset(key, shape, dtype)
if naux == 0:
feri.close()
return erifile
if j_only and aosym[:2] == 's2':
assert (shls_slice[2] == 0)
nao_pair = nii
else:
nao_pair = nij
if gamma_point(kptij_lst):
dtype = numpy.double
else:
dtype = numpy.complex128
buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * ni * nj * comp)))
auxdims = aux_loc[shls_slice[4] + 1:shls_slice[5] +
1] - aux_loc[shls_slice[4]:shls_slice[5]]
auxranges = balance_segs(auxdims, buflen)
buflen = max([x[2] for x in auxranges])
buf = numpy.empty(nkptij * comp * ni * nj * buflen, dtype=dtype)
buf1 = numpy.empty(ni * nj * buflen, dtype=dtype)
int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)
naux0 = 0
for istep, auxrange in enumerate(auxranges):
sh0, sh1, nrow = auxrange
sub_slice = (shls_slice[0], shls_slice[1], shls_slice[2],
shls_slice[3], shls_slice[4] + sh0, shls_slice[4] + sh1)
mat = numpy.ndarray((nkptij, comp, nao_pair, nrow),
dtype=dtype,
buffer=buf)
mat = int3c(sub_slice, mat)
for k, kptij in enumerate(kptij_lst):
h5dat = feri['%s/%d' % (dataname, k)]
for icomp, v in enumerate(mat[k]):
v = lib.transpose(v, out=buf1)
if gamma_point(kptij):
v = v.real
if aosym_ks2[k] and v.shape[1] == ni**2:
v = lib.pack_tril(v.reshape(-1, ni, ni))
if comp == 1:
h5dat[naux0:naux0 + nrow] = v
else:
h5dat[icomp, naux0:naux0 + nrow] = v
naux0 += nrow
if not isinstance(erifile, h5py.Group):
feri.close()
return erifile
def _aux_e2(cell,
auxcell,
erifile,
intor='int3c2e',
aosym='s2ij',
comp=None,
kptij_lst=None,
dataname='eri_mo',
shls_slice=None,
max_memory=2000,
verbose=0):
r'''3-center AO integrals (ij|L) with double lattice sum:
\sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
Three-index integral tensor (kptij_idx, nao_pair, naux) or four-index
integral tensor (kptij_idx, comp, nao_pair, naux) are stored on disk.
**This function should be only used by df and mdf initialization function
_make_j3c**
Args:
kptij_lst : (*,2,3) array
A list of (kpti, kptj)
'''
intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor),
comp)
if isinstance(erifile, h5py.Group):
feri = erifile
elif h5py.is_hdf5(erifile):
feri = h5py.File(erifile, 'a')
else:
feri = h5py.File(erifile, 'w')
if dataname in feri:
del (feri[dataname])
if dataname + '-kptij' in feri:
del (feri[dataname + '-kptij'])
if kptij_lst is None:
kptij_lst = numpy.zeros((1, 2, 3))
feri[dataname + '-kptij'] = kptij_lst
if shls_slice is None:
shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)
ao_loc = cell.ao_loc_nr()
aux_loc = auxcell.ao_loc_nr(auxcell.cart
or 'ssc' in intor)[:shls_slice[5] + 1]
ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
nkptij = len(kptij_lst)
nii = (ao_loc[shls_slice[1]] * (ao_loc[shls_slice[1]] + 1) // 2 -
ao_loc[shls_slice[0]] * (ao_loc[shls_slice[0]] + 1) // 2)
nij = ni * nj
kpti = kptij_lst[:, 0]
kptj = kptij_lst[:, 1]
aosym_ks2 = abs(kpti - kptj).sum(axis=1) < KPT_DIFF_TOL
j_only = numpy.all(aosym_ks2)
#aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
aosym_ks2 &= aosym[:2] == 's2'
if j_only and aosym[:2] == 's2':
assert (shls_slice[2] == 0)
nao_pair = nii
else:
nao_pair = nij
if gamma_point(kptij_lst):
dtype = numpy.double
else:
dtype = numpy.complex128
buflen = max(8, int(max_memory * .47e6 / 16 / (nkptij * ni * nj * comp)))
auxdims = aux_loc[shls_slice[4] + 1:shls_slice[5] +
1] - aux_loc[shls_slice[4]:shls_slice[5]]
auxranges = balance_segs(auxdims, buflen)
buflen = max([x[2] for x in auxranges])
buf = numpy.empty(nkptij * comp * ni * nj * buflen, dtype=dtype)
buf1 = numpy.empty_like(buf)
int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)
kptis = kptij_lst[:, 0]
kptjs = kptij_lst[:, 1]
kpt_ji = kptjs - kptis
uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
# sorted_ij_idx: Sort and group the kptij_lst according to the ordering in
# df._make_j3c to reduce the data fragment in the hdf5 file. When datasets
# are written to hdf5, they are saved sequentially. If the integral data are
# saved as the order of kptij_lst, removing the datasets in df._make_j3c will
# lead to holes that can not be reused.
sorted_ij_idx = numpy.hstack(
[numpy.where(uniq_inverse == k)[0] for k, kpt in enumerate(uniq_kpts)])
tril_idx = numpy.tril_indices(ni)
tril_idx = tril_idx[0] * ni + tril_idx[1]
def save(istep, mat):
for k in sorted_ij_idx:
v = mat[k]
if gamma_point(kptij_lst[k]):
v = v.real
if aosym_ks2[k] and nao_pair == ni**2:
v = v[:, tril_idx]
feri['%s/%d/%d' % (dataname, k, istep)] = v
with lib.call_in_background(save) as bsave:
for istep, auxrange in enumerate(auxranges):
sh0, sh1, nrow = auxrange
sub_slice = (shls_slice[0], shls_slice[1], shls_slice[2],
shls_slice[3], shls_slice[4] + sh0,
shls_slice[4] + sh1)
mat = numpy.ndarray((nkptij, comp, nao_pair, nrow),
dtype=dtype,
buffer=buf)
bsave(istep, int3c(sub_slice, mat))
buf, buf1 = buf1, buf
if not isinstance(erifile, h5py.Group):
feri.close()
return erifile
def aux_e2(cell,
auxcell,
erifile,
intor='cint3c2e_sph',
aosym='s1',
comp=1,
kptij_lst=None,
dataname='eri_mo',
max_memory=2000,
verbose=0):
'''3-center AO integrals (ij|L) with double lattice sum:
\sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
On diks, the integrals are stored as (kptij_idx, naux, nao_pair)
Args:
kptij_lst : (*,2,3) array
A list of (kpti, kptj)
'''
if comp > 1:
raise NotImplementedError('comp = %d' % comp)
if h5py.is_hdf5(erifile):
feri = h5py.File(erifile)
if dataname in feri:
del (feri[dataname])
if dataname + '-kptij' in feri:
del (feri[dataname + '-kptij'])
else:
feri = h5py.File(erifile, 'w')
if kptij_lst is None:
kptij_lst = numpy.zeros((1, 2, 3))
feri[dataname + '-kptij'] = kptij_lst
nkptij = len(kptij_lst)
# sum over largest number of images in either cell or auxcell
nimgs = numpy.max((cell.nimgs, auxcell.nimgs), axis=0)
Ls = cell.get_lattice_Ls(nimgs)
logger.debug1(cell, "pbc.df.outcore.Images %s", nimgs)
logger.debug3(cell, "Ls = %s", Ls)
nao = cell.nao_nr()
#naux = auxcell.nao_nr('ssc' in intor)
naux = auxcell.nao_nr()
aosym_s2 = numpy.zeros(nkptij, dtype=bool)
for k, kptij in enumerate(kptij_lst):
key = '%s/%d' % (dataname, k)
if abs(kptij).sum() < 1e-9: # gamma_point:
dtype = 'f8'
else:
dtype = 'c16'
aosym_s2[k] = abs(kptij[0] - kptij[1]).sum() < 1e-9
if aosym_s2[k]:
nao_pair = nao * (nao + 1) // 2
else:
nao_pair = nao * nao
if comp == 1:
shape = (naux, nao_pair)
else:
shape = (comp, naux, nao_pair)
chunks = (min(256, naux), min(256, nao_pair)) # 512 KB
feri.create_dataset(key, shape, dtype, chunks=chunks)
if naux == 0:
feri.close()
return erifile
aux_loc = auxcell.ao_loc_nr('ssc' in intor)
buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * nao**2 * comp)))
auxranges = balance_segs(aux_loc[1:] - aux_loc[:-1], buflen)
buflen = max([x[2] for x in auxranges])
buf = [
numpy.zeros(nao * nao * buflen * comp, dtype=numpy.complex128)
for k in range(nkptij)
]
ints = incore._wrap_int3c(cell, auxcell, intor, comp, Ls, buf)
atm, bas, env = ints._envs[:3]
xyz = numpy.asarray(cell.atom_coords(), order='C')
ptr_coordL = atm[:cell.natm, PTR_COORD]
ptr_coordL = numpy.vstack(
(ptr_coordL, ptr_coordL + 1, ptr_coordL + 2)).T.copy('C')
kpti = kptij_lst[:, 0]
kptj = kptij_lst[:, 1]
if numpy.all(aosym_s2):
def ccsum_or_reorder(Lpq):
tmp = numpy.asarray(Lpq.transpose(0, 2, 1).conj(), order='C')
tmp += Lpq
return tmp
else:
def ccsum_or_reorder(Lpq):
return numpy.asarray(Lpq, order='C')
naux0 = 0
for istep, auxrange in enumerate(auxranges):
sh0, sh1, nrow = auxrange
c_shls_slice = (ctypes.c_int * 6)(0, cell.nbas, cell.nbas,
cell.nbas * 2, cell.nbas * 2 + sh0,
cell.nbas * 2 + sh1)
if numpy.all(aosym_s2):
for l, L1 in enumerate(Ls):
env[ptr_coordL] = xyz + L1
e = numpy.dot(
Ls[:l + 1] - L1,
kptj.T) # Lattice sum over half of the images {1..l}
exp_Lk = numpy.exp(1j * numpy.asarray(e, order='C'))
exp_Lk[l] = .5
ints(exp_Lk, c_shls_slice)
else:
for l, L1 in enumerate(Ls):
env[ptr_coordL] = xyz + L1
e = numpy.dot(Ls, kptj.T) - numpy.dot(L1, kpti.T)
exp_Lk = numpy.exp(1j * numpy.asarray(e, order='C'))
ints(exp_Lk, c_shls_slice)
for k, kptij in enumerate(kptij_lst):
h5dat = feri['%s/%d' % (dataname, k)]
# transpose 3201 as (comp,L,i,j)
mat = numpy.ndarray((nao, nao, nrow, comp),
order='F',
dtype=numpy.complex128,
buffer=buf[k])
for icomp, vi in enumerate(mat.transpose(3, 2, 0, 1)):
v = ccsum_or_reorder(vi)
if abs(kptij).sum() < 1e-9: # gamma_point:
v = v.real
if aosym_s2[k]:
v = lib.pack_tril(v)
else:
v = v.reshape(nrow, -1)
if comp == 1:
h5dat[naux0:naux0 + nrow] = v
else:
h5dat[icomp, naux0:naux0 + nrow] = v
mat[:] = 0
naux0 += nrow
feri.close()
return erifile
def aux_e1(cell, auxcell, erifile, intor='int3c2e', aosym='s2ij', comp=None,
kptij_lst=None, dataname='eri_mo', shls_slice=None, max_memory=2000,
verbose=0):
r'''3-center AO integrals (L|ij) with double lattice sum:
\sum_{lm} (L[0]|i[l]j[m]), where L is the auxiliary basis.
Three-index integral tensor (kptij_idx, naux, nao_pair) or four-index
integral tensor (kptij_idx, comp, naux, nao_pair) are stored on disk.
Args:
kptij_lst : (*,2,3) array
A list of (kpti, kptj)
'''
intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor), comp)
if isinstance(erifile, h5py.Group):
feri = erifile
elif h5py.is_hdf5(erifile):
feri = h5py.File(erifile)
else:
feri = h5py.File(erifile, 'w')
if dataname in feri:
del(feri[dataname])
if dataname+'-kptij' in feri:
del(feri[dataname+'-kptij'])
if kptij_lst is None:
kptij_lst = numpy.zeros((1,2,3))
feri[dataname+'-kptij'] = kptij_lst
if shls_slice is None:
shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)
ao_loc = cell.ao_loc_nr()
aux_loc = auxcell.ao_loc_nr(auxcell.cart or 'ssc' in intor)[:shls_slice[5]+1]
ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
nkptij = len(kptij_lst)
nii = (ao_loc[shls_slice[1]]*(ao_loc[shls_slice[1]]+1)//2 -
ao_loc[shls_slice[0]]*(ao_loc[shls_slice[0]]+1)//2)
nij = ni * nj
kpti = kptij_lst[:,0]
kptj = kptij_lst[:,1]
aosym_ks2 = abs(kpti-kptj).sum(axis=1) < KPT_DIFF_TOL
j_only = numpy.all(aosym_ks2)
#aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
aosym_ks2 &= aosym[:2] == 's2'
for k, kptij in enumerate(kptij_lst):
key = '%s/%d' % (dataname, k)
if gamma_point(kptij):
dtype = 'f8'
else:
dtype = 'c16'
if aosym_ks2[k]:
nao_pair = nii
else:
nao_pair = nij
if comp == 1:
shape = (naux,nao_pair)
else:
shape = (comp,naux,nao_pair)
feri.create_dataset(key, shape, dtype)
if naux == 0:
feri.close()
return erifile
if j_only and aosym[:2] == 's2':
assert(shls_slice[2] == 0)
nao_pair = nii
else:
nao_pair = nij
if gamma_point(kptij_lst):
dtype = numpy.double
else:
dtype = numpy.complex128
buflen = max(8, int(max_memory*1e6/16/(nkptij*ni*nj*comp)))
auxdims = aux_loc[shls_slice[4]+1:shls_slice[5]+1] - aux_loc[shls_slice[4]:shls_slice[5]]
auxranges = balance_segs(auxdims, buflen)
buflen = max([x[2] for x in auxranges])
buf = numpy.empty(nkptij*comp*ni*nj*buflen, dtype=dtype)
buf1 = numpy.empty(ni*nj*buflen, dtype=dtype)
int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)
naux0 = 0
for istep, auxrange in enumerate(auxranges):
sh0, sh1, nrow = auxrange
sub_slice = (shls_slice[0], shls_slice[1],
shls_slice[2], shls_slice[3],
shls_slice[4]+sh0, shls_slice[4]+sh1)
mat = numpy.ndarray((nkptij,comp,nao_pair,nrow), dtype=dtype, buffer=buf)
mat = int3c(sub_slice, mat)
for k, kptij in enumerate(kptij_lst):
h5dat = feri['%s/%d'%(dataname,k)]
for icomp, v in enumerate(mat[k]):
v = lib.transpose(v, out=buf1)
if gamma_point(kptij):
v = v.real
if aosym_ks2[k] and v.shape[1] == ni**2:
v = lib.pack_tril(v.reshape(-1,ni,ni))
if comp == 1:
h5dat[naux0:naux0+nrow] = v
else:
h5dat[icomp,naux0:naux0+nrow] = v
naux0 += nrow
if not isinstance(erifile, h5py.Group):
feri.close()
return erifile
def _aux_e2(cell, auxcell, erifile, intor='int3c2e', aosym='s2ij', comp=None,
kptij_lst=None, dataname='eri_mo', shls_slice=None, max_memory=2000,
verbose=0):
r'''3-center AO integrals (ij|L) with double lattice sum:
\sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
Three-index integral tensor (kptij_idx, nao_pair, naux) or four-index
integral tensor (kptij_idx, comp, nao_pair, naux) are stored on disk.
**This function should be only used by df and mdf initialization function
_make_j3c**
Args:
kptij_lst : (*,2,3) array
A list of (kpti, kptj)
'''
intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor), comp)
if isinstance(erifile, h5py.Group):
feri = erifile
elif h5py.is_hdf5(erifile):
feri = h5py.File(erifile)
else:
feri = h5py.File(erifile, 'w')
if dataname in feri:
del(feri[dataname])
if dataname+'-kptij' in feri:
del(feri[dataname+'-kptij'])
if kptij_lst is None:
kptij_lst = numpy.zeros((1,2,3))
feri[dataname+'-kptij'] = kptij_lst
if shls_slice is None:
shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)
ao_loc = cell.ao_loc_nr()
aux_loc = auxcell.ao_loc_nr(auxcell.cart or 'ssc' in intor)[:shls_slice[5]+1]
ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
nkptij = len(kptij_lst)
nii = (ao_loc[shls_slice[1]]*(ao_loc[shls_slice[1]]+1)//2 -
ao_loc[shls_slice[0]]*(ao_loc[shls_slice[0]]+1)//2)
nij = ni * nj
kpti = kptij_lst[:,0]
kptj = kptij_lst[:,1]
aosym_ks2 = abs(kpti-kptj).sum(axis=1) < KPT_DIFF_TOL
j_only = numpy.all(aosym_ks2)
#aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
aosym_ks2 &= aosym[:2] == 's2'
if j_only and aosym[:2] == 's2':
assert(shls_slice[2] == 0)
nao_pair = nii
else:
nao_pair = nij
if gamma_point(kptij_lst):
dtype = numpy.double
else:
dtype = numpy.complex128
buflen = max(8, int(max_memory*.47e6/16/(nkptij*ni*nj*comp)))
auxdims = aux_loc[shls_slice[4]+1:shls_slice[5]+1] - aux_loc[shls_slice[4]:shls_slice[5]]
auxranges = balance_segs(auxdims, buflen)
buflen = max([x[2] for x in auxranges])
buf = numpy.empty(nkptij*comp*ni*nj*buflen, dtype=dtype)
buf1 = numpy.empty_like(buf)
int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)
kptis = kptij_lst[:,0]
kptjs = kptij_lst[:,1]
kpt_ji = kptjs - kptis
uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
# sorted_ij_idx: Sort and group the kptij_lst according to the ordering in
# df._make_j3c to reduce the data fragment in the hdf5 file. When datasets
# are written to hdf5, they are saved sequentially. If the integral data are
# saved as the order of kptij_lst, removing the datasets in df._make_j3c will
# lead to holes that can not be reused.
sorted_ij_idx = numpy.hstack([numpy.where(uniq_inverse == k)[0]
for k, kpt in enumerate(uniq_kpts)])
tril_idx = numpy.tril_indices(ni)
tril_idx = tril_idx[0] * ni + tril_idx[1]
def save(istep, mat):
for k in sorted_ij_idx:
v = mat[k]
if gamma_point(kptij_lst[k]):
v = v.real
if aosym_ks2[k] and nao_pair == ni**2:
v = v[:,tril_idx]
feri['%s/%d/%d' % (dataname,k,istep)] = v
with lib.call_in_background(save) as bsave:
for istep, auxrange in enumerate(auxranges):
sh0, sh1, nrow = auxrange
sub_slice = (shls_slice[0], shls_slice[1],
shls_slice[2], shls_slice[3],
shls_slice[4]+sh0, shls_slice[4]+sh1)
mat = numpy.ndarray((nkptij,comp,nao_pair,nrow), dtype=dtype, buffer=buf)
bsave(istep, int3c(sub_slice, mat))
buf, buf1 = buf1, buf
if not isinstance(erifile, h5py.Group):
feri.close()
return erifile
def aux_e2(cell, auxcell, erifile, intor='cint3c2e_sph', aosym='s1', comp=1,
kptij_lst=None, dataname='eri_mo', max_memory=2000, verbose=0):
'''3-center AO integrals (ij|L) with double lattice sum:
\sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
On diks, the integrals are stored as (kptij_idx, naux, nao_pair)
Args:
kptij_lst : (*,2,3) array
A list of (kpti, kptj)
'''
if comp > 1:
raise NotImplementedError('comp = %d' % comp)
if h5py.is_hdf5(erifile):
feri = h5py.File(erifile)
if dataname in feri:
del(feri[dataname])
if dataname+'-kptij' in feri:
del(feri[dataname+'-kptij'])
else:
feri = h5py.File(erifile, 'w')
if kptij_lst is None:
kptij_lst = numpy.zeros((1,2,3))
feri[dataname+'-kptij'] = kptij_lst
nkptij = len(kptij_lst)
# sum over largest number of images in either cell or auxcell
nimgs = numpy.max((cell.nimgs, auxcell.nimgs), axis=0)
Ls = cell.get_lattice_Ls(nimgs)
logger.debug1(cell, "pbc.df.outcore.Images %s", nimgs)
logger.debug3(cell, "Ls = %s", Ls)
nao = cell.nao_nr()
#naux = auxcell.nao_nr('ssc' in intor)
naux = auxcell.nao_nr()
aosym_s2 = numpy.zeros(nkptij, dtype=bool)
for k, kptij in enumerate(kptij_lst):
key = '%s/%d' % (dataname, k)
if abs(kptij).sum() < 1e-9: # gamma_point:
dtype = 'f8'
else:
dtype = 'c16'
aosym_s2[k] = abs(kptij[0]-kptij[1]).sum() < 1e-9
if aosym_s2[k]:
nao_pair = nao * (nao+1) // 2
else:
nao_pair = nao * nao
if comp == 1:
shape = (naux,nao_pair)
else:
shape = (comp,naux,nao_pair)
chunks = (min(256,naux), min(256,nao_pair)) # 512 KB
feri.create_dataset(key, shape, dtype, chunks=chunks)
if naux == 0:
feri.close()
return erifile
aux_loc = auxcell.ao_loc_nr('ssc' in intor)
buflen = max(8, int(max_memory*1e6/16/(nkptij*nao**2*comp)))
auxranges = balance_segs(aux_loc[1:]-aux_loc[:-1], buflen)
buflen = max([x[2] for x in auxranges])
buf = [numpy.zeros(nao*nao*buflen*comp, dtype=numpy.complex128)
for k in range(nkptij)]
ints = incore._wrap_int3c(cell, auxcell, intor, comp, Ls, buf)
atm, bas, env = ints._envs[:3]
xyz = numpy.asarray(cell.atom_coords(), order='C')
ptr_coordL = atm[:cell.natm,PTR_COORD]
ptr_coordL = numpy.vstack((ptr_coordL,ptr_coordL+1,ptr_coordL+2)).T.copy('C')
kpti = kptij_lst[:,0]
kptj = kptij_lst[:,1]
if numpy.all(aosym_s2):
def ccsum_or_reorder(Lpq):
tmp = numpy.asarray(Lpq.transpose(0,2,1).conj(), order='C')
tmp += Lpq
return tmp
else:
def ccsum_or_reorder(Lpq):
return numpy.asarray(Lpq, order='C')
naux0 = 0
for istep, auxrange in enumerate(auxranges):
sh0, sh1, nrow = auxrange
c_shls_slice = (ctypes.c_int*6)(0, cell.nbas, cell.nbas, cell.nbas*2,
cell.nbas*2+sh0, cell.nbas*2+sh1)
if numpy.all(aosym_s2):
for l, L1 in enumerate(Ls):
env[ptr_coordL] = xyz + L1
e = numpy.dot(Ls[:l+1]-L1, kptj.T) # Lattice sum over half of the images {1..l}
exp_Lk = numpy.exp(1j * numpy.asarray(e, order='C'))
exp_Lk[l] = .5
ints(exp_Lk, c_shls_slice)
else:
for l, L1 in enumerate(Ls):
env[ptr_coordL] = xyz + L1
e = numpy.dot(Ls, kptj.T) - numpy.dot(L1, kpti.T)
exp_Lk = numpy.exp(1j * numpy.asarray(e, order='C'))
ints(exp_Lk, c_shls_slice)
for k, kptij in enumerate(kptij_lst):
h5dat = feri['%s/%d'%(dataname,k)]
# transpose 3201 as (comp,L,i,j)
mat = numpy.ndarray((nao,nao,nrow,comp), order='F',
dtype=numpy.complex128, buffer=buf[k])
for icomp, vi in enumerate(mat.transpose(3,2,0,1)):
v = ccsum_or_reorder(vi)
if abs(kptij).sum() < 1e-9: # gamma_point:
v = v.real
if aosym_s2[k]:
v = lib.pack_tril(v)
else:
v = v.reshape(nrow,-1)
if comp == 1:
h5dat[naux0:naux0+nrow] = v
else:
h5dat[icomp,naux0:naux0+nrow] = v
mat[:] = 0
naux0 += nrow
feri.close()
return erifile