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
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
def _update_vk_(vk, aoaoks, dms, coulG, weight, kpti_idx, kptj_idx, swap_2e):
dmsR, dmsI = dms
vkR, vkI = vk
nG = len(coulG)
n_dm = vkR.shape[0]
nao = vkR.shape[-1]
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(n_dm):
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 *= coulG.reshape(1, nG, 1)
iLkI *= coulG.reshape(1, nG, 1)
zdotNC(iLkR.reshape(nao, -1), iLkI.reshape(nao, -1),
pLqR.reshape(nao, -1).T,
pLqI.reshape(nao, -1).T, weight, 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)
# <r|-G+k_rs|s> = conj(<s|G-k_rs|r>) = conj(<s|G+k_sr|r>)
if swap_2e:
for i in range(n_dm):
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 *= coulG.reshape(1, nG, 1)
iLkI *= coulG.reshape(1, nG, 1)
zdotCN(
pLqR.reshape(-1, nao).T,
pLqI.reshape(-1, nao).T, iLkR.reshape(-1, nao),
iLkI.reshape(-1, nao), weight, vkR[i, kj], vkI[i, kj], 1)
def contract_k(pLqR, pLqI, pjqR, pjqI):
# K ~ 'iLj,lLk*,li->kj' + 'lLk*,iLj,li->kj'
#:Lpq = LpqR + LpqI*1j
#:j3c = j3cR + j3cI*1j
#:for i in range(nset):
#: tmp = numpy.dot(dms[i], j3c.reshape(nao,-1))
#: vk1 = numpy.dot(Lpq.reshape(-1,nao).conj().T, tmp.reshape(-1,nao))
#: tmp = numpy.dot(dms[i], Lpq.reshape(nao,-1))
#: vk1+= numpy.dot(j3c.reshape(-1,nao).conj().T, tmp.reshape(-1,nao))
#: vkR[i] += vk1.real
#: vkI[i] += vk1.imag
nrow = pLqR.shape[1]
tmpR = numpy.ndarray((nao,nrow*nao), buffer=buf3R)
if k_real:
for i in range(nset):
tmpR = lib.ddot(dmsR[i], pjqR.reshape(nao,-1), 1, tmpR)
vk1R = lib.ddot(pLqR.reshape(-1,nao).T, tmpR.reshape(-1,nao))
vkR[i] += vk1R
if hermi:
vkR[i] += vk1R.T
else:
tmpR = lib.ddot(dmsR[i], pLqR.reshape(nao,-1), 1, tmpR)
lib.ddot(pjqR.reshape(-1,nao).T, tmpR.reshape(-1,nao),
1, vkR[i], 1)
else:
tmpI = numpy.ndarray((nao,nrow*nao), buffer=buf3I)
for i in range(nset):
tmpR, tmpI = zdotNN(dmsR[i], dmsI[i], pjqR.reshape(nao,-1),
pjqI.reshape(nao,-1), 1, tmpR, tmpI, 0)
vk1R, vk1I = zdotCN(pLqR.reshape(-1,nao).T, pLqI.reshape(-1,nao).T,
tmpR.reshape(-1,nao), tmpI.reshape(-1,nao))
vkR[i] += vk1R
vkI[i] += vk1I
if hermi:
vkR[i] += vk1R.T
vkI[i] -= vk1I.T
else:
tmpR, tmpI = zdotNN(dmsR[i], dmsI[i], pLqR.reshape(nao,-1),
pLqI.reshape(nao,-1), 1, tmpR, tmpI, 0)
zdotCN(pjqR.reshape(-1,nao).T, pjqI.reshape(-1,nao).T,
tmpR.reshape(-1,nao), tmpI.reshape(-1,nao),
1, vkR[i], vkI[i], 1)
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 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)
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 get_jk(mydf, dm, hermi=1, kpt=numpy.zeros(3),
kpt_band=None, with_j=True, with_k=True, exxdiv=None):
'''JK for given k-point'''
from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0
vj = vk = None
if kpt_band is not None and abs(kpt-kpt_band).sum() > 1e-9:
kpt = numpy.reshape(kpt, (1,3))
if with_k:
vk = get_k_kpts(mydf, dm, hermi, kpt, kpt_band, exxdiv)
if with_j:
vj = get_j_kpts(mydf, dm, hermi, kpt, kpt_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)
kptii = numpy.asarray((kpt,kpt))
kpt_allow = numpy.zeros(3)
if with_j:
vjcoulG = mydf.weighted_coulG(kpt_allow, False, mydf.gs)
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, mydf.gs)
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)
bufR = numpy.empty(blksize*nao**2)
bufI = numpy.empty(blksize*nao**2)
for pqkR, pqkI, p0, p1 in mydf.pw_loop(mydf.gs, 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:
coulG = numpy.sqrt(vkcoulG[p0:p1])
pqkR *= coulG
pqkI *= coulG
#: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=bufR).reshape(-1,nao)
pLqI = lib.transpose(pqkI, axes=(0,2,1), out=bufI).reshape(-1,nao)
iLkR = numpy.ndarray((nao*(p1-p0),nao), buffer=pqkR)
iLkI = numpy.ndarray((nao*(p1-p0),nao), buffer=pqkI)
for i in range(nset):
if k_real:
lib.dot(pLqR, dmsR[i], 1, iLkR)
lib.dot(pLqI, dmsR[i], 1, iLkI)
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, iLkI)
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
if cell.dimension != 3 and exxdiv:
assert(exxdiv.lower() == 'ewald')
_ewald_exxdiv_for_G0(cell, kpt, dms, vk)
vk = vk.reshape(dm.shape)
return vj, vk
def get_eri(mydf, kpts=None, compact=True):
if mydf._cderi is None:
mydf.build()
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
nao = cell.nao_nr()
nao_pair = nao * (nao+1) // 2
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*8/1e6)
####################
# gamma point, the integral is real and with s4 symmetry
if abs(kptijkl).sum() < KPT_DIFF_TOL:
eriR = numpy.zeros((nao_pair,nao_pair))
for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, True):
lib.ddot(LpqR.T, LpqR, 1, eriR, 1)
LpqR = LpqI = None
if not compact:
eriR = ao2mo.restore(1, eriR, nao).reshape(nao**2,-1)
return eriR
elif (abs(kpti-kptk).sum() < KPT_DIFF_TOL) and (abs(kptj-kptl).sum() < KPT_DIFF_TOL):
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNN(LpqR.T, LpqI.T, LpqR, LpqI, 1, 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 (abs(kpti-kptl).sum() < KPT_DIFF_TOL) and (abs(kptj-kptk).sum() < KPT_DIFF_TOL):
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNC(LpqR.T, LpqI.T, LpqR, LpqI, 1, 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))
for (LpqR, LpqI), (LrsR, LrsI) in \
lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False),
mydf.sr_loop(kptijkl[2:], max_memory, False)):
zdotNN(LpqR.T, LpqI.T, LrsR, LrsI, 1, eriR, eriI, 1)
LpqR = LpqI = LrsR = LrsI = None
return eriR + eriI*1j
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 = (logger.process_clock(), logger.perf_counter())
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:
vkcoulG = mydf.weighted_coulG(kpt_allow, exxdiv, 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 = pLqR = pLqI = iLkR = iLkI = None
#t2 = log.timer_debug1('%d:%d'%(p0,p1), *t2)
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 get_eri(mydf, kpts=None, compact=True):
if mydf._cderi is None:
mydf.build()
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
nao = cell.nao_nr()
nao_pair = nao * (nao + 1) // 2
max_memory = max(2000, mydf.max_memory - lib.current_memory()[0] - nao ** 4 * 8 / 1e6)
####################
# gamma point, the integral is real and with s4 symmetry
if abs(kptijkl).sum() < KPT_DIFF_TOL:
eriR = numpy.zeros((nao_pair, nao_pair))
for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, True):
lib.ddot(LpqR.T, LpqR, 1, eriR, 1)
LpqR = LpqI = None
if not compact:
eriR = ao2mo.restore(1, eriR, nao).reshape(nao ** 2, -1)
return eriR
elif (abs(kpti - kptk).sum() < KPT_DIFF_TOL) and (abs(kptj - kptl).sum() < KPT_DIFF_TOL):
eriR = numpy.zeros((nao * nao, nao * nao))
eriI = numpy.zeros((nao * nao, nao * nao))
for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNN(LpqR.T, LpqI.T, LpqR, LpqI, 1, 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 (abs(kpti - kptl).sum() < KPT_DIFF_TOL) and (abs(kptj - kptk).sum() < KPT_DIFF_TOL):
eriR = numpy.zeros((nao * nao, nao * nao))
eriI = numpy.zeros((nao * nao, nao * nao))
for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNC(LpqR.T, LpqI.T, LpqR, LpqI, 1, 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))
for (LpqR, LpqI), (LrsR, LrsI) in lib.izip(
mydf.sr_loop(kptijkl[:2], max_memory, False), mydf.sr_loop(kptijkl[2:], max_memory, False)
):
zdotNN(LpqR.T, LpqI.T, LrsR, LrsI, 1, eriR, eriI, 1)
LpqR = LpqI = LrsR = LrsI = None
return eriR + eriI * 1j
def get_jk(mydf, dm, hermi=1, kpt=numpy.zeros(3),
kpt_band=None, with_j=True, with_k=True, exxdiv=None):
'''JK for given k-point'''
from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0
vj = vk = None
if kpt_band is not None and abs(kpt-kpt_band).sum() > 1e-9:
kpt = numpy.reshape(kpt, (1,3))
if with_k:
vk = get_k_kpts(mydf, [dm], hermi, kpt, kpt_band, exxdiv)
if with_j:
vj = get_j_kpts(mydf, [dm], hermi, kpt, kpt_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)
kptii = numpy.asarray((kpt,kpt))
kpt_allow = numpy.zeros(3)
if with_j:
vjcoulG = mydf.weighted_coulG(kpt_allow, False, mydf.gs)
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, mydf.gs)
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)
bufR = numpy.empty(blksize*nao**2)
bufI = numpy.empty(blksize*nao**2)
for pqkR, pqkI, p0, p1 in mydf.pw_loop(mydf.gs, 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:
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:
coulG = numpy.sqrt(vkcoulG[p0:p1])
pqkR *= coulG
pqkI *= coulG
#: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=bufR).reshape(-1,nao)
pLqI = lib.transpose(pqkI, axes=(0,2,1), out=bufI).reshape(-1,nao)
iLkR = numpy.ndarray((nao*(p1-p0),nao), buffer=pqkR)
iLkI = numpy.ndarray((nao*(p1-p0),nao), buffer=pqkI)
for i in range(nset):
if k_real:
lib.dot(pLqR, dmsR[i], 1, iLkR)
lib.dot(pLqI, dmsR[i], 1, iLkI)
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, iLkI)
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('pwdf_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
if cell.dimension != 3 and exxdiv is not None:
assert(exxdiv.lower() == 'ewald')
_ewald_exxdiv_for_G0(cell, kpt, dms, vk)
vk = vk.reshape(dm.shape)
return vj, vk
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)
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)
kk_todo[kpti_idx,kptj_idx] = False
if swap_2e and not is_zero(kpt):
kk_todo[kptj_idx,kpti_idx] = False
# Note: kj-ki for electorn 1 and ki-kj for electron 2
# j2c ~ ({kj-ki}|{ks-kr}) ~ ({kj-ki}|-{kj-ki}) ~ ({kj-ki}|{ki-kj})
# j3c ~ (Q|kj,ki) = j3c{ji} = (Q|ki,kj)* = conj(transpose(j3c{ij}, (0,2,1)))
bufR = numpy.empty((mydf.blockdim*nao**2))
bufI = numpy.empty((mydf.blockdim*nao**2))
for ki,kj in zip(kpti_idx,kptj_idx):
kpti = kpts_band[ki]
kptj = kpts[kj]
kptij = numpy.asarray((kpti,kptj))
for LpqR, LpqI, j3cR, j3cI in mydf.sr_loop(kptij, max_memory, False):
nrow = LpqR.shape[0]
pLqR = numpy.ndarray((nao,nrow,nao), buffer=bufR)
pLqI = numpy.ndarray((nao,nrow,nao), buffer=bufI)
pjqR = numpy.ndarray((nao,nrow,nao), buffer=LpqR)
pjqI = numpy.ndarray((nao,nrow,nao), buffer=LpqI)
tmpR = numpy.ndarray((nao,nrow*nao), buffer=j3cR)
tmpI = numpy.ndarray((nao,nrow*nao), buffer=j3cI)
pLqR[:] = LpqR.reshape(-1,nao,nao).transpose(1,0,2)
pLqI[:] = LpqI.reshape(-1,nao,nao).transpose(1,0,2)
pjqR[:] = j3cR.reshape(-1,nao,nao).transpose(1,0,2)
pjqI[:] = j3cI.reshape(-1,nao,nao).transpose(1,0,2)
#:Lpq = LpqR + LpqI*1j
#:j3c = j3cR + j3cI*1j
#:for i in range(nset):
#: dm = dms[i,ki]
#: tmp = numpy.dot(dm, j3c.reshape(nao,-1))
#: vk1 = numpy.dot(Lpq.reshape(-1,nao).conj().T, tmp.reshape(-1,nao))
#: tmp = numpy.dot(dm, Lpq.reshape(nao,-1))
#: vk1+= numpy.dot(j3c.reshape(-1,nao).conj().T, tmp.reshape(-1,nao))
#: vkR[i,kj] += vk1.real
#: vkI[i,kj] += vk1.imag
#:if swap_2e and not is_zero(kpt):
#: # K ~ 'Lij,Llk*,jk->il' + 'Llk*,Lij,jk->il'
#: for i in range(nset):
#: dm = dms[i,kj]
#: tmp = numpy.dot(j3c.reshape(-1,nao), dm)
#: vk1 = numpy.dot(tmp.reshape(nao,-1), Lpq.reshape(nao,-1).conj().T)
#: tmp = numpy.dot(Lpq.reshape(-1,nao), dm)
#: vk1+= numpy.dot(tmp.reshape(nao,-1), j3c.reshape(nao,-1).conj().T)
#: vkR[i,ki] += vk1.real
#: vkI[i,ki] += vk1.imag
# K ~ 'iLj,lLk*,li->kj' + 'lLk*,iLj,li->kj'
for i in range(nset):
tmpR, tmpI = zdotNN(dmsR[i,ki], dmsI[i,ki], pjqR.reshape(nao,-1),
pjqI.reshape(nao,-1), 1, tmpR, tmpI)
vk1R, vk1I = zdotCN(pLqR.reshape(-1,nao).T, pLqI.reshape(-1,nao).T,
tmpR.reshape(-1,nao), tmpI.reshape(-1,nao))
vkR[i,kj] += vk1R
vkI[i,kj] += vk1I
if hermi:
vkR[i,kj] += vk1R.T
vkI[i,kj] -= vk1I.T
else:
tmpR, tmpI = zdotNN(dmsR[i,ki], dmsI[i,ki], pLqR.reshape(nao,-1),
pLqI.reshape(nao,-1), 1, tmpR, tmpI)
zdotCN(pjqR.reshape(-1,nao).T, pjqI.reshape(-1,nao).T,
tmpR.reshape(-1,nao), tmpI.reshape(-1,nao),
1, vkR[i,kj], vkI[i,kj], 1)
if swap_2e and not is_zero(kpt):
tmpR = tmpR.reshape(nao*nrow,nao)
tmpI = tmpI.reshape(nao*nrow,nao)
# K ~ 'iLj,lLk*,jk->il' + 'lLk*,iLj,jk->il'
for i in range(nset):
tmpR, tmpI = zdotNN(pjqR.reshape(-1,nao), pjqI.reshape(-1,nao),
dmsR[i,kj], dmsI[i,kj], 1, tmpR, tmpI)
vk1R, vk1I = zdotNC(tmpR.reshape(nao,-1), tmpI.reshape(nao,-1),
pLqR.reshape(nao,-1).T, pLqI.reshape(nao,-1).T)
vkR[i,ki] += vk1R
vkI[i,ki] += vk1I
if hermi:
vkR[i,ki] += vk1R.T
vkI[i,ki] -= vk1I.T
else:
tmpR, tmpI = 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),
pjqR.reshape(nao,-1).T, pjqI.reshape(nao,-1).T,
1, vkR[i,ki], vkI[i,ki], 1)
LpqR = LpqI = j3cR = j3cI = tmpR = tmpI = None
return None
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 get_eri(mydf, kpts=None, compact=True):
if mydf._cderi is None:
mydf.build()
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
eri = pwdf_ao2mo.get_eri(mydf, kptijkl, compact=True)
nao = cell.nao_nr()
max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0] - nao**4*8/1e6) * .8)
####################
# gamma point, the integral is real and with s4 symmetry
if abs(kptijkl).sum() < KPT_DIFF_TOL:
eri *= .5 # because we'll do +cc later
for LpqR, LpqI, j3cR, j3cI in mydf.sr_loop(kptijkl[:2], max_memory, True):
lib.ddot(j3cR.T, LpqR, 1, eri, 1)
LpqR = LpqI = j3cR = j3cI = None
eri = lib.transpose_sum(eri, inplace=True)
if not compact:
eri = ao2mo.restore(1, eri, nao).reshape(nao**2,-1)
return eri
####################
# (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 (abs(kpti-kptl).sum() < KPT_DIFF_TOL) and (abs(kptj-kptk).sum() < KPT_DIFF_TOL):
eriR = numpy.zeros((nao*nao,nao*nao))
eriI = numpy.zeros((nao*nao,nao*nao))
for LpqR, LpqI, j3cR, j3cI in mydf.sr_loop(kptijkl[:2], max_memory, False):
zdotNC(j3cR.T, j3cI.T, LpqR, LpqI, 1, eriR, eriI, 1)
# eri == eri.transpose(3,2,1,0).conj()
# zdotNC(LpqR.T, LpqI.T, j3cR, j3cI, 1, eriR, eriI, 1)
LpqR = LpqI = j3cR = j3cI = None
# eri == eri.transpose(3,2,1,0).conj()
eriR = lib.transpose_sum(eriR, inplace=True)
buf = lib.transpose(eriI)
eriI -= buf
eriR = lib.transpose(eriR.reshape(-1,nao,nao), axes=(0,2,1), out=buf)
eri += eriR.reshape(eri.shape)
eriI = lib.transpose(eriI.reshape(-1,nao,nao), axes=(0,2,1), out=buf)
eri += eriI.reshape(eri.shape)*1j
return eri
####################
# 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))
max_memory *= .5
for (LpqR, LpqI, jpqR, jpqI), (LrsR, LrsI, jrsR, jrsI) in \
lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False),
mydf.sr_loop(kptijkl[2:], max_memory, False)):
zdotNN(jpqR.T, jpqI.T, LrsR, LrsI, 1, eriR, eriI, 1)
zdotNN(LpqR.T, LpqI.T, jrsR, jrsI, 1, eriR, eriI, 1)
LpqR = LpqI = jpqR = jpqI = LrsR = LrsI = jrsR = jrsI = None
eri += eriR
eri += eriI*1j
return eri