def get_jk(mf,
cell=None,
dm_kpts=None,
hermi=0,
kpts=None,
kpts_band=None,
with_j=True,
with_k=True,
**kwargs):
if cell is None: cell = mf.cell
if dm_kpts is None: dm_kpts = mf.make_rdm1()
if kpts is None: kpts = mf.kpts
if kpts_band is None: kpts_band = kpts
nkpts = len(kpts)
nband = len(kpts_band)
dm_kpts = np.asarray(dm_kpts)
nso = dm_kpts.shape[-1]
nao = nso // 2
dms = dm_kpts.reshape(-1, nkpts, nso, nso)
n_dm = dms.shape[0]
dmaa = dms[:, :, :nao, :nao]
dmab = dms[:, :, nao:, :nao]
dmbb = dms[:, :, nao:, nao:]
dms = np.vstack((dmaa, dmbb, dmab))
j1, k1 = mf.with_df.get_jk(dms,
hermi,
kpts,
kpts_band,
with_j,
with_k,
exxdiv=mf.exxdiv)
j1 = j1.reshape(3, n_dm, nband, nao, nao)
k1 = k1.reshape(3, n_dm, nband, nao, nao)
vj = vk = None
if with_j:
vj = np.zeros((n_dm, nband, nso, nso), j1.dtype)
vj[:, :, :nao, :nao] = vj[:, :, nao:, nao:] = j1[0] + j1[1]
vj = _format_jks(vj, dm_kpts, kpts_band, kpts)
if with_k:
vk = np.zeros((n_dm, nband, nso, nso), k1.dtype)
vk[:, :, :nao, :nao] = k1[0]
vk[:, :, nao:, nao:] = k1[1]
vk[:, :, :nao, nao:] = k1[2]
vk[:, :, nao:, :nao] = k1[2].transpose(0, 1, 3, 2).conj()
vk = _format_jks(vk, dm_kpts, kpts_band, kpts)
return vj, vk
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)
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):
cell = mydf.cell
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)
coulG = mydf.weighted_coulG(kpt_allow, False, mydf.gs)
ngs = len(coulG)
vR = numpy.zeros((nset,ngs))
vI = numpy.zeros((nset,ngs))
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
for k, pqkR, pqkI, p0, p1 \
in mydf.ft_loop(mydf.gs, kpt_allow, kpts, max_memory=max_memory):
#:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
for i in range(nset):
rhoR = numpy.dot(dmsR[i,k], pqkR)
rhoR+= numpy.dot(dmsI[i,k], pqkI)
rhoI = numpy.dot(dmsI[i,k], pqkR)
rhoI-= numpy.dot(dmsR[i,k], pqkI)
vR[i,p0:p1] += rhoR * coulG[p0:p1]
vI[i,p0:p1] += rhoI * coulG[p0:p1]
pqkR = pqkI = coulG = None
weight = 1./len(kpts)
vR *= weight
vI *= weight
t1 = log.timer_debug1('get_j pass 1 to compute J(G)', *t1)
kpts_band, single_kpt_band = _format_kpts_band(kpts_band, kpts)
gamma_point = abs(kpts_band).sum() < 1e-9
nband = len(kpts_band)
vjR = numpy.zeros((nset,nband,nao*nao))
vjI = numpy.zeros((nset,nband,nao*nao))
for k, pqkR, pqkI, p0, p1 \
in mydf.ft_loop(mydf.gs, kpt_allow, kpts_band,
max_memory=max_memory):
for i in range(nset):
vjR[i,k] += numpy.dot(pqkR, vR[i,p0:p1])
vjR[i,k] -= numpy.dot(pqkI, vI[i,p0:p1])
if not gamma_point:
for i in range(nset):
vjI[i,k] += numpy.dot(pqkI, vR[i,p0:p1])
vjI[i,k] += numpy.dot(pqkR, vI[i,p0:p1])
pqkR = pqkI = coulG = None
if gamma_point:
vj_kpts = vjR
else:
vj_kpts = vjR + vjI*1j
vj_kpts = vj_kpts.reshape(nset,nband,nao,nao)
t1 = log.timer_debug1('get_j pass 2', *t1)
return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts, single_kpt_band)
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_j_kpts(mydf,
dm_kpts,
hermi=1,
kpts=numpy.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
rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0)
rhoG = rhoG[:, 0]
coulG = tools.get_coulG(cell,
mesh=cell.mesh,
low_dim_ft_type=mydf.low_dim_ft_type)
ngrids = coulG.size
vG = numpy.einsum('ng,g->ng', rhoG.reshape(-1, ngrids), coulG)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
vj_kpts = _get_j_pass2(mydf, vG, kpts_band)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
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 get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray
Density matrix at each k-point
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
vk : (nkpts, nao, nao) ndarray
or list of vj and vk if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
gs = mydf.gs
coords = cell.gen_uniform_grids(gs)
ngs = coords.shape[0]
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngs)
kpts_band, single_kpt_band = _format_kpts_band(kpts_band, kpts)
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
for k2, ao_k2 in mydf.aoR_loop(gs, kpts):
kpt2 = kpts[k2]
aoR_dms = [lib.dot(ao_k2, dms[i, k2]) for i in range(nset)]
for k1, ao_k1 in mydf.aoR_loop(gs, kpts_band):
kpt1 = kpts_band[k1]
vkR_k1k2 = get_vkR(mydf, cell, ao_k1, ao_k2, kpt1, kpt2, coords,
gs, exxdiv)
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i, k1] += weight * lib.dot(ao_k1.T.conj(), tmp_Rq)
vkR_k1k2 = aoR_dms = tmp_Rq = None
return _format_jks(vk_kpts, dm_kpts, kpts_band, kpts, single_kpt_band)
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_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 get_jk(mf, cell=None, dm_kpts=None, hermi=0, kpts=None, kpts_band=None,
with_j=True, with_k=True):
if cell is None: cell = mf.cell
if dm_kpts is None: dm_kpts = mf.make_rdm1()
if kpts is None: kpts = mf.kpts
if kpts_band is None: kpts_band = kpts
nkpts = len(kpts)
nband = len(kpts_band)
dm_kpts = np.asarray(dm_kpts)
nso = dm_kpts.shape[-1]
nao = nso // 2
dms = dm_kpts.reshape(-1,nkpts,nso,nso)
n_dm = dms.shape[0]
dmaa = dms[:,:,:nao,:nao]
dmab = dms[:,:,nao:,:nao]
dmbb = dms[:,:,nao:,nao:]
dms = np.vstack((dmaa, dmbb, dmab))
j1, k1 = mf.with_df.get_jk(dms, hermi, kpts, kpts_band, with_j, with_k,
exxdiv=mf.exxdiv)
j1 = j1.reshape(3,n_dm,nband,nao,nao)
k1 = k1.reshape(3,n_dm,nband,nao,nao)
vj = vk = None
if with_j:
vj = np.zeros((n_dm,nband,nso,nso), j1.dtype)
vj[:,:,:nao,:nao] = vj[:,:,nao:,nao:] = j1[0] + j1[1]
vj = _format_jks(vj, dm_kpts, kpts_band, kpts)
if with_k:
vk = np.zeros((n_dm,nband,nso,nso), k1.dtype)
vk[:,:,:nao,:nao] = k1[0]
vk[:,:,nao:,nao:] = k1[1]
vk[:,:,:nao,nao:] = k1[2]
vk[:,:,nao:,:nao] = k1[2].transpose(0,1,3,2).conj()
vk = _format_jks(vk, dm_kpts, kpts_band, kpts)
return vj, vk
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, single_kpt_band = _format_kpts_band(kpts_band, kpts)
nband = len(kpts_band)
vj_kpts = []
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, kpts_band, kpts, single_kpt_band)
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 get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)),
kpts_band=None):
mydf = _sync_mydf(mydf)
cell = mydf.cell
mesh = mydf.mesh
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)
ngrids = len(coulG)
vR = rhoR = numpy.zeros((nset,ngrids))
for ao_ks_etc, p0, p1 in mydf.mpi_aoR_loop(mydf.grids, kpts):
ao_ks = ao_ks_etc[0]
for k, ao in enumerate(ao_ks):
for i in range(nset):
rhoR[i,p0:p1] += numint.eval_rho(cell, ao, dms[i,k])
ao = ao_ks = None
rhoR = mpi.allreduce(rhoR)
for i in range(nset):
rhoR[i] *= 1./nkpts
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 = numpy.zeros((nset,nband,nao,nao))
else:
vj_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128)
for ao_ks_etc, p0, p1 in mydf.mpi_aoR_loop(mydf.grids, kpts_band):
ao_ks = ao_ks_etc[0]
for k, ao in enumerate(ao_ks):
for i in range(nset):
vj_kpts[i,k] += lib.dot(ao.T.conj()*vR[i,p0:p1], ao)
vj_kpts = mpi.reduce(vj_kpts)
if gamma_point(kpts_band):
vj_kpts = vj_kpts.real
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
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 uks_j_xc(mydf,
dm_kpts,
xc_code,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None,
with_j=WITH_J,
j_in_xc=J_IN_XC):
log = lib.logger.Logger(mydf.stdout, mydf.verbose)
cell = mydf.cell
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
dms = None
#TODO: Handle multiple sets of KUKS density matrices (2,nset,nkpts,nao,nao)
assert (nset == 2) # alpha and beta density matrices in KUKS
ni = mydf._numint
xctype = ni._xc_type(xc_code)
if xctype == 'LDA':
deriv = 0
rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0)
def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
for k in range(nkpts):
aow = numpy.einsum('pi,p->pi', ao_l[k], vR[0])
v[0, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k])
aow = numpy.einsum('pi,p->pi', ao_l[k], vR[1])
v[1, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k])
def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
add_j_(v, ao_l, ao_h, idx_l, idx_h, wv[:, 0])
elif xctype == 'GGA':
deriv = 1
if RHOG_HIGH_DERIV:
rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv)
else:
Gv = cell.Gv
ngrids = Gv.shape[0]
rhoG = numpy.empty((2, 4, ngrids), dtype=numpy.complex128)
rhoG[:, :1] = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0)
rhoG[:, 1:] = numpy.einsum('np,px->nxp', 1j * rhoG[:, 0], Gv)
def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
for k in range(nkpts):
aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[0])
v[0, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[1])
v[1, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
wva, wvb = wv
for k in range(nkpts):
aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wva)
v1 = lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wva[1:4])
v1 += lib.dot(ao_l[k][0].conj().T, aow)
v[0, k, idx_l[:, None], idx_h] += v1
aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wvb)
v1 = lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wvb[1:4])
v1 += lib.dot(ao_l[k][0].conj().T, aow)
v[1, k, idx_l[:, None], idx_h] += v1
else: # MGGA
deriv = 2
#TODO: RHOG_HIGH_DERIV:
rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv)
def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
for k in range(nkpts):
aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[0])
v[0, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[1])
v[1, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
wva, wvb = wv
for k in range(nkpts):
aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wva[:4])
v1 = lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wva[1:4])
v1 += lib.dot(ao_l[k][0].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][1], wva[4], out=aow)
v1 += lib.dot(ao_l[k][1].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][2], wva[4], out=aow)
v1 += lib.dot(ao_l[k][2].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][3], wva[4], out=aow)
v1 += lib.dot(ao_l[k][3].conj().T, aow)
v[0, k, idx_l[:, None], idx_h] += v1
aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wvb[:4])
v1 = lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wvb[1:4])
v1 += lib.dot(ao_l[k][0].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][1], wvb[4], out=aow)
v1 += lib.dot(ao_l[k][1].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][2], wvb[4], out=aow)
v1 += lib.dot(ao_l[k][2].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][3], wvb[4], out=aow)
v1 += lib.dot(ao_l[k][3].conj().T, aow)
v[1, k, idx_l[:, None], idx_h] += v1
mesh = cell.mesh
coulG = tools.get_coulG(cell,
mesh=mesh,
low_dim_ft_type=mydf.low_dim_ft_type)
ngrids = coulG.size
vG = numpy.einsum('ng,g->ng', rhoG[:, 0].reshape(-1, ngrids), coulG)
vG = vG.reshape(2, *mesh)
weight = cell.vol / ngrids
# *(1./weight) because rhoR is scaled by weight in _eval_rhoG. When
# computing rhoR with IFFT, the weight factor is not needed.
rhoR = tools.ifft(rhoG.reshape(-1, ngrids), mesh) * (1. / weight)
rhoR = rhoR.real.reshape(2, -1, ngrids)
nelec = numpy.zeros(2)
excsum = 0
exc, vxc = ni.eval_xc(xc_code, rhoR, 1, deriv=1)[:2]
if xctype == 'LDA':
vrho = vxc[0]
wva = vrho[:, 0].reshape(1, ngrids)
wvb = vrho[:, 1].reshape(1, ngrids)
elif xctype == 'GGA':
vrho, vsigma = vxc[:2]
wva = numpy.empty((4, ngrids))
wvb = numpy.empty((4, ngrids))
wva[0] = vrho[:, 0]
wva[1:4] = rhoR[0, 1:4] * (vsigma[:, 0] * 2) # sigma_uu
wva[1:4] += rhoR[1, 1:4] * vsigma[:, 1] # sigma_ud
wvb[0] = vrho[:, 1]
wvb[1:4] = rhoR[1, 1:4] * (vsigma[:, 2] * 2) # sigma_dd
wvb[1:4] += rhoR[0, 1:4] * vsigma[:, 1] # sigma_ud
else:
vrho, vsigma, vlapl, vtau = vxc
wva = numpy.empty((5, ngrids))
wvb = numpy.empty((5, ngrids))
wva[0] = vrho[:, 0]
wva[1:4] = rhoR[0, 1:4] * (vsigma[:, 0] * 2) # sigma_uu
wva[1:4] += rhoR[1, 1:4] * vsigma[:, 1] # sigma_ud
wvb[0] = vrho[:, 1]
wvb[1:4] = rhoR[1, 1:4] * (vsigma[:, 2] * 2) # sigma_dd
wvb[1:4] += rhoR[0, 1:4] * vsigma[:, 1] # sigma_ud
if vlapl is None:
wvb[4] = .5 * vtau[:, 1]
wva[4] = .5 * vtau[:, 0]
else:
wva[4] = (.5 * vtau[:, 0] + 2 * vlapl[:, 0])
wvb[4] = (.5 * vtau[:, 1] + 2 * vlapl[:, 1])
nelec[0] += rhoR[0, 0].sum() * weight
nelec[1] += rhoR[1, 0].sum() * weight
excsum += (rhoR[0, 0] * exc).sum() * weight
excsum += (rhoR[1, 0] * exc).sum() * weight
wv_freq = tools.fft(numpy.vstack((wva, wvb)), mesh) * weight
wv_freq = wv_freq.reshape(2, -1, *mesh)
if j_in_xc:
wv_freq[:, 0] += vG
vR = tools.ifft(vG.reshape(-1, ngrids), mesh)
ecoul = numpy.einsum('ng,ng->', rhoR[:, 0].real, vR.real) * .5
log.debug('Coulomb energy %s', ecoul)
excsum += ecoul
rhoR = rhoG = None
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
if gamma_point(kpts_band):
veff = numpy.zeros((2, nkpts, nao, nao))
vj = numpy.zeros((2, nkpts, nao, nao))
else:
veff = numpy.zeros((2, nkpts, nao, nao), dtype=numpy.complex128)
vj = numpy.zeros((2, nkpts, nao, nao), dtype=numpy.complex128)
for grids_high, grids_low in mydf.tasks:
cell_high = grids_high.cell
mesh = grids_high.mesh
coords_idx = grids_high.coords_idx
ngrids0 = numpy.prod(mesh)
ngrids1 = grids_high.coords.shape[0]
log.debug('mesh %s, ngrids %s/%s', mesh, ngrids1, ngrids0)
gx = numpy.fft.fftfreq(mesh[0], 1. / mesh[0]).astype(int)
gy = numpy.fft.fftfreq(mesh[1], 1. / mesh[1]).astype(int)
gz = numpy.fft.fftfreq(mesh[2], 1. / mesh[2]).astype(int)
sub_wvG = wv_freq[:, :, gx[:, None, None], gy[:, None],
gz].reshape(-1, ngrids0)
wv = tools.ifft(sub_wvG, mesh).real.reshape(2, -1, ngrids0)
wv = wv[:, :, coords_idx]
if with_j:
sub_vG = vG[:, gx[:, None, None], gy[:, None],
gz].reshape(-1, ngrids0)
vR = tools.ifft(sub_vG, mesh).real.reshape(2, ngrids0)
vR = vR[:, coords_idx]
idx_h = grids_high.ao_idx
if grids_low is None:
for ao_h_etc, p0, p1 in mydf.aoR_loop(grids_high, kpts, deriv):
ao_h = ao_h_etc[0]
add_xc_(veff, ao_h, ao_h, idx_h, idx_h, wv[:, :, p0:p1])
if with_j:
add_j_(vj, ao_h, ao_h, idx_h, idx_h, vR[:, p0:p1])
ao_h = ao_h_etc = None
else:
idx_l = grids_low.ao_idx
for ao_h_etc, ao_l_etc in zip(
mydf.aoR_loop(grids_high, kpts, deriv),
mydf.aoR_loop(grids_low, kpts, deriv)):
p0, p1 = ao_h_etc[1:3]
ao_h = ao_h_etc[0][0]
ao_l = ao_l_etc[0][0]
add_xc_(veff, ao_h, ao_h, idx_h, idx_h, wv[:, :, p0:p1])
add_xc_(veff, ao_h, ao_l, idx_h, idx_l, wv[:, :, p0:p1])
add_xc_(veff, ao_l, ao_h, idx_l, idx_h, wv[:, :, p0:p1])
if with_j:
add_j_(vj, ao_h, ao_h, idx_h, idx_h, vR[:, p0:p1])
add_j_(vj, ao_h, ao_l, idx_h, idx_l, vR[:, p0:p1])
add_j_(vj, ao_l, ao_h, idx_l, idx_h, vR[:, p0:p1])
ao_h = ao_l = ao_h_etc = ao_l_etc = None
vj = _format_jks(vj, dm_kpts, input_band, kpts)
veff = _format_jks(veff, dm_kpts, input_band, kpts)
return nelec, excsum, veff, vj
def get_j_for_bands(self,
dm_kpts,
hermi=1,
kpts=np.zeros((1, 3)),
kpts_band=None):
log = logger.Logger(self.stdout, self.verbose)
t1 = (time.clock(), time.time())
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]
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
ngrids = len(coulG)
vG = np.zeros((n_dm, ngrids), dtype=np.complex128)
vG_SR = np.zeros((n_dm, ngrids), dtype=np.complex128)
max_memory = (self.max_memory - lib.current_memory()[0]) * .8
for aoaoks, p0, p1 in self.ft_loop(mesh,
kpt_allow,
kpts,
max_memory=max_memory):
#:rho = np.einsum('lkL,lk->L', pqk.conj(), dm)
for k, aoao in enumerate(aoaoks):
aoao = aoao.reshape(-1, nao, nao)
if nao_compact < nao:
for i in range(n_dm):
rho = np.einsum('ij,Lji->L', dms[i, k], aoao.conj())
vG[i, p0:p1] += rho * coulG[p0:p1]
aoao[:, nao_compact:, nao_compact:] = 0
for i in range(n_dm):
rho = np.einsum('ij,Lji->L', dms[i, k], aoao.conj())
vG_SR[i, p0:p1] += rho * coulG_SR[p0:p1]
else:
for i in range(n_dm):
rho = np.einsum('ij,Lji->L', dms[i, k], aoao.conj())
vG[i, p0:p1] += rho * coulG_LR[p0:p1]
aoao = aoaoks = p0 = p1 = None
weight = 1. / len(kpts)
vG *= weight
vG_SR *= 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 = np.zeros((n_dm, nband, nao, nao), dtype=np.complex128)
for aoaoks, p0, p1 in self.ft_loop(mesh,
kpt_allow,
kpts_band,
max_memory=max_memory):
for k, aoao in enumerate(aoaoks):
aoao = aoao.reshape(-1, nao, nao)
if nao_compact < nao:
for i in range(n_dm):
vj_kpts[i, k] += np.einsum('L,Lij->ij', vG[i, p0:p1],
aoao)
aoao[:, nao_compact:, nao_compact:] = 0
for i in range(n_dm):
vj_kpts[i, k] -= np.einsum('L,Lij->ij',
vG_SR[i, p0:p1], aoao)
else:
for i in range(n_dm):
vj_kpts[i, k] += np.einsum('L,Lij->ij', vG[i, p0:p1],
aoao)
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))
ovlp[:, nao_compact:, nao_compact:] = 0
ovlp_b = np.asarray(
cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts_band))
ovlp_b[:, 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_b)
vj_kpts -= np.pi / self.omega**2 * weight * G0_weight * vj_G0
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_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())
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, single_kpt_band = _format_kpts_band(kpts_band, kpts)
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, mydf.gs)
kptjs = kpts[kptj_idx]
# <r|-G+k_rs|s> = conj(<s|G-k_rs|r>) = conj(<s|G+k_sr|r>)
for k, pqkR, pqkI, p0, p1 \
in mydf.ft_loop(mydf.gs, kpt, kptjs, max_memory=max_memory1):
ki = kpti_idx[k]
kj = kptj_idx[k]
coulG = numpy.sqrt(vkcoulG[p0:p1])
# case 1: k_pq = (pi|iq)
#:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj())
#:vk += numpy.einsum('ijkl,jk->il', v4, dm)
pqkR *= coulG
pqkI *= coulG
pLqR = lib.transpose(pqkR.reshape(nao,nao,-1), axes=(0,2,1), out=bufR)
pLqI = lib.transpose(pqkI.reshape(nao,nao,-1), axes=(0,2,1), out=bufI)
iLkR = numpy.empty((nao*(p1-p0),nao))
iLkI = numpy.empty((nao*(p1-p0),nao))
for i in range(nset):
iLkR, iLkI = zdotNN(pLqR.reshape(-1,nao), pLqI.reshape(-1,nao),
dmsR[i,kj], dmsI[i,kj], 1, iLkR, iLkI)
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):
iLkR = iLkR.reshape(nao,-1)
iLkI = iLkI.reshape(nao,-1)
for i in range(nset):
iLkR, iLkI = zdotNN(dmsR[i,ki], dmsI[i,ki], pLqR.reshape(nao,-1),
pLqI.reshape(nao,-1), 1, iLkR, iLkI)
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)
pqkR = pqkI = coulG = pLqR = pLqI = iLkR = iLkI = None
for ki, kpti in enumerate(kpts_band):
for kj, kptj in enumerate(kpts):
if kk_todo[ki,kj]:
make_kpt(kptj-kpti)
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
# G=0 was not included in the non-uniform grids
if cell.dimension != 3 and exxdiv:
assert(exxdiv.lower() == 'ewald')
_ewald_exxdiv_for_G0(cell, kpts_band, dms, vk_kpts, kpts_band)
return _format_jks(vk_kpts, dm_kpts, kpts_band, kpts, single_kpt_band)
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 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 = (logger.process_clock(), logger.perf_counter())
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
vkcoulG = mydf.weighted_coulG(kpt, exxdiv, mesh)
kptjs = kpts[kptj_idx]
weight = 1./len(kpts)
perm_sym = swap_2e and not is_zero(kpt)
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt, kptjs, max_memory=max_memory1):
_update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI), vkcoulG[p0:p1],
weight, kpti_idx, kptj_idx, perm_sym)
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
# 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 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 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
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)
ngrids = len(coulG)
if hermi == 1 or gamma_point(kpts):
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
else: # vR may be complex if the underlying density is complex
vR = rhoR = np.zeros((nset, ngrids), dtype=np.complex128)
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):
for k, ao in enumerate(ao_ks):
ao_dm = lib.dot(ao, dms[i, k])
rhoR[i, p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj())
rhoR *= 1. / nkpts
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh)
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)
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):
# ni.eval_mat can handle real vR only
# vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA')
for k, ao in enumerate(ao_ks):
aow = np.einsum('xi,x->xi', ao, vR[i, p0:p1])
vj_kpts[i, k] += lib.dot(ao.conj().T, aow)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def get_k_e1_kpts(mydf,
dm_kpts,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''Derivatives of exchange (K) AO matrix at sampled k-points.
'''
cell = mydf.cell
mesh = mydf.mesh
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if getattr(dm_kpts, 'mo_coeff', None) is not None:
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128)
coords = mydf.grids.coords
if input_band is None:
ao2_kpts = [
np.asarray(ao.transpose(0, 2, 1), order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts, deriv=1)
]
ao1_kpts = ao2_kpts
ao2_kpts = [ao2_kpt[0] for ao2_kpt in ao2_kpts]
else:
ao2_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)
]
ao1_kpts = [
np.asarray(ao.transpose(0, 2, 1), order='C') for ao in
mydf._numint.eval_ao(cell, coords, kpts=kpts_band, deriv=1)
]
if mo_coeff is not None and nset == 1:
mo_coeff = [
mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
for k, occ in enumerate(mo_occ)
]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(
min(nao,
max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / 3 / ngrids / nao)))
lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d',
max_memory, blksize)
ao1_dtype = np.result_type(*ao1_kpts)
ao2_dtype = np.result_type(*ao2_kpts)
vR_dm = np.empty((3, nset, nao, ngrids), dtype=vk_kpts.dtype)
t1 = (time.clock(), time.time())
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
naoj = ao2T.shape[0]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
mydf.exxdiv = exxdiv
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh)
if is_zero(kpt1 - kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('aig,jg->aijg',
ao1T[1:, p0:p1].conj() * expmikr, ao2T)
vG = tools.fft(rho1.reshape(-1, ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(3, p1 - p0, naoj, ngrids)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('aijg,jg->aig',
vR,
ao_dms[i],
out=vR_dm[:, i, p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[:, i, k1] -= weight * np.einsum(
'aig,jg->aij', vR_dm[:, i], ao1T[0])
t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2,
*t1)
# Ewald correction has no contribution to nuclear gradient unless range separted Coulomb is used
# The gradient correction part is not added in the vk matrix
if exxdiv == 'ewald' and cell.omega != 0:
raise NotImplementedError("Range Separated Coulomb")
# when cell.omega !=0: madelung constant will have a non-zero derivative
vk_kpts = np.asarray(
[_format_jks(vk, dm_kpts, input_band, kpts) for vk in vk_kpts])
return vk_kpts
def get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray
Density matrix at each k-point
kpts : (nkpts, 3) ndarray
Kwargs:
hermi : int
Whether K matrix is hermitian
| 0 : not hermitian and not symmetric
| 1 : hermitian
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
vk : (nkpts, nao, nao) ndarray
or list of vj and vk if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
mesh = mydf.mesh
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if getattr(dm_kpts, 'mo_coeff', None) is not None:
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
coords = mydf.grids.coords
ao2_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)
]
if input_band is None:
ao1_kpts = ao2_kpts
else:
ao1_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band)
]
if mo_coeff is not None and nset == 1:
mo_coeff = [
mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
for k, occ in enumerate(mo_occ)
]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(
min(nao, max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / ngrids / nao)))
lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d',
max_memory, blksize)
#ao1_dtype = np.result_type(*ao1_kpts)
#ao2_dtype = np.result_type(*ao2_kpts)
vR_dm = np.empty((nset, nao, ngrids), dtype=vk_kpts.dtype)
t1 = (time.clock(), time.time())
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
naoj = ao2T.shape[0]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
mydf.exxdiv = exxdiv
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh)
if is_zero(kpt1 - kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('ig,jg->ijg', ao1T[p0:p1].conj() * expmikr,
ao2T)
vG = tools.fft(rho1.reshape(-1, ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(p1 - p0, naoj, ngrids)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[i, k1] += weight * lib.dot(vR_dm[i], ao1T.T)
t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2,
*t1)
# Function _ewald_exxdiv_for_G0 to add back in the G=0 component to vk_kpts
# Note in the _ewald_exxdiv_for_G0 implementation, the G=0 treatments are
# different for 1D/2D and 3D systems. The special treatments for 1D and 2D
# can only be used with AFTDF/GDF/MDF method. In the FFTDF method, 1D, 2D
# and 3D should use the ewald probe charge correction.
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_j_e1_kpts(mydf, dm_kpts, kpts=np.zeros((1, 3)), kpts_band=None):
'''Derivatives of Coulomb (J) AO matrix at sampled k-points.
'''
cell = mydf.cell
mesh = mydf.mesh
ni = mydf._numint
make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi=1)
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)
ngrids = len(coulG)
if gamma_point(kpts):
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
else: # vR may be complex if the underlying density is complex
vR = rhoR = np.zeros((nset, ngrids), dtype=np.complex128)
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):
for k, ao in enumerate(ao_ks):
ao_dm = lib.dot(ao, dms[i, k])
rhoR[i, p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj())
rhoR *= 1. / nkpts
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh)
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((3, nset, nband, nao, nao))
else:
vj_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128)
rho = None
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band, deriv=1):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
# ni.eval_mat can handle real vR only
# vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA')
for k, ao in enumerate(ao_ks):
aow = np.einsum('xi,x->xi', ao[0], vR[i, p0:p1])
vj_kpts[:, i, k] -= lib.einsum('axi,xj->aij', ao[1:].conj(),
aow)
vj_kpts = np.asarray(
[_format_jks(vj, dm_kpts, input_band, kpts) for vj in vj_kpts])
return vj_kpts
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)),
kpts_band=None, exxdiv=None):
mydf = _sync_mydf(mydf)
cell = mydf.cell
mesh = mydf.mesh
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if hasattr(dm_kpts, 'mo_coeff'):
if dm_kpts.ndim == 3: # KRHF
mo_coeff = [dm_kpts.mo_coeff]
mo_occ = [dm_kpts.mo_occ ]
else: # KUHF
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
elif hasattr(dm_kpts[0], 'mo_coeff'):
mo_coeff = [dm.mo_coeff for dm in dm_kpts]
mo_occ = [dm.mo_occ for dm in dm_kpts]
else:
mo_coeff = None
kpts = numpy.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1./nkpts * (cell.vol/ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = numpy.zeros((nset,nband,nao,nao), dtype=dms.dtype)
else:
vk_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128)
coords = mydf.grids.coords
ao2_kpts = [numpy.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)]
if input_band is None:
ao1_kpts = ao2_kpts
else:
ao1_kpts = [numpy.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(min(nao, max(1, (max_memory-mem_now)*1e6/16/4/ngrids/nao)))
lib.logger.debug1(mydf, 'max_memory %s blksize %d', max_memory, blksize)
ao1_dtype = numpy.result_type(*ao1_kpts)
ao2_dtype = numpy.result_type(*ao2_kpts)
vR_dm = numpy.empty((nset,nao,ngrids), dtype=vk_kpts.dtype)
ao_dms_buf = [None] * nkpts
tasks = [(k1,k2) for k2 in range(nkpts) for k1 in range(nband)]
for k1, k2 in mpi.static_partition(tasks):
ao1T = ao1_kpts[k1]
ao2T = ao2_kpts[k2]
kpt1 = kpts_band[k1]
kpt2 = kpts[k2]
if ao2T.size == 0 or ao1T.size == 0:
continue
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
mydf.exxdiv = exxdiv
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2-kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2-kpt1, True, mydf, mesh)
if is_zero(kpt1-kpt2):
expmikr = numpy.array(1.)
else:
expmikr = numpy.exp(-1j * numpy.dot(coords, kpt2-kpt1))
if ao_dms_buf[k2] is None:
if mo_coeff is None:
ao_dms = [lib.dot(dm[k2], ao2T.conj()) for dm in dms]
else:
ao_dms = []
for i, dm in enumerate(dms):
occ = mo_occ[i][k2]
mo_scaled = mo_coeff[i][k2][:,occ>0] * numpy.sqrt(occ[occ>0])
ao_dms.append(lib.dot(mo_scaled.T, ao2T).conj())
ao_dms_buf[k2] = ao_dms
else:
ao_dms = ao_dms_buf[k2]
if mo_coeff is None:
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = numpy.einsum('ig,jg->ijg', ao1T[p0:p1].conj()*expmikr, ao2T)
vG = tools.fft(rho1.reshape(-1,ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(p1-p0,nao,ngrids)
vG = None
if vR_dm.dtype == numpy.double:
vR = vR.real
for i in range(nset):
numpy.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i,p0:p1])
vR = None
else:
for p0, p1 in lib.prange(0, nao, blksize):
for i in range(nset):
rho1 = numpy.einsum('ig,jg->ijg',
ao1T[p0:p1].conj()*expmikr,
ao_dms[i].conj())
vG = tools.fft(rho1.reshape(-1,ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(p1-p0,-1,ngrids)
vG = None
if vR_dm.dtype == numpy.double:
vR = vR.real
numpy.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i,p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[i,k1] += weight * lib.dot(vR_dm[i], ao1T.T)
vk_kpts = mpi.reduce(lib.asarray(vk_kpts))
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = vk_kpts.real
if rank == 0:
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray
Density matrix at each k-point
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
vk : (nkpts, nao, nao) ndarray
or list of vj and vk if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
gs = mydf.gs
coords = cell.gen_uniform_grids(gs)
ngs = coords.shape[0]
if hasattr(dm_kpts, 'mo_coeff'):
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngs)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
ao2_kpts = mydf._numint.eval_ao(cell, coords, kpts, non0tab=mydf.non0tab)
ao2_kpts = [np.asarray(ao.T, order='C') for ao in ao2_kpts]
if input_band is None:
ao1_kpts = ao2_kpts
else:
ao1_kpts = mydf._numint.eval_ao(cell,
coords,
kpts_band,
non0tab=mydf.non0tab)
ao1_kpts = [np.asarray(ao.T, order='C') for ao in ao1_kpts]
if mo_coeff is not None and nset == 1:
mo_coeff = [
mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
for k, occ in enumerate(mo_occ)
]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
naoj = ao2_kpts[0].shape[0]
else:
naoj = nao
max_memory = mydf.max_memory - lib.current_memory()[0]
blksize = int(max(max_memory * 1e6 / 16 / 2 / ngs / nao, 1))
ao1_dtype = np.result_type(*ao1_kpts)
ao2_dtype = np.result_type(*ao2_kpts)
buf = np.empty((blksize, naoj, ngs),
dtype=np.result_type(ao1_dtype, ao2_dtype))
vR_dm = np.empty((nset, nao, ngs), dtype=vk_kpts.dtype)
ao_dms = np.empty((nset, naoj, ngs), dtype=np.result_type(dms, ao2_dtype))
for k2, ao2T in enumerate(ao2_kpts):
kpt2 = kpts[k2]
if mo_coeff is None or nset > 1:
for i in range(nset):
lib.dot(dms[i, k2], ao2T.conj(), c=ao_dms[i])
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
mydf.exxdiv = exxdiv
coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, gs)
if is_zero(kpt1 - kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('ig,jg->ijg',
ao1T[p0:p1].conj() * expmikr,
ao2T,
out=buf[:p1 - p0])
vG = tools.fft(rho1.reshape(-1, ngs), gs)
vG *= coulG
vR = tools.ifft(vG, gs).reshape(p1 - p0, naoj, ngs)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[i, k1] += weight * lib.dot(vR_dm[i], ao1T.T)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_k_kpts(self,
dm_kpts,
hermi=1,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''
C ~ compact basis, D ~ diffused basis
Compute K matrix with coulG_LR:
(CC|CC) (CC|CD) (CC|DC) (CD|CC) (CD|CD) (CD|DC) (DC|CC) (DC|CD) (DC|DC)
Compute K matrix with full coulG:
(CC|DD) (CD|DD) (DC|DD) (DD|CC) (DD|CD) (DD|DC) (DD|DD)
'''
cell = self.cell
log = logger.Logger(self.stdout, self.verbose)
t1 = (time.clock(), time.time())
mesh = self.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 = np.ones(kk_table.shape[:2], dtype=bool)
vkR = np.zeros((nset, nband, nao, nao))
vkI = np.zeros((nset, nband, nao, nao))
dmsR = np.asarray(dms.real, order='C')
dmsI = np.asarray(dms.imag, order='C')
weight = 1. / nkpts
n_diffused = cell._nbas_each_set[2]
nao_compact = cell.ao_loc[cell.nbas - n_diffused]
mem_now = lib.current_memory()[0]
max_memory = max(2000, (self.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 = np.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
kpti_idx, kptj_idx = np.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 = np.empty((blksize*nao**2))
#bufI = np.empty((blksize*nao**2))
# Use DF object to mimic KRHF/KUHF object in function get_coulG
vkcoulG = self.weighted_coulG(kpt, exxdiv, mesh)
coulG_SR = self.weighted_coulG_SR(kpt, False, mesh)
coulG_LR = vkcoulG - coulG_SR
kptjs = kpts[kptj_idx]
perm_sym = swap_2e and not is_zero(kpt)
for aoaoks, p0, p1 in self.ft_loop(mesh,
kpt,
kptjs,
max_memory=max_memory1):
if nao_compact < nao:
aoaoks = [aoao.reshape(-1, nao, nao) for aoao in aoaoks]
aft_jk._update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI),
vkcoulG[p0:p1], weight, kpti_idx,
kptj_idx, perm_sym)
for aoao in aoaoks:
aoao[:, nao_compact:, nao_compact:] = 0
aft_jk._update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI),
coulG_SR[p0:p1], -weight, kpti_idx,
kptj_idx, perm_sym)
else:
aft_jk._update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI),
coulG_LR[p0:p1], weight, kpti_idx,
kptj_idx, perm_sym)
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 np.iscomplexobj(dm_kpts)):
vk_kpts = vkR
else:
vk_kpts = vkR + vkI * 1j
# G=0 associated to 2e integrals in real-space
if cell.dimension >= 2:
ovlp = np.asarray(cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts))
ovlp[:, nao_compact:, nao_compact:] = 0
kws = cell.get_Gv_weights(mesh)[2]
G0_weight = kws[0] if isinstance(kws, np.ndarray) else kws
vk_G0 = lib.einsum('kpq,nkqr,krs->nkps', ovlp, dm_kpts, ovlp)
vk_kpts -= np.pi / self.omega**2 * weight * G0_weight * vk_G0
# 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)
