def _make_mo_eris_incore(agf2, mo_coeff=None):
''' Returns _ChemistsERIs
'''
cput0 = (logger.process_clock(), logger.perf_counter())
log = logger.Logger(agf2.stdout, agf2.verbose)
eris = _ChemistsERIs()
eris._common_init_(agf2, mo_coeff)
with_df = agf2.with_df
nmo = eris.fock.shape[0]
npair = nmo * (nmo + 1) // 2
naux = with_df.get_naoaux()
qxy = np.zeros((naux, npair))
mo = np.asarray(eris.mo_coeff, order='F')
sij = (0, nmo, 0, nmo)
sym = dict(aosym='s2', mosym='s2')
for p0, p1 in with_df.prange():
eri0 = with_df._cderi[p0:p1]
qxy[p0:p1] = ao2mo._ao2mo.nr_e2(eri0, mo, sij, out=qxy[p0:p1], **sym)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(qxy)
eris.eri = eris.qxy = qxy
log.timer('MO integral transformation', *cput0)
return eris
def build_mats_dfuagf2_lowmem(qxi,
qja,
e_occ,
e_vir,
os_factor=1.0,
ss_factor=1.0):
''' Wrapper for AGF2udf_vv_vev_islice_lowmem
'''
fdrv = getattr(libagf2, 'AGF2udf_vv_vev_islice_lowmem')
naux = qxi[0].shape[0]
noa, nob = e_occ[0].size, e_occ[1].size
nva, nvb = e_vir[0].size, e_vir[1].size
nmo = qxi[0].size // (naux * noa)
assert qxi[0].size == (naux * nmo * noa)
assert qja[0].size == (naux * noa * nva)
assert qja[1].size == (naux * nob * nvb)
qxi_a, qxi_b = qxi
qja_a, qja_b = qja
qxi = np.asarray(qxi_a, order='C')
qja = np.asarray(qja_a, order='C')
qJA = np.asarray(qja_b, order='C')
e_i = np.asarray(e_occ[0], order='C')
e_I = np.asarray(e_occ[1], order='C')
e_a = np.asarray(e_vir[0], order='C')
e_A = np.asarray(e_vir[1], order='C')
vv = np.zeros((nmo * nmo))
vev = np.zeros((nmo * nmo))
rank, size = mpi_helper.rank, mpi_helper.size
nomax = max(noa, nob)
start = rank * (noa * nomax) // size
end = (noa * nomax) if rank == (size -
1) else (rank + 1) * (noa * nomax) // size
fdrv(qxi.ctypes.data_as(ctypes.c_void_p),
qja.ctypes.data_as(ctypes.c_void_p),
qJA.ctypes.data_as(ctypes.c_void_p),
e_i.ctypes.data_as(ctypes.c_void_p),
e_I.ctypes.data_as(ctypes.c_void_p),
e_a.ctypes.data_as(ctypes.c_void_p),
e_A.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(os_factor),
ctypes.c_double(ss_factor), ctypes.c_int(nmo), ctypes.c_int(noa),
ctypes.c_int(nob), ctypes.c_int(nva), ctypes.c_int(nvb),
ctypes.c_int(naux), ctypes.c_int(start), ctypes.c_int(end),
vv.ctypes.data_as(ctypes.c_void_p),
vev.ctypes.data_as(ctypes.c_void_p))
vv = vv.reshape(nmo, nmo)
vev = vev.reshape(nmo, nmo)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vv)
mpi_helper.allreduce_safe_inplace(vev)
return vv, vev
def load_agf2(chkfile):
''' Load the AGF2 data from the chkfile.
'''
if mpi_helper.rank == 0:
dic = chkutil.load(chkfile, 'agf2')
else:
dic = None
mpi_helper.barrier()
dic = mpi_helper.bcast_dict(dic)
if 'gf' in dic:
gf = dic['gf']
dic['gf'] = GreensFunction(gf['energy'], gf['coupling'], chempot=gf['chempot'])
elif 'gfa' in dic:
gfa, gfb = dic['gfa'], dic['gfb']
dic['gf'] = (GreensFunction(gfa['energy'], gfa['coupling'], chempot=gfa['chempot']),
GreensFunction(gfb['energy'], gfb['coupling'], chempot=gfb['chempot']))
del(dic['gfa'], dic['gfb'])
if 'se' in dic:
se = dic['se']
dic['se'] = SelfEnergy(se['energy'], se['coupling'], chempot=se['chempot'])
elif 'sea' in dic:
sea, seb = dic['sea'], dic['seb']
dic['se'] = (SelfEnergy(sea['energy'], sea['coupling'], chempot=sea['chempot']),
SelfEnergy(seb['energy'], seb['coupling'], chempot=seb['chempot']))
del(dic['sea'], dic['seb'])
if 'ngf' in dic:
dic['nmom'] = (dic.get('ngf', None), dic.get('nse', None))
del(dic['ngf'], dic['nse'])
return load_mol(chkfile), dic
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 build_mats_dfragf2_outcore(qxi,
qja,
e_occ,
e_vir,
os_factor=1.0,
ss_factor=1.0):
''' Python version of AGF2df_vv_vev_islice to support outcore
'''
naux = qxi.shape[0]
nocc = e_occ.size
nvir = e_vir.size
nmo = qxi.size // (naux * nocc)
assert qxi.size == (naux * nmo * nocc)
assert qja.size == (naux * nocc * nvir)
vv = np.zeros((nmo, nmo))
vev = np.zeros((nmo, nmo))
fpos = os_factor + ss_factor
fneg = -ss_factor
eja = lib.direct_sum('j,a->ja', e_occ, -e_vir)
eja = eja.ravel()
buf = (np.zeros((nmo, nocc * nvir)), np.zeros((nmo * nocc, nvir)))
for i in mpi_helper.nrange(nocc):
qx = qxi.reshape(naux, nmo, nocc)[:, :, i]
xija = lib.dot(qx.T, qja, c=buf[0])
xjia = lib.dot(qxi.T, qja[:, i * nvir:(i + 1) * nvir], c=buf[1])
xjia = xjia.reshape(nmo, nocc * nvir)
eija = eja + e_occ[i]
vv = lib.dot(xija, xija.T, alpha=fpos, beta=1, c=vv)
vv = lib.dot(xija, xjia.T, alpha=fneg, beta=1, c=vv)
exija = xija * eija[None]
vev = lib.dot(exija, xija.T, alpha=fpos, beta=1, c=vev)
vev = lib.dot(exija, xjia.T, alpha=fneg, beta=1, c=vev)
vv = vv.reshape(nmo, nmo)
vev = vev.reshape(nmo, nmo)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vv)
mpi_helper.allreduce_safe_inplace(vev)
return vv, vev
def _make_mo_eris_incore(agf2, mo_coeff=None):
''' Returns _ChemistsERIs
'''
cput0 = (logger.process_clock(), logger.perf_counter())
log = logger.Logger(agf2.stdout, agf2.verbose)
eris = _ChemistsERIs()
eris._common_init_(agf2, mo_coeff)
with_df = agf2.with_df
moa, mob = eris.mo_coeff
nmoa, nmob = moa.shape[1], mob.shape[1]
npaira, npairb = nmoa * (nmoa + 1) // 2, nmob * (nmob + 1) // 2
naux = with_df.get_naoaux()
qxy_a = np.zeros((naux, npaira))
qxy_b = np.zeros((naux, npairb))
moa = np.asarray(moa, order='F')
mob = np.asarray(mob, order='F')
sija = (0, nmoa, 0, nmoa)
sijb = (0, nmob, 0, nmob)
sym = dict(aosym='s2', mosym='s2')
for p0, p1 in with_df.prange():
eri0 = with_df._cderi[p0:p1]
qxy_a[p0:p1] = ao2mo._ao2mo.nr_e2(eri0,
moa,
sija,
out=qxy_a[p0:p1],
**sym)
qxy_b[p0:p1] = ao2mo._ao2mo.nr_e2(eri0,
mob,
sijb,
out=qxy_b[p0:p1],
**sym)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(qxy_a)
mpi_helper.allreduce_safe_inplace(qxy_b)
eris.eri_a = qxy_a
eris.eri_b = qxy_b
eris.eri_aa = (eris.eri_a, eris.eri_a)
eris.eri_ab = (eris.eri_a, eris.eri_b)
eris.eri_ba = (eris.eri_b, eris.eri_a)
eris.eri_bb = (eris.eri_b, eris.eri_b)
eris.eri = (eris.eri_a, eris.eri_b)
log.timer('MO integral transformation', *cput0)
return eris
def build_mats_dfragf2_lowmem(qxi,
qja,
e_occ,
e_vir,
os_factor=1.0,
ss_factor=1.0):
''' Wrapper for AGF2df_vv_vev_islice_lowmem
'''
fdrv = getattr(libagf2, 'AGF2df_vv_vev_islice_lowmem')
naux = qxi.shape[0]
nocc = e_occ.size
nvir = e_vir.size
nmo = qxi.size // (naux * nocc)
assert qxi.size == (naux * nmo * nocc)
assert qja.size == (naux * nocc * nvir)
qxi = np.asarray(qxi, order='C')
qja = np.asarray(qja, order='C')
e_i = np.asarray(e_occ, order='C')
e_a = np.asarray(e_vir, order='C')
rank, size = mpi_helper.rank, mpi_helper.size
vv = np.zeros((nmo * nmo))
vev = np.zeros((nmo * nmo))
start = rank * (nocc * nocc) // size
end = nocc * nocc if rank == (size -
1) else (rank + 1) * (nocc * nocc) // size
fdrv(qxi.ctypes.data_as(ctypes.c_void_p),
qja.ctypes.data_as(ctypes.c_void_p),
e_i.ctypes.data_as(ctypes.c_void_p),
e_a.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(os_factor),
ctypes.c_double(ss_factor), ctypes.c_int(nmo), ctypes.c_int(nocc),
ctypes.c_int(nvir), ctypes.c_int(naux), ctypes.c_int(start),
ctypes.c_int(end), vv.ctypes.data_as(ctypes.c_void_p),
vev.ctypes.data_as(ctypes.c_void_p))
vv = vv.reshape(nmo, nmo)
vev = vev.reshape(nmo, nmo)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vv)
mpi_helper.allreduce_safe_inplace(vev)
return vv, vev
def load(chkfile, key):
''' Load array(s) from chkfile.
See pyscf.lib.chkfile
'''
if mpi_helper.rank == 0:
vals = chkutil.load(chkfile, key)
else:
vals = None
mpi_helper.barrier()
vals = mpi_helper.bcast_dict(vals)
return vals
def build_mats_uagf2_outcore(qeri, e_occ, e_vir, os_factor=1.0, ss_factor=1.0):
''' Python version of AGF2uee_vv_vev_islice to support outcore
'''
assert qeri[0].ndim == qeri[1].ndim == 4
nmo = qeri[0].shape[0]
# noa, nob = e_occ[0].size, e_occ[1].size
# nva, nvb = e_vir[0].size, e_vir[1].size
noa = e_occ[0].size
vv = np.zeros((nmo, nmo))
vev = np.zeros((nmo, nmo))
fposa = ss_factor
fnega = -ss_factor
fposb = os_factor
eja_a = lib.direct_sum('j,a->ja', e_occ[0], -e_vir[0]).ravel()
eja_b = lib.direct_sum('j,a->ja', e_occ[1], -e_vir[1]).ravel()
for i in mpi_helper.nrange(noa):
xija_aa = qeri[0][:, i].reshape(nmo, -1)
xija_ab = qeri[1][:, i].reshape(nmo, -1)
xjia_aa = qeri[0][:, :, i].reshape(nmo, -1)
eija_aa = eja_a + e_occ[0][i]
eija_ab = eja_b + e_occ[0][i]
vv = lib.dot(xija_aa, xija_aa.T, alpha=fposa, beta=1, c=vv)
vv = lib.dot(xija_aa, xjia_aa.T, alpha=fnega, beta=1, c=vv)
vv = lib.dot(xija_ab, xija_ab.T, alpha=fposb, beta=1, c=vv)
exija_aa = xija_aa * eija_aa[None]
exija_ab = xija_ab * eija_ab[None]
vev = lib.dot(exija_aa, xija_aa.T, alpha=fposa, beta=1, c=vev)
vev = lib.dot(exija_aa, xjia_aa.T, alpha=fnega, beta=1, c=vev)
vev = lib.dot(exija_ab, xija_ab.T, alpha=fposb, beta=1, c=vev)
vv = vv.reshape(nmo, nmo)
vev = vev.reshape(nmo, nmo)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vv)
mpi_helper.allreduce_safe_inplace(vev)
return vv, vev
def load_mol(chkfile):
''' Load the mol from the chkfile.
See pyscf.lib.chkfile
'''
if mpi_helper.rank == 0:
mol = chkutil.load_mol(chkfile)
dumps = mol.dumps()
else:
dumps = None
mpi_helper.barrier()
dumps = mpi_helper.bcast_dict(dumps)
mol = gto.loads(dumps)
return mol
def dfragf2_slow_fast_build(qxi,
qja,
e_occ,
e_vir,
nmom_max,
os_factor=1.0,
ss_factor=1.0):
lib = np.ctypeslib.load_library('agf2_ext/lib/dfragf2_slow_fast.so',
'dfragf2_slow_fast')
fdrv = lib.AGF2df_build_lowmem_islice
naux = qxi.shape[0]
nocc = e_occ.size
nvir = e_vir.size
nmo = qxi.size // (naux * nocc)
assert qxi.size == (naux * nmo * nocc)
assert qja.size == (naux * nocc * nvir)
qxi = np.asarray(qxi, order='C')
qja = np.asarray(qja, order='C')
e_i = np.asarray(e_occ, order='C')
e_a = np.asarray(e_vir, order='C')
rank, size = mpi_helper.rank, mpi_helper.size
out = np.zeros(((nmom_max + 1) * nmo * nmo))
start = rank * (nocc * nocc) // size
end = nocc * nocc if rank == (size -
1) else (rank + 1) * (nocc * nocc) // size
fdrv(qxi.ctypes.data_as(ctypes.c_void_p),
qja.ctypes.data_as(ctypes.c_void_p),
e_i.ctypes.data_as(ctypes.c_void_p),
e_a.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(os_factor),
ctypes.c_double(ss_factor), ctypes.c_int(nmo), ctypes.c_int(nocc),
ctypes.c_int(nvir), ctypes.c_int(naux), ctypes.c_int(nmom_max),
ctypes.c_int(start), ctypes.c_int(end),
out.ctypes.data_as(ctypes.c_void_p))
out = out.reshape(nmom_max + 1, nmo, nmo)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(out)
return out
def build_mats_uagf2_incore(qeri, e_occ, e_vir, os_factor=1.0, ss_factor=1.0):
''' Wraps AGF2uee_vv_vev_islice
'''
fdrv = getattr(libagf2, 'AGF2uee_vv_vev_islice')
assert qeri[0].ndim == qeri[1].ndim == 4
nmo = qeri[0].shape[0]
noa, nob = e_occ[0].size, e_occ[1].size
nva, nvb = e_vir[0].size, e_vir[1].size
qeri_a = np.asarray(qeri[0], order='C')
qeri_b = np.asarray(qeri[1], order='C')
e_i = np.asarray(e_occ[0], order='C')
e_I = np.asarray(e_occ[1], order='C')
e_a = np.asarray(e_vir[0], order='C')
e_A = np.asarray(e_vir[1], order='C')
vv = np.zeros((nmo * nmo))
vev = np.zeros((nmo * nmo))
rank, size = mpi_helper.rank, mpi_helper.size
istart = rank * noa // size
iend = noa if rank == (size - 1) else (rank + 1) * noa // size
fdrv(qeri_a.ctypes.data_as(ctypes.c_void_p),
qeri_b.ctypes.data_as(ctypes.c_void_p),
e_i.ctypes.data_as(ctypes.c_void_p),
e_I.ctypes.data_as(ctypes.c_void_p),
e_a.ctypes.data_as(ctypes.c_void_p),
e_A.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(os_factor),
ctypes.c_double(ss_factor), ctypes.c_int(nmo), ctypes.c_int(noa),
ctypes.c_int(nob), ctypes.c_int(nva), ctypes.c_int(nvb),
ctypes.c_int(istart), ctypes.c_int(iend),
vv.ctypes.data_as(ctypes.c_void_p),
vev.ctypes.data_as(ctypes.c_void_p))
vv = vv.reshape(nmo, nmo)
vev = vev.reshape(nmo, nmo)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vv)
mpi_helper.allreduce_safe_inplace(vev)
return vv, vev
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 build_mats_ragf2_outcore(qeri, e_occ, e_vir, os_factor=1.0, ss_factor=1.0):
''' Python version of AGF2ee_vv_vev_islice to support outcore
'''
assert qeri.ndim == 4
nmo = qeri.shape[0]
nocc = e_occ.size
nvir = e_vir.size
vv = np.zeros((nmo, nmo))
vev = np.zeros((nmo, nmo))
fpos = os_factor + ss_factor
fneg = -ss_factor
eja = lib.direct_sum('j,a->ja', e_occ, -e_vir)
eja = eja.ravel()
for i in mpi_helper.nrange(nocc):
xija = qeri[:, i].reshape(nmo, -1)
xjia = qeri[:, :, i].reshape(nmo, -1)
eija = eja + e_occ[i]
vv = lib.dot(xija, xija.T, alpha=fpos, beta=1, c=vv)
vv = lib.dot(xija, xjia.T, alpha=fneg, beta=1, c=vv)
exija = xija * eija[None]
vev = lib.dot(exija, xija.T, alpha=fpos, beta=1, c=vev)
vev = lib.dot(exija, xjia.T, alpha=fneg, beta=1, c=vev)
vv = vv.reshape(nmo, nmo)
vev = vev.reshape(nmo, nmo)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vv)
mpi_helper.allreduce_safe_inplace(vev)
return vv, vev
def energy_2body(agf2, gf, se):
''' Calculates the two-body energy using analytically integrated
Galitskii-Migdal formula. The formula is symmetric and only
one side needs to be calculated.
Args:
gf : GreensFunction
Auxiliaries of the Green's function
se : SelfEnergy
Auxiliaries of the self-energy
Returns
Two-body energy
'''
assert type(gf) is aux.GreensFunction
assert type(se) is aux.SelfEnergy
gf_occ = gf.get_occupied()
se_vir = se.get_virtual()
e2b = 0.0
for l in mpi_helper.nrange(gf_occ.naux):
vxl = gf_occ.coupling[:, l]
vxk = se_vir.coupling
dlk = gf_occ.energy[l] - se_vir.energy
vv = vxk * vxl[:, None]
e2b += lib.einsum('xk,yk,k->', vv, vv.conj(), 1. / dlk)
e2b *= 2
mpi_helper.barrier()
e2b = mpi_helper.allreduce(e2b)
return np.ravel(e2b.real)[0]
def build_se_part(gf2,
eri,
gf_occ,
gf_vir,
nmom,
os_factor=1.0,
ss_factor=1.0):
''' Builds the auxiliaries.
'''
cput0 = (time.clock(), time.time())
log = logger.Logger(gf2.stdout, gf2.verbose)
assert type(gf_occ) is agf2.GreensFunction
assert type(gf_vir) is agf2.GreensFunction
nmo = gf2.nmo
nocc = gf_occ.naux
nvir = gf_vir.naux
tol = gf2.weight_tol
if not (gf2.frozen is None or gf2.frozen == 0):
mask = agf2.ragf2.get_frozen_mask(agf2)
nmo -= np.sum(~mask)
ei, ci = gf_occ.energy, gf_occ.coupling
ea, ca = gf_vir.energy, gf_vir.coupling
qxi, qja = agf2.dfragf2._make_qmo_eris_incore(gf2, eri, (ci, ci, ca))
t = np.zeros((2 * nmom + 2, nmo, nmo))
eija = lib.direct_sum('i,j,a->ija', ei, ei, -ea)
naux = qxi.shape[0]
buf = [np.empty((nmo, nocc * nvir)), np.empty((nmo * nocc, nvir))]
for i in mpi_helper.nrange(nocc):
qx = qxi.reshape(naux, nmo, nocc)[:, :, i]
xija = lib.dot(qx.T, qja, c=buf[0])
xjia = lib.dot(qxi.T, qja[:, i * nvir:(i + 1) * nvir], c=buf[1])
xjia = xjia.reshape(nmo, nocc * nvir)
xjia = 2.0 * xija - xjia #TODO: oo/ss factors
eja = eija[i].ravel()
for n in range(2 * nmom + 2):
t[n] = lib.dot(xija, xjia.T, beta=1, c=t[n])
xjia *= eja[None]
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(t)
#t = _agf2.dfragf2_slow_fast_build(qxi, qja, ei, ea, 2*nmom+2, os_factor=os_factor, ss_factor=ss_factor)
if not np.all(np.isfinite(t)):
raise ValueError('Overflow from large moments')
m, b = block_lanczos.block_lanczos(t, nmom + 1)
e, v = block_lanczos.build_from_tridiag(m, b)
se = agf2.SelfEnergy(e, v, chempot=gf_occ.chempot)
se.remove_uncoupled(tol=tol)
if not (gf2.frozen is None or gf2.frozen == 0):
coupling = np.zeros((nmo, se.naux))
coupling[mask] = se.coupling
se = agf2.SelfEnergy(se.energy, coupling, chempot=se.chempot)
log.timer('se part', *cput0)
return se
def build_mats_dfuagf2_outcore(qxi,
qja,
e_occ,
e_vir,
os_factor=1.0,
ss_factor=1.0):
''' Python version of AGF2udf_vv_vev_islice to support outcore
'''
naux = qxi[0].shape[0]
noa, nob = e_occ[0].size, e_occ[1].size
nva, nvb = e_vir[0].size, e_vir[1].size
nmo = qxi[0].size // (naux * noa)
assert qxi[0].size == (naux * nmo * noa)
assert qja[0].size == (naux * noa * nva)
assert qja[1].size == (naux * nob * nvb)
qxi_a, qxi_b = qxi
qja_a, qja_b = qja
vv = np.zeros((nmo, nmo))
vev = np.zeros((nmo, nmo))
fposa = ss_factor
fnega = -ss_factor
fposb = os_factor
eja_a = lib.direct_sum('j,a->ja', e_occ[0], -e_vir[0]).ravel()
eja_b = lib.direct_sum('j,a->ja', e_occ[1], -e_vir[1]).ravel()
buf = (np.zeros((nmo, noa * nva)), np.zeros(
(nmo, nob * nvb)), np.zeros((nmo * noa, nva)))
for i in mpi_helper.nrange(noa):
qx_a = qxi_a.reshape(naux, nmo, noa)[:, :, i]
xija_aa = lib.dot(qx_a.T, qja_a, c=buf[0])
xija_ab = lib.dot(qx_a.T, qja_b, c=buf[1])
xjia_aa = lib.dot(qxi_a.T, qja_a[:, i * nva:(i + 1) * nva], c=buf[2])
xjia_aa = xjia_aa.reshape(nmo, -1)
eija_aa = eja_a + e_occ[0][i]
eija_ab = eja_b + e_occ[0][i]
vv = lib.dot(xija_aa, xija_aa.T, alpha=fposa, beta=1, c=vv)
vv = lib.dot(xija_aa, xjia_aa.T, alpha=fnega, beta=1, c=vv)
vv = lib.dot(xija_ab, xija_ab.T, alpha=fposb, beta=1, c=vv)
exija_aa = xija_aa * eija_aa[None]
exija_ab = xija_ab * eija_ab[None]
vev = lib.dot(exija_aa, xija_aa.T, alpha=fposa, beta=1, c=vev)
vev = lib.dot(exija_aa, xjia_aa.T, alpha=fnega, beta=1, c=vev)
vev = lib.dot(exija_ab, xija_ab.T, alpha=fposb, beta=1, c=vev)
vv = vv.reshape(nmo, nmo)
vev = vev.reshape(nmo, nmo)
mpi_helper.barrier()
mpi_helper.allreduce_safe_inplace(vv)
mpi_helper.allreduce_safe_inplace(vev)
return vv, vev
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
