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 cc_Woooo(t1T, t2T, eris, tauT=None, vlocs=None):
nvir = t1T.shape[0]
if vlocs is None:
ntasks = mpi.pool.size
vlocs = [_task_location(nvir, task_id) for task_id in range(ntasks)]
vloc0, vloc1 = vlocs[rank]
if tauT is None:
tauT = make_tauT(t1T, t2T, vlocs=vlocs)
Wmnij = einsum('efmn, efij -> mnij', eris.xvoo, tauT)
Wmnij *= 0.25
tauT = None
#tmp = einsum('mnie, ej -> mnij', eris.ooox, t1T[vloc0:vloc1])
#Wmnij += tmp
#Wmnij -= tmp.transpose(0, 1, 3, 2)
#tmp = None
Wmnij = mpi.allreduce_inplace(Wmnij)
Wmnij += eris.oooo
return Wmnij
def cc_Foo(t1T, t2T, eris, tauT_tilde=None, vlocs=None):
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]
Fmi = 0.5 * einsum('efmn, efin -> mi', eris.xvoo, tauT_tilde)
tauT_tilde = None
#Fmi += np.einsum('mnie, en -> mi', eris.ooox, t1T[vloc0:vloc1], optimize=True)
Fmi = mpi.allreduce_inplace(Fmi)
fov = eris.fock[:nocc, nocc:]
foo = eris.fock[:nocc, :nocc]
Fmi += foo
#Fmi += 0.5 * np.dot(fov, t1T)
return Fmi
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 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)
# 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)
l1Tnew = np.zeros_like(l1T)
l2Tnew = m3
l2Tnew += vvoo
#fvo1 = fvo #+ mpi.allreduce(einsum('cbkj, ck -> bj', vvoo, t1T[vloc0:vloc1]))
#tmp = np.einsum('ai, bj -> abij', l1T[vloc0:vloc1], fvo1, optimize=True)
tmp = 0.0
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 -= tmp_2
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('abki, jk -> abij', l2T, v2)
tmp1oo = mij #+ np.dot(l1T.T, t1T) # ik
tmp -= einsum('ik, abkj -> abij', tmp1oo, vvoo)
vvoo = None
l2Tnew += tmp
l2Tnew -= tmp.transpose(0, 1, 3, 2)
tmp = None
#l1Tnew += fvo
#tmp = einsum('bj, ibja -> ai', -l1T[vloc0:vloc1], eris.oxov)
#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])
#l1Tnew += mpi.allreduce(tmp)
#l1Tnew -= mpi.allgather(einsum('jika, kj -> ai', eris.ooox, tmp1oo))
#tmp = fvo - mpi.allreduce(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 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)