def allgather(n, m): import numpy from mpi4pyscf.tools import mpi mpi.INT_MAX = 7 arrs = numpy.ones((n - mpi.rank, m)) res = mpi.allgather(arrs) print(res.shape)
def cc_Fov(t1T, eris, vlocs=None): """ Fov: me. """ nvir, nocc = t1T.shape if vlocs is None: ntasks = mpi.pool.size vlocs = [_task_location(nvir, task_id) for task_id in range(ntasks)] fov = eris.fock[:nocc, nocc:] Fme = einsum_mv('efmn, fn -> em', eris.xvoo, t1T) Fme = mpi.allgather(Fme).T Fme += fov return Fme
def _int_nuc_vloc(mydf, nuccell, kpts, intor='int3c2e_sph'): '''Vnuc - Vloc''' cell = mydf.cell nkpts = len(kpts) # Use the 3c2e code with steep s gaussians to mimic nuclear density fakenuc = aft._fake_nuc(cell) fakenuc._atm, fakenuc._bas, fakenuc._env = \ gto.conc_env(nuccell._atm, nuccell._bas, nuccell._env, fakenuc._atm, fakenuc._bas, fakenuc._env) kptij_lst = numpy.hstack((kpts, kpts)).reshape(-1, 2, 3) ishs = mpi.work_balanced_partition(numpy.arange(cell.nbas), costs=numpy.arange(1, cell.nbas + 1)) if len(ishs) > 0: ish0, ish1 = ishs[0], ishs[-1] + 1 buf = incore.aux_e2(cell, fakenuc, intor, aosym='s2', kptij_lst=kptij_lst, shls_slice=(ish0, ish1, 0, cell.nbas, 0, fakenuc.nbas)) else: buf = numpy.zeros(0) charge = cell.atom_charges() charge = numpy.append(charge, -charge) # (charge-of-nuccell, charge-of-fakenuc) nao = cell.nao_nr() nchg = len(charge) nao_pair = nao * (nao + 1) // 2 buf = buf.reshape(nkpts, -1, nchg) # scaled by 1./mpi.pool.size because nuc is mpi.reduced in get_nuc function buf = numpy.einsum('kxz,z->kx', buf, 1. / mpi.pool.size * charge) mat = numpy.empty((nkpts, nao_pair), dtype=numpy.complex128) for k in range(nkpts): mat[k] = mpi.allgather(buf[k]) if rank == 0 and cell.dimension == 3: nucbar = sum([ z / nuccell.bas_exp(i)[0] for i, z in enumerate(cell.atom_charges()) ]) nucbar *= numpy.pi / cell.vol ovlp = cell.pbc_intor('int1e_ovlp_sph', 1, lib.HERMITIAN, kpts) for k in range(nkpts): s = lib.pack_tril(ovlp[k]) mat[k] += nucbar * s return mat
def _int_nuc_vloc(mydf, nuccell, kpts, intor='int3c2e_sph', aosym='s2', comp=1): '''Vnuc - Vloc''' cell = mydf.cell nkpts = len(kpts) # Use the 3c2e code with steep s gaussians to mimic nuclear density fakenuc = aft._fake_nuc(cell) fakenuc._atm, fakenuc._bas, fakenuc._env = \ gto.conc_env(nuccell._atm, nuccell._bas, nuccell._env, fakenuc._atm, fakenuc._bas, fakenuc._env) kptij_lst = numpy.hstack((kpts,kpts)).reshape(-1,2,3) ishs = mpi.work_balanced_partition(numpy.arange(cell.nbas), costs=numpy.arange(1, cell.nbas+1)) if len(ishs) > 0: ish0, ish1 = ishs[0], ishs[-1]+1 buf = incore.aux_e2(cell, fakenuc, intor, aosym='s2', kptij_lst=kptij_lst, shls_slice=(ish0,ish1,0,cell.nbas,0,fakenuc.nbas)) else: buf = numpy.zeros(0) charge = cell.atom_charges() charge = numpy.append(charge, -charge) # (charge-of-nuccell, charge-of-fakenuc) nao = cell.nao_nr() nchg = len(charge) nao_pair = nao*(nao+1)//2 buf = buf.reshape(nkpts,-1,nchg) # scaled by 1./mpi.pool.size because nuc is mpi.reduced in get_nuc function buf = numpy.einsum('kxz,z->kx', buf, 1./mpi.pool.size*charge) mat = numpy.empty((nkpts,nao_pair), dtype=numpy.complex128) for k in range(nkpts): mat[k] = mpi.allgather(buf[k]) if (rank == 0 and cell.dimension == 3 and intor in ('int3c2e', 'int3c2e_sph', 'int3c2e_cart')): assert(comp == 1) charges = cell.atom_charges() nucbar = sum([z/nuccell.bas_exp(i)[0] for i,z in enumerate(charges)]) nucbar *= numpy.pi/cell.vol ovlp = cell.pbc_intor('int1e_ovlp_sph', 1, lib.HERMITIAN, kpts) for k in range(nkpts): if aosym == 's1': mat[k] += nucbar * ovlp[k].reshape(nao_pair) else: mat[k] += nucbar * lib.pack_tril(ovlp[k]) return mat
def cc_Fvv(t1T, t2T, eris, tauT_tilde=None, vlocs=None): """ Fvv: ae. """ nvir, nocc = t1T.shape if vlocs is None: ntasks = mpi.pool.size vlocs = [_task_location(nvir, task_id) for task_id in range(ntasks)] if tauT_tilde is None: tauT_tilde = make_tauT(t1T, t2T, fac=0.5, vlocs=vlocs) vloc0, vloc1 = vlocs[rank] #fvo = eris.fock[nocc:, :nocc] fvv = eris.fock[nocc + vloc0:nocc + vloc1, nocc:] Fea = fvv #- 0.5 * np.dot(fvo[vloc0:vloc1], t1T.T) Fae = (-0.5) * einsum('femn, famn -> ae', eris.xvoo, tauT_tilde) Fae = mpi.allreduce_inplace(Fae) tauT_tilde = None #Fea += np.einsum('mafe, fm -> ea', eris.ovvx, t1T, optimize=True) Fae += mpi.allgather(Fea).T return Fae
def _gamma1_intermediates(mycc, t1, t2, l1, l2): t1T = t1.T t2T = t2.transpose(2, 3, 0, 1) l1T = l1.T l2T = l2.transpose(2, 3, 0, 1) t1 = t2 = l1 = l2 = None doo = -np.dot(l1T.T, t1T) doo -= mpi.allreduce_inplace(einsum('efim, efjm -> ij', l2T, t2T) * 0.5) dvv = np.dot(t1T, l1T.T) dvv += mpi.allreduce_inplace(einsum('eamn, ebmn -> ab', t2T, l2T) * 0.5) xt1 = mpi.allreduce_inplace(einsum('efmn, efin -> mi', l2T, t2T) * 0.5) xt2 = mpi.allreduce_inplace(einsum('famn, femn -> ae', t2T, l2T) * 0.5) xt2 += np.dot(t1T, l1T.T) dvo = mpi.allgather(np.einsum('aeim, em -> ai', t2T, l1T, optimize=True)) dvo -= np.dot(t1T, xt1) dvo -= np.dot(xt2, t1T) dvo += t1T dov = l1T.T return doo, dov, dvo, dvv
def make_intermediates(mycc, t1, t2, eris): t1T = t1.T t2T = np.asarray(t2.transpose(2, 3, 0, 1), order='C') nvir_seg, nvir, nocc = t2T.shape[:3] t1 = t2 = None ntasks = mpi.pool.size vlocs = [_task_location(nvir, task_id) for task_id in range(ntasks)] vloc0, vloc1 = vlocs[rank] assert vloc1 - vloc0 == nvir_seg class _IMDS: pass imds = _IMDS() imds.ftmp = lib.H5TmpFile() dtype = t1T.dtype imds.woooo = imds.ftmp.create_dataset('woooo', (nocc, nocc, nocc, nocc), dtype) imds.wovvo = imds.ftmp.create_dataset('wovvo', (nocc, nvir_seg, nvir, nocc), dtype) imds.wovoo = imds.ftmp.create_dataset('wovoo', (nocc, nvir_seg, nocc, nocc), dtype) imds.wvvvo = imds.ftmp.create_dataset('wvvvo', (nvir_seg, nvir, nvir, nocc), dtype) foo = eris.fock[:nocc, :nocc] fov = eris.fock[:nocc, nocc:] fvo = eris.fock[nocc:, :nocc] fvv = eris.fock[nocc:, nocc:] #tauT = np.einsum('ai, bj -> abij', t1T[vloc0:vloc1] * 2.0, t1T, optimize=True) tauT = t2T v1 = np.array(fvv, copy=True) #- np.dot(t1T, fov) #tmp = einsum('jbac, cj -> ba', eris.oxvv, t1T) tmp = 0.0 v4 = 0.0 eris_voov = eris.xvoo.transpose(0, 2, 3, 1) for task_id, eri_tmp, p0, p1 in _rotate_vir_block(eris_voov, vlocs=vlocs): tmp += einsum('cjka, bcjk -> ba', eri_tmp, tauT[:, p0:p1]) v4 += einsum('dljb, cdkl -> jcbk', eri_tmp, t2T[:, p0:p1]) eri_tmp = None eris_voov = None tmp *= 0.5 v1 += mpi.allgather(tmp) v2 = np.array(foo, copy=True) #+ np.dot(fov, t1T) #tmp = einsum('kijb, bk -> ij', eris.ooox, t1T[vloc0:vloc1]) tmp = einsum('bcik, bcjk -> ij', eris.xvoo, tauT) tmp *= 0.5 v2 += mpi.allreduce_inplace(tmp) #v4 -= np.asarray(eris.oxov).transpose(0, 1, 3, 2) v4 -= np.asarray(eris.xovo).transpose(1, 0, 2, 3) v5 = fvo + mpi.allgather(einsum('kc, bcjk -> bj', fov, t2T)) #tmp = fvo[vloc0:vloc1] #+ einsum('cdkl, dl -> ck', eris.xvoo, t1T) #v5 += mpi.allreduce(np.einsum('ck, bk, cj -> bj', tmp, t1T, t1T[vloc0:vloc1], optimize=True)) #v5 += mpi.allreduce(einsum('kljc, cbkl -> bj', eris.ooox, t2T)) * 0.5 v5 += mpi.allreduce_inplace( einsum('cjlk, cbkl -> bj', eris.xooo, t2T) * 0.5) tmp = 0.0 # ZHC NOTE FIXME it seems that the tmp does not contribute to rdm for task_id, t2T_tmp, p0, p1 in _rotate_vir_block(t2T, vlocs=vlocs): #tmp += einsum('kbcd, cdjk -> bj', eris.oxvv[:, :, p0:p1], t2T_tmp) tmp -= einsum('bkcd, cdjk -> bj', eris.xovv[:, :, p0:p1], t2T_tmp) t2T_tmp = None tmp *= 0.5 tmp = mpi.allgather(tmp) v5 -= tmp #w3 = np.array(v5[vloc0:vloc1], copy=True) #+ einsum('jcbk, bj -> ck', v4, t1T) #w3 += np.dot(v1[vloc0:vloc1], t1T) #w3 -= np.dot(t1T[vloc0:vloc1], v2) #w3 = mpi.allgather(w3) w3 = v5 woooo = einsum('cdij, cdkl -> ijkl', eris.xvoo, tauT) woooo *= 0.25 #woooo += einsum('jilc, ck -> jilk', eris.ooox, t1T[vloc0:vloc1]) woooo = mpi.allreduce_inplace(woooo) woooo += np.asarray(eris.oooo) * 0.5 imds.woooo[:] = woooo woooo = None # ZHC NOTE: wovvo, v4 has shape j[c]bk wovvo = v4 #+ einsum('jcbd, dk -> jcbk', eris.oxvv, t1T) #tmp = einsum('bdlj, dk -> bklj', eris.xvoo, t1T) #for task_id, tmp_2, p0, p1 in _rotate_vir_block(tmp, vlocs=vlocs): # wovvo[:, :, p0:p1] += einsum('bklj, cl -> jcbk', tmp_2, t1T[vloc0:vloc1]) # tmp_2 = None #tmp = None #eris_vooo = eris.ooox.transpose(3, 2, 1, 0) #for task_id, eri_tmp, p0, p1 in _rotate_vir_block(eris_vooo, vlocs=vlocs): # wovvo[:, :, p0:p1] -= einsum('bkjl, cl -> jcbk', eri_tmp, t1T[vloc0:vloc1]) # eri_tmp = None #eris_vooo = None imds.wovvo[:] = wovvo wovvo = None wovoo = 0.0 for task_id, tauT_tmp, p0, p1 in _rotate_vir_block(tauT, vlocs=vlocs): #wovoo += einsum('icdb, dbjk -> icjk', eris.oxvv[:, :, p0:p1], tauT_tmp) wovoo -= einsum('cidb, dbjk -> icjk', eris.xovv[:, :, p0:p1], tauT_tmp) tauT_tmp = None wovoo *= 0.25 #wovoo += np.asarray(eris.ooox.transpose(2, 3, 0, 1)) * 0.5 wovoo += np.asarray(eris.xooo.transpose(1, 0, 2, 3)) * (-0.5) #wovoo += einsum('icbk, bj -> icjk', v4, t1T) tauT = tauT * 0.25 #eris_vooo = eris.ooox.transpose(3, 0, 1, 2) eris_vooo = eris.xooo.transpose(0, 3, 2, 1) for task_id, eri_tmp, p0, p1 in _rotate_vir_block(eris_vooo, vlocs=vlocs): wovoo -= einsum('blij, cbkl -> icjk', eri_tmp, t2T[:, p0:p1]) imds.wvvvo[:, :, p0:p1] = einsum('bcjl, ajlk -> bcak', tauT, eri_tmp) eri_tmp = None eris_vooo = None imds.wovoo[:] = wovoo wovoo = None tauT = None #v4 = v4.transpose(1, 0, 2, 3) #for task_id, v4_tmp, p0, p1 in _rotate_vir_block(v4, vlocs=vlocs): # imds.wvvvo[:, p0:p1] += einsum('bj, cjak -> bcak', t1T[vloc0:vloc1], v4_tmp) # v4_tmp = None #v4 = None #wvvvo = np.asarray(eris.ovvx).conj().transpose(3, 2, 1, 0) * 0.5 wvvvo = np.asarray(eris.xvvo) * 0.5 #eris_ovvv = eris.oxvv #for task_id, t2T_tmp, p0, p1 in _rotate_vir_block(t2T, vlocs=vlocs): # wvvvo[:, p0:p1] -= einsum('kbad, cdjk -> bcaj', eris_ovvv, t2T_tmp) # t2T_tmp = None eris_vovv = eris.xovv for task_id, t2T_tmp, p0, p1 in _rotate_vir_block(t2T, vlocs=vlocs): wvvvo[:, p0:p1] -= einsum('bkda, cdjk -> bcaj', eris_vovv, t2T_tmp) t2T_tmp = None imds.wvvvo -= wvvvo wvvvo = None imds.v1 = v1 imds.v2 = v2 imds.w3 = w3 imds.ftmp.flush() return imds
def update_amps(mycc, t1, t2, eris): """ Update GCCD amplitudes. """ time0 = logger.process_clock(), logger.perf_counter() log = logger.Logger(mycc.stdout, mycc.verbose) t1T = t1.T t2T = np.asarray(t2.transpose(2, 3, 0, 1), order='C') nvir_seg, nvir, nocc = t2T.shape[:3] t2 = None ntasks = mpi.pool.size vlocs = [_task_location(nvir, task_id) for task_id in range(ntasks)] vloc0, vloc1 = vlocs[rank] log.debug2('vlocs %s', vlocs) assert vloc1 - vloc0 == nvir_seg fock = eris.fock fvo = fock[nocc:, :nocc] mo_e_o = eris.mo_energy[:nocc] mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift tauT_tilde = make_tauT(t1T, t2T, fac=0.5, vlocs=vlocs) Fvv = cc_Fvv(t1T, t2T, eris, tauT_tilde=tauT_tilde, vlocs=vlocs) Foo = cc_Foo(t1T, t2T, eris, tauT_tilde=tauT_tilde, vlocs=vlocs) tauT_tilde = None Fov = cc_Fov(t1T, eris, vlocs=vlocs) # Move energy terms to the other side Fvv[np.diag_indices(nvir)] -= mo_e_v Foo[np.diag_indices(nocc)] -= mo_e_o # T1 equation t1Tnew = np.zeros_like(t1T) #t1Tnew = np.dot(Fvv, t1T) #t1Tnew -= np.dot(t1T, Foo) tmp = einsum('aeim, me -> ai', t2T, Fov) #tmp -= np.einsum('fn, naif -> ai', t1T, eris.oxov, optimize=True) tmp = mpi.allgather(tmp) #tmp2 = einsum('eamn, mnie -> ai', t2T, eris.ooox) tmp2 = einsum('eamn, einm -> ai', t2T, eris.xooo) #tmp2 += einsum('efim, mafe -> ai', t2T, eris.ovvx) tmp2 += einsum('efim, efam -> ai', t2T, eris.xvvo) tmp2 *= 0.5 tmp2 = mpi.allreduce_inplace(tmp2) tmp += tmp2 tmp2 = None #t1Tnew += tmp #t1Tnew += fvo # T2 equation Ftmp = Fvv #- 0.5 * np.dot(t1T, Fov) t2Tnew = einsum('aeij, be -> abij', t2T, Ftmp) t2T_tmp = mpi.alltoall_new([t2Tnew[:, p0:p1] for p0, p1 in vlocs], split_recvbuf=True) for task_id, (p0, p1) in enumerate(vlocs): tmp = t2T_tmp[task_id].reshape(p1 - p0, nvir_seg, nocc, nocc) t2Tnew[:, p0:p1] -= tmp.transpose(1, 0, 2, 3) tmp = None t2T_tmp = None Ftmp = Foo #+ 0.5 * np.dot(Fov, t1T) tmp = einsum('abim, mj -> abij', t2T, Ftmp) t2Tnew -= tmp t2Tnew += tmp.transpose(0, 1, 3, 2) tmp = None t2Tnew += np.asarray(eris.xvoo) tauT = make_tauT(t1T, t2T, vlocs=vlocs) Woooo = cc_Woooo(t1T, t2T, eris, tauT=tauT, vlocs=vlocs) Woooo *= 0.5 t2Tnew += einsum('abmn, mnij -> abij', tauT, Woooo) Woooo = None Wvvvv = cc_Wvvvv(t1T, t2T, eris, tauT=tauT, vlocs=vlocs) for task_id, tauT_tmp, p0, p1 in _rotate_vir_block(tauT, vlocs=vlocs): tmp = einsum('abef, efij -> abij', Wvvvv[:, :, p0:p1], tauT_tmp) tmp *= 0.5 t2Tnew += tmp tmp = tauT_tmp = None Wvvvv = None tauT = None #tmp = einsum('mbje, ei -> bmij', eris.oxov, t1T) # [b]mij #tmp = mpi.allgather(tmp) # bmij #tmp = einsum('am, bmij -> abij', t1T[vloc0:vloc1], tmp) # [a]bij tmp = 0.0 Wvovo = cc_Wovvo(t1T, t2T, eris, vlocs=vlocs).transpose(2, 0, 1, 3) for task_id, w_tmp, p0, p1 in _rotate_vir_block(Wvovo, vlocs=vlocs): tmp += einsum('aeim, embj -> abij', t2T[:, p0:p1], w_tmp) w_tmp = None Wvovo = None tmp = tmp - tmp.transpose(0, 1, 3, 2) t2Tnew += tmp tmpT = mpi.alltoall_new([tmp[:, p0:p1] for p0, p1 in vlocs], split_recvbuf=True) for task_id, (p0, p1) in enumerate(vlocs): tmp = tmpT[task_id].reshape(p1 - p0, nvir_seg, nocc, nocc) t2Tnew[:, p0:p1] -= tmp.transpose(1, 0, 2, 3) tmp = None tmpT = None #tmp = einsum('ei, jeba -> abij', t1T, eris.ovvx) #t2Tnew += tmp #t2Tnew -= tmp.transpose(0, 1, 3, 2) #tmp = einsum('am, ijmb -> baij', t1T, eris.ooox.conj()) #t2Tnew += tmp #tmpT = mpi.alltoall([tmp[:, p0:p1] for p0, p1 in vlocs], # split_recvbuf=True) #for task_id, (p0, p1) in enumerate(vlocs): # tmp = tmpT[task_id].reshape(p1-p0, nvir_seg, nocc, nocc) # t2Tnew[:, p0:p1] -= tmp.transpose(1, 0, 2, 3) # tmp = None #tmpT = None eia = mo_e_o[:, None] - mo_e_v #t1Tnew /= eia.T for i in range(vloc0, vloc1): t2Tnew[i - vloc0] /= lib.direct_sum('i + jb -> bij', eia[:, i], eia) time0 = log.timer_debug1('update t1 t2', *time0) return t1Tnew.T, t2Tnew.transpose(2, 3, 0, 1)
def update_lambda(mycc, t1, t2, l1, l2, eris, imds): """ Update GCCSD lambda. """ time0 = logger.process_clock(), logger.perf_counter() log = logger.Logger(mycc.stdout, mycc.verbose) t1T = t1.T t2T = np.asarray(t2.transpose(2, 3, 0, 1), order='C') t1 = t2 = None nvir_seg, nvir, nocc = t2T.shape[:3] l1T = l1.T l2T = np.asarray(l2.transpose(2, 3, 0, 1), order='C') l1 = l2 = None ntasks = mpi.pool.size vlocs = [_task_location(nvir, task_id) for task_id in range(ntasks)] vloc0, vloc1 = vlocs[rank] log.debug2('vlocs %s', vlocs) assert vloc1 - vloc0 == nvir_seg fvo = eris.fock[nocc:, :nocc] mo_e_o = eris.mo_energy[:nocc] mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift v1 = imds.v1 - np.diag(mo_e_v) v2 = imds.v2 - np.diag(mo_e_o) mba = einsum('cakl, cbkl -> ba', l2T, t2T) * 0.5 mba = mpi.allreduce_inplace(mba) mij = einsum('cdki, cdkj -> ij', l2T, t2T) * 0.5 mij = mpi.allreduce_inplace(mij) # m3 [a]bij m3 = einsum('abkl, ijkl -> abij', l2T, np.asarray(imds.woooo)) tauT = t2T + np.einsum( 'ai, bj -> abij', t1T[vloc0:vloc1] * 2.0, t1T, optimize=True) tmp = einsum('cdij, cdkl -> ijkl', l2T, tauT) tmp = mpi.allreduce_inplace(tmp) tauT = None vvoo = np.asarray(eris.xvoo) tmp = einsum('abkl, ijkl -> abij', vvoo, tmp) tmp *= 0.25 m3 += tmp tmp = None tmp = einsum('cdij, dk -> ckij', l2T, t1T) for task_id, tmp, p0, p1 in _rotate_vir_block(tmp, vlocs=vlocs): #m3 -= einsum('kcba, ckij -> abij', eris.ovvx[:, p0:p1], tmp) m3 -= einsum('abck, ckij -> abij', eris.xvvo[:, :, p0:p1], tmp) tmp = None eris_vvvv = eris.xvvv.transpose(2, 3, 0, 1) tmp_2 = np.empty_like(l2T) # used for line 387 for task_id, l2T_tmp, p0, p1 in _rotate_vir_block(l2T, vlocs=vlocs): tmp = einsum('cdij, cdab -> abij', l2T_tmp, eris_vvvv[p0:p1]) tmp *= 0.5 m3 += tmp tmp_2[:, p0:p1] = einsum('acij, cb -> baij', l2T_tmp, v1[:, vloc0:vloc1]) tmp = l2T_tmp = None eris_vvvv = None l1Tnew = einsum('abij, bj -> ai', m3, t1T) l1Tnew = mpi.allgather(l1Tnew) l2Tnew = m3 l2Tnew += vvoo fvo1 = fvo + mpi.allreduce_inplace( einsum('cbkj, ck -> bj', vvoo, t1T[vloc0:vloc1])) tmp = np.einsum('ai, bj -> abij', l1T[vloc0:vloc1], fvo1, optimize=True) wvovo = np.asarray(imds.wovvo).transpose(1, 0, 2, 3) for task_id, w_tmp, p0, p1 in _rotate_vir_block(wvovo, vlocs=vlocs): tmp -= einsum('acki, cjbk -> abij', l2T[:, p0:p1], w_tmp) w_tmp = None wvovo = None tmp = tmp - tmp.transpose(0, 1, 3, 2) l2Tnew += tmp tmpT = mpi.alltoall_new([tmp[:, p0:p1] for p0, p1 in vlocs], split_recvbuf=True) for task_id, (p0, p1) in enumerate(vlocs): tmp = tmpT[task_id].reshape(p1 - p0, nvir_seg, nocc, nocc) l2Tnew[:, p0:p1] -= tmp.transpose(1, 0, 2, 3) tmp = None #tmp = einsum('ak, ijkb -> baij', l1T, eris.ooox) tmp = einsum('ak, bkji -> baij', l1T, eris.xooo) tmp -= tmp_2 tmp1vv = mba + np.dot(t1T, l1T.T) # ba tmp -= einsum('ca, bcij -> baij', tmp1vv, vvoo) l2Tnew += tmp tmpT = mpi.alltoall_new([tmp[:, p0:p1] for p0, p1 in vlocs], split_recvbuf=True) for task_id, (p0, p1) in enumerate(vlocs): tmp = tmpT[task_id].reshape(p1 - p0, nvir_seg, nocc, nocc) l2Tnew[:, p0:p1] -= tmp.transpose(1, 0, 2, 3) tmp = None #tmp = einsum('jcab, ci -> baji', eris.ovvx, -l1T) tmp = einsum('bacj, ci -> baji', eris.xvvo, -l1T) tmp += einsum('abki, jk -> abij', l2T, v2) tmp1oo = mij + np.dot(l1T.T, t1T) # ik tmp -= einsum('ik, abkj -> abij', tmp1oo, vvoo) l2Tnew += tmp l2Tnew -= tmp.transpose(0, 1, 3, 2) tmp = None l1Tnew += fvo #tmp = einsum('bj, ibja -> ai', -l1T[vloc0:vloc1], eris.oxov) tmp = einsum('bj, biaj -> ai', -l1T[vloc0:vloc1], eris.xovo) l1Tnew += np.dot(v1.T, l1T) l1Tnew -= np.dot(l1T, v2.T) tmp -= einsum('cakj, icjk -> ai', l2T, imds.wovoo) tmp -= einsum('bcak, bcik -> ai', imds.wvvvo, l2T) tmp += einsum('baji, bj -> ai', l2T, imds.w3[vloc0:vloc1]) tmp_2 = t1T[vloc0:vloc1] - np.dot(tmp1vv[vloc0:vloc1], t1T) tmp_2 -= np.dot(t1T[vloc0:vloc1], mij) tmp_2 += einsum('bcjk, ck -> bj', t2T, l1T) tmp += einsum('baji, bj -> ai', vvoo, tmp_2) tmp_2 = None #tmp += einsum('icab, bc -> ai', eris.oxvv, tmp1vv[:, vloc0:vloc1]) tmp += einsum('ciba, bc -> ai', eris.xovv, tmp1vv[:, vloc0:vloc1]) l1Tnew += mpi.allreduce_inplace(tmp) #l1Tnew -= mpi.allgather(einsum('jika, kj -> ai', eris.ooox, tmp1oo)) l1Tnew -= mpi.allgather(einsum('akij, kj -> ai', eris.xooo, tmp1oo)) tmp = fvo - mpi.allreduce_inplace( einsum('bakj, bj -> ak', vvoo, t1T[vloc0:vloc1])) vvoo = None l1Tnew -= np.dot(tmp, mij.T) l1Tnew -= np.dot(mba.T, tmp) eia = mo_e_o[:, None] - mo_e_v l1Tnew /= eia.T for i in range(vloc0, vloc1): l2Tnew[i - vloc0] /= lib.direct_sum('i + jb -> bij', eia[:, i], eia) time0 = log.timer_debug1('update l1 l2', *time0) return l1Tnew.T, l2Tnew.transpose(2, 3, 0, 1)
def get_jkgrd(mol_or_mf, dm, mo_coeff=None, mo_occ=None): '''MPI version of scf.hf.get_jk function''' if rank == 0: print('uuuu00', lib.current_memory()[0]) if isinstance(mol_or_mf, gto.mole.Mole): mf = hf.SCF(mol_or_mf).view(SCF) else: mf = mol_or_mf # dm may be too big for mpi4py library to serialize. Broadcast dm here. if any(comm.allgather(isinstance(dm, str) and dm == 'SKIPPED_ARG')): dm = mpi.bcast_tagged_array_occdf(dm) mf.unpack_(comm.bcast(mf.pack())) if mf.opt is None: mf.opt = mf.init_direct_scf() mf.with_df = mf mol = mf.mol auxbasis = mf.auxbasis auxbasis = comm.bcast(auxbasis) mf.auxbasis = comm.bcast(mf.auxbasis) auxmol = df.addons.make_auxmol(mol, auxbasis) nao = mol.nao_nr() naux = auxmol.nao_nr() if rank == 0: print('number of AOs', nao) print('number of auxiliary basis functions', naux) # (d/dX i,j|P) # int3c_e1 = df.incore.aux_e2(mol, auxmol, intor='int3c2e_ip1', aosym='s1', comp=3) int3c_e1 = loop_aux(mf, intor='int3c2e_ip1', aosym='s1', comp=3) # (i,j|d/dX P) # int3c_e2 = df.incore.aux_e2(mol, auxmol, intor='int3c2e_ip2', aosym='s1', comp=3) int3c_e2 = loop_aux(mf, intor='int3c2e_ip2', aosym='s1', comp=3) # (d/dX P|Q) int2c_e1 = auxmol.intor('int2c2e_ip1', aosym='s1', comp=3) fmmm = _ao2mo.libao2mo.AO2MOmmm_bra_nr_s2 fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv ftrans = _ao2mo.libao2mo.AO2MOtranse2_nr_s2 null = lib.c_null_ptr() dms = numpy.asarray(dm) dm_shape = dms.shape nao = dm_shape[-1] dms = dms.reshape(-1, nao, nao) nset = dms.shape[0] vj = [0] * nset vk = [0] * nset if 0 == 0: nmo = mo_occ.shape[-1] mo_coeff = mo_coeff.reshape(-1, nao, nmo) mo_occ = mo_occ.reshape(-1, nmo) if mo_occ.shape[0] * 2 == nset: # handle ROHF DM mo_coeff = numpy.vstack((mo_coeff, mo_coeff)) mo_occa = numpy.array(mo_occ > 0, dtype=numpy.double) mo_occb = numpy.array(mo_occ == 2, dtype=numpy.double) assert (mo_occa.sum() + mo_occb.sum() == mo_occ.sum()) mo_occ = numpy.vstack((mo_occa, mo_occb)) dmtril = [] orbo = [] orbo0 = [] kiv = [] for k in range(nset): c = numpy.einsum('pi,i->pi', mo_coeff[k][:, mo_occ[k] > 0], numpy.sqrt(mo_occ[k][mo_occ[k] > 0])) orbo.append(numpy.asarray(c, order='F')) orbo0.append( numpy.asarray(mo_coeff[k][:, mo_occ[k] > 0], order='F')) nocc = orbo[k].shape[1] kiv.append(numpy.zeros((nocc, nao))) rho = [] split = [] buf1 = [] iokr = mpi_occ_df_incore.iokr rec = mpi_occ_df_incore.rec k = 0 dmtril = [] for k in range(nset): dmtril.append(lib.pack_tril(dms[k] + dms[k].T)) i = numpy.arange(nao) dmtril[k][i * (i + 1) // 2 + i] *= .5 tmp0 = int3c_e2.swapaxes(1, 3).reshape(-1, nao * nao) tmp = numpy.dot(tmp0, dm.reshape(-1)) tmp0 = numpy.einsum('xp,p->xp', tmp.reshape(3, -1), rec[k]) ec1_3cp = mpi.gather(tmp0.reshape(-1, order='F')).reshape((3, -1), order='F') ec1_3cu_b = numpy.einsum('xuvp,p,uv->xu', int3c_e1, rec[k], dm) ec1_3cu = comm.reduce(ec1_3cu_b) if rank == 0: j0 = 0 else: j0 = pauxz[rank - 1] j1 = pauxz[rank] tmp0 = int2c_e1[:, :, j0:j1].dot(rec[k]) tmp = comm.allreduce(tmp0) tmp0 = numpy.einsum('xp,p->xp', tmp[:, j0:j1], rec[k]) ec1_2c = mpi.gather(tmp0.reshape(-1, order='F')).reshape((3, -1), order='F') coeff3mo = numpy.sqrt(2.0e0) * (iokr[k].reshape(-1, nocc, nocc)) tmp = numpy.tensordot(coeff3mo, orbo[k], axes=([2], [1])) coeffb = numpy.tensordot(tmp, orbo[k], axes=([1], [1])) ex1_3cp_b = numpy.einsum('puv,xuvp->xp', coeffb, int3c_e2) ex1_3cp = mpi.gather(ex1_3cp_b.reshape(-1, order='F')).reshape( (3, -1), order='F') ex1_3cu_b = numpy.einsum('puv,xuvp->xu', coeffb, int3c_e1) ex1_3cu = comm.reduce(ex1_3cu_b) tmp0 = mpi.allgather(coeff3mo) tmp = numpy.tensordot(tmp0, coeff3mo, axes=([1, 2], [1, 2])) ex1_2c_b = numpy.einsum('xpq,pq->xp', int2c_e1[:, :, j0:j1], tmp) ex1_2c = comm.reduce(ex1_2c_b) eaux1p = numpy.empty((3, naux)) eaux1u = numpy.empty((3, naux)) if rank == 0: eaux1p = -0.5 * (2 * ec1_3cp - 2 * ec1_2c - 0.5 * (2 * ex1_3cp - 2 * ex1_2c)) eaux1u = -0.5 * (4 * ec1_3cu - 0.5 * (4 * ex1_3cu)) return eaux1p, eaux1u