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)
else:
v = lib.dot(j2c, v)
feri['j3c/%d/%d'%(ji,istep)] = v
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)
if (cell.dimension == 1 or cell.dimension == 2) and is_zero(kpt):
G0idx, SI_on_z = pbcgto.cell._SI_for_uniform_model_charge(cell, Gv[p0:p1])
if SI_on_z.size > 0:
for k, aoao in enumerate(dat):
aoao[G0idx] -= numpy.einsum('g,i->gi', SI_on_z, ovlp[k])
aux = fuse(ft_ao.ft_ao(fused_cell, Gv[p0:p1][G0idx]).T)
vG_mod = numpy.einsum('ig,g,g->i', aux.conj(),
wcoulG[p0:p1][G0idx], SI_on_z)
if gamma_point(adapted_kptjs[k]):
j3cR[k][:naux] -= vG_mod[:,None].real * ovlp[k]
else:
tmp = vG_mod[:,None] * ovlp[k]
j3cR[k][:naux] -= tmp.real
j3cI[k][:naux] -= tmp.imag
tmp = aux = vG_mod
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][naux:], 1)
lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k][naux:], 1)
if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k][naux:], 1)
lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k][naux:], 1)
for k, ji 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)
else:
v = lib.dot(j2c, v)
feri['j3c/%d/%d'%(ji,istep)] = v
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
if kpts_band is not None:
return get_j_for_bands(mydf, dm_kpts, hermi, kpts, kpts_band)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
vj_kpts = numpy.zeros((nset,nkpts,nao,nao), dtype=numpy.complex128)
kpt_allow = numpy.zeros(3)
mesh = mydf.mesh
coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
weight = 1./len(kpts)
dmsC = dms.conj()
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
vG = [0] * nset
#:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
for k, aoao in enumerate(aoaoks):
for i in range(nset):
rho = numpy.einsum('ij,Lij->L', dmsC[i,k],
aoao.reshape(-1,nao,nao)).conj()
vG[i] += rho * coulG[p0:p1]
for i in range(nset):
vG[i] *= weight
for k, aoao in enumerate(aoaoks):
for i in range(nset):
vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i],
aoao.reshape(-1,nao,nao))
aoao = aoaoks = p0 = p1 = None
if gamma_point(kpts):
vj_kpts = vj_kpts.real.copy()
return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)
def ft_ao(mol, Gv, shls_slice=None, b=None,
gxyz=None, Gvbase=None, kpt=numpy.zeros(3), verbose=None):
if gamma_point(kpt):
return mol_ft_ao(mol, Gv, shls_slice, b, gxyz, Gvbase, verbose)
else:
kG = Gv + kpt
return mol_ft_ao(mol, kG, shls_slice, None, None, None, verbose)
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
if kpts_band is not None:
return get_j_for_bands(mydf, dm_kpts, hermi, kpts, kpts_band)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
vj_kpts = numpy.zeros((nset,nkpts,nao,nao), dtype=numpy.complex128)
kpt_allow = numpy.zeros(3)
mesh = mydf.mesh
coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
weight = 1./len(kpts)
dmsC = dms.conj()
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
vG = [0] * nset
#:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
for k, aoao in enumerate(aoaoks):
for i in range(nset):
rho = numpy.einsum('ij,Lij->L', dmsC[i,k],
aoao.reshape(-1,nao,nao)).conj()
vG[i] += rho * coulG[p0:p1]
for i in range(nset):
vG[i] *= weight
for k, aoao in enumerate(aoaoks):
for i in range(nset):
vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i],
aoao.reshape(-1,nao,nao))
aoao = aoaoks = p0 = p1 = None
if gamma_point(kpts):
vj_kpts = vj_kpts.real.copy()
return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)
def get_j_kpts(mydf,
dm_kpts,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None):
if kpts_band is not None:
return get_j_for_bands(mydf, dm_kpts, hermi, kpts, kpts_band)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
n_dm, nkpts, nao = dms.shape[:3]
vj_kpts = numpy.zeros((n_dm, nkpts, nao, nao), dtype=numpy.complex128)
kpt_allow = numpy.zeros(3)
mesh = mydf.mesh
coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
weight = 1. / len(kpts)
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt_allow,
kpts,
max_memory=max_memory):
_update_vj_(vj_kpts, aoaoks, dms, coulG[p0:p1], weight)
aoaoks = None
if gamma_point(kpts):
vj_kpts = vj_kpts.real.copy()
return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)
def get_mo_pairs_G(mydf, mo_coeffs, kpts=numpy.zeros((2,3)), q=None,
compact=getattr(__config__, 'pbc_df_mo_pairs_compact', False)):
'''Calculate forward (G|ij) FFT of all MO pairs.
Args:
mo_coeff: length-2 list of (nao,nmo) ndarrays
The two sets of MO coefficients to use in calculating the
product |ij).
Returns:
mo_pairs_G : (ngrids, nmoi*nmoj) ndarray
The FFT of the real-space MO pairs.
'''
if kpts is None: kpts = numpy.zeros((2,3))
cell = mydf.cell
kpts = numpy.asarray(kpts)
coords = cell.gen_uniform_grids(mydf.mesh)
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
ngrids = len(coords)
def trans(aoi, aoj, fac=1):
if id(aoi) == id(aoj) and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
moi = moj = numpy.asarray(lib.dot(mo_coeffs[0].T,aoi.T), order='C')
else:
moi = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T), order='C')
moj = numpy.asarray(lib.dot(mo_coeffs[1].T, aoj.T), order='C')
mo_pairs_G = numpy.empty((nmoi,nmoj,ngrids), dtype=numpy.complex128)
for i in range(nmoi):
mo_pairs_G[i] = tools.fft(fac * moi[i].conj() * moj, mydf.mesh)
mo_pairs_G = mo_pairs_G.reshape(-1,ngrids).T
return mo_pairs_G
if gamma_point(kpts): # gamma point, real
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
if compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
mo = numpy.asarray(lib.dot(mo_coeffs[0].T, ao.T), order='C')
npair = nmoi*(nmoi+1)//2
mo_pairs_G = numpy.empty((npair,ngrids), dtype=numpy.complex128)
ij = 0
for i in range(nmoi):
mo_pairs_G[ij:ij+i+1] = tools.fft(mo[i].conj() * mo[:i+1], mydf.mesh)
ij += i + 1
mo_pairs_G = mo_pairs_G.T
else:
mo_pairs_G = trans(ao, ao)
elif is_zero(kpts[0]-kpts[1]):
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
mo_pairs_G = trans(ao, ao)
else:
if q is None:
q = kpts[1] - kpts[0]
aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts)
fac = numpy.exp(-1j * numpy.dot(coords, q))
mo_pairs_G = trans(aoi, aoj, fac)
return mo_pairs_G
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpts_band=None):
'''Get the Coulomb (J) AO matrix at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray
Density matrix at each k-point. If a list of k-point DMs, eg,
UHF alpha and beta DM, the alpha and beta DMs are contracted
separately.
kpts : (nkpts, 3) ndarray
Kwargs:
kpts_band : (3,) ndarray or (*,3) ndarray
A list of arbitrary "band" k-points at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
or list of vj if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
mesh = mydf.mesh
low_dim_ft_type = mydf.low_dim_ft_type
ni = mydf._numint
make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
coulG = tools.get_coulG(cell, mesh=mesh, low_dim_ft_type=low_dim_ft_type)
ngrids = len(coulG)
vR = rhoR = np.zeros((nset, ngrids))
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
rhoR[i, p0:p1] += make_rho(i, ao_ks, mask, 'LDA')
ao = ao_ks = None
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh).real
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
weight = cell.vol / ngrids
vR *= weight
if gamma_point(kpts_band):
vj_kpts = np.zeros((nset, nband, nao, nao))
else:
vj_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
rho = None
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i, p0:p1],
mask, 'LDA')
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
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
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
def get_j_for_bands(mydf,
dm_kpts,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None):
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
dmsR = dms.real.reshape(nset, nkpts, nao**2)
dmsI = dms.imag.reshape(nset, nkpts, nao**2)
kpt_allow = numpy.zeros(3)
mesh = mydf.mesh
coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
ngrids = len(coulG)
vG = numpy.zeros((nset, ngrids), dtype=numpy.complex128)
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
dmsC = dms.conj()
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt_allow,
kpts,
max_memory=max_memory):
#:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
for k, aoao in enumerate(aoaoks):
for i in range(nset):
rho = numpy.einsum('ij,Lij->L', dmsC[i, k],
aoao.reshape(-1, nao, nao)).conj()
vG[i, p0:p1] += rho * coulG[p0:p1]
aoao = aoaoks = p0 = p1 = None
weight = 1. / len(kpts)
vG *= weight
t1 = log.timer_debug1('get_j pass 1 to compute J(G)', *t1)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
vj_kpts = numpy.zeros((nset, nband, nao, nao), dtype=numpy.complex128)
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt_allow,
kpts_band,
max_memory=max_memory):
for k, aoao in enumerate(aoaoks):
for i in range(nset):
vj_kpts[i, k] += numpy.einsum('L,Lij->ij', vG[i, p0:p1],
aoao.reshape(-1, nao, nao))
aoao = aoaoks = p0 = p1 = None
if gamma_point(kpts_band):
vj_kpts = vj_kpts.real.copy()
t1 = log.timer_debug1('get_j pass 2', *t1)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def get_eri(mydf, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_get_eri_compact', True)):
cell = mydf.cell
nao = cell.nao_nr()
kptijkl = _format_kpts(kpts)
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'fft_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros((nao,nao,nao,nao))
kpti, kptj, kptk, kptl = kptijkl
q = kptj - kpti
coulG = tools.get_coulG(cell, q, mesh=mydf.mesh)
coords = cell.gen_uniform_grids(mydf.mesh)
max_memory = mydf.max_memory - lib.current_memory()[0]
####################
# gamma point, the integral is real and with s4 symmetry
if gamma_point(kptijkl):
#:ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=compact)
#:ao_pairs_G *= numpy.sqrt(coulG).reshape(-1,1)
#:eri = lib.dot(ao_pairs_G.T, ao_pairs_G, cell.vol/ngrids**2)
ao = mydf._numint.eval_ao(cell, coords, kpti)[0]
ao = numpy.asarray(ao.T, order='C')
eri = _contract_compact(mydf, (ao,ao), coulG, max_memory=max_memory)
if not compact:
eri = ao2mo.restore(1, eri, nao).reshape(nao**2,nao**2)
return eri
####################
# aosym = s1, complex integrals
else:
#:ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=False)
#:# ao_pairs_invG = rho_kl(-(G+k_ij)) = conj(rho_lk(G+k_ij)).swap(r,s)
#:#=get_ao_pairs_G(mydf, [kptl,kptk], q, compact=False).transpose(0,2,1).conj()
#:ao_pairs_invG = get_ao_pairs_G(mydf, -kptijkl[2:], q, compact=False).conj()
#:ao_pairs_G *= coulG.reshape(-1,1)
#:eri = lib.dot(ao_pairs_G.T, ao_pairs_invG, cell.vol/ngrids**2)
if is_zero(kpti-kptl) and is_zero(kptj-kptk):
if is_zero(kpti-kptj):
aoi = mydf._numint.eval_ao(cell, coords, kpti)[0]
aoi = aoj = numpy.asarray(aoi.T, order='C')
else:
aoi, aoj = mydf._numint.eval_ao(cell, coords, kptijkl[:2])
aoi = numpy.asarray(aoi.T, order='C')
aoj = numpy.asarray(aoj.T, order='C')
aos = (aoi, aoj, aoj, aoi)
else:
aos = mydf._numint.eval_ao(cell, coords, kptijkl)
aos = [numpy.asarray(x.T, order='C') for x in aos]
fac = numpy.exp(-1j * numpy.dot(coords, q))
max_memory = max_memory - aos[0].nbytes*4*1e-6
eri = _contract_plain(mydf, aos, coulG, fac, max_memory=max_memory)
return eri
def get_j_kpts(self, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None):
'''
C ~ compact basis, D ~ diffused basis
Compute J matrix with coulG_LR:
(CC|CC) (CC|CD) (CC|DC) (CD|CC) (CD|CD) (CD|DC) (DC|CC) (DC|CD) (DC|DC)
Compute J matrix with full coulG:
(CC|DD) (CD|DD) (DC|DD) (DD|CC) (DD|CD) (DD|DC) (DD|DD)
'''
if kpts_band is not None:
return self.get_j_for_bands(dm_kpts, hermi, kpts, kpts_band)
cell = self.cell
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
n_dm, nkpts, nao = dms.shape[:3]
n_diffused = cell._nbas_each_set[2]
nao_compact = cell.ao_loc[cell.nbas-n_diffused]
vj_kpts = np.zeros((n_dm,nkpts,nao,nao), dtype=np.complex128)
kpt_allow = np.zeros(3)
mesh = self.mesh
coulG = self.weighted_coulG(kpt_allow, False, mesh)
coulG_LR = self.weighted_coulG_LR(kpt_allow, False, mesh)
coulG_SR = coulG - coulG_LR
max_memory = (self.max_memory - lib.current_memory()[0]) * .8
weight = 1./len(kpts)
for aoaoks, p0, p1 in self.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
if nao_compact < nao:
aoaoks = [aoao.reshape(-1,nao,nao) for aoao in aoaoks]
aft_jk._update_vj_(vj_kpts, aoaoks, dms, coulG[p0:p1], weight)
for aoao in aoaoks:
aoao[:,nao_compact:,nao_compact:] = 0
aft_jk._update_vj_(vj_kpts, aoaoks, dms, coulG_SR[p0:p1], -weight)
else:
aft_jk._update_vj_(vj_kpts, aoaoks, dms, coulG_LR[p0:p1], weight)
aoao = aoaoks = p0 = p1 = None
# G=0 contribution, associated to 2e integrals in real-space
if cell.dimension >= 2:
ovlp = np.asarray(cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts))
if nao_compact < nao:
ovlp[:,nao_compact:,nao_compact:] = 0
kws = cell.get_Gv_weights(mesh)[2]
G0_weight = kws[0] if isinstance(kws, np.ndarray) else kws
vj_G0 = lib.einsum('kpq,nkqp,lrs->nlrs', ovlp, dm_kpts, ovlp)
vj_kpts -= np.pi/self.omega**2 * weight * G0_weight * vj_G0
if gamma_point(kpts):
vj_kpts = vj_kpts.real.copy()
return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)
def _init_df_eris(cc, eris):
from pyscf.pbc.df import df
from pyscf.ao2mo import _ao2mo
if cc._scf.with_df._cderi is None:
cc._scf.with_df.build()
cell = cc._scf.cell
if cell.dimension == 2:
# 2D ERIs are not positive definite. The 3-index tensors are stored in
# two part. One corresponds to the positive part and one corresponds
# to the negative part. The negative part is not considered in the
# DF-driven CCSD implementation.
raise NotImplementedError
nocc = cc.nocc
nmo = cc.nmo
nvir = nmo - nocc
nao = cell.nao_nr()
kpts = cc.kpts
nkpts = len(kpts)
naux = cc._scf.with_df.get_naoaux()
if gamma_point(kpts):
dtype = np.double
else:
dtype = np.complex128
dtype = np.result_type(dtype, *eris.mo_coeff)
eris.Lpv = Lpv = np.empty((nkpts, nkpts), dtype=object)
with h5py.File(cc._scf.with_df._cderi, 'r') as f:
kptij_lst = f['j3c-kptij'].value
tao = []
ao_loc = None
for ki, kpti in enumerate(kpts):
for kj, kptj in enumerate(kpts):
kpti_kptj = np.array((kpti, kptj))
Lpq = np.asarray(df._getitem(f, 'j3c', kpti_kptj, kptij_lst))
mo = np.hstack((eris.mo_coeff[ki], eris.mo_coeff[kj][:,
nocc:]))
mo = np.asarray(mo, dtype=dtype, order='F')
if dtype == np.double:
out = _ao2mo.nr_e2(Lpq,
mo, (0, nmo, nmo, nmo + nvir),
aosym='s2')
else:
#Note: Lpq.shape[0] != naux if linear dependency is found in auxbasis
if Lpq[0].size != nao**2: # aosym = 's2'
Lpq = lib.unpack_tril(Lpq).astype(np.complex128)
out = _ao2mo.r_e2(Lpq, mo, (0, nmo, nmo, nmo + nvir), tao,
ao_loc)
Lpv[ki, kj] = out.reshape(-1, nmo, nvir)
return eris
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][naux:], 1)
lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k][naux:], 1)
if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k][naux:], 1)
lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k][naux:], 1)
for k, ji 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)
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 kernel(self, x0=None):
'''TDHF diagonalization with non-Hermitian eigenvalue solver
'''
logger.warn(
self, 'PBC-TDDFT is an experimental feature. '
'It is numerically sensitive to the accuracy of integrals '
'(relating to cell.precision).')
self.check_sanity()
self.dump_flags()
vind, hdiag = self.gen_vind(self._scf)
precond = self.get_precond(hdiag)
if x0 is None:
x0 = self.init_guess(self._scf, self.nstates)
real_system = (gamma_point(self._scf.kpts)
and self._scf.mo_coeff[0].dtype == numpy.double)
# We only need positive eigenvalues
def pickeig(w, v, nroots, envs):
realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD)
& (w.real > POSTIVE_EIG_THRESHOLD))[0]
return lib.linalg_helper._eigs_cmplx2real(w, v, realidx,
real_system)
log = logger.Logger(self.stdout, self.verbose)
precision = self.cell.precision * 1e-2
hermi = 0
self.converged, w, x1 = \
lib.davidson_nosym1(vind, x0, precond,
tol=self.conv_tol,
nroots=self.nstates, lindep=self.lindep,
max_space=self.max_space, pick=pickeig,
fill_heff=purify_krlyov_heff(precision, hermi, log),
verbose=self.verbose)
mo_occ = self._scf.mo_occ
self.e = w
def norm_xy(z):
x, y = z.reshape(2, -1)
norm = 2 * (lib.norm(x)**2 - lib.norm(y)**2)
norm = 1 / numpy.sqrt(norm)
x *= norm
y *= norm
return _unpack(x, mo_occ), _unpack(y, mo_occ)
self.xy = [norm_xy(z) for z in x1]
return self.e, self.xy
def general(mydf, mo_coeffs, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_general_compact', True)):
'''General MO integral transformation'''
from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
warn_pbc2d_eri(mydf)
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
mo_coeffs = (mo_coeffs,) * 4
mo_coeffs = [numpy.asarray(mo, order='F') for mo in mo_coeffs]
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'fft_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros([mo.shape[1] for mo in mo_coeffs])
allreal = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs)
q = kptj - kpti
coulG = tools.get_coulG(cell, q, mesh=mydf.mesh)
coords = cell.gen_uniform_grids(mydf.mesh)
max_memory = mydf.max_memory - lib.current_memory()[0]
if gamma_point(kptijkl) and allreal:
ao = mydf._numint.eval_ao(cell, coords, kpti)[0]
if ((iden_coeffs(mo_coeffs[0], mo_coeffs[1]) and
iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
iden_coeffs(mo_coeffs[0], mo_coeffs[3]))):
moiT = mojT = numpy.asarray(lib.dot(mo_coeffs[0].T,ao.T), order='C')
ao = None
max_memory = max_memory - moiT.nbytes*1e-6
eri = _contract_compact(mydf, (moiT,mojT), coulG, max_memory=max_memory)
if not compact:
nmo = moiT.shape[0]
eri = ao2mo.restore(1, eri, nmo).reshape(nmo**2,nmo**2)
else:
mos = [numpy.asarray(lib.dot(c.T, ao.T), order='C') for c in mo_coeffs]
ao = None
fac = numpy.array(1.)
max_memory = max_memory - sum([x.nbytes for x in mos])*1e-6
eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory).real
return eri
else:
aos = mydf._numint.eval_ao(cell, coords, kptijkl)
mos = [numpy.asarray(lib.dot(c.T, aos[i].T), order='C')
for i,c in enumerate(mo_coeffs)]
aos = None
fac = numpy.exp(-1j * numpy.dot(coords, q))
max_memory = max_memory - sum([x.nbytes for x in mos])*1e-6
eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory)
return eri
def __init__(self,
the_mf,
mp_grid_dim,
seedname,
bands=None,
plot=False,
reorder=False,
max_memory=None):
self.mp_grid_dim = np.asarray(mp_grid_dim, dtype=np.int32)
self.num_kpts = mp_grid_dim[0] * mp_grid_dim[1] * mp_grid_dim[2]
self.seedname = seedname
self.plot = plot
self.reorder = reorder
self.mf = the_mf
self.cell = self.mf.cell
self.real_lattice = self.cell.lattice_vectors()
self.recip_lattice = self.cell.reciprocal_vectors()
self.max_memory = max_memory
if self.max_memory is None: self.max_memory = self.cell.max_memory
self.num_bands_tot = self.cell.nao_nr()
self.bands = bands
if self.bands is None: self.bands = np.arange(self.num_bands_tot)
self.num_atoms = self.cell.natm
self.atom_symbols = [x[0] for x in self.cell._atom]
self.atoms_cart = np.asarray([x[1] for x in self.cell._atom])
self.kpts = None
if isinstance(self.mf, pbc.scf.khf.KSCF): self.kpts = self.mf.kpts
elif isinstance(self.mf, pbc.scf.hf.SCF): self.kpts = self.mf.kpt
self.kpts = self.kpts.reshape((-1, 3))
self.kpt_latt = self.cell.get_scaled_kpts(self.kpts)
assert (len(self.kpt_latt) == self.num_kpts)
self.gamma_only = False
if gamma_point(self.kpts): self.gamma_only = True
eig = self.mf.mo_energy
eig *= param.HARTREE2EV #in eV
self.eig = eig.reshape((-1, self.num_bands_tot))
self.nnlist,self.nncell,self.nntot,self.num_bands,self.num_wann,\
self.proj_site,self.proj_l,self.proj_m,self.proj_radial,\
self.proj_z,self.proj_x,self.proj_zona,self.exclude_bands,\
self.proj_s,self.proj_s_qaxis = self.wannier90_setup()
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpts_band=None):
'''Get the Coulomb (J) AO matrix at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray
Density matrix at each k-point. If a list of k-point DMs, eg,
UHF alpha and beta DM, the alpha and beta DMs are contracted
separately.
kpts : (nkpts, 3) ndarray
Kwargs:
kpts_band : (3,) ndarray or (*,3) ndarray
A list of arbitrary "band" k-points at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
or list of vj if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
gs = mydf.gs
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
coulG = tools.get_coulG(cell, gs=gs)
ngs = len(coulG)
vR = rhoR = np.zeros((nset, ngs))
for k, aoR in mydf.aoR_loop(gs, kpts):
for i in range(nset):
rhoR[i] += numint.eval_rho(cell, aoR, dms[i, k])
for i in range(nset):
rhoR[i] *= 1. / nkpts
rhoG = tools.fft(rhoR[i], gs)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, gs).real
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
weight = cell.vol / ngs
if gamma_point(kpts_band):
vj_kpts = np.empty((nset, nband, nao, nao))
else:
vj_kpts = np.empty((nset, nband, nao, nao), dtype=np.complex128)
for k, aoR in mydf.aoR_loop(gs, kpts_band):
for i in range(nset):
vj_kpts[i, k] = weight * lib.dot(aoR.T.conj() * vR[i], aoR)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def aux_e2(cell,
auxcell,
intor='int3c2e_sph',
aosym='s1',
comp=1,
kptij_lst=numpy.zeros((1, 2, 3)),
shls_slice=None):
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.
Returns:
(nao_pair, naux) array
'''
intor = gto.moleintor.ascint3(intor)
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('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)
kpti = kptij_lst[:, 0]
kptj = kptij_lst[:, 1]
j_only = is_zero(kpti - kptj)
if j_only and aosym[:2] == 's2':
assert (shls_slice[2] == 0)
nao_pair = (ao_loc[shls_slice[1]] * (ao_loc[shls_slice[1]] + 1) // 2 -
ao_loc[shls_slice[0]] * (ao_loc[shls_slice[0]] + 1) // 2)
else:
nao_pair = ni * nj
if gamma_point(kptij_lst):
dtype = numpy.double
else:
dtype = numpy.complex128
int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)
out = numpy.empty((nkptij, comp, nao_pair, naux), dtype=dtype)
out = int3c(shls_slice, out)
if comp == 1:
out = out[:, 0]
if nkptij == 1:
out = out[0]
return out
def kernel(self, x0=None):
'''TDHF diagonalization with non-Hermitian eigenvalue solver
'''
self.check_sanity()
self.dump_flags()
vind, hdiag = self.gen_vind(self._scf)
precond = self.get_precond(hdiag)
if x0 is None:
x0 = self.init_guess(self._scf, self.nstates)
real_system = (gamma_point(self._scf.kpts)
and self._scf.mo_coeff[0][0].dtype == numpy.double)
# We only need positive eigenvalues
def pickeig(w, v, nroots, envs):
realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD)
& (w.real > POSTIVE_EIG_THRESHOLD))[0]
return lib.linalg_helper._eigs_cmplx2real(w, v, realidx,
real_system)
log = logger.Logger(self.stdout, self.verbose)
precision = self.cell.precision * 1e-2
self.converged, w, x1 = \
lib.davidson_nosym1(vind, x0, precond,
tol=self.conv_tol,
nroots=self.nstates, lindep=self.lindep,
max_space=self.max_space, pick=pickeig,
fill_heff=purify_krlyov_heff(precision, 0, log),
verbose=self.verbose)
mo_occ = self._scf.mo_occ
e = []
xy = []
for i, z in enumerate(x1):
xs, ys = z.reshape(2, -1)
norm = lib.norm(xs)**2 - lib.norm(ys)**2
if norm > 0:
norm = 1 / numpy.sqrt(norm)
xs *= norm
ys *= norm
e.append(w[i])
xy.append((_unpack(xs, mo_occ), _unpack(ys, mo_occ)))
self.e = numpy.array(e)
self.xy = xy
return self.e, self.xy
def get_j_for_bands(mydf,
dm_kpts,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None):
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (logger.process_clock(), logger.perf_counter())
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
kpt_allow = numpy.zeros(3)
mesh = mydf.mesh
coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
ngrids = len(coulG)
rhoG = numpy.zeros((nset, ngrids), dtype=numpy.complex128)
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt_allow,
kpts,
max_memory=max_memory):
for k, aoao in enumerate(aoaoks):
rhoG[:, p0:p1] += numpy.einsum('nij,Lij->nL', dms[:, k].conj(),
aoao.reshape(-1, nao, nao)).conj()
aoao = aoaoks = p0 = p1 = None
weight = 1. / len(kpts)
vG = rhoG * coulG * weight
t1 = log.timer_debug1('get_j pass 1 to compute J(G)', *t1)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
vj_kpts = numpy.zeros((nset, nband, nao, nao), dtype=numpy.complex128)
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt_allow,
kpts_band,
max_memory=max_memory):
for k, aoao in enumerate(aoaoks):
vj_kpts[:, k] += numpy.einsum('nL,Lij->nij', vG[:, p0:p1],
aoao.reshape(-1, nao, nao))
aoao = aoaoks = p0 = p1 = None
if gamma_point(kpts_band):
vj_kpts = vj_kpts.real.copy()
t1 = log.timer_debug1('get_j pass 2', *t1)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
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
def get_j_for_bands(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
dmsR = dms.real.reshape(nset,nkpts,nao**2)
dmsI = dms.imag.reshape(nset,nkpts,nao**2)
kpt_allow = numpy.zeros(3)
mesh = mydf.mesh
coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
ngrids = len(coulG)
vG = numpy.zeros((nset,ngrids), dtype=numpy.complex128)
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
dmsC = dms.conj()
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
#:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
for k, aoao in enumerate(aoaoks):
for i in range(nset):
rho = numpy.einsum('ij,Lij->L', dmsC[i,k],
aoao.reshape(-1,nao,nao)).conj()
vG[i,p0:p1] += rho * coulG[p0:p1]
aoao = aoaoks = p0 = p1 = None
weight = 1./len(kpts)
vG *= weight
t1 = log.timer_debug1('get_j pass 1 to compute J(G)', *t1)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
vj_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128)
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts_band,
max_memory=max_memory):
for k, aoao in enumerate(aoaoks):
for i in range(nset):
vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i,p0:p1],
aoao.reshape(-1,nao,nao))
aoao = aoaoks = p0 = p1 = None
if gamma_point(kpts_band):
vj_kpts = vj_kpts.real.copy()
t1 = log.timer_debug1('get_j pass 2', *t1)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def load(aux_slice):
col0, col1 = aux_slice
j3cR = []
j3cI = []
for k, idx in enumerate(adapted_ji_idx):
v = numpy.vstack([fswap['j3c-junk/%d/%d'%(idx,i)][0,col0:col1].T
for i in range(nsegs)])
# vbar is the interaction between the background charge
# and the auxiliary basis. 0D, 1D, 2D do not have vbar.
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
def get_ao_pairs_G(mydf,
kpts=numpy.zeros((2, 3)),
q=None,
shls_slice=None,
compact=getattr(__config__, 'pbc_df_ao_pairs_compact',
False)):
'''Calculate forward Fourier tranform (G|ij) of all AO pairs.
Returns:
ao_pairs_G : 2D complex array
For gamma point, the shape is (ngrids, nao*(nao+1)/2); otherwise the
shape is (ngrids, nao*nao)
'''
if kpts is None: kpts = numpy.zeros((2, 3))
cell = mydf.cell
kpts = numpy.asarray(kpts)
q = kpts[1] - kpts[0]
coords = cell.gen_uniform_grids(mydf.mesh)
ngrids = len(coords)
max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .5)
if shls_slice is None:
shls_slice = (0, cell.nbas, 0, cell.nbas)
ish0, ish1, jsh0, jsh1 = shls_slice
ao_loc = cell.ao_loc_nr()
i0 = ao_loc[ish0]
i1 = ao_loc[ish1]
j0 = ao_loc[jsh0]
j1 = ao_loc[jsh1]
compact = compact and (i0 == j0) and (i1 == j1)
if compact and gamma_point(kpts): # gamma point
aosym = 's2'
else:
aosym = 's1'
ao_pairs_G = numpy.empty((ngrids, (i1 - i0) * (j1 - j0)),
dtype=numpy.complex)
for pqkR, pqkI, p0, p1 \
in mydf.pw_loop(mydf.mesh, kpts, q, shls_slice,
max_memory=max_memory, aosym=aosym):
ao_pairs_G[p0:p1] = pqkR.T + pqkI.T * 1j
return ao_pairs_G
def general(mydf, mo_coeffs, kpts=None, compact=False):
'''General MO integral transformation'''
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
mo_coeffs = (mo_coeffs,) * 4
mo_coeffs = [numpy.asarray(mo, order='F') for mo in mo_coeffs]
allreal = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs)
q = kptj - kpti
coulG = tools.get_coulG(cell, q, gs=mydf.gs)
coords = cell.gen_uniform_grids(mydf.gs)
max_memory = mydf.max_memory - lib.current_memory()[0]
if gamma_point(kptijkl) and allreal:
ao = mydf._numint.eval_ao(cell, coords, kpti)[0]
if ((iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
iden_coeffs(mo_coeffs[1], mo_coeffs[3]))):
moiT = mojT = numpy.asarray(lib.dot(mo_coeffs[0].T,ao.T), order='C')
ao = None
max_memory = max_memory - moiT.nbytes*1e-6
eri = _contract_compact(mydf, (moiT,mojT), coulG, max_memory=max_memory)
if not compact:
nmo = moiT.shape[0]
eri = ao2mo.restore(1, eri, nmo).reshape(nmo**2,nmo**2)
else:
mos = [numpy.asarray(lib.dot(c.T, ao.T), order='C') for c in mo_coeffs]
ao = None
fac = numpy.array(1.)
max_memory = max_memory - sum([x.nbytes for x in mos])*1e-6
eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory).real
return eri
else:
aos = mydf._numint.eval_ao(cell, coords, kptijkl)
mos = [numpy.asarray(lib.dot(c.T, aos[i].T), order='C')
for i,c in enumerate(mo_coeffs)]
aos = None
fac = numpy.exp(-1j * numpy.dot(coords, q))
max_memory = max_memory - sum([x.nbytes for x in mos])*1e-6
eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory)
return eri
def get_ao_pairs_G(mydf, kpts=numpy.zeros((2,3)), q=None, shls_slice=None,
compact=getattr(__config__, 'pbc_df_ao_pairs_compact', False)):
'''Calculate forward Fourier tranform (G|ij) of all AO pairs.
Returns:
ao_pairs_G : 2D complex array
For gamma point, the shape is (ngrids, nao*(nao+1)/2); otherwise the
shape is (ngrids, nao*nao)
'''
if kpts is None: kpts = numpy.zeros((2,3))
cell = mydf.cell
kpts = numpy.asarray(kpts)
q = kpts[1] - kpts[0]
coords = cell.gen_uniform_grids(mydf.mesh)
ngrids = len(coords)
if shls_slice is None:
shls_slice = (0, cell.nbas, 0, cell.nbas)
ish0, ish1, jsh0, jsh1 = shls_slice
ao_loc = cell.ao_loc_nr()
i0 = ao_loc[ish0]
i1 = ao_loc[ish1]
j0 = ao_loc[jsh0]
j1 = ao_loc[jsh1]
compact = compact and (i0 == j0) and (i1 == j1)
if compact and gamma_point(kpts): # gamma point
aosym = 's2'
else:
aosym = 's1'
ao_pairs_G = numpy.empty((ngrids,(i1-i0)*(j1-j0)), dtype=numpy.complex)
for pqkR, pqkI, p0, p1 \
in mydf.pw_loop(mydf.mesh, kptijkl[:2], q, shls_slice,
max_memory=max_memory, aosym=aosym):
ao_pairs_G[p0:p1] = pqkR.T + pqkI.T * 1j
return ao_pairs_G
def jk_method(self, J='FFTDF', K=None):
'''
Set up the schemes to evaluate Coulomb and exchange matrix
FFTDF: planewave density fitting using Fast Fourier Transform
AFTDF: planewave density fitting using analytic Fourier Transform
GDF: Gaussian density fitting
MDF: Gaussian and planewave mix density fitting
RS: range-separation JK builder
RSDF: range-separation density fitting
'''
if K is None:
K = J
if J != K:
raise NotImplementedError('J != K')
if 'DF' in J or 'DF' in K:
if 'DF' in J and 'DF' in K:
assert J == K
else:
df_method = J if 'DF' in J else K
self.with_df = getattr(df, df_method)(self.cell, self.kpt)
if 'RS' in J or 'RS' in K:
if not gamma_point(self.kpt):
raise NotImplementedError('Single k-point must be gamma point')
self.rsjk = RangeSeparationJKBuilder(self.cell, self.kpt)
self.rsjk.verbose = self.verbose
# For nuclear attraction
if J == 'RS' and K == 'RS' and not isinstance(self.with_df, df.GDF):
self.with_df = df.GDF(self.cell, self.kpt)
nuc = self.with_df.__class__.__name__
logger.debug1(self, 'Apply %s for J, %s for K, %s for nuc', J, K, nuc)
return self
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None):
cell = mydf.cell
if kpts is None:
kpts = mydf.kpts
nkpts = len(kpts)
if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3:
mo_coeff_kpts = [mo_coeff_kpts] * 4
else:
mo_coeff_kpts = list(mo_coeff_kpts)
# Shape of the orbitals can be different on different k-points. The
# orbital coefficients must be formatted (padded by zeros) so that the
# shape of the orbital coefficients are the same on all k-points. This can
# be achieved by calling pbc.mp.kmp2.padded_mo_coeff function
nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts]
eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol)
if gamma_point(kpts):
dtype = numpy.result_type(*mo_coeff_kpts)
else:
dtype = numpy.complex128
if out is None:
out = numpy.empty(eri_shape, dtype=dtype)
else:
assert(out.shape == eri_shape)
kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts])
kptis_lst = kptij_lst[:,0]
kptjs_lst = kptij_lst[:,1]
kpt_ji = kptjs_lst - kptis_lst
uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
ngrids = numpy.prod(mydf.mesh)
nao = cell.nao_nr()
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6) * .5
tao = []
ao_loc = None
kconserv = kpts_helper.get_kconserv(cell, kpts)
for uniq_id, kpt in enumerate(uniq_kpts):
q = uniq_kpts[uniq_id]
adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0]
for ji, ji_idx in enumerate(adapted_ji_idx):
ki = ji_idx // nkpts
kj = ji_idx % nkpts
moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:]
zij = []
for LpqR, LpqI, sign in mydf.sr_loop(kpts[[ki,kj]], max_memory, False, mydf.blockdim):
zij.append(_ao2mo.r_e2(LpqR+LpqI*1j, moij, ijslice, tao, ao_loc))
for kk in range(nkpts):
kl = kconserv[ki, kj, kk]
mokl, klslice = _conc_mos(mo_coeff_kpts[2][kk], mo_coeff_kpts[3][kl])[2:]
eri_mo = numpy.zeros((nmoi*nmoj,nmok*nmol), dtype=numpy.complex128)
for i, (LrsR, LrsI, sign) in \
enumerate(mydf.sr_loop(kpts[[kk,kl]], max_memory, False, mydf.blockdim)):
zkl = _ao2mo.r_e2(LrsR+LrsI*1j, mokl, klslice, tao, ao_loc)
lib.dot(zij[i].T, zkl, sign*factor, eri_mo, 1)
if dtype == numpy.double:
eri_mo = eri_mo.real
out[ki,kj,kk] = eri_mo.reshape(eri_shape[3:])
return out
def general(mydf, mo_coeffs, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_general_compact', True)):
warn_pbc2d_eri(mydf)
if mydf._cderi is None:
mydf.build()
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
mo_coeffs = (mo_coeffs,) * 4
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'df_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros([mo.shape[1] for mo in mo_coeffs])
all_real = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs)
max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]))
####################
# gamma point, the integral is real and with s4 symmetry
if gamma_point(kptijkl) and all_real:
ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact)
klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact)
eri_mo = numpy.zeros((nij_pair,nkl_pair))
sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
iden_coeffs(mo_coeffs[1], mo_coeffs[3]))
ijR = klR = None
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, True):
ijR, klR = _dtrans(LpqR, ijR, ijmosym, moij, ijslice,
LpqR, klR, klmosym, mokl, klslice, sym)
lib.ddot(ijR.T, klR, sign, eri_mo, 1)
LpqR = LpqI = None
return eri_mo
elif is_zero(kpti-kptk) and is_zero(kptj-kptl):
mo_coeffs = _mo_as_complex(mo_coeffs)
nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:]
nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3])[1:]
eri_mo = numpy.zeros((nij_pair,nkl_pair), dtype=numpy.complex)
sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
iden_coeffs(mo_coeffs[1], mo_coeffs[3]))
zij = zkl = None
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False):
buf = LpqR+LpqI*1j
zij, zkl = _ztrans(buf, zij, moij, ijslice,
buf, zkl, mokl, klslice, sym)
lib.dot(zij.T, zkl, sign, eri_mo, 1)
LpqR = LpqI = buf = None
return eri_mo
####################
# (kpt) i == j == k == l != 0
# (kpt) i == l && j == k && i != j && j != k =>
#
elif is_zero(kpti-kptl) and is_zero(kptj-kptk):
mo_coeffs = _mo_as_complex(mo_coeffs)
nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:]
nlk_pair, molk, lkslice = _conc_mos(mo_coeffs[3], mo_coeffs[2])[1:]
eri_mo = numpy.zeros((nij_pair,nlk_pair), dtype=numpy.complex)
sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[3]) and
iden_coeffs(mo_coeffs[1], mo_coeffs[2]))
zij = zlk = None
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False):
buf = LpqR+LpqI*1j
zij, zlk = _ztrans(buf, zij, moij, ijslice,
buf, zlk, molk, lkslice, sym)
lib.dot(zij.T, zlk.conj(), sign, eri_mo, 1)
LpqR = LpqI = buf = None
nmok = mo_coeffs[2].shape[1]
nmol = mo_coeffs[3].shape[1]
eri_mo = lib.transpose(eri_mo.reshape(-1,nmol,nmok), axes=(0,2,1))
return eri_mo.reshape(nij_pair,nlk_pair)
####################
# aosym = s1, complex integrals
#
# If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave
# vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition.
# So kptl/b - kptk/b must be -1 < k/b < 1. => kptl == kptk
#
else:
mo_coeffs = _mo_as_complex(mo_coeffs)
nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:]
nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3])[1:]
nao = mo_coeffs[0].shape[0]
eri_mo = numpy.zeros((nij_pair,nkl_pair), dtype=numpy.complex)
blksize = int(min(max_memory*.3e6/16/nij_pair,
max_memory*.3e6/16/nkl_pair,
max_memory*.3e6/16/nao**2))
zij = zkl = None
for (LpqR, LpqI, sign), (LrsR, LrsI, sign1) in \
lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False, blksize),
mydf.sr_loop(kptijkl[2:], max_memory, False, blksize)):
zij, zkl = _ztrans(LpqR+LpqI*1j, zij, moij, ijslice,
LrsR+LrsI*1j, zkl, mokl, klslice, False)
lib.dot(zij.T, zkl, sign, eri_mo, 1)
LpqR = LpqI = LrsR = LrsI = None
return eri_mo
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 get_jk(mydf, dm, hermi=1, kpt=numpy.zeros(3),
kpts_band=None, with_j=True, with_k=True, exxdiv=None):
'''JK for given k-point'''
vj = vk = None
if kpts_band is not None and abs(kpt-kpts_band).sum() > 1e-9:
kpt = numpy.reshape(kpt, (1,3))
if with_k:
vk = get_k_kpts(mydf, dm, hermi, kpt, kpts_band, exxdiv)
if with_j:
vj = get_j_kpts(mydf, dm, hermi, kpt, kpts_band)
return vj, vk
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
dm = numpy.asarray(dm, order='C')
dms = _format_dms(dm, [kpt])
nset, _, nao = dms.shape[:3]
dms = dms.reshape(nset,nao,nao)
j_real = gamma_point(kpt)
k_real = gamma_point(kpt) and not numpy.iscomplexobj(dms)
mesh = mydf.mesh
kptii = numpy.asarray((kpt,kpt))
kpt_allow = numpy.zeros(3)
if with_j:
vjcoulG = mydf.weighted_coulG(kpt_allow, False, mesh)
vjR = numpy.zeros((nset,nao,nao))
vjI = numpy.zeros((nset,nao,nao))
if with_k:
mydf.exxdiv = exxdiv
vkcoulG = mydf.weighted_coulG(kpt_allow, True, mesh)
vkR = numpy.zeros((nset,nao,nao))
vkI = numpy.zeros((nset,nao,nao))
dmsR = numpy.asarray(dms.real.reshape(nset,nao,nao), order='C')
dmsI = numpy.asarray(dms.imag.reshape(nset,nao,nao), order='C')
mem_now = lib.current_memory()[0]
max_memory = max(2000, (mydf.max_memory - mem_now)) * .8
log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)
t2 = t1
# rho_rs(-G+k_rs) is computed as conj(rho_{rs^*}(G-k_rs))
# == conj(transpose(rho_sr(G+k_sr), (0,2,1)))
blksize = max(int(max_memory*.25e6/16/nao**2), 16)
pLqR = pLqI = None
for pqkR, pqkI, p0, p1 in mydf.pw_loop(mesh, kptii, max_memory=max_memory):
t2 = log.timer_debug1('%d:%d ft_aopair'%(p0,p1), *t2)
pqkR = pqkR.reshape(nao,nao,-1)
pqkI = pqkI.reshape(nao,nao,-1)
if with_j:
#:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj())
#:vj += numpy.einsum('ijkl,lk->ij', v4, dm)
for i in range(nset):
rhoR = numpy.einsum('pq,pqk->k', dmsR[i], pqkR)
rhoR+= numpy.einsum('pq,pqk->k', dmsI[i], pqkI)
rhoI = numpy.einsum('pq,pqk->k', dmsI[i], pqkR)
rhoI-= numpy.einsum('pq,pqk->k', dmsR[i], pqkI)
rhoR *= vjcoulG[p0:p1]
rhoI *= vjcoulG[p0:p1]
vjR[i] += numpy.einsum('pqk,k->pq', pqkR, rhoR)
vjR[i] -= numpy.einsum('pqk,k->pq', pqkI, rhoI)
if not j_real:
vjI[i] += numpy.einsum('pqk,k->pq', pqkR, rhoI)
vjI[i] += numpy.einsum('pqk,k->pq', pqkI, rhoR)
#t2 = log.timer_debug1(' with_j', *t2)
if with_k:
#:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj())
#:vk += numpy.einsum('ijkl,jk->il', v4, dm)
pLqR = lib.transpose(pqkR, axes=(0,2,1), out=pLqR).reshape(-1,nao)
pLqI = lib.transpose(pqkI, axes=(0,2,1), out=pLqI).reshape(-1,nao)
nG = p1 - p0
iLkR = numpy.ndarray((nao,nG,nao), buffer=pqkR)
iLkI = numpy.ndarray((nao,nG,nao), buffer=pqkI)
for i in range(nset):
if k_real:
lib.dot(pLqR, dmsR[i], 1, iLkR.reshape(nao*nG,nao))
lib.dot(pLqI, dmsR[i], 1, iLkI.reshape(nao*nG,nao))
iLkR *= vkcoulG[p0:p1].reshape(1,nG,1)
iLkI *= vkcoulG[p0:p1].reshape(1,nG,1)
lib.dot(iLkR.reshape(nao,-1), pLqR.reshape(nao,-1).T, 1, vkR[i], 1)
lib.dot(iLkI.reshape(nao,-1), pLqI.reshape(nao,-1).T, 1, vkR[i], 1)
else:
zdotNN(pLqR, pLqI, dmsR[i], dmsI[i], 1,
iLkR.reshape(-1,nao), iLkI.reshape(-1,nao))
iLkR *= vkcoulG[p0:p1].reshape(1,nG,1)
iLkI *= vkcoulG[p0:p1].reshape(1,nG,1)
zdotNC(iLkR.reshape(nao,-1), iLkI.reshape(nao,-1),
pLqR.reshape(nao,-1).T, pLqI.reshape(nao,-1).T,
1, vkR[i], vkI[i])
#t2 = log.timer_debug1(' with_k', *t2)
pqkR = pqkI = coulG = pLqR = pLqI = iLkR = iLkI = None
#t2 = log.timer_debug1('%d:%d'%(p0,p1), *t2)
bufR = bufI = None
t1 = log.timer_debug1('aft_jk.get_jk', *t1)
if with_j:
if j_real:
vj = vjR
else:
vj = vjR + vjI * 1j
vj = vj.reshape(dm.shape)
if with_k:
if k_real:
vk = vkR
else:
vk = vkR + vkI * 1j
# Add ewald_exxdiv contribution because G=0 was not included in the
# non-uniform grids
if (exxdiv == 'ewald' and
(cell.dimension < 2 or # 0D and 1D are computed with inf_vacuum
(cell.dimension == 2 and cell.low_dim_ft_type == 'inf_vacuum'))):
_ewald_exxdiv_for_G0(cell, kpt, dms, vk)
vk = vk.reshape(dm.shape)
return vj, vk
def _init_cis_df_eris(cis, eris):
"""Add 3-center electron repulsion integrals, i.e. (L|pq), in `eris`,
where `L` denotes DF auxiliary basis functions and `p` and `q` canonical
crystalline orbitals. Note that `p` and `q` contain kpt indices `kp` and `kq`,
and the third kpt index `kL` is determined by the conservation of momentum.
Arguments:
cis {KCIS} -- A KCIS instance
eris {_CIS_ERIS} -- A _CIS_ERIS instance to which we want to add 3c ints
Returns:
_CIS_ERIS -- A _CIS_ERIS instance with 3c ints
"""
from pyscf.pbc.df import df
from pyscf.ao2mo import _ao2mo
from pyscf.pbc.lib.kpts_helper import gamma_point
log = logger.Logger(cis.stdout, cis.verbose)
if cis._scf.with_df._cderi is None:
cis._scf.with_df.build()
cell = cis._scf.cell
if cell.dimension == 2:
# 2D ERIs are not positive definite. The 3-index tensors are stored in
# two part. One corresponds to the positive part and one corresponds
# to the negative part. The negative part is not considered in the
# DF-driven CCSD implementation.
raise NotImplementedError
nocc = cis.nocc
nmo = cis.nmo
nao = cell.nao_nr()
kpts = cis.kpts
nkpts = len(kpts)
if gamma_point(kpts):
dtype = np.double
else:
dtype = np.complex128
eris.dtype = dtype = np.result_type(dtype, *eris.mo_coeff)
eris.Lpq_mo = Lpq_mo = np.empty((nkpts, nkpts), dtype=object)
cput0 = (time.clock(), time.time())
with h5py.File(cis._scf.with_df._cderi, 'r') as f:
kptij_lst = f['j3c-kptij'].value
tao = []
ao_loc = None
for ki, kpti in enumerate(kpts):
for kj, kptj in enumerate(kpts):
kpti_kptj = np.array((kpti, kptj))
Lpq_ao = np.asarray(df._getitem(f, 'j3c', kpti_kptj,
kptij_lst))
mo = np.hstack((eris.mo_coeff[ki], eris.mo_coeff[kj]))
mo = np.asarray(mo, dtype=dtype, order='F')
if dtype == np.double:
out = _ao2mo.nr_e2(Lpq_ao,
mo, (0, nmo, nmo, nmo + nmo),
aosym='s2')
else:
#Note: Lpq.shape[0] != naux if linear dependency is found in auxbasis
if Lpq_ao[0].size != nao**2: # aosym = 's2'
Lpq_ao = lib.unpack_tril(Lpq_ao).astype(np.complex128)
out = _ao2mo.r_e2(Lpq_ao, mo, (0, nmo, nmo, nmo + nmo),
tao, ao_loc)
Lpq_mo[ki, kj] = out.reshape(-1, nmo, nmo)
log.timer_debug1("transforming DF-CIS integrals", *cput0)
return eris
def make_kpt(uniq_kptji_id, cholesky_j2c):
kpt = uniq_kpts[uniq_kptji_id] # kpt = kptj - kpti
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
shls_slice = (auxcell.nbas, fused_cell.nbas)
Gaux = ft_ao.ft_ao(fused_cell, Gv, shls_slice, b, gxyz, Gvbase, kpt)
wcoulG = mydf.weighted_coulG(kpt, False, mesh)
Gaux *= wcoulG.reshape(-1, 1)
kLR = Gaux.real.copy('C')
kLI = Gaux.imag.copy('C')
Gaux = None
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 = numpy.vstack([
fswap['j3c-junk/%d/%d' % (idx, i)][0, col0:col1].T
for i in range(nsegs)
])
# vbar is the interaction between the background charge
# and the auxiliary basis. 0D, 1D, 2D do not have vbar.
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][naux:], 1)
lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k][naux:], 1)
if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k][naux:], 1)
lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k][naux:], 1)
for k, ji 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)
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 get_jk(mydf, dm, hermi=1, kpt=numpy.zeros(3),
kpts_band=None, with_j=True, with_k=True, exxdiv=None):
'''JK for given k-point'''
vj = vk = None
if kpts_band is not None and abs(kpt-kpts_band).sum() > 1e-9:
kpt = numpy.reshape(kpt, (1,3))
if with_k:
vk = get_k_kpts(mydf, dm, hermi, kpt, kpts_band, exxdiv)
if with_j:
vj = get_j_kpts(mydf, dm, hermi, kpt, kpts_band)
return vj, vk
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
if mydf._cderi is None or not mydf.has_kpts(kpts_band):
if mydf._cderi is not None:
log.warn('DF integrals for band k-points were not found %s. '
'DF integrals will be rebuilt to include band k-points.',
mydf._cderi)
mydf.build(kpts_band=kpts_band)
t1 = log.timer_debug1('Init get_jk', *t1)
dm = numpy.asarray(dm, order='C')
dms = _format_dms(dm, [kpt])
nset, _, nao = dms.shape[:3]
dms = dms.reshape(nset,nao,nao)
j_real = gamma_point(kpt)
k_real = gamma_point(kpt) and not numpy.iscomplexobj(dms)
kptii = numpy.asarray((kpt,kpt))
dmsR = dms.real.reshape(nset,nao,nao)
dmsI = dms.imag.reshape(nset,nao,nao)
mem_now = lib.current_memory()[0]
max_memory = max(2000, (mydf.max_memory - mem_now))
if with_j:
vjR = numpy.zeros((nset,nao,nao))
vjI = numpy.zeros((nset,nao,nao))
if with_k:
vkR = numpy.zeros((nset,nao,nao))
vkI = numpy.zeros((nset,nao,nao))
buf1R = numpy.empty((mydf.blockdim*nao**2))
buf2R = numpy.empty((mydf.blockdim*nao**2))
buf1I = numpy.zeros((mydf.blockdim*nao**2))
buf2I = numpy.empty((mydf.blockdim*nao**2))
max_memory *= .5
log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)
def contract_k(pLqR, pLqI):
# K ~ 'iLj,lLk*,li->kj' + 'lLk*,iLj,li->kj'
#:pLq = (LpqR + LpqI.reshape(-1,nao,nao)*1j).transpose(1,0,2)
#:tmp = numpy.dot(dm, pLq.reshape(nao,-1))
#:vk += numpy.dot(pLq.reshape(-1,nao).conj().T, tmp.reshape(-1,nao))
nrow = pLqR.shape[1]
tmpR = numpy.ndarray((nao,nrow*nao), buffer=buf2R)
if k_real:
for i in range(nset):
lib.ddot(dmsR[i], pLqR.reshape(nao,-1), 1, tmpR)
lib.ddot(pLqR.reshape(-1,nao).T, tmpR.reshape(-1,nao), 1, vkR[i], 1)
else:
tmpI = numpy.ndarray((nao,nrow*nao), buffer=buf2I)
for i in range(nset):
zdotNN(dmsR[i], dmsI[i], pLqR.reshape(nao,-1),
pLqI.reshape(nao,-1), 1, tmpR, tmpI, 0)
zdotCN(pLqR.reshape(-1,nao).T, pLqI.reshape(-1,nao).T,
tmpR.reshape(-1,nao), tmpI.reshape(-1,nao),
1, vkR[i], vkI[i], 1)
pLqI = None
thread_k = None
for LpqR, LpqI, sign in mydf.sr_loop(kptii, max_memory, False):
LpqR = LpqR.reshape(-1,nao,nao)
t1 = log.timer_debug1(' load', *t1)
if thread_k is not None:
thread_k.join()
if with_j:
#:rho_coeff = numpy.einsum('Lpq,xqp->xL', Lpq, dms)
#:vj += numpy.dot(rho_coeff, Lpq.reshape(-1,nao**2))
rhoR = numpy.einsum('Lpq,xpq->xL', LpqR, dmsR)
if not j_real:
LpqI = LpqI.reshape(-1,nao,nao)
rhoR -= numpy.einsum('Lpq,xpq->xL', LpqI, dmsI)
rhoI = numpy.einsum('Lpq,xpq->xL', LpqR, dmsI)
rhoI += numpy.einsum('Lpq,xpq->xL', LpqI, dmsR)
vjR += sign * numpy.einsum('xL,Lpq->xpq', rhoR, LpqR)
if not j_real:
vjR -= sign * numpy.einsum('xL,Lpq->xpq', rhoI, LpqI)
vjI += sign * numpy.einsum('xL,Lpq->xpq', rhoR, LpqI)
vjI += sign * numpy.einsum('xL,Lpq->xpq', rhoI, LpqR)
t1 = log.timer_debug1(' with_j', *t1)
if with_k:
nrow = LpqR.shape[0]
pLqR = numpy.ndarray((nao,nrow,nao), buffer=buf1R)
pLqR[:] = LpqR.transpose(1,0,2)
if not k_real:
pLqI = numpy.ndarray((nao,nrow,nao), buffer=buf1I)
if LpqI is not None:
pLqI[:] = LpqI.reshape(-1,nao,nao).transpose(1,0,2)
thread_k = lib.background_thread(contract_k, pLqR, pLqI)
t1 = log.timer_debug1(' with_k', *t1)
LpqR = LpqI = pLqR = pLqI = None
if thread_k is not None:
thread_k.join()
thread_k = None
if with_j:
if j_real:
vj = vjR
else:
vj = vjR + vjI * 1j
vj = vj.reshape(dm.shape)
if with_k:
if k_real:
vk = vkR
else:
vk = vkR + vkI * 1j
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpt, dms, vk)
vk = vk.reshape(dm.shape)
t1 = log.timer('sr jk', *t1)
return vj, vk
def general(mydf, mo_coeffs, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_general_compact', True)):
warn_pbc2d_eri(mydf)
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
mo_coeffs = (mo_coeffs,) * 4
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'aft_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros([mo.shape[1] for mo in mo_coeffs])
q = kptj - kpti
mesh = mydf.mesh
coulG = mydf.weighted_coulG(q, False, mesh)
all_real = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs)
max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .5)
####################
# gamma point, the integral is real and with s4 symmetry
if gamma_point(kptijkl) and all_real:
ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact)
klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact)
eri_mo = numpy.zeros((nij_pair,nkl_pair))
sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
iden_coeffs(mo_coeffs[1], mo_coeffs[3]))
ijR = ijI = klR = klI = buf = None
for pqkR, pqkI, p0, p1 \
in mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory,
aosym='s2'):
buf = lib.transpose(pqkR, out=buf)
ijR, klR = _dtrans(buf, ijR, ijmosym, moij, ijslice,
buf, klR, klmosym, mokl, klslice, sym)
lib.ddot(ijR.T, klR*coulG[p0:p1,None], 1, eri_mo, 1)
buf = lib.transpose(pqkI, out=buf)
ijI, klI = _dtrans(buf, ijI, ijmosym, moij, ijslice,
buf, klI, klmosym, mokl, klslice, sym)
lib.ddot(ijI.T, klI*coulG[p0:p1,None], 1, eri_mo, 1)
pqkR = pqkI = None
return eri_mo
####################
# (kpt) i == j == k == l != 0
# (kpt) i == l && j == k && i != j && j != k =>
#
elif is_zero(kpti-kptl) and is_zero(kptj-kptk):
mo_coeffs = _mo_as_complex(mo_coeffs)
nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:]
nlk_pair, molk, lkslice = _conc_mos(mo_coeffs[3], mo_coeffs[2])[1:]
eri_mo = numpy.zeros((nij_pair,nlk_pair), dtype=numpy.complex)
sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[3]) and
iden_coeffs(mo_coeffs[1], mo_coeffs[2]))
zij = zlk = buf = None
for pqkR, pqkI, p0, p1 \
in mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory):
buf = lib.transpose(pqkR+pqkI*1j, out=buf)
zij, zlk = _ztrans(buf, zij, moij, ijslice,
buf, zlk, molk, lkslice, sym)
lib.dot(zij.T, zlk.conj()*coulG[p0:p1,None], 1, eri_mo, 1)
pqkR = pqkI = None
nmok = mo_coeffs[2].shape[1]
nmol = mo_coeffs[3].shape[1]
eri_mo = lib.transpose(eri_mo.reshape(-1,nmol,nmok), axes=(0,2,1))
return eri_mo.reshape(nij_pair,nlk_pair)
####################
# aosym = s1, complex integrals
#
# If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave
# vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition.
# So kptl/b - kptk/b must be -1 < k/b < 1. => kptl == kptk
#
else:
mo_coeffs = _mo_as_complex(mo_coeffs)
nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:]
nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3])[1:]
eri_mo = numpy.zeros((nij_pair,nkl_pair), dtype=numpy.complex)
tao = []
ao_loc = None
zij = zkl = buf = None
for (pqkR, pqkI, p0, p1), (rskR, rskI, q0, q1) in \
lib.izip(mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory*.5),
mydf.pw_loop(mesh,-kptijkl[2:], q, max_memory=max_memory*.5)):
buf = lib.transpose(pqkR+pqkI*1j, out=buf)
zij = _ao2mo.r_e2(buf, moij, ijslice, tao, ao_loc, out=zij)
buf = lib.transpose(rskR-rskI*1j, out=buf)
zkl = _ao2mo.r_e2(buf, mokl, klslice, tao, ao_loc, out=zkl)
zij *= coulG[p0:p1,None]
lib.dot(zij.T, zkl, 1, eri_mo, 1)
pqkR = pqkI = rskR = rskI = None
return eri_mo
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
if mydf._cderi is None or not mydf.has_kpts(kpts_band):
if mydf._cderi is not None:
log.warn('DF integrals for band k-points were not found %s. '
'DF integrals will be rebuilt to include band k-points.',
mydf._cderi)
mydf.build(kpts_band=kpts_band)
t1 = log.timer_debug1('Init get_j_kpts', *t1)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
if mydf.auxcell is None:
# If mydf._cderi is the file that generated from another calculation,
# guess naux based on the contents of the integral file.
naux = mydf.get_naoaux()
else:
naux = mydf.auxcell.nao_nr()
nao_pair = nao * (nao+1) // 2
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
j_real = gamma_point(kpts_band) and not numpy.iscomplexobj(dms)
dmsR = dms.real.transpose(0,1,3,2).reshape(nset,nkpts,nao**2)
dmsI = dms.imag.transpose(0,1,3,2).reshape(nset,nkpts,nao**2)
rhoR = numpy.zeros((nset,naux))
rhoI = numpy.zeros((nset,naux))
max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]))
for k, kpt in enumerate(kpts):
kptii = numpy.asarray((kpt,kpt))
p1 = 0
for LpqR, LpqI, sign in mydf.sr_loop(kptii, max_memory, False):
p0, p1 = p1, p1+LpqR.shape[0]
#:Lpq = (LpqR + LpqI*1j).reshape(-1,nao,nao)
#:rhoR[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k]).real
#:rhoI[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k]).imag
rhoR[:,p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqR, dmsR[:,k])
rhoI[:,p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqR, dmsI[:,k])
if LpqI is not None:
rhoR[:,p0:p1] -= sign * numpy.einsum('Lp,xp->xL', LpqI, dmsI[:,k])
rhoI[:,p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqI, dmsR[:,k])
LpqR = LpqI = None
t1 = log.timer_debug1('get_j pass 1', *t1)
weight = 1./nkpts
rhoR *= weight
rhoI *= weight
vjR = numpy.zeros((nset,nband,nao_pair))
vjI = numpy.zeros((nset,nband,nao_pair))
for k, kpt in enumerate(kpts_band):
kptii = numpy.asarray((kpt,kpt))
p1 = 0
for LpqR, LpqI, sign in mydf.sr_loop(kptii, max_memory, True):
p0, p1 = p1, p1+LpqR.shape[0]
#:Lpq = (LpqR + LpqI*1j)#.reshape(-1,nao,nao)
#:vjR[:,k] += numpy.dot(rho[:,p0:p1], Lpq).real
#:vjI[:,k] += numpy.dot(rho[:,p0:p1], Lpq).imag
vjR[:,k] += numpy.dot(rhoR[:,p0:p1], LpqR)
if not j_real:
vjI[:,k] += numpy.dot(rhoI[:,p0:p1], LpqR)
if LpqI is not None:
vjR[:,k] -= numpy.dot(rhoI[:,p0:p1], LpqI)
vjI[:,k] += numpy.dot(rhoR[:,p0:p1], LpqI)
LpqR = LpqI = None
t1 = log.timer_debug1('get_j pass 2', *t1)
if j_real:
vj_kpts = vjR
else:
vj_kpts = vjR + vjI*1j
vj_kpts = lib.unpack_tril(vj_kpts.reshape(-1,nao_pair))
vj_kpts = vj_kpts.reshape(nset,nband,nao,nao)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None):
cell = mydf.cell
if kpts is None:
kpts = mydf.kpts
nkpts = len(kpts)
if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3:
mo_coeff_kpts = [mo_coeff_kpts] * 4
else:
mo_coeff_kpts = list(mo_coeff_kpts)
mo_ids = [id(x) for x in mo_coeff_kpts]
moTs = []
coords = cell.gen_uniform_grids(mydf.mesh)
aos = mydf._numint.eval_ao(cell, coords, kpts)
for n, mo_id in enumerate(mo_ids):
if mo_id in mo_ids[:n]:
moTs.append(moTs[mo_ids[:n].index(mo_id)])
else:
moTs.append([lib.dot(mo.T, aos[k].T) for k,mo in enumerate(mo_coeff_kpts[n])])
# Shape of the orbitals can be different on different k-points. The
# orbital coefficients must be formatted (padded by zeros) so that the
# shape of the orbital coefficients are the same on all k-points. This can
# be achieved by calling pbc.mp.kmp2.padded_mo_coeff function
nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts]
eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol)
if gamma_point(kpts):
dtype = numpy.result_type(*mo_coeff_kpts)
else:
dtype = numpy.complex128
if out is None:
out = numpy.empty(eri_shape, dtype=dtype)
else:
assert(out.shape == eri_shape)
kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts])
kptis_lst = kptij_lst[:,0]
kptjs_lst = kptij_lst[:,1]
kpt_ji = kptjs_lst - kptis_lst
uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
ngrids = numpy.prod(mydf.mesh)
# To hold intermediates
fswap = lib.H5TmpFile()
kconserv = kpts_helper.get_kconserv(cell, kpts)
for uniq_id, kpt in enumerate(uniq_kpts):
q = uniq_kpts[uniq_id]
adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0]
ki = adapted_ji_idx[0] // nkpts
kj = adapted_ji_idx[0] % nkpts
coulG = tools.get_coulG(cell, q, mesh=mydf.mesh)
coulG *= (cell.vol/ngrids) * factor
phase = numpy.exp(-1j * numpy.dot(coords, q))
for kk in range(nkpts):
kl = kconserv[ki, kj, kk]
mokT = moTs[2][kk]
molT = moTs[3][kl]
mo_pairs = numpy.einsum('ig,g,jg->ijg', mokT.conj(), phase.conj(), molT)
v = tools.ifft(mo_pairs.reshape(-1,ngrids), mydf.mesh)
v *= coulG
v = tools.fft(v.reshape(-1,ngrids), mydf.mesh)
v *= phase
fswap['zkl/'+str(kk)] = v
for ji_idx in adapted_ji_idx:
ki = ji_idx // nkpts
kj = ji_idx % nkpts
for kk in range(nkpts):
moiT = moTs[0][ki]
mojT = moTs[1][kj]
mo_pairs = numpy.einsum('ig,jg->ijg', moiT.conj(), mojT)
tmp = lib.dot(mo_pairs.reshape(-1,ngrids),
numpy.asarray(fswap['zkl/'+str(kk)]).T)
if dtype == numpy.double:
tmp = tmp.real
out[ki,kj,kk] = tmp.reshape(eri_shape[3:])
del(fswap['zkl'])
return out
def get_ao_pairs_G(mydf, kpts=numpy.zeros((2,3)), q=None, shls_slice=None,
compact=getattr(__config__, 'pbc_df_ao_pairs_compact', False)):
'''Calculate forward (G|ij) FFT of all AO pairs.
Returns:
ao_pairs_G : 2D complex array
For gamma point, the shape is (ngrids, nao*(nao+1)/2); otherwise the
shape is (ngrids, nao*nao)
'''
if kpts is None: kpts = numpy.zeros((2,3))
cell = mydf.cell
kpts = numpy.asarray(kpts)
coords = cell.gen_uniform_grids(mydf.mesh)
ngrids = len(coords)
if shls_slice is None:
i0, i1 = j0, j1 = (0, cell.nao_nr())
else:
ish0, ish1, jsh0, jsh1 = shls_slice
ao_loc = cell.ao_loc_nr()
i0 = ao_loc[ish0]
i1 = ao_loc[ish1]
j0 = ao_loc[jsh0]
j1 = ao_loc[jsh1]
def trans(aoi, aoj, fac=1):
if id(aoi) == id(aoj):
aoi = aoj = numpy.asarray(aoi.T, order='C')
else:
aoi = numpy.asarray(aoi.T, order='C')
aoj = numpy.asarray(aoj.T, order='C')
ni = aoi.shape[0]
nj = aoj.shape[0]
ao_pairs_G = numpy.empty((ni,nj,ngrids), dtype=numpy.complex128)
for i in range(ni):
ao_pairs_G[i] = tools.fft(fac * aoi[i].conj() * aoj, mydf.mesh)
ao_pairs_G = ao_pairs_G.reshape(-1,ngrids).T
return ao_pairs_G
if compact and gamma_point(kpts): # gamma point
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
ao = numpy.asarray(ao.T, order='C')
npair = i1*(i1+1)//2 - i0*(i0+1)//2
ao_pairs_G = numpy.empty((npair,ngrids), dtype=numpy.complex128)
ij = 0
for i in range(i0, i1):
ao_pairs_G[ij:ij+i+1] = tools.fft(ao[i] * ao[:i+1], mydf.mesh)
ij += i + 1
ao_pairs_G = ao_pairs_G.T
elif is_zero(kpts[0]-kpts[1]):
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
ao_pairs_G = trans(ao[:,i0:i1], ao[:,j0:j1])
else:
if q is None:
q = kpts[1] - kpts[0]
aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts[:2])
fac = numpy.exp(-1j * numpy.dot(coords, q))
ao_pairs_G = trans(aoi[:,i0:i1], aoj[:,j0:j1], fac)
return ao_pairs_G
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None):
cell = mydf.cell
if kpts is None:
kpts = mydf.kpts
nkpts = len(kpts)
if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3:
mo_coeff_kpts = [mo_coeff_kpts] * 4
else:
mo_coeff_kpts = list(mo_coeff_kpts)
# Shape of the orbitals can be different on different k-points. The
# orbital coefficients must be formatted (padded by zeros) so that the
# shape of the orbital coefficients are the same on all k-points. This can
# be achieved by calling pbc.mp.kmp2.padded_mo_coeff function
nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts]
eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol)
if gamma_point(kpts):
dtype = numpy.result_type(*mo_coeff_kpts)
else:
dtype = numpy.complex128
if out is None:
out = numpy.empty(eri_shape, dtype=dtype)
else:
assert(out.shape == eri_shape)
kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts])
kptis_lst = kptij_lst[:,0]
kptjs_lst = kptij_lst[:,1]
kpt_ji = kptjs_lst - kptis_lst
uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
ngrids = numpy.prod(mydf.mesh)
nao = cell.nao_nr()
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6) * .5
fswap = lib.H5TmpFile()
tao = []
ao_loc = None
kconserv = kpts_helper.get_kconserv(cell, kpts)
for uniq_id, kpt in enumerate(uniq_kpts):
q = uniq_kpts[uniq_id]
adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0]
kptjs = kptjs_lst[adapted_ji_idx]
coulG = mydf.weighted_coulG(q, False, mydf.mesh)
coulG *= factor
moij_list = []
ijslice_list = []
for ji, ji_idx in enumerate(adapted_ji_idx):
ki = ji_idx // nkpts
kj = ji_idx % nkpts
moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:]
moij_list.append(moij)
ijslice_list.append(ijslice)
fswap.create_dataset('zij/'+str(ji), (ngrids,nmoi*nmoj), 'D')
for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, q, kptjs):
for ji, aoao in enumerate(aoaoks):
ki = adapted_ji_idx[ji] // nkpts
kj = adapted_ji_idx[ji] % nkpts
buf = aoao.transpose(1,2,0).reshape(nao**2,ngrids)
zij = _ao2mo.r_e2(lib.transpose(buf), moij_list[ji],
ijslice_list[ji], tao, ao_loc)
zij *= coulG[p0:p1,None]
fswap['zij/'+str(ji)][p0:p1] = zij
mokl_list = []
klslice_list = []
for kk in range(nkpts):
kl = kconserv[ki, kj, kk]
mokl, klslice = _conc_mos(mo_coeff_kpts[2][kk], mo_coeff_kpts[3][kl])[2:]
mokl_list.append(mokl)
klslice_list.append(klslice)
fswap.create_dataset('zkl/'+str(kk), (ngrids,nmok*nmol), 'D')
ki = adapted_ji_idx[0] // nkpts
kj = adapted_ji_idx[0] % nkpts
kptls = kpts[kconserv[ki, kj, :]]
for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, q, -kptls):
for kk, aoao in enumerate(aoaoks):
buf = aoao.conj().transpose(1,2,0).reshape(nao**2,ngrids)
zkl = _ao2mo.r_e2(lib.transpose(buf), mokl_list[kk],
klslice_list[kk], tao, ao_loc)
fswap['zkl/'+str(kk)][p0:p1] = zkl
for ji, ji_idx in enumerate(adapted_ji_idx):
ki = ji_idx // nkpts
kj = ji_idx % nkpts
moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:]
zij = []
for LpqR, LpqI, sign in mydf.sr_loop(kpts[[ki,kj]], max_memory, False, mydf.blockdim):
zij.append(_ao2mo.r_e2(LpqR+LpqI*1j, moij, ijslice, tao, ao_loc))
for kk in range(nkpts):
kl = kconserv[ki, kj, kk]
eri_mo = lib.dot(numpy.asarray(fswap['zij/'+str(ji)]).T,
numpy.asarray(fswap['zkl/'+str(kk)]))
for i, (LrsR, LrsI, sign) in \
enumerate(mydf.sr_loop(kpts[[kk,kl]], max_memory, False, mydf.blockdim)):
zkl = _ao2mo.r_e2(LrsR+LrsI*1j, mokl_list[kk],
klslice_list[kk], tao, ao_loc)
lib.dot(zij[i].T, zkl, sign*factor, eri_mo, 1)
if dtype == numpy.double:
eri_mo = eri_mo.real
out[ki,kj,kk] = eri_mo.reshape(eri_shape[3:])
del(fswap['zij'])
del(fswap['zkl'])
return out
def get_pp_loc_part1(mydf, kpts=None):
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1,3))
else:
kpts_lst = numpy.reshape(kpts, (-1,3))
log = logger.Logger(mydf.stdout, mydf.verbose)
t0 = t1 = (time.clock(), time.time())
mesh = numpy.asarray(mydf.mesh)
nkpts = len(kpts_lst)
nao = cell.nao_nr()
nao_pair = nao * (nao+1) // 2
charges = cell.atom_charges()
kpt_allow = numpy.zeros(3)
if mydf.eta == 0:
if cell.dimension > 0:
ke_guess = estimate_ke_cutoff(cell, cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if numpy.any(mesh[:cell.dimension] < mesh_guess[:cell.dimension]*.8):
logger.warn(mydf, 'mesh %s is not enough for AFTDF.get_nuc function '
'to get integral accuracy %g.\nRecommended mesh is %s.',
mesh, cell.precision, mesh_guess)
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
vpplocG = -numpy.einsum('ij,ij->j', cell.get_SI(Gv), vpplocG)
vpplocG *= kws
vG = vpplocG
vj = numpy.zeros((nkpts,nao_pair), dtype=numpy.complex128)
else:
if cell.dimension > 0:
ke_guess = estimate_ke_cutoff_for_eta(cell, mydf.eta, cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if numpy.any(mesh < mesh_guess*.8):
logger.warn(mydf, 'mesh %s is not enough for AFTDF.get_nuc function '
'to get integral accuracy %g.\nRecommended mesh is %s.',
mesh, cell.precision, mesh_guess)
mesh_min = numpy.min((mesh_guess, mesh), axis=0)
if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
mesh[:cell.dimension] = mesh_min[:cell.dimension]
else:
mesh = mesh_min
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
nuccell = _compensate_nuccell(mydf)
# PP-loc part1 is handled by fakenuc in _int_nuc_vloc
vj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst))
t0 = t1 = log.timer_debug1('vnuc pass1: analytic int', *t0)
coulG = tools.get_coulG(cell, kpt_allow, mesh=mesh, Gv=Gv) * kws
aoaux = ft_ao.ft_ao(nuccell, Gv)
vG = numpy.einsum('i,xi->x', -charges, aoaux) * coulG
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts_lst,
max_memory=max_memory, aosym='s2'):
for k, aoao in enumerate(aoaoks):
# rho_ij(G) nuc(-G) / G^2
# = [Re(rho_ij(G)) + Im(rho_ij(G))*1j] [Re(nuc(G)) - Im(nuc(G))*1j] / G^2
if gamma_point(kpts_lst[k]):
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real)
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag)
else:
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao)
t1 = log.timer_debug1('contracting Vnuc [%s:%s]'%(p0, p1), *t1)
log.timer_debug1('contracting Vnuc', *t0)
vj_kpts = []
for k, kpt in enumerate(kpts_lst):
if gamma_point(kpt):
vj_kpts.append(lib.unpack_tril(vj[k].real.copy()))
else:
vj_kpts.append(lib.unpack_tril(vj[k]))
if kpts is None or numpy.shape(kpts) == (3,):
vj_kpts = vj_kpts[0]
return numpy.asarray(vj_kpts)
def get_eri(mydf, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_get_eri_compact', True)):
cell = mydf.cell
nao = cell.nao_nr()
kptijkl = _format_kpts(kpts)
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'aft_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros((nao,nao,nao,nao))
kpti, kptj, kptk, kptl = kptijkl
q = kptj - kpti
mesh = mydf.mesh
coulG = mydf.weighted_coulG(q, False, mesh)
nao_pair = nao * (nao+1) // 2
max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .8)
####################
# gamma point, the integral is real and with s4 symmetry
if gamma_point(kptijkl):
eriR = numpy.zeros((nao_pair,nao_pair))
for pqkR, pqkI, p0, p1 \
in mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory,
aosym='s2'):
lib.ddot(pqkR*coulG[p0:p1], pqkR.T, 1, eriR, 1)
lib.ddot(pqkI*coulG[p0:p1], pqkI.T, 1, eriR, 1)
pqkR = pqkI = None
if not compact:
eriR = ao2mo.restore(1, eriR, nao).reshape(nao**2,-1)
return eriR
####################
# (kpt) i == j == k == l != 0
# (kpt) i == l && j == k && i != j && j != k =>
#
# complex integrals, N^4 elements
elif is_zero(kpti-kptl) and is_zero(kptj-kptk):
eriR = numpy.zeros((nao**2,nao**2))
eriI = numpy.zeros((nao**2,nao**2))
for pqkR, pqkI, p0, p1 \
in mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory):
# rho_pq(G+k_pq) * conj(rho_rs(G-k_rs))
zdotNC(pqkR*coulG[p0:p1], pqkI*coulG[p0:p1], pqkR.T, pqkI.T,
1, eriR, eriI, 1)
pqkR = pqkI = None
pqkR = pqkI = coulG = None
# transpose(0,1,3,2) because
# j == k && i == l =>
# (L|ij).transpose(0,2,1).conj() = (L^*|ji) = (L^*|kl) => (M|kl)
# rho_rs(-G+k_rs) = conj(transpose(rho_sr(G+k_sr), (0,2,1)))
eri = lib.transpose((eriR+eriI*1j).reshape(-1,nao,nao), axes=(0,2,1))
return eri.reshape(nao**2,-1)
####################
# aosym = s1, complex integrals
#
# If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave
# vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition.
# So kptl/b - kptk/b must be -1 < k/b < 1. => kptl == kptk
#
else:
eriR = numpy.zeros((nao**2,nao**2))
eriI = numpy.zeros((nao**2,nao**2))
#
# (pq|rs) = \sum_G 4\pi rho_pq rho_rs / |G+k_{pq}|^2
# rho_pq = 1/N \sum_{Tp,Tq} \int exp(-i(G+k_{pq})*r) p(r-Tp) q(r-Tq) dr
# = \sum_{Tq} exp(i k_q*Tq) \int exp(-i(G+k_{pq})*r) p(r) q(r-Tq) dr
# Note the k-point wrap-around for rho_rs, which leads to G+k_{pq} in FT
# rho_rs = 1/N \sum_{Tr,Ts} \int exp( i(G+k_{pq})*r) r(r-Tr) s(r-Ts) dr
# = \sum_{Ts} exp(i k_s*Ts) \int exp( i(G+k_{pq})*r) r(r) s(r-Ts) dr
# rho_pq can be directly evaluated by AFT (function pw_loop)
# rho_pq = pw_loop(k_q, G+k_{pq})
# Assuming r(r) and s(r) are real functions, rho_rs is evaluated
# rho_rs = 1/N \sum_{Tr,Ts} \int exp( i(G+k_{pq})*r) r(r-Tr) s(r-Ts) dr
# = conj(\sum_{Ts} exp(-i k_s*Ts) \int exp(-i(G+k_{pq})*r) r(r) s(r-Ts) dr)
# = conj( pw_loop(-k_s, G+k_{pq}) )
#
# TODO: For complex AO function r(r) and s(r), pw_loop function needs to be
# extended to include Gv vector in the arguments
for (pqkR, pqkI, p0, p1), (rskR, rskI, q0, q1) in \
lib.izip(mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory*.5),
mydf.pw_loop(mesh,-kptijkl[2:], q, max_memory=max_memory*.5)):
pqkR *= coulG[p0:p1]
pqkI *= coulG[p0:p1]
zdotNC(pqkR, pqkI, rskR.T, rskI.T, 1, eriR, eriI, 1)
pqkR = pqkI = rskR = rskI = None
return (eriR+eriI*1j)
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None,
exxdiv=None):
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
if mydf._cderi is None or not mydf.has_kpts(kpts_band):
if mydf._cderi is not None:
log.warn('DF integrals for band k-points were not found %s. '
'DF integrals will be rebuilt to include band k-points.',
mydf._cderi)
mydf.build(kpts_band=kpts_band)
t1 = log.timer_debug1('Init get_k_kpts', *t1)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
vkR = numpy.zeros((nset,nband,nao,nao))
vkI = numpy.zeros((nset,nband,nao,nao))
dmsR = numpy.asarray(dms.real, order='C')
dmsI = numpy.asarray(dms.imag, order='C')
# K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
bufR = numpy.empty((mydf.blockdim*nao**2))
bufI = numpy.empty((mydf.blockdim*nao**2))
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
def make_kpt(ki, kj, swap_2e):
kpti = kpts[ki]
kptj = kpts_band[kj]
for LpqR, LpqI, sign in mydf.sr_loop((kpti,kptj), max_memory, False):
nrow = LpqR.shape[0]
pLqR = numpy.ndarray((nao,nrow,nao), buffer=bufR)
pLqI = numpy.ndarray((nao,nrow,nao), buffer=bufI)
tmpR = numpy.ndarray((nao,nrow*nao), buffer=LpqR)
tmpI = numpy.ndarray((nao,nrow*nao), buffer=LpqI)
pLqR[:] = LpqR.reshape(-1,nao,nao).transpose(1,0,2)
pLqI[:] = LpqI.reshape(-1,nao,nao).transpose(1,0,2)
for i in range(nset):
zdotNN(dmsR[i,ki], dmsI[i,ki], pLqR.reshape(nao,-1),
pLqI.reshape(nao,-1), 1, tmpR, tmpI)
zdotCN(pLqR.reshape(-1,nao).T, pLqI.reshape(-1,nao).T,
tmpR.reshape(-1,nao), tmpI.reshape(-1,nao),
sign, vkR[i,kj], vkI[i,kj], 1)
if swap_2e:
tmpR = tmpR.reshape(nao*nrow,nao)
tmpI = tmpI.reshape(nao*nrow,nao)
for i in range(nset):
zdotNN(pLqR.reshape(-1,nao), pLqI.reshape(-1,nao),
dmsR[i,kj], dmsI[i,kj], 1, tmpR, tmpI)
zdotNC(tmpR.reshape(nao,-1), tmpI.reshape(nao,-1),
pLqR.reshape(nao,-1).T, pLqI.reshape(nao,-1).T,
sign, vkR[i,ki], vkI[i,ki], 1)
if kpts_band is kpts: # normal k-points HF/DFT
for ki in range(nkpts):
for kj in range(ki):
make_kpt(ki, kj, True)
make_kpt(ki, ki, False)
t1 = log.timer_debug1('get_k_kpts: make_kpt ki>=kj (%d,*)'%ki, *t1)
else:
for ki in range(nkpts):
for kj in range(nband):
make_kpt(ki, kj, False)
t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)'%ki, *t1)
if (gamma_point(kpts) and gamma_point(kpts_band) and
not numpy.iscomplexobj(dm_kpts)):
vk_kpts = vkR
else:
vk_kpts = vkR + vkI * 1j
vk_kpts *= 1./nkpts
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
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 get_eri(mydf, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_get_eri_compact', True)):
if mydf._cderi is None:
mydf.build()
cell = mydf.cell
nao = cell.nao_nr()
kptijkl = _format_kpts(kpts)
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'df_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros((nao,nao,nao,nao))
kpti, kptj, kptk, kptl = kptijkl
nao_pair = nao * (nao+1) // 2
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6)
####################
# gamma point, the integral is real and with s4 symmetry
if gamma_point(kptijkl):
eriR = numpy.zeros((nao_pair,nao_pair))
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, True):
lib.ddot(LpqR.T, LpqR, sign, eriR, 1)
LpqR = LpqI = None
if not compact:
eriR = ao2mo.restore(1, eriR, nao).reshape(nao**2,-1)
return eriR
elif is_zero(kpti-kptk) and is_zero(kptj-kptl):
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNN(LpqR.T, LpqI.T, LpqR, LpqI, sign, eriR, eriI, 1)
LpqR = LpqI = None
return eriR + eriI*1j
####################
# (kpt) i == j == k == l != 0
#
# (kpt) i == l && j == k && i != j && j != k =>
# both vbar and ovlp are zero. It corresponds to the exchange integral.
#
# complex integrals, N^4 elements
elif is_zero(kpti-kptl) and is_zero(kptj-kptk):
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNC(LpqR.T, LpqI.T, LpqR, LpqI, sign, eriR, eriI, 1)
LpqR = LpqI = None
# transpose(0,1,3,2) because
# j == k && i == l =>
# (L|ij).transpose(0,2,1).conj() = (L^*|ji) = (L^*|kl) => (M|kl)
eri = lib.transpose((eriR+eriI*1j).reshape(-1,nao,nao), axes=(0,2,1))
return eri.reshape(nao**2,-1)
####################
# aosym = s1, complex integrals
#
# kpti == kptj => kptl == kptk
# If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave
# vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition.
# So kptl/b - kptk/b must be -1 < k/b < 1.
#
else:
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
blksize = int(max_memory*.4e6/16/nao**2)
for (LpqR, LpqI, sign), (LrsR, LrsI, sign1) in \
lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False, blksize),
mydf.sr_loop(kptijkl[2:], max_memory, False, blksize)):
zdotNN(LpqR.T, LpqI.T, LrsR, LrsI, sign, eriR, eriI, 1)
LpqR = LpqI = LrsR = LrsI = None
return eriR + eriI*1j
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 get_eri(mydf, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_get_eri_compact', True)):
if mydf._cderi is None:
mydf.build()
cell = mydf.cell
nao = cell.nao_nr()
kptijkl = _format_kpts(kpts)
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'df_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros((nao,nao,nao,nao))
kpti, kptj, kptk, kptl = kptijkl
nao_pair = nao * (nao+1) // 2
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6)
####################
# gamma point, the integral is real and with s4 symmetry
if gamma_point(kptijkl):
eriR = numpy.zeros((nao_pair,nao_pair))
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, True):
lib.ddot(LpqR.T, LpqR, sign, eriR, 1)
LpqR = LpqI = None
if not compact:
eriR = ao2mo.restore(1, eriR, nao).reshape(nao**2,-1)
return eriR
elif is_zero(kpti-kptk) and is_zero(kptj-kptl):
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNN(LpqR.T, LpqI.T, LpqR, LpqI, sign, eriR, eriI, 1)
LpqR = LpqI = None
return eriR + eriI*1j
####################
# (kpt) i == j == k == l != 0
#
# (kpt) i == l && j == k && i != j && j != k =>
# both vbar and ovlp are zero. It corresponds to the exchange integral.
#
# complex integrals, N^4 elements
elif is_zero(kpti-kptl) and is_zero(kptj-kptk):
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNC(LpqR.T, LpqI.T, LpqR, LpqI, sign, eriR, eriI, 1)
LpqR = LpqI = None
# transpose(0,1,3,2) because
# j == k && i == l =>
# (L|ij).transpose(0,2,1).conj() = (L^*|ji) = (L^*|kl) => (M|kl)
eri = lib.transpose((eriR+eriI*1j).reshape(-1,nao,nao), axes=(0,2,1))
return eri.reshape(nao**2,-1)
####################
# aosym = s1, complex integrals
#
# kpti == kptj => kptl == kptk
# If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave
# vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition.
# So kptl/b - kptk/b must be -1 < k/b < 1.
#
else:
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
blksize = int(max_memory*.4e6/16/nao**2)
for (LpqR, LpqI, sign), (LrsR, LrsI, sign1) in \
lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False, blksize),
mydf.sr_loop(kptijkl[2:], max_memory, False, blksize)):
zdotNN(LpqR.T, LpqI.T, LrsR, LrsI, sign, eriR, eriI, 1)
LpqR = LpqI = LrsR = LrsI = None
return eriR + eriI*1j
def make_kpt(uniq_kptji_id): # 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)
Gaux = ft_ao.ft_ao(fused_cell, Gv, None, b, gxyz, Gvbase, kpt).T
Gaux = fuse(Gaux)
Gaux *= mydf.weighted_coulG(kpt, False, gs)
kLR = Gaux.T.real.copy('C')
kLI = Gaux.T.imag.copy('C')
j2c = numpy.asarray(feri['j2c/%d' % uniq_kptji_id])
# Note large difference may be found in results between the CD/eig treatments.
# In some systems, small integral errors can lead to different treatments of
# linear dependency which can be observed in the total energy/orbital energy
# around 4th decimal place.
# try:
# j2c = scipy.linalg.cholesky(j2c, lower=True)
# j2ctag = 'CD'
# except scipy.linalg.LinAlgError as e:
#
# Abandon CD treatment for better numerical stablity
w, v = scipy.linalg.eigh(j2c)
log.debug('MDF metric for kpt %s cond = %.4g, drop %d bfns',
uniq_kptji_id, w[-1] / w[0],
numpy.count_nonzero(w < mydf.linear_dep_threshold))
v = v[:, w > mydf.linear_dep_threshold].T.conj()
v /= numpy.sqrt(w[w > mydf.linear_dep_threshold]).reshape(-1, 1)
j2c = v
j2ctag = 'eig'
naux0 = j2c.shape[0]
if is_zero(kpt): # kpti == kptj
aosym = 's2'
nao_pair = nao * (nao + 1) // 2
vbar = fuse(mydf.auxbar(fused_cell))
ovlp = cell.pbc_intor('int1e_ovlp_sph',
hermi=1,
kpts=adapted_kptjs)
for k, ji in enumerate(adapted_ji_idx):
ovlp[k] = lib.pack_tril(ovlp[k])
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 * .6 * 1e6 / 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 * .2 * 1e6 / 16 / buflen / (nkptj + 1)))
else:
Gblksize = max(
16, int(max_memory * .4 * 1e6 / 16 / buflen / (nkptj + 1)))
Gblksize = min(Gblksize, ngs, 16384)
pqkRbuf = numpy.empty(buflen * Gblksize)
pqkIbuf = numpy.empty(buflen * Gblksize)
# buf for ft_aopair
buf = numpy.empty((nkptj, buflen * Gblksize), dtype=numpy.complex128)
col1 = 0
for istep, sh_range in enumerate(shranges):
log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \
istep+1, len(shranges), *sh_range)
bstart, bend, ncol = sh_range
col0, col1 = col1, col1 + ncol
j3cR = []
j3cI = []
for k, idx in enumerate(adapted_ji_idx):
v = fuse(numpy.asarray(feri['j3c/%d' % idx][:, col0:col1]))
if is_zero(kpt):
for i, c in enumerate(vbar):
if c != 0:
v[i] -= c * 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
if aosym == 's2':
shls_slice = (bstart, bend, 0, bend)
else:
shls_slice = (bstart, bend, 0, cell.nbas)
for p0, p1 in lib.prange(0, ngs, 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)
else:
v = lib.dot(j2c, v)
feri['j3c/%d' % ji][:naux0, col0:col1] = v
del (feri['j2c/%d' % uniq_kptji_id])
for k, ji in enumerate(adapted_ji_idx):
v = feri['j3c/%d' % ji][:naux0]
del (feri['j3c/%d' % ji])
feri['j3c/%d' % ji] = v
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None,
exxdiv=None):
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
mesh = mydf.mesh
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
swap_2e = (kpts_band is None)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
kk_table = kpts_band.reshape(-1,1,3) - kpts.reshape(1,-1,3)
kk_todo = numpy.ones(kk_table.shape[:2], dtype=bool)
vkR = numpy.zeros((nset,nband,nao,nao))
vkI = numpy.zeros((nset,nband,nao,nao))
dmsR = numpy.asarray(dms.real, order='C')
dmsI = numpy.asarray(dms.imag, order='C')
mem_now = lib.current_memory()[0]
max_memory = max(2000, (mydf.max_memory - mem_now)) * .8
log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)
# K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
def make_kpt(kpt): # kpt = kptj - kpti
# search for all possible ki and kj that has ki-kj+kpt=0
kk_match = numpy.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
kpti_idx, kptj_idx = numpy.where(kk_todo & kk_match)
nkptj = len(kptj_idx)
log.debug1('kpt = %s', kpt)
log.debug2('kpti_idx = %s', kpti_idx)
log.debug2('kptj_idx = %s', kptj_idx)
kk_todo[kpti_idx,kptj_idx] = False
if swap_2e and not is_zero(kpt):
kk_todo[kptj_idx,kpti_idx] = False
max_memory1 = max_memory * (nkptj+1)/(nkptj+5)
#blksize = max(int(max_memory1*4e6/(nkptj+5)/16/nao**2), 16)
#bufR = numpy.empty((blksize*nao**2))
#bufI = numpy.empty((blksize*nao**2))
# Use DF object to mimic KRHF/KUHF object in function get_coulG
mydf.exxdiv = exxdiv
vkcoulG = mydf.weighted_coulG(kpt, True, mesh)
kptjs = kpts[kptj_idx]
# <r|-G+k_rs|s> = conj(<s|G-k_rs|r>) = conj(<s|G+k_sr|r>)
#buf1R = numpy.empty((blksize*nao**2))
#buf1I = numpy.empty((blksize*nao**2))
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt, kptjs, max_memory=max_memory1):
nG = p1 - p0
bufR = numpy.empty((nG*nao**2))
bufI = numpy.empty((nG*nao**2))
buf1R = numpy.empty((nG*nao**2))
buf1I = numpy.empty((nG*nao**2))
for k, aoao in enumerate(aoaoks):
ki = kpti_idx[k]
kj = kptj_idx[k]
# case 1: k_pq = (pi|iq)
#:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj())
#:vk += numpy.einsum('ijkl,jk->il', v4, dm)
pLqR = numpy.ndarray((nao,nG,nao), buffer=bufR)
pLqI = numpy.ndarray((nao,nG,nao), buffer=bufI)
pLqR[:] = aoao.real.reshape(nG,nao,nao).transpose(1,0,2)
pLqI[:] = aoao.imag.reshape(nG,nao,nao).transpose(1,0,2)
iLkR = numpy.ndarray((nao,nG,nao), buffer=buf1R)
iLkI = numpy.ndarray((nao,nG,nao), buffer=buf1I)
for i in range(nset):
zdotNN(pLqR.reshape(-1,nao), pLqI.reshape(-1,nao),
dmsR[i,kj], dmsI[i,kj], 1,
iLkR.reshape(-1,nao), iLkI.reshape(-1,nao))
iLkR *= vkcoulG[p0:p1].reshape(1,nG,1)
iLkI *= vkcoulG[p0:p1].reshape(1,nG,1)
zdotNC(iLkR.reshape(nao,-1), iLkI.reshape(nao,-1),
pLqR.reshape(nao,-1).T, pLqI.reshape(nao,-1).T,
1, vkR[i,ki], vkI[i,ki], 1)
# case 2: k_pq = (iq|pi)
#:v4 = numpy.einsum('iLj,lLk->ijkl', pqk, pqk.conj())
#:vk += numpy.einsum('ijkl,li->kj', v4, dm)
if swap_2e and not is_zero(kpt):
for i in range(nset):
zdotNN(dmsR[i,ki], dmsI[i,ki], pLqR.reshape(nao,-1),
pLqI.reshape(nao,-1), 1,
iLkR.reshape(nao,-1), iLkI.reshape(nao,-1))
iLkR *= vkcoulG[p0:p1].reshape(1,nG,1)
iLkI *= vkcoulG[p0:p1].reshape(1,nG,1)
zdotCN(pLqR.reshape(-1,nao).T, pLqI.reshape(-1,nao).T,
iLkR.reshape(-1,nao), iLkI.reshape(-1,nao),
1, vkR[i,kj], vkI[i,kj], 1)
for ki, kpti in enumerate(kpts_band):
for kj, kptj in enumerate(kpts):
if kk_todo[ki,kj]:
make_kpt(kptj-kpti)
t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)'%ki, *t1)
if (gamma_point(kpts) and gamma_point(kpts_band) and
not numpy.iscomplexobj(dm_kpts)):
vk_kpts = vkR
else:
vk_kpts = vkR + vkI * 1j
vk_kpts *= 1./nkpts
# Add ewald_exxdiv contribution because G=0 was not included in the
# non-uniform grids
if (exxdiv == 'ewald' and
(cell.dimension < 2 or # 0D and 1D are computed with inf_vacuum
(cell.dimension == 2 and cell.low_dim_ft_type == 'inf_vacuum'))):
_ewald_exxdiv_for_G0(cell, kpts_band, dms, vk_kpts, kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
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 make_kpt(uniq_kptji_id, cholesky_j2c):
kpt = uniq_kpts[uniq_kptji_id] # kpt = kptj - kpti
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
shls_slice = (auxcell.nbas, fused_cell.nbas)
Gaux = ft_ao.ft_ao(fused_cell, Gv, shls_slice, b, gxyz, Gvbase, kpt)
wcoulG = mydf.weighted_coulG(kpt, False, mesh)
Gaux *= wcoulG.reshape(-1,1)
kLR = Gaux.real.copy('C')
kLI = Gaux.imag.copy('C')
Gaux = None
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)
pqkRbuf = numpy.empty(buflen*Gblksize)
pqkIbuf = numpy.empty(buflen*Gblksize)
# buf for ft_aopair
buf = numpy.empty(nkptj*buflen*Gblksize, dtype=numpy.complex128)
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][naux:], 1)
lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k][naux:], 1)
if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k][naux:], 1)
lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k][naux:], 1)
for k, ji 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)
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)
with lib.call_in_background(pw_contract) as compute:
col1 = 0
for istep, sh_range in enumerate(shranges):
log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \
istep+1, len(shranges), *sh_range)
bstart, bend, ncol = sh_range
col0, col1 = col1, col1+ncol
j3cR = []
j3cI = []
for k, idx in enumerate(adapted_ji_idx):
v = numpy.vstack([fswap['j3c-junk/%d/%d'%(idx,i)][0,col0:col1].T
for i in range(nsegs)])
# vbar is the interaction between the background charge
# and the auxiliary basis. 0D, 1D, 2D do not have vbar.
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
compute(istep, sh_range, j3cR, j3cI)
for ji in adapted_ji_idx:
del(fswap['j3c-junk/%d'%ji])
