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)
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.
'''
if isinstance(eri, np.ndarray):
vj, vk = _vhf.incore(eri, rdm1, with_j=with_j, with_k=with_k)
else:
nmo = rdm1.shape[0]
npair = nmo * (nmo + 1) // 2
vj = vk = None
if with_j:
rdm1_tril = lib.pack_tril(rdm1 + np.tril(rdm1, k=-1))
vj = np.zeros_like(rdm1_tril)
if with_k:
vk = np.zeros_like(rdm1)
blksize = _agf2.get_blksize(agf2.max_memory, (nmo * npair, nmo**3))
blksize = min(1, max(BLKMIN, blksize))
logger.debug1(agf2, 'blksize (ragf2.get_jk) = %d' % blksize)
tril2sq = lib.square_mat_in_trilu_indices(nmo)
for p0, p1 in lib.prange(0, nmo, blksize):
idx = list(np.concatenate(tril2sq[p0:p1]))
eri0 = eri[idx]
# vj built in tril layout with scaled rdm1_tril
if with_j:
vj[idx] = np.dot(eri0, rdm1_tril)
if with_k:
eri0 = lib.unpack_tril(eri0, axis=-1)
eri0 = eri0.reshape(p1 - p0, nmo, nmo, nmo)
vk[p0:p1] = lib.einsum('ijkl,jk->il', eri0, rdm1)
if with_j:
vj = lib.unpack_tril(vj)
return vj, vk
def _make_qmo_eris_outcore(agf2, eri, coeffs):
''' Returns H5 dataset
'''
cput0 = (logger.process_clock(), logger.perf_counter())
log = logger.Logger(agf2.stdout, agf2.verbose)
nmo = eri.nmo
ci, cj, ca = coeffs
ni = ci.shape[1]
nj = cj.shape[1]
na = ca.shape[1]
npair = nmo * (nmo + 1) // 2
mask = get_frozen_mask(agf2)
frozen = np.sum(~mask)
# possible to have incore MO, outcore QMO
if getattr(eri, 'feri', None) is None:
eri.feri = lib.H5TmpFile()
elif 'qmo' in eri.feri:
del eri.feri['qmo']
eri.feri.create_dataset('qmo', (nmo - frozen, ni, nj, na), 'f8')
blksize = _agf2.get_blksize(agf2.max_memory, (nmo * npair, nj * na, npair),
(nmo * ni, nj * na))
blksize = min(nmo, max(BLKMIN, blksize))
log.debug1('blksize (ragf2._make_qmo_eris_outcore) = %d', blksize)
tril2sq = lib.square_mat_in_trilu_indices(nmo)
q1 = 0
for p0, p1 in lib.prange(0, nmo, blksize):
if not np.any(mask[p0:p1]):
# block is fully frozen
continue
inds = np.arange(p0, p1)[mask[p0:p1]]
q0, q1 = q1, q1 + len(inds)
idx = list(np.concatenate(tril2sq[inds]))
buf = eri.eri[idx] # (blk, nmo, npair)
buf = buf.reshape((q1 - q0) * nmo, -1) # (blk*nmo, npair)
jasym, nja, cja, sja = ao2mo.incore._conc_mos(cj, ca, compact=True)
buf = ao2mo._ao2mo.nr_e2(buf, cja, sja, 's2kl', 's1')
buf = buf.reshape(q1 - q0, nmo, nj, na)
buf = lib.einsum('xpja,pi->xija', buf, ci)
eri.feri['qmo'][q0:q1] = np.asarray(buf, order='C')
log.timer('QMO integral transformation', *cput0)
return eri.feri['qmo']
def _make_qmo_eris_outcore(agf2, eri, coeffs_a, coeffs_b, spin=None):
''' Returns nested tuple of H5 dataset
spin = None: ((aaaa, aabb), (bbaa, bbbb))
spin = 0: (aaaa, aabb)
spin = 1: (bbbb, bbaa)
'''
cput0 = (time.clock(), time.time())
log = logger.Logger(agf2.stdout, agf2.verbose)
nmo = eri.nmo
nmoa, nmob = nmo
cxa = np.eye(nmoa)
cxb = np.eye(nmob)
mask = get_frozen_mask(agf2)
frozena = np.sum(~mask[0])
frozenb = np.sum(~mask[1])
npaira, npairb = nmoa*(nmoa+1)//2, nmob*(nmob+1)//2
cia, cja, caa = coeffs_a
cib, cjb, cab = coeffs_b
nia, nja, naa = [x.shape[1] for x in coeffs_a]
nib, njb, nab = [x.shape[1] for x in coeffs_b]
# possible to have incore MO, outcore QMO
if getattr(eri, 'feri', None) is None:
eri.feri = lib.H5TmpFile()
else:
for key in ['aa', 'ab', 'ba', 'bb']:
if 'qmo/%s'%key in eri.feri:
del eri.feri['qmo/%s'%key]
if spin is None or spin == 0:
eri.feri.create_dataset('qmo/aa', (nmoa-frozena, nia, nja, naa), 'f8')
eri.feri.create_dataset('qmo/ab', (nmoa-frozena, nia, njb, nab), 'f8')
blksize = _agf2.get_blksize(agf2.max_memory, (nmoa**3, nmoa*nja*naa),
(nmoa*nmob**2, nmoa*njb*nab))
blksize = min(nmoa, max(BLKMIN, blksize))
log.debug1('blksize (uagf2._make_qmo_eris_outcore) = %d', blksize)
tril2sq = lib.square_mat_in_trilu_indices(nmoa)
q1 = 0
for p0, p1 in lib.prange(0, nmoa, blksize):
if not np.any(mask[0][p0:p1]):
# block is fully frozen
continue
inds = np.arange(p0, p1)[mask[0][p0:p1]]
q0, q1 = q1, q1 + len(inds)
idx = list(np.concatenate(tril2sq[inds]))
# aa
buf = eri.eri_aa[idx] # (blk, nmoa, npaira)
buf = buf.reshape((q1-q0)*nmoa, -1) # (blk*nmoa, npaira)
jasym_aa, nja_aa, cja_aa, sja_aa = ao2mo.incore._conc_mos(cja, caa)
buf = ao2mo._ao2mo.nr_e2(buf, cja_aa, sja_aa, 's2kl', 's1')
buf = buf.reshape(q1-q0, nmoa, nja, naa)
buf = lib.einsum('xpja,pi->xija', buf, cia)
eri.feri['qmo/aa'][q0:q1] = np.asarray(buf, order='C')
# ab
buf = eri.eri_ab[idx] # (blk, nmoa, npairb)
buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmoa, npairb)
jasym_ab, nja_ab, cja_ab, sja_ab = ao2mo.incore._conc_mos(cjb, cab)
buf = ao2mo._ao2mo.nr_e2(buf, cja_ab, sja_ab, 's2kl', 's1')
buf = buf.reshape(q1-q0, nmoa, njb, nab)
buf = lib.einsum('xpja,pi->xija', buf, cia)
eri.feri['qmo/ab'][q0:q1] = np.asarray(buf, order='C')
if spin is None or spin == 1:
eri.feri.create_dataset('qmo/ba', (nmob-frozenb, nib, nja, naa), 'f8')
eri.feri.create_dataset('qmo/bb', (nmob-frozenb, nib, njb, nab), 'f8')
max_memory = agf2.max_memory - lib.current_memory()[0]
blksize = int((max_memory/8e-6) / max(nmob**3+nmob*njb*nab,
nmob*nmoa**2*nja*naa))
blksize = min(nmob, max(BLKMIN, blksize))
log.debug1('blksize (uagf2._make_qmo_eris_outcore) = %d', blksize)
tril2sq = lib.square_mat_in_trilu_indices(nmob)
q1 = 0
for p0, p1 in lib.prange(0, nmob, blksize):
if not np.any(mask[1][p0:p1]):
# block is fully frozen
continue
inds = np.arange(p0, p1)[mask[1][p0:p1]]
q0, q1 = q1, q1 + len(inds)
idx = list(np.concatenate(tril2sq[inds]))
# ba
buf = eri.eri_ba[idx] # (blk, nmob, npaira)
buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmob, npaira)
jasym_ba, nja_ba, cja_ba, sja_ba = ao2mo.incore._conc_mos(cja, caa)
buf = ao2mo._ao2mo.nr_e2(buf, cja_ba, sja_ba, 's2kl', 's1')
buf = buf.reshape(q1-q0, nmob, nja, naa)
buf = lib.einsum('xpja,pi->xija', buf, cib)
eri.feri['qmo/ba'][q0:q1] = np.asarray(buf, order='C')
# bb
buf = eri.eri_bb[idx] # (blk, nmob, npairb)
buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmob, npairb)
jasym_bb, nja_bb, cja_bb, sja_bb = ao2mo.incore._conc_mos(cjb, cab)
buf = ao2mo._ao2mo.nr_e2(buf, cja_bb, sja_bb, 's2kl', 's1')
buf = buf.reshape(q1-q0, nmob, njb, nab)
buf = lib.einsum('xpja,pi->xija', buf, cib)
eri.feri['qmo/bb'][q0:q1] = np.asarray(buf, order='C')
if spin is None:
qeri = ((eri.feri['qmo/aa'], eri.feri['qmo/ab']),
(eri.feri['qmo/ba'], eri.feri['qmo/bb']))
elif spin == 0:
qeri = (eri.feri['qmo/aa'], eri.feri['qmo/ab'])
elif spin == 1:
qeri = (eri.feri['qmo/bb'], eri.feri['qmo/ba'])
log.timer('QMO integral transformation', *cput0)
return qeri
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
