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