def fjk(dm):
fmmm = _ri._fpointer('RIhalfmmm_r_s2_bra_noconj')
fdrv = _ao2mo.libao2mo.AO2MOr_e2_drv
ftrans = _ri._fpointer('RItranse2_r_s2')
vj = numpy.zeros_like(dm)
vk = numpy.zeros_like(dm)
fcopy = _ri._fpointer('RImmm_r_s2_transpose')
rargs = (ctypes.c_int(n2c), ctypes.c_int(0), ctypes.c_int(n2c),
ctypes.c_int(0), ctypes.c_int(0))
dmll = numpy.asarray(dm[:n2c, :n2c], order='C')
dmls = numpy.asarray(dm[:n2c, n2c:], order='C') * c1
dmsl = numpy.asarray(dm[n2c:, :n2c], order='C') * c1
dmss = numpy.asarray(dm[n2c:, n2c:], order='C') * c1**2
for erill, eriss in dfobj.loop():
naux, nao_pair = erill.shape
buf = numpy.empty((naux, n2c, n2c), dtype=numpy.complex)
buf1 = numpy.empty((naux, n2c, n2c), dtype=numpy.complex)
fdrv(ftrans, fmmm, buf.ctypes.data_as(ctypes.c_void_p),
erill.ctypes.data_as(ctypes.c_void_p),
dmll.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux),
*rargs) # buf == (P|LL)
rho = numpy.einsum('kii->k', buf)
fdrv(ftrans, fcopy, buf1.ctypes.data_as(ctypes.c_void_p),
erill.ctypes.data_as(ctypes.c_void_p),
dmll.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux),
*rargs) # buf1 == (P|LL)
vk[:n2c, :n2c] += numpy.dot(
buf1.reshape(-1, n2c).T, buf.reshape(-1, n2c))
fdrv(ftrans, fmmm, buf.ctypes.data_as(ctypes.c_void_p),
eriss.ctypes.data_as(ctypes.c_void_p),
dmls.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux),
*rargs) # buf == (P|LS)
vk[:n2c, n2c:] += numpy.dot(
buf1.reshape(-1, n2c).T, buf.reshape(-1, n2c)) * c1
fdrv(ftrans, fmmm, buf.ctypes.data_as(ctypes.c_void_p),
eriss.ctypes.data_as(ctypes.c_void_p),
dmss.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux),
*rargs) # buf == (P|SS)
rho += numpy.einsum('kii->k', buf)
vj[:n2c, :n2c] += pyscf.lib.unpack_tril(numpy.dot(rho, erill), 1)
vj[n2c:,
n2c:] += pyscf.lib.unpack_tril(numpy.dot(rho, eriss), 1) * c1**2
fdrv(ftrans, fcopy, buf1.ctypes.data_as(ctypes.c_void_p),
eriss.ctypes.data_as(ctypes.c_void_p),
dmss.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux),
*rargs) # buf == (P|SS)
vk[n2c:, n2c:] += numpy.dot(
buf1.reshape(-1, n2c).T, buf.reshape(-1, n2c)) * c1**2
if not hermi == 1:
fdrv(ftrans, fmmm, buf.ctypes.data_as(ctypes.c_void_p),
erill.ctypes.data_as(ctypes.c_void_p),
dmsl.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux),
*rargs) # buf == (P|SL)
vk[n2c:, :n2c] += numpy.dot(
buf1.reshape(-1, n2c).T, buf.reshape(-1, n2c)) * c1
if hermi == 1:
vk[n2c:, :n2c] = vk[:n2c, n2c:].T.conj()
return vj, vk
def fjk(dm):
fmmm = _ri._fpointer('RIhalfmmm_r_s2_bra_noconj')
fdrv = _ao2mo.libao2mo.AO2MOr_e2_drv
ftrans = _ri._fpointer('RItranse2_r_s2')
vj = numpy.zeros_like(dm)
vk = numpy.zeros_like(dm)
fcopy = _ri._fpointer('RImmm_r_s2_transpose')
rargs = (ctypes.c_int(n2c),
ctypes.c_int(0), ctypes.c_int(n2c),
ctypes.c_int(0), ctypes.c_int(0))
dmll = numpy.asarray(dm[:n2c,:n2c], order='C')
dmls = numpy.asarray(dm[:n2c,n2c:], order='C') * c1
dmsl = numpy.asarray(dm[n2c:,:n2c], order='C') * c1
dmss = numpy.asarray(dm[n2c:,n2c:], order='C') * c1**2
for erill, eriss in dfobj.loop():
naux, nao_pair = erill.shape
buf = numpy.empty((naux,n2c,n2c), dtype=numpy.complex)
buf1 = numpy.empty((naux,n2c,n2c), dtype=numpy.complex)
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
erill.ctypes.data_as(ctypes.c_void_p),
dmll.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs) # buf == (P|LL)
rho = numpy.einsum('kii->k', buf)
fdrv(ftrans, fcopy,
buf1.ctypes.data_as(ctypes.c_void_p),
erill.ctypes.data_as(ctypes.c_void_p),
dmll.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs) # buf1 == (P|LL)
vk[:n2c,:n2c] += numpy.dot(buf1.reshape(-1,n2c).T,
buf.reshape(-1,n2c))
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
eriss.ctypes.data_as(ctypes.c_void_p),
dmls.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs) # buf == (P|LS)
vk[:n2c,n2c:] += numpy.dot(buf1.reshape(-1,n2c).T,
buf.reshape(-1,n2c)) * c1
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
eriss.ctypes.data_as(ctypes.c_void_p),
dmss.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs) # buf == (P|SS)
rho += numpy.einsum('kii->k', buf)
vj[:n2c,:n2c] += pyscf.lib.unpack_tril(numpy.dot(rho, erill), 1)
vj[n2c:,n2c:] += pyscf.lib.unpack_tril(numpy.dot(rho, eriss), 1) * c1**2
fdrv(ftrans, fcopy,
buf1.ctypes.data_as(ctypes.c_void_p),
eriss.ctypes.data_as(ctypes.c_void_p),
dmss.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs) # buf == (P|SS)
vk[n2c:,n2c:] += numpy.dot(buf1.reshape(-1,n2c).T,
buf.reshape(-1,n2c)) * c1**2
if not hermi == 1:
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
erill.ctypes.data_as(ctypes.c_void_p),
dmsl.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs) # buf == (P|SL)
vk[n2c:,:n2c] += numpy.dot(buf1.reshape(-1,n2c).T,
buf.reshape(-1,n2c)) * c1
if hermi == 1:
vk[n2c:,:n2c] = vk[:n2c,n2c:].T.conj()
return vj, vk
def get_jk(dfobj, mol, dms, hermi=1, vhfopt=None, with_j=True, with_k=True):
t0 = t1 = (time.clock(), time.time())
log = logger.Logger(dfobj.stdout, dfobj.verbose)
if len(dms) == 0:
return [], []
elif isinstance(dms, numpy.ndarray) and dms.ndim == 2:
nset = 1
dms = [dms]
else:
nset = len(dms)
nao = dms[0].shape[0]
fmmm = _ri._fpointer('RIhalfmmm_nr_s2_bra')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2')
vj = numpy.zeros((nset, nao, nao))
vk = numpy.zeros((nset, nao, nao))
null = pyscf.lib.c_null_ptr()
#:vj = reduce(numpy.dot, (cderi.reshape(-1,nao*nao), dm.reshape(-1),
#: cderi.reshape(-1,nao*nao))).reshape(nao,nao)
if hermi == 1: # and numpy.einsum('ij,ij->', dm, ovlp) > 0.1
# I cannot assume dm is positive definite because it might be the density
# matrix difference when the mf.direct_scf flag is set.
dmtril = []
cpos = []
cneg = []
for k, dm in enumerate(dms):
if with_j:
dmtril.append(pyscf.lib.pack_tril(dm + dm.T))
i = numpy.arange(nao)
dmtril[k][i * (i + 1) // 2 + i] *= .5
if with_k:
e, c = scipy.linalg.eigh(dm)
pos = e > OCCDROP
neg = e < -OCCDROP
#:vk = numpy.einsum('pij,jk->kpi', cderi, c[:,abs(e)>OCCDROP])
#:vk = numpy.einsum('kpi,kpj->ij', vk, vk)
tmp = numpy.einsum('ij,j->ij', c[:, pos], numpy.sqrt(e[pos]))
cpos.append(numpy.asarray(tmp, order='F'))
tmp = numpy.einsum('ij,j->ij', c[:, neg], numpy.sqrt(-e[neg]))
cneg.append(numpy.asarray(tmp, order='F'))
buf = numpy.empty((dfobj.blockdim * nao, nao))
for eri1 in dfobj.loop():
naux, nao_pair = eri1.shape
assert (nao_pair == nao * (nao + 1) // 2)
for k in range(nset):
if with_j:
buf1 = reduce(numpy.dot, (eri1, dmtril[k], eri1))
vj[k] += pyscf.lib.unpack_tril(buf1, hermi)
if with_k and cpos[k].shape[1] > 0:
buf1 = buf[:naux * cpos[k].shape[1]]
fdrv(ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
cpos[k].ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux),
ctypes.c_int(nao), ctypes.c_int(0),
ctypes.c_int(cpos[k].shape[1]), ctypes.c_int(0),
ctypes.c_int(0), null, ctypes.c_int(0))
vk[k] += pyscf.lib.dot(buf1.T, buf1)
if with_k and cneg[k].shape[1] > 0:
buf1 = buf[:naux * cneg[k].shape[1]]
fdrv(ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
cneg[k].ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux),
ctypes.c_int(nao), ctypes.c_int(0),
ctypes.c_int(cneg[k].shape[1]), ctypes.c_int(0),
ctypes.c_int(0), null, ctypes.c_int(0))
vk[k] -= pyscf.lib.dot(buf1.T, buf1)
t1 = log.timer_debug1('jk', *t1)
else:
#:vk = numpy.einsum('pij,jk->pki', cderi, dm)
#:vk = numpy.einsum('pki,pkj->ij', cderi, vk)
fcopy = _ri._fpointer('RImmm_nr_s2_copy')
rargs = (ctypes.c_int(nao), ctypes.c_int(0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(0), null, ctypes.c_int(0))
dms = [numpy.asarray(dm, order='F') for dm in dms]
buf = numpy.empty((2, dfobj.blockdim, nao, nao))
for eri1 in dfobj.loop():
naux, nao_pair = eri1.shape
for k in range(nset):
buf1 = buf[0, :naux]
fdrv(ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
dms[k].ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs)
rho = numpy.einsum('kii->k', buf1)
vj[k] += pyscf.lib.unpack_tril(numpy.dot(rho, eri1), 1)
if with_k:
buf2 = buf[1, :naux]
fdrv(ftrans, fcopy, buf2.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
dms[k].ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs)
vk[k] += pyscf.lib.dot(
buf1.reshape(-1, nao).T, buf2.reshape(-1, nao))
t1 = log.timer_debug1('jk', *t1)
if len(dms) == 1:
vj = vj[0]
vk = vk[0]
logger.timer(dfobj, 'vj and vk', *t0)
return vj, vk
def get_jk(dfobj, mol, dms, hermi=1, vhfopt=None, with_j=True, with_k=True):
t0 = t1 = (time.clock(), time.time())
log = logger.Logger(dfobj.stdout, dfobj.verbose)
if len(dms) == 0:
return [], []
elif isinstance(dms, numpy.ndarray) and dms.ndim == 2:
nset = 1
dms = [dms]
else:
nset = len(dms)
nao = dms[0].shape[0]
fmmm = _ri._fpointer('RIhalfmmm_nr_s2_bra')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2')
vj = numpy.zeros((nset,nao,nao))
vk = numpy.zeros((nset,nao,nao))
null = pyscf.lib.c_null_ptr()
#:vj = reduce(numpy.dot, (cderi.reshape(-1,nao*nao), dm.reshape(-1),
#: cderi.reshape(-1,nao*nao))).reshape(nao,nao)
if hermi == 1: # and numpy.einsum('ij,ij->', dm, ovlp) > 0.1
# I cannot assume dm is positive definite because it might be the density
# matrix difference when the mf.direct_scf flag is set.
dmtril = []
cpos = []
cneg = []
for k, dm in enumerate(dms):
if with_j:
dmtril.append(pyscf.lib.pack_tril(dm+dm.T))
i = numpy.arange(nao)
dmtril[k][i*(i+1)//2+i] *= .5
if with_k:
e, c = scipy.linalg.eigh(dm)
pos = e > OCCDROP
neg = e < -OCCDROP
#:vk = numpy.einsum('pij,jk->kpi', cderi, c[:,abs(e)>OCCDROP])
#:vk = numpy.einsum('kpi,kpj->ij', vk, vk)
tmp = numpy.einsum('ij,j->ij', c[:,pos], numpy.sqrt(e[pos]))
cpos.append(numpy.asarray(tmp, order='F'))
tmp = numpy.einsum('ij,j->ij', c[:,neg], numpy.sqrt(-e[neg]))
cneg.append(numpy.asarray(tmp, order='F'))
buf = numpy.empty((dfobj.blockdim*nao,nao))
for eri1 in dfobj.loop():
naux, nao_pair = eri1.shape
assert(nao_pair == nao*(nao+1)//2)
for k in range(nset):
if with_j:
buf1 = reduce(numpy.dot, (eri1, dmtril[k], eri1))
vj[k] += pyscf.lib.unpack_tril(buf1, hermi)
if with_k and cpos[k].shape[1] > 0:
buf1 = buf[:naux*cpos[k].shape[1]]
fdrv(ftrans, fmmm,
buf1.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
cpos[k].ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(cpos[k].shape[1]),
ctypes.c_int(0), ctypes.c_int(0), null, ctypes.c_int(0))
vk[k] += pyscf.lib.dot(buf1.T, buf1)
if with_k and cneg[k].shape[1] > 0:
buf1 = buf[:naux*cneg[k].shape[1]]
fdrv(ftrans, fmmm,
buf1.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
cneg[k].ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(cneg[k].shape[1]),
ctypes.c_int(0), ctypes.c_int(0), null, ctypes.c_int(0))
vk[k] -= pyscf.lib.dot(buf1.T, buf1)
t1 = log.timer_debug1('jk', *t1)
else:
#:vk = numpy.einsum('pij,jk->pki', cderi, dm)
#:vk = numpy.einsum('pki,pkj->ij', cderi, vk)
fcopy = _ri._fpointer('RImmm_nr_s2_copy')
rargs = (ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(0), null, ctypes.c_int(0))
dms = [numpy.asarray(dm, order='F') for dm in dms]
buf = numpy.empty((2,dfobj.blockdim,nao,nao))
for eri1 in dfobj.loop():
naux, nao_pair = eri1.shape
for k in range(nset):
buf1 = buf[0,:naux]
fdrv(ftrans, fmmm,
buf1.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
dms[k].ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs)
rho = numpy.einsum('kii->k', buf1)
vj[k] += pyscf.lib.unpack_tril(numpy.dot(rho, eri1), 1)
if with_k:
buf2 = buf[1,:naux]
fdrv(ftrans, fcopy,
buf2.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
dms[k].ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(naux), *rargs)
vk[k] += pyscf.lib.dot(buf1.reshape(-1,nao).T,
buf2.reshape(-1,nao))
t1 = log.timer_debug1('jk', *t1)
if len(dms) == 1:
vj = vj[0]
vk = vk[0]
logger.timer(dfobj, 'vj and vk', *t0)
return vj, vk