def get_jk(self,
dm,
hermi=1,
kpts=None,
kpt_band=None,
with_j=True,
with_k=True,
exxdiv='ewald'):
if kpts is None:
if numpy.all(self.kpts == 0):
# Gamma-point calculation by default
kpts = numpy.zeros(3)
else:
kpts = self.kpts
if kpts.shape == (3, ):
return pwdf_jk.get_jk(self, dm, hermi, kpts, kpt_band, with_j,
with_k, exxdiv)
vj = vk = None
if with_k:
vk = pwdf_jk.get_k_kpts(self, dm, hermi, kpts, kpt_band, exxdiv)
if with_j:
vj = pwdf_jk.get_j_kpts(self, dm, hermi, kpts, kpt_band)
return vj, vk
def test_pwdf_j(self):
numpy.random.seed(1)
nao = cell.nao_nr()
dm = numpy.random.random((4,nao,nao))
dm = dm + dm.transpose(0,2,1)
mydf = pwdf.PWDF(cell)
mydf.kpts = numpy.random.random((4,3))
mydf.gs = numpy.asarray((5,)*3)
mydf.auxbasis = 'weigend'
vj = pwdf_jk.get_j_kpts(mydf, dm, 1, mydf.kpts)
self.assertAlmostEqual(finger(vj[0]), (0.93946193432413905+0.00010862804196223034j)/4, 9)
self.assertAlmostEqual(finger(vj[1]), (0.94866073525335626+0.005571199307452865j) /4, 9)
self.assertAlmostEqual(finger(vj[2]), (1.1492194255929766+0.0093705761598793739j) /4, 9)
self.assertAlmostEqual(finger(vj[3]), (1.1397493412770023+0.010731970529096637j) /4, 9)
def test_pwdf_j(self):
numpy.random.seed(1)
nao = cell.nao_nr()
dm = numpy.random.random((4, nao, nao))
dm = dm + dm.transpose(0, 2, 1)
mydf = pwdf.PWDF(cell)
mydf.kpts = numpy.random.random((4, 3))
mydf.gs = numpy.asarray((5, ) * 3)
mydf.auxbasis = 'weigend'
vj = pwdf_jk.get_j_kpts(mydf, dm, 1, mydf.kpts)
self.assertAlmostEqual(finger(
vj[0]), (0.93946193432413905 + 0.00010862804196223034j) / 4, 9)
self.assertAlmostEqual(finger(
vj[1]), (0.94866073525335626 + 0.005571199307452865j) / 4, 9)
self.assertAlmostEqual(finger(
vj[2]), (1.1492194255929766 + 0.0093705761598793739j) / 4, 9)
self.assertAlmostEqual(finger(
vj[3]), (1.1397493412770023 + 0.010731970529096637j) / 4, 9)
def get_jk(self, dm, hermi=1, kpts=None, kpt_band=None,
with_j=True, with_k=True, exxdiv='ewald'):
if kpts is None:
if numpy.all(self.kpts == 0):
# Gamma-point calculation by default
kpts = numpy.zeros(3)
else:
kpts = self.kpts
if kpts.shape == (3,):
return pwdf_jk.get_jk(self, dm, hermi, kpts, kpt_band, with_j,
with_k, exxdiv)
vj = vk = None
if with_k:
vk = pwdf_jk.get_k_kpts(self, dm, hermi, kpts, kpt_band, exxdiv)
if with_j:
vj = pwdf_jk.get_j_kpts(self, dm, hermi, kpts, kpt_band)
return vj, vk
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpt_band=None):
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
if mydf._cderi is None:
mydf.build()
t1 = log.timer_debug1('Init get_j_kpts', *t1)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
if kpt_band is None:
kpts_band = kpts
else:
kpts_band = numpy.reshape(kpt_band, (-1,3))
nband = len(kpts_band)
j_real = gamma_point(kpts_band)
mem_now = lib.current_memory()[0]
max_memory = max(2000, (mydf.max_memory - mem_now)) * .9
log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)
naux = mydf.auxcell.nao_nr()
dmsR = dms.real.transpose(0,1,3,2).reshape(nset,nkpts,nao**2)
dmsI = dms.imag.transpose(0,1,3,2).reshape(nset,nkpts,nao**2)
rhoR = numpy.zeros((nset,naux))
rhoI = numpy.zeros((nset,naux))
jauxR = numpy.zeros((nset,naux))
jauxI = numpy.zeros((nset,naux))
for k, kpt in enumerate(kpts_band):
kptii = numpy.asarray((kpt,kpt))
p1 = 0
for LpqR, LpqI, j3cR, j3cI in mydf.sr_loop(kptii, max_memory, False):
p0, p1 = p1, p1+LpqR.shape[0]
#:Lpq = LpqR + LpqI*1j
#:j3c = j3cR + j3cI*1j
#:rho [:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k])
#:jaux[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', j3c, dms[:,k])
rhoR [:,p0:p1]+= numpy.einsum('Lp,xp->xL', LpqR, dmsR[:,k])
rhoR [:,p0:p1]-= numpy.einsum('Lp,xp->xL', LpqI, dmsI[:,k])
rhoI [:,p0:p1]+= numpy.einsum('Lp,xp->xL', LpqR, dmsI[:,k])
rhoI [:,p0:p1]+= numpy.einsum('Lp,xp->xL', LpqI, dmsR[:,k])
jauxR[:,p0:p1]+= numpy.einsum('Lp,xp->xL', j3cR, dmsR[:,k])
jauxR[:,p0:p1]-= numpy.einsum('Lp,xp->xL', j3cI, dmsI[:,k])
jauxI[:,p0:p1]+= numpy.einsum('Lp,xp->xL', j3cR, dmsI[:,k])
jauxI[:,p0:p1]+= numpy.einsum('Lp,xp->xL', j3cI, dmsR[:,k])
LpqR = LpqI = j3cR = j3cI = None
weight = 1./nkpts
jauxR *= weight
jauxI *= weight
rhoR *= weight
rhoI *= weight
vjR = numpy.zeros((nset,nband,nao,nao))
vjI = numpy.zeros((nset,nband,nao,nao))
for k, kpt in enumerate(kpts_band):
kptii = numpy.asarray((kpt,kpt))
p1 = 0
for LpqR, LpqI, j3cR, j3cI in mydf.sr_loop(kptii, max_memory, True):
p0, p1 = p1, p1+LpqR.shape[0]
#:v = numpy.dot(jaux, Lpq) + numpy.dot(rho, j3c)
#:vj_kpts[:,k] += lib.unpack_tril(v)
v = numpy.dot(jauxR[:,p0:p1], LpqR)
v -= numpy.dot(jauxI[:,p0:p1], LpqI)
v += numpy.dot(rhoR [:,p0:p1], j3cR)
v -= numpy.dot(rhoI [:,p0:p1], j3cI)
vjR[:,k] += lib.unpack_tril(v)
if not j_real:
v = numpy.dot(jauxR[:,p0:p1], LpqI)
v += numpy.dot(jauxI[:,p0:p1], LpqR)
v += numpy.dot(rhoR [:,p0:p1], j3cI)
v += numpy.dot(rhoI [:,p0:p1], j3cR)
vjI[:,k] += lib.unpack_tril(v, lib.ANTIHERMI)
LpqR = LpqI = j3cR = j3cI = None
t1 = log.timer_debug1('get_j pass 2', *t1)
vj_kpts = pwdf_jk.get_j_kpts(mydf, dm_kpts, hermi, kpts, kpt_band)
if j_real:
vj_kpts += vjR.reshape(vj_kpts.shape)
else:
vj_kpts += (vjR+vjI*1j).reshape(vj_kpts.shape)
return vj_kpts
