def prange(self, start=None, stop=None, step=None):
if start is None: start = 0
if stop is None: stop = self.get_naoaux()
if step is None: step = self.blockdim
for p0, p1 in mpi_helper.prange(start, stop, step):
yield p0, p1
def _make_qmo_eris_incore(agf2, eri, coeffs_a, coeffs_b):
''' Returns nested tuple of ndarray
'''
cput0 = (time.clock(), time.time())
log = logger.Logger(agf2.stdout, agf2.verbose)
cxa, cxb = np.eye(agf2.nmo[0]), np.eye(agf2.nmo[1])
if not (agf2.frozen is None or agf2.frozen == 0):
mask = uagf2.get_frozen_mask(agf2)
cxa = cxa[:,mask[0]]
cxb = cxb[:,mask[1]]
nmoa, nmob = agf2.nmo
npaira, npairb = nmoa*(nmoa+1)//2, nmob*(nmob+1)//2
with_df = agf2.with_df
naux = with_df.get_naoaux()
cia, cja, caa = coeffs_a
cib, cjb, cab = coeffs_b
xisym_a, nxi_a, cxi_a, sxi_a = ao2mo.incore._conc_mos(cxa, cia, compact=False)
jasym_a, nja_a, cja_a, sja_a = ao2mo.incore._conc_mos(cja, caa, compact=False)
xisym_b, nxi_b, cxi_b, sxi_b = ao2mo.incore._conc_mos(cxb, cib, compact=False)
jasym_b, nja_b, cja_b, sja_b = ao2mo.incore._conc_mos(cjb, cab, compact=False)
sym = dict(aosym='s2', mosym='s1')
qxi_a = np.zeros((naux, nxi_a))
qxi_b = np.zeros((naux, nxi_b))
qja_a = np.zeros((naux, nja_a))
qja_b = np.zeros((naux, nja_b))
buf = (np.zeros((with_df.blockdim, npaira)), np.zeros((with_df.blockdim, npairb)))
for p0, p1 in mpi_helper.prange(0, naux, with_df.blockdim):
naux0 = p1 - p0
bufa0 = buf[0][:naux0]
bufb0 = buf[1][:naux0]
bufa0[:] = eri.eri[0][p0:p1]
bufb0[:] = eri.eri[1][p0:p1]
qxi_a[p0:p1] = ao2mo._ao2mo.nr_e2(bufa0, cxi_a, sxi_a, out=qxi_a[p0:p1], **sym)
qxi_b[p0:p1] = ao2mo._ao2mo.nr_e2(bufb0, cxi_b, sxi_b, out=qxi_b[p0:p1], **sym)
qja_a[p0:p1] = ao2mo._ao2mo.nr_e2(bufa0, cja_a, sja_a, out=qja_a[p0:p1], **sym)
qja_b[p0:p1] = ao2mo._ao2mo.nr_e2(bufb0, cja_b, sja_b, out=qja_b[p0:p1], **sym)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(qxi_a)
mpi_helper.allreduce_safe_inplace(qxi_b)
mpi_helper.allreduce_safe_inplace(qja_a)
mpi_helper.allreduce_safe_inplace(qja_b)
qxi_a = qxi_a.reshape(naux, -1)
qxi_b = qxi_b.reshape(naux, -1)
qja_a = qja_a.reshape(naux, -1)
qja_b = qja_b.reshape(naux, -1)
log.timer('QMO integral transformation', *cput0)
return ((qxi_a, qja_a), (qxi_b, qja_b))
def _make_qmo_eris_incore(agf2, eri, coeffs):
''' Returns tuple of ndarray
'''
cput0 = (logger.process_clock(), logger.perf_counter())
log = logger.Logger(agf2.stdout, agf2.verbose)
cx = np.eye(agf2.nmo)
if not (agf2.frozen is None or agf2.frozen == 0):
mask = ragf2.get_frozen_mask(agf2)
cx = cx[:, mask]
nmo = eri.fock.shape[0]
npair = nmo * (nmo + 1) // 2
with_df = agf2.with_df
naux = with_df.get_naoaux()
ci, cj, ca = coeffs
xisym, nxi, cxi, sxi = ao2mo.incore._conc_mos(cx, ci, compact=False)
jasym, nja, cja, sja = ao2mo.incore._conc_mos(cj, ca, compact=False)
sym = dict(aosym='s2', mosym='s1')
qxi = np.zeros((naux, nxi))
qja = np.zeros((naux, nja))
buf = np.zeros((with_df.blockdim, npair))
for p0, p1 in mpi_helper.prange(0, naux, with_df.blockdim):
naux0 = p1 - p0
buf0 = buf[:naux0]
buf0[:] = eri.eri[p0:p1]
qxi[p0:p1] = ao2mo._ao2mo.nr_e2(buf0, cxi, sxi, out=qxi[p0:p1], **sym)
qja[p0:p1] = ao2mo._ao2mo.nr_e2(buf0, cja, sja, out=qja[p0:p1], **sym)
qxi = qxi.reshape(naux, -1)
qja = qja.reshape(naux, -1)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(qxi)
mpi_helper.allreduce_safe_inplace(qja)
log.timer('QMO integral transformation', *cput0)
return (qxi, qja)
def get_jk(agf2, eri, rdm1, with_j=True, with_k=True):
''' Get the J/K matrices.
Args:
eri : ndarray or H5 dataset
Electronic repulsion integrals (NOT as _ChemistsERIs). In
the case of no bra/ket symmetry, a tuple can be passed.
rdm1 : 2D array
Reduced density matrix
Kwargs:
with_j : bool
Whether to compute J. Default value is True
with_k : bool
Whether to compute K. Default value is True
Returns:
tuple of ndarrays corresponding to J and K, if either are
not requested then they are set to None.
'''
nmo = rdm1.shape[0]
npair = nmo * (nmo + 1) // 2
naux = agf2.with_df.get_naoaux()
vj = vk = None
if with_j:
rdm1_tril = lib.pack_tril(rdm1 + np.tril(rdm1, k=-1))
vj = np.zeros((npair, ))
if with_k:
vk = np.zeros((nmo, nmo))
fdrv = ao2mo._ao2mo.libao2mo.AO2MOnr_e2_drv
fmmm = ao2mo._ao2mo.libao2mo.AO2MOmmm_bra_nr_s2
ftrans = ao2mo._ao2mo.libao2mo.AO2MOtranse2_nr_s2
if isinstance(eri, tuple):
bra, ket = eri
else:
bra = ket = eri
blksize = _agf2.get_blksize(agf2.max_memory,
(npair, npair, 1, nmo**2, nmo**2))
blksize = min(nmo, max(BLKMIN, blksize))
logger.debug1(agf2, 'blksize (dfragf2.get_jk) = %d' % blksize)
buf = (np.empty((blksize, nmo, nmo)), np.empty((blksize, nmo, nmo)))
for p0, p1 in mpi_helper.prange(0, naux, blksize):
bra0 = bra[p0:p1]
ket0 = ket[p0:p1]
rho = np.dot(ket0, rdm1_tril)
if with_j:
vj += np.dot(rho, bra0)
if with_k:
buf1 = buf[0][:p1 - p0]
fdrv(ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p),
bra0.ctypes.data_as(ctypes.c_void_p),
rdm1.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(p1 - p0),
ctypes.c_int(nmo), (ctypes.c_int * 4)(0, nmo, 0, nmo),
lib.c_null_ptr(), ctypes.c_int(0))
buf2 = lib.unpack_tril(ket0, out=buf[1])
buf1 = buf1.reshape(-1, nmo)
buf2 = buf2.reshape(-1, nmo)
vk = lib.dot(buf1.T, buf2, c=vk, beta=1)
if with_j:
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vj)
mpi_helper.barrier()
vj = lib.unpack_tril(vj)
if with_k:
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vk)
return vj, vk