def test_ft_aopair_bvk(self):
from pyscf.pbc.tools import k2gamma
n = 2
cell = pgto.Cell()
cell.a = numpy.eye(3) * 4
cell.mesh = numpy.array([n, n, n])
cell.atom = '''C 1.3 .2 .3
C .1 .1 1.1
'''
cell.basis = 'ccpvdz'
cell.unit = 'B'
cell.build()
kpts = cell.make_kpts([2, 2, 2])
Gv, Gvbase, kws = cell.get_Gv_weights()
b = cell.reciprocal_vectors()
gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
bvk_kmesh = k2gamma.kpts_to_kmesh(cell, kpts)
ref = ft_ao.ft_aopair_kpts(cell,
Gv,
b=b,
gxyz=gxyz,
Gvbase=Gvbase,
kptjs=kpts)
aopair = ft_ao.ft_aopair_kpts(cell,
Gv,
b=b,
gxyz=gxyz,
Gvbase=Gvbase,
kptjs=kpts,
bvk_kmesh=bvk_kmesh)
self.assertAlmostEqual(abs(ref - aopair).max(), 0, 8)
self.assertAlmostEqual(lib.fp(aopair),
(-5.735639500461687 - 12.425151458809875j), 8)
def ft_loop(self,
mesh=None,
q=numpy.zeros(3),
kpts=None,
shls_slice=None,
max_memory=4000,
aosym='s1',
intor='GTO_ft_ovlp',
comp=1,
bvk_kmesh=None):
'''
Fourier transform iterator for all kpti which satisfy
2pi*N = (kpts - kpti - q)*a, N = -1, 0, 1
'''
cell = self.cell
if mesh is None:
mesh = self.mesh
if kpts is None:
assert (is_zero(q))
kpts = self.kpts
kpts = numpy.asarray(kpts)
nkpts = len(kpts)
ao_loc = cell.ao_loc_nr()
b = cell.reciprocal_vectors()
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
ngrids = gxyz.shape[0]
if shls_slice is None:
shls_slice = (0, cell.nbas, 0, cell.nbas)
if aosym == 's2':
assert (shls_slice[2] == 0)
i0 = ao_loc[shls_slice[0]]
i1 = ao_loc[shls_slice[1]]
nij = i1 * (i1 + 1) // 2 - i0 * (i0 + 1) // 2
else:
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 = max(16, int(max_memory * .9e6 / (nij * nkpts * 16 * comp)))
blksize = min(blksize, ngrids, 16384)
buf = numpy.empty(nkpts * nij * blksize * comp, dtype=numpy.complex128)
for p0, p1 in self.prange(0, ngrids, blksize):
dat = ft_ao.ft_aopair_kpts(cell,
Gv[p0:p1],
shls_slice,
aosym,
b,
gxyz[p0:p1],
Gvbase,
q,
kpts,
intor,
comp,
bvk_kmesh=bvk_kmesh,
out=buf)
yield dat, p0, p1
def test_ft_aoao_with_kpts1(self):
numpy.random.seed(1)
kpti, kptj = kpts = numpy.random.random((2, 3))
Gv = cell.get_Gv([11] * 3)
q = numpy.random.random(3)
dat = ft_ao.ft_aopair_kpts(cell, Gv, q=q, kptjs=kpts)
self.assertAlmostEqual(finger(dat[0]),
(2.3753953914129382 - 2.5365192689115088j), 9)
self.assertAlmostEqual(finger(dat[1]),
(2.4951510097641840 - 3.1990956672116355j), 9)
dat = ft_ao.ft_aopair(cell, Gv)
self.assertAlmostEqual(finger(dat),
(1.2534723618134684 + 1.830086071817564j), 9)
def make_kpt(uniq_kptji_id, cholesky_j2c): # kpt = kptj - kpti
kpt = uniq_kpts[uniq_kptji_id]
log.debug1('kpt = %s', kpt)
adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0]
adapted_kptjs = kptjs[adapted_ji_idx]
nkptj = len(adapted_kptjs)
log.debug1('adapted_ji_idx = %s', adapted_ji_idx)
j2c, j2c_negative, j2ctag = cholesky_j2c
Gaux = ft_ao.ft_ao(fused_cell, Gv, None, b, gxyz, Gvbase, kpt).T
Gaux = fuse(Gaux)
Gaux *= mydf.weighted_coulG(kpt, False, mesh)
kLR = Gaux.T.real.copy('C')
kLI = Gaux.T.imag.copy('C')
if is_zero(kpt): # kpti == kptj
aosym = 's2'
nao_pair = nao*(nao+1)//2
if 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]
else:
aosym = 's1'
nao_pair = nao**2
mem_now = lib.current_memory()[0]
log.debug2('memory = %s', mem_now)
max_memory = max(2000, mydf.max_memory-mem_now)
# nkptj for 3c-coulomb arrays plus 1 Lpq array
buflen = min(max(int(max_memory*.38e6/16/naux/(nkptj+1)), 1), nao_pair)
shranges = _guess_shell_ranges(cell, buflen, aosym)
buflen = max([x[2] for x in shranges])
# +1 for a pqkbuf
if aosym == 's2':
Gblksize = max(16, int(max_memory*.1e6/16/buflen/(nkptj+1)))
else:
Gblksize = max(16, int(max_memory*.2e6/16/buflen/(nkptj+1)))
Gblksize = min(Gblksize, ngrids, 16384)
def load(aux_slice):
col0, col1 = aux_slice
j3cR = []
j3cI = []
for k, idx in enumerate(adapted_ji_idx):
v = [fswap['j3c-junk/%d/%d'%(idx,i)][0,col0:col1].T for i in range(nsegs)]
v = fuse(numpy.vstack(v))
if is_zero(kpt) and cell.dimension == 3:
for i in numpy.where(vbar != 0)[0]:
v[i] -= vbar[i] * ovlp[k][col0:col1]
j3cR.append(numpy.asarray(v.real, order='C'))
if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
j3cI.append(None)
else:
j3cI.append(numpy.asarray(v.imag, order='C'))
v = None
return j3cR, j3cI
pqkRbuf = numpy.empty(buflen*Gblksize)
pqkIbuf = numpy.empty(buflen*Gblksize)
# buf for ft_aopair
buf = numpy.empty((nkptj,buflen*Gblksize), dtype=numpy.complex128)
cols = [sh_range[2] for sh_range in shranges]
locs = numpy.append(0, numpy.cumsum(cols))
tasks = zip(locs[:-1], locs[1:])
for istep, (j3cR, j3cI) in enumerate(lib.map_with_prefetch(load, tasks)):
bstart, bend, ncol = shranges[istep]
log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d',
istep+1, len(shranges), bstart, bend, ncol)
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)
j3cR = j3cI = None
for ji in adapted_ji_idx:
del(fswap['j3c-junk/%d'%ji])
def pw_loop(self,
mesh=None,
kpti_kptj=None,
q=None,
shls_slice=None,
max_memory=2000,
aosym='s1',
blksize=None,
intor='GTO_ft_ovlp',
comp=1,
bvk_kmesh=None):
'''
Fourier transform iterator for AO pair
'''
cell = self.cell
if mesh is None:
mesh = self.mesh
if kpti_kptj is None:
kpti = kptj = numpy.zeros(3)
else:
kpti, kptj = kpti_kptj
if q is None:
q = kptj - kpti
ao_loc = cell.ao_loc_nr()
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
b = cell.reciprocal_vectors()
gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
ngrids = gxyz.shape[0]
if shls_slice is None:
shls_slice = (0, cell.nbas, 0, cell.nbas)
if aosym == 's2':
assert (shls_slice[2] == 0)
i0 = ao_loc[shls_slice[0]]
i1 = ao_loc[shls_slice[1]]
nij = i1 * (i1 + 1) // 2 - i0 * (i0 + 1) // 2
else:
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
if blksize is None:
blksize = min(
max(64, int(max_memory * 1e6 * .75 / (nij * 16 * comp))),
16384)
sublk = int(blksize // 4)
else:
sublk = blksize
buf = numpy.empty(nij * blksize * comp, dtype=numpy.complex128)
pqkRbuf = numpy.empty(nij * sublk * comp)
pqkIbuf = numpy.empty(nij * sublk * comp)
for p0, p1 in self.prange(0, ngrids, blksize):
#aoao = ft_ao.ft_aopair(cell, Gv[p0:p1], shls_slice, aosym,
# b, Gvbase, gxyz[p0:p1], mesh, (kpti, kptj), q)
aoao = ft_ao.ft_aopair_kpts(cell,
Gv[p0:p1],
shls_slice,
aosym,
b,
gxyz[p0:p1],
Gvbase,
q,
kptj.reshape(1, 3),
intor,
comp,
bvk_kmesh=bvk_kmesh,
out=buf)[0]
aoao = aoao.reshape(p1 - p0, nij)
for i0, i1 in lib.prange(0, p1 - p0, sublk):
nG = i1 - i0
if comp == 1:
pqkR = numpy.ndarray((nij, nG), buffer=pqkRbuf)
pqkI = numpy.ndarray((nij, nG), buffer=pqkIbuf)
pqkR[:] = aoao[i0:i1].real.T
pqkI[:] = aoao[i0:i1].imag.T
else:
pqkR = numpy.ndarray((comp, nij, nG), buffer=pqkRbuf)
pqkI = numpy.ndarray((comp, nij, nG), buffer=pqkIbuf)
pqkR[:] = aoao[i0:i1].real.transpose(0, 2, 1)
pqkI[:] = aoao[i0:i1].imag.transpose(0, 2, 1)
yield (pqkR, pqkI, p0 + i0, p0 + i1)