def ao2mo(self, mo_coeff, nocc):
time0 = (time.clock(), time.time())
log = logger.Logger(self.stdout, self.verbose)
cderi_file = tempfile.NamedTemporaryFile()
df.outcore.general(self.mol, (mo_coeff[:,:nocc], mo_coeff[:,nocc:]),
cderi_file.name, auxbasis=self.auxbasis, verbose=log)
time1 = log.timer('Integral transformation (P|ia)', *time0)
return df.load(cderi_file)
def ao2mo(self, mo_coeff, nocc):
time0 = (time.clock(), time.time())
log = logger.Logger(self.stdout, self.verbose)
cderi_file = tempfile.NamedTemporaryFile()
df.outcore.general(self.mol, (mo_coeff[:, :nocc], mo_coeff[:, nocc:]),
cderi_file.name,
auxbasis=self.auxbasis,
verbose=log)
time1 = log.timer('Integral transformation (P|ia)', *time0)
return df.load(cderi_file)
def ao2mo_(casscf, mo):
t0 = (time.clock(), time.time())
log = logger.Logger(casscf.stdout, casscf.verbose)
# using dm=[], a hacky call to dfhf.get_jk, to generate casscf._cderi
dfhf.get_jk_(casscf, casscf.mol, [])
if log.verbose >= logger.DEBUG1:
t1 = log.timer('Generate density fitting integrals', *t0)
if hasattr(casscf._scf, '_tag_df') and casscf._scf._tag_df:
eris = _ERIS(casscf, mo)
else:
# Only approximate the orbital rotation, call the 4-center integral
# transformation. CASSCF is exact.
eris = mc_ao2mo._ERIS(casscf, mo, 'incore', level=2)
t0 = (time.clock(), time.time())
mo = numpy.asarray(mo, order='F')
nao, nmo = mo.shape
ncore = casscf.ncore
eris.j_pc = numpy.zeros((nmo,ncore))
k_cp = numpy.zeros((ncore,nmo))
fmmm = _ao2mo._fpointer('AO2MOmmm_nr_s2_iltj')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2')
bufs1 = numpy.empty((dfhf.BLOCKDIM,nmo,nmo))
with df.load(casscf._cderi) as feri:
for b0, b1 in dfhf.prange(0, casscf._naoaux, dfhf.BLOCKDIM):
eri1 = numpy.asarray(feri[b0:b1], order='C')
buf = bufs1[:b1-b0]
if log.verbose >= logger.DEBUG1:
t1 = log.timer('load buf %d:%d'%(b0,b1), *t1)
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
mo.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(b1-b0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_void_p(0), ctypes.c_int(0))
if log.verbose >= logger.DEBUG1:
t1 = log.timer('transform [%d:%d]'%(b0,b1), *t1)
bufd = numpy.einsum('kii->ki', buf)
eris.j_pc += numpy.einsum('ki,kj->ij', bufd, bufd[:,:ncore])
k_cp += numpy.einsum('kij,kij->ij', buf[:,:ncore], buf[:,:ncore])
if log.verbose >= logger.DEBUG1:
t1 = log.timer('j_pc and k_pc', *t1)
eri1 = None
eris.k_pc = k_cp.T.copy()
log.timer('ao2mo density fit part', *t0)
return eris
def test_df_ao2mo(self):
mf = scf.density_fit(msym)
mf.max_memory = 100
mf.kernel()
mc = mcscf.DFCASSCF(mf, 4, 4)
with df.load(mf._cderi) as feri:
cderi = numpy.asarray(feri)
eri0 = numpy.dot(cderi.T, cderi)
nmo = mc.mo_coeff.shape[1]
ncore = mc.ncore
nocc = ncore + mc.ncas
eri0 = ao2mo.restore(1, ao2mo.kernel(eri0, mc.mo_coeff), nmo)
eris = mc.ao2mo(mc.mo_coeff)
self.assertTrue(numpy.allclose(eri0[:,:,ncore:nocc,ncore:nocc], eris.ppaa))
self.assertTrue(numpy.allclose(eri0[:,ncore:nocc,:,ncore:nocc], eris.papa))
def ao2mo(self, mo):
t0 = (time.clock(), time.time())
ncore = self.ncore
log = pyscf.lib.logger.Logger(self.stdout, self.verbose)
# using dm=[], a hacky call to dfhf.get_jk, to generate self._cderi
self.get_jk(self.mol, [])
if log.verbose >= pyscf.lib.logger.DEBUG1:
t1 = log.timer('Generate density fitting integrals', *t0)
eris = mc_ao2mo._ERIS(self, mo, 'incore', level=2)
t0 = (time.clock(), time.time())
mo = numpy.asarray(mo, order='F')
nao, nmo = mo.shape
eris.j_pc = numpy.zeros((nmo,ncore))
k_cp = numpy.zeros((ncore,nmo))
fmmm = _ao2mo._fpointer('AO2MOmmm_nr_s2_iltj')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2kl')
bufs1 = numpy.empty((dfhf.BLOCKDIM,nmo,nmo))
with df.load(self._cderi) as feri:
for b0, b1 in dfhf.prange(0, self._naoaux, dfhf.BLOCKDIM):
eri1 = numpy.array(feri[b0:b1], copy=False)
buf = bufs1[:b1-b0]
if log.verbose >= pyscf.lib.logger.DEBUG1:
t1 = log.timer('load buf %d:%d'%(b0,b1), *t1)
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
mo.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(b1-b0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_void_p(0), ctypes.c_int(0))
if log.verbose >= pyscf.lib.logger.DEBUG1:
t1 = log.timer('transform [%d:%d]'%(b0,b1), *t1)
bufd = numpy.einsum('kii->ki', buf).copy()
#:eris.j_pc += numpy.einsum('ki,kj->ij', bufd, bufd[:,:ncore])
pyscf.lib.dot(bufd.T, numpy.asarray(bufd[:,:ncore],order='C'),
1, eris.j_pc, 1)
k_cp += numpy.einsum('kij,kij->ij', buf[:,:ncore], buf[:,:ncore])
if log.verbose >= pyscf.lib.logger.DEBUG1:
t1 = log.timer('j_pc and k_pc', *t1)
eri1 = None
eris.k_pc = k_cp.T.copy()
log.timer('ao2mo density fit part', *t0)
return eris
def ao2mo(self, mo_coeff):
log = logger.Logger(self.stdout, self.verbose)
# the exact integral transformation
eris = casscf_class.ao2mo(self, mo_coeff)
# using dm=[], a hacky call to dfhf.get_jk, to generate self._cderi
dfhf.get_jk_(self, self.mol, [])
# Add the approximate diagonal term for orbital hessian
t0 = (time.clock(), time.time())
mo = numpy.asarray(mo_coeff, order='F')
nao, nmo = mo.shape
ncore = self.ncore
eris.j_pc = numpy.zeros((nmo,ncore))
k_cp = numpy.zeros((ncore,nmo))
fmmm = _ao2mo._fpointer('AO2MOmmm_nr_s2_iltj')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2')
bufs1 = numpy.empty((dfhf.BLOCKDIM,nmo,nmo))
with df.load(self._cderi) as feri:
for b0, b1 in dfhf.prange(0, self._naoaux, dfhf.BLOCKDIM):
eri1 = numpy.asarray(feri[b0:b1], order='C')
buf = bufs1[:b1-b0]
if log.verbose >= logger.DEBUG1:
t1 = log.timer('load buf %d:%d'%(b0,b1), *t1)
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
mo.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(b1-b0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_void_p(0), ctypes.c_int(0))
if log.verbose >= logger.DEBUG1:
t1 = log.timer('transform [%d:%d]'%(b0,b1), *t1)
bufd = numpy.einsum('kii->ki', buf)
eris.j_pc += numpy.einsum('ki,kj->ij', bufd, bufd[:,:ncore])
k_cp += numpy.einsum('kij,kij->ij', buf[:,:ncore], buf[:,:ncore])
if log.verbose >= logger.DEBUG1:
t1 = log.timer('j_pc and k_pc', *t1)
eri1 = None
eris.k_pc = k_cp.T.copy()
log.timer('ao2mo density fit part', *t0)
return eris
def ao2mo(self, mo):
ncore = self.ncore
#self._cderi = None # FIXME? leave as much memory as possible for mc_ao2mo
eris = mc_ao2mo._ERIS(self, mo, 'incore', 2)
# using dm=[], a hacky call to dfhf.get_jk, to generate self._cderi
t0 = (time.clock(), time.time())
log = pyscf.lib.logger.Logger(self.stdout, self.verbose)
self.get_jk(self.mol, [])
if log.verbose >= pyscf.lib.logger.DEBUG1:
t1 = log.timer('Generate density fitting integrals', *t0)
mo = numpy.asarray(mo, order='F')
nao, nmo = mo.shape
eris.j_cp = numpy.zeros((ncore,nmo))
eris.k_cp = numpy.zeros((ncore,nmo))
fmmm = _ao2mo._fpointer('AO2MOmmm_nr_s2_iltj')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2kl')
with df.load(self._cderi) as feri:
for b0, b1 in dfhf.prange(0, self._naoaux, dfhf.BLOCKDIM):
eri1 = numpy.array(feri[b0:b1], copy=False)
buf = numpy.empty((b1-b0,nmo,nmo))
if log.verbose >= pyscf.lib.logger.DEBUG1:
t1 = log.timer('load buf %d:%d'%(b0,b1), *t1)
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
mo.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(b1-b0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_void_p(0), ctypes.c_int(0))
if log.verbose >= pyscf.lib.logger.DEBUG1:
t1 = log.timer('transform [%d:%d]'%(b0,b1), *t1)
bufd = numpy.einsum('kii->ki', buf).copy()
#:eris.j_cp += numpy.einsum('ki,kj->ij', bufd[:,:ncore], bufd)
pyscf.lib.dot(bufd[:,:ncore].T.copy(), bufd, 1, eris.j_cp, 1)
eris.k_cp += numpy.einsum('kij,kij->ij', buf[:,:ncore], buf[:,:ncore])
if log.verbose >= pyscf.lib.logger.DEBUG1:
t1 = log.timer('j_cp and k_cp', *t1)
return eris
def ao2mo_aaaa(casscf, mo):
dfhf.get_jk_(casscf, casscf.mol, [])
nao, nmo = mo.shape
buf = numpy.empty((casscf._naoaux,nmo*(nmo+1)//2))
mo = numpy.asarray(mo, order='F')
fmmm = _ao2mo._fpointer('AO2MOmmm_nr_s2_s2')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2')
with df.load(casscf._cderi) as feri:
for b0, b1 in dfhf.prange(0, casscf._naoaux, dfhf.BLOCKDIM):
eri1 = numpy.asarray(feri[b0:b1], order='C')
fdrv(ftrans, fmmm,
buf[b0:b1].ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
mo.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(b1-b0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_void_p(0), ctypes.c_int(0))
eri1 = None
eri = pyscf.lib.dot(buf.T, buf)
return eri
def get_jk_(mf, mol, dms, hermi=1, with_j=True, with_k=True):
t0 = (time.clock(), time.time())
log = logger.Logger(mf.stdout, mf.verbose)
if not hasattr(mf, '_cderi') or mf._cderi is None:
nao = mol.nao_nr()
nao_pair = nao*(nao+1)//2
auxmol = df.incore.format_aux_basis(mol, mf.auxbasis)
mf._naoaux = auxmol.nao_nr()
if (nao_pair*mf._naoaux*8/1e6*2+pyscf.lib.current_memory()[0]
< mf.max_memory*.8):
mf._cderi = df.incore.cholesky_eri(mol, auxbasis=mf.auxbasis,
verbose=log)
else:
mf._cderi = tempfile.NamedTemporaryFile()
df.outcore.cholesky_eri(mol, mf._cderi.name, auxbasis=mf.auxbasis,
verbose=log)
# if (nao_pair*mf._naoaux*8/1e6+pyscf.lib.current_memory()[0]
# < mf.max_memory*.9):
# with df.load(mf._cderi) as feri:
# cderi = numpy.asarray(feri)
# mf._cderi = cderi
if mf._naoaux is None:
# By overwriting mf._cderi, one can provide the Cholesky integrals for "DF/RI" calculation
with df.load(mf._cderi) as feri:
mf._naoaux = feri.shape[0]
if len(dms) == 0:
return [], []
cderi = mf._cderi
nao = mol.nao_nr()
fmmm = df.incore._fpointer('RIhalfmmm_nr_s2_bra')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2')
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
dms = [dms]
nset = 1
else:
nset = len(dms)
vj = numpy.zeros((nset,nao,nao))
vk = numpy.zeros((nset,nao,nao))
#:vj = reduce(numpy.dot, (cderi.reshape(-1,nao*nao), dm.reshape(-1),
#: cderi.reshape(-1,nao*nao))).reshape(nao,nao)
if hermi == 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))
for i in range(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'))
if mf.verbose >= logger.DEBUG1:
t1 = log.timer('Initialization', *t0)
with df.load(cderi) as feri:
buf = numpy.empty((BLOCKDIM*nao,nao))
for b0, b1 in prange(0, mf._naoaux, BLOCKDIM):
eri1 = numpy.array(feri[b0:b1], copy=False)
if mf.verbose >= logger.DEBUG1:
t1 = log.timer('load buf %d:%d'%(b0,b1), *t1)
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[:(b1-b0)*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(b1-b0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(cpos[k].shape[1]),
ctypes.c_int(0), ctypes.c_int(0))
vk[k] += pyscf.lib.dot(buf1.T, buf1)
if with_k and cneg[k].shape[1] > 0:
buf1 = buf[:(b1-b0)*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(b1-b0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(cneg[k].shape[1]),
ctypes.c_int(0), ctypes.c_int(0))
vk[k] -= pyscf.lib.dot(buf1.T, buf1)
if mf.verbose >= logger.DEBUG1:
t1 = log.timer('jk', *t1)
else:
#:vk = numpy.einsum('pij,jk->pki', cderi, dm)
#:vk = numpy.einsum('pki,pkj->ij', cderi, vk)
fcopy = df.incore._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))
dms = [numpy.asarray(dm, order='F') for dm in dms]
if mf.verbose >= logger.DEBUG1:
t1 = log.timer('Initialization', *t0)
with df.load(cderi) as feri:
buf = numpy.empty((2,BLOCKDIM,nao,nao))
for b0, b1 in prange(0, mf._naoaux, BLOCKDIM):
eri1 = numpy.array(feri[b0:b1], copy=False)
if mf.verbose >= logger.DEBUG1:
t1 = log.timer('load buf %d:%d'%(b0,b1), *t1)
for k in range(nset):
buf1 = buf[0,:b1-b0]
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(b1-b0), *rargs)
rho = numpy.einsum('kii->k', buf1)
vj[k] += pyscf.lib.unpack_tril(numpy.dot(rho, eri1), 1)
if with_k:
buf2 = buf[1,:b1-b0]
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(b1-b0), *rargs)
vk[k] += pyscf.lib.dot(buf1.reshape(-1,nao).T,
buf2.reshape(-1,nao))
if mf.verbose >= logger.DEBUG1:
t1 = log.timer('jk', *t1)
if len(dms) == 1:
vj = vj[0]
vk = vk[0]
logger.timer(mf, 'vj and vk', *t0)
return vj, vk
def fjk(dm):
fmmm = df.r_incore._fpointer('RIhalfmmm_r_s2_bra_noconj')
fdrv = _ao2mo.libao2mo.AO2MOr_e2_drv
ftrans = df.r_incore._fpointer('RItranse2_r_s2')
vj = numpy.zeros_like(dm)
vk = numpy.zeros_like(dm)
fcopy = df.incore._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
with df.load(mf._cderi[0]) as ferill:
with df.load(mf._cderi[1]) as feriss: # python2.6 not support multiple with
for b0, b1 in prange(0, mf._naoaux, BLOCKDIM):
erill = numpy.array(ferill[b0:b1], copy=False)
eriss = numpy.array(feriss[b0:b1], copy=False)
buf = numpy.empty((b1-b0,n2c,n2c), dtype=numpy.complex)
buf1 = numpy.empty((b1-b0,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(b1-b0), *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(b1-b0), *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(b1-b0), *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(b1-b0), *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(b1-b0), *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(b1-b0), *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 __init__(self, casscf, mo):
assert(casscf._scf._tag_df)
import gc
gc.collect()
log = logger.Logger(casscf.stdout, casscf.verbose)
mol = casscf.mol
nao, nmo = mo.shape
ncore = casscf.ncore
ncas = casscf.ncas
nocc = ncore + ncas
naoaux = casscf._naoaux
mem_incore, mem_outcore, mem_basic = _mem_usage(ncore, ncas, nmo)
mem_now = pyscf.lib.current_memory()[0]
max_memory = max(3000, casscf.max_memory*.9-mem_now)
if max_memory < mem_basic:
log.warn('Calculation needs %d MB memory, over CASSCF.max_memory (%d MB) limit',
(mem_basic+mem_now)/.9, casscf.max_memory)
t0 = (time.clock(), time.time())
self._tmpfile = tempfile.NamedTemporaryFile()
self.feri = h5py.File(self._tmpfile.name, 'w')
self.ppaa = self.feri.create_dataset('ppaa', (nmo,nmo,ncas,ncas), 'f8')
self.papa = self.feri.create_dataset('papa', (nmo,ncas,nmo,ncas), 'f8')
self.j_pc = numpy.zeros((nmo,ncore))
k_cp = numpy.zeros((ncore,nmo))
mo = numpy.asarray(mo, order='F')
_tmpfile1 = tempfile.NamedTemporaryFile()
fxpp = h5py.File(_tmpfile1.name)
bufpa = numpy.empty((naoaux,nmo,ncas))
bufs1 = numpy.empty((dfhf.BLOCKDIM,nmo,nmo))
fmmm = _ao2mo._fpointer('AO2MOmmm_nr_s2_iltj')
fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv
ftrans = _ao2mo._fpointer('AO2MOtranse2_nr_s2')
t2 = t1 = t0
fxpp_keys = []
with df.load(casscf._cderi) as feri:
for b0, b1 in dfhf.prange(0, naoaux, dfhf.BLOCKDIM):
eri1 = numpy.asarray(feri[b0:b1], order='C')
if log.verbose >= logger.DEBUG1:
t2 = log.timer('load buf %d:%d'%(b0,b1), *t2)
bufpp = bufs1[:b1-b0]
fdrv(ftrans, fmmm,
bufpp.ctypes.data_as(ctypes.c_void_p),
eri1.ctypes.data_as(ctypes.c_void_p),
mo.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(b1-b0), ctypes.c_int(nao),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_int(0), ctypes.c_int(nmo),
ctypes.c_void_p(0), ctypes.c_int(0))
fxpp_keys.append([str(b0), b0, b1])
fxpp[str(b0)] = bufpp.transpose(1,2,0)
bufpa[b0:b1] = bufpp[:,:,ncore:nocc]
bufd = numpy.einsum('kii->ki', bufpp)
self.j_pc += numpy.einsum('ki,kj->ij', bufd, bufd[:,:ncore])
k_cp += numpy.einsum('kij,kij->ij', bufpp[:,:ncore], bufpp[:,:ncore])
if log.verbose >= logger.DEBUG1:
t1 = log.timer('j_pc and k_pc', *t1)
eri1 = None
self.k_pc = k_cp.T.copy()
bufs1 = bufpp = None
t1 = log.timer('density fitting ao2mo pass1', *t0)
mem_now = pyscf.lib.current_memory()[0]
nblk = int(max(8, min(nmo, ((max_memory-mem_now)*1e6/8-bufpa.size)/(ncas**2*nmo))))
bufs1 = numpy.empty((nblk,ncas,nmo,ncas))
dgemm = pyscf.lib.numpy_helper._dgemm
for p0, p1 in prange(0, nmo, nblk):
#tmp = numpy.dot(bufpa[:,p0:p1].reshape(naoaux,-1).T,
# bufpa.reshape(naoaux,-1))
tmp = bufs1[:p1-p0]
dgemm('T', 'N', (p1-p0)*ncas, nmo*ncas, naoaux,
bufpa.reshape(naoaux,-1), bufpa.reshape(naoaux,-1),
tmp.reshape(-1,nmo*ncas), 1, 0, p0*ncas, 0, 0)
self.papa[p0:p1] = tmp.reshape(p1-p0,ncas,nmo,ncas)
bufaa = bufpa[:,ncore:nocc,:].copy().reshape(-1,ncas**2)
bufs1 = bufpa = None
t1 = log.timer('density fitting papa pass2', *t1)
mem_now = pyscf.lib.current_memory()[0]
nblk = int(max(8, min(nmo, (max_memory-mem_now)*1e6/8/(nmo*naoaux+ncas**2*nmo))))
bufs1 = numpy.empty((nblk,nmo,naoaux))
bufs2 = numpy.empty((nblk,nmo,ncas,ncas))
for p0, p1 in prange(0, nmo, nblk):
nrow = p1 - p0
buf = bufs1[:nrow]
tmp = bufs2[:nrow].reshape(-1,ncas**2)
col0 = 0
for key, col0, col1 in fxpp_keys:
buf[:nrow,:,col0:col1] = fxpp[key][p0:p1]
pyscf.lib.dot(buf.reshape(-1,naoaux), bufaa, 1, tmp)
self.ppaa[p0:p1] = tmp.reshape(p1-p0,nmo,ncas,ncas)
bufs1 = bufs2 = buf = None
t1 = log.timer('density fitting ppaa pass2', *t1)
fxpp.close()
self.feri.flush()
dm_core = numpy.dot(mo[:,:ncore], mo[:,:ncore].T)
vj, vk = casscf.get_jk(mol, dm_core)
self.vhf_c = reduce(numpy.dot, (mo.T, vj*2-vk, mo))
t0 = log.timer('density fitting ao2mo', *t0)