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