def matvec(y):
y = np.asarray(y, order='C')
r = np.zeros_like(y)
yi = y[:nocc]
ri = r[:nocc]
yija = y[nocc:].reshape(nocc, -1)
rija = r[nocc:].reshape(nocc, -1)
for i in mpi_helper.distr_iter(range(nocc)):
kja = np.dot(Loo[:, i].T, Lov.reshape(-1, nocc * nvir))
kia = np.dot(Loo.reshape(-1, nocc * nocc).T, Lov[:, i])
kia = kia.reshape(nocc, -1)
v = 2.0 * kja - kia
ri += np.dot(kja, yija[i])
rija[i] += np.dot(v.T, yi)
rija[i] += eija[i] * yija[i]
mpi_helper.barrier()
mpi_helper.allreduce_inplace(ri)
mpi_helper.allreduce_inplace(rija)
ri += np.dot(h1, yi)
return r
def matvec(y):
y = np.asarray(y, order='C')
r = np.zeros_like(y)
yi = y[:nocc]
ri = r[:nocc]
yija = y[nocc:]
rija = r[nocc:]
yija_block = yija[q0:q1]
rija_block = rija[q0:q1]
ri += np.dot(ooov_as_block, yija_block)
rija_block += np.dot(yi, ooov_block)
rija_block += eija_block * yija_block
#TODO: really not a fan of this... mixing and matching MPI strategies
p0, p1 = mpi_helper.distr_blocks(nocc)
yija = yija.reshape(nocc, nocc, nvir)
yija_as = 2.0 * yija - yija.swapaxes(0, 1)
rija_block = rija.reshape(nocc, nocc, nvir)[p0:p1]
rija_block -= utils.einsum('ikjl,kla->ija', oooo[p0:p1], yija) * sign
rija_block += utils.einsum('ilba,ljb->ija', oovv[p0:p1], yija) * sign
rija_block -= utils.einsum('jalb,ilb->ija', ovov,
yija_as[p0:p1]) * sign
rija_block += utils.einsum('jlba,ilb->ija', oovv, yija[p0:p1]) * sign
mpi_helper.barrier()
mpi_helper.allreduce_inplace(ri)
mpi_helper.allreduce_inplace(rija)
ri += np.dot(h1, yi)
return r
def matvec(y):
y = np.asarray(y, order='C')
r = np.zeros_like(y)
yi = y[:nocc]
ri = r[:nocc]
yija = y[nocc:].reshape(nocc, nocc, nvir)
rija = r[nocc:].reshape(nocc, nocc, nvir)
ri += np.dot(h1, yi)
ri += mpi_helper.einsum('kija,ija->k', as2(ooov, (1, 2)), yija)
rija += mpi_helper.einsum('k,kija->ija', yi, ooov)
rija += eija * yija
tmp1 = np.zeros_like(rija)
for i in mpi_helper.distr_iter(range(nocc)):
v = np.dot(helper.Loo[:, i].T, helper.Loo.reshape(-1, nocc * nocc))
v = v.reshape(nocc, nocc, nocc).swapaxes(0, 1).reshape(nocc, -1)
tmp1[i] -= np.dot(v, yija.reshape(nocc * nocc, -1)) * sign
mpi_helper.barrier()
mpi_helper.allreduce_inplace(tmp1)
rija += tmp1
rija += mpi_helper.einsum('ilba,ljb->ija', oovv, yija) * sign
rija += mpi_helper.einsum('jlba,ilb->ija', oovv, yija) * sign
rija -= mpi_helper.einsum('jalb,ilb->ija', ovov, as2(yija,
(0, 1))) * sign
tmp1 = mpi_helper.einsum('ibjc,jia->cab', t2a, as1(yija,
(0, 1))) * 0.25
ri += mpi_helper.einsum('cab,icab->i', tmp1, as1(ovvv)) * sign
tmp1 = mpi_helper.einsum('ibjc,jia->cab', t2, yija)
ri += mpi_helper.einsum('cab,icab->i', tmp1, ovvv) * sign
tmp1 = mpi_helper.einsum('i,ibac->cab', yi, ovvv)
rija += mpi_helper.einsum('cab,jcib->ija', tmp1, t2) * sign
tmp1 = mpi_helper.einsum('jalb,kja->blk', t2a, as1(yija, (0, 1)))
tmp1 += mpi_helper.einsum('jalb,kja->blk', t2, yija)
ri += mpi_helper.einsum('blk,iklb->i', tmp1, as1(ooov, (0, 2))) * sign
tmp1 = mpi_helper.einsum('jalb,kja->blk', t2, as1(yija, (0, 1)))
tmp1 += mpi_helper.einsum('jalb,kja->blk', t2a, yija)
ri += mpi_helper.einsum('blk,iklb->i', tmp1, ooov) * sign
tmp1 = mpi_helper.einsum('jbla,jka->blk', t2, yija)
ri -= mpi_helper.einsum('blk,lkib->i', tmp1, ooov) * sign
tmp1 = mpi_helper.einsum('i,iklb->kbl', yi, as1(ooov, (0, 2)))
rija += mpi_helper.einsum('kbl,jalb->kja', tmp1, t2) * sign
tmp1 = mpi_helper.einsum('i,iklb->kbl', yi, ooov)
rija += mpi_helper.einsum('kbl,jalb->kja', tmp1, t2a) * sign
tmp1 = mpi_helper.einsum('i,ljib->ljb', yi, ooov)
rija -= mpi_helper.einsum('ljb,kbla->kja', tmp1, t2) * sign
return mpi_helper.mean(r) # FIXME: robust enough?
def get_1h(helper):
t2, ovov, ooov, oooo, oovv, eija = helper.unpack()
nocc, nvir = helper.nocc, helper.nvir
p0, p1 = mpi_helper.distr_blocks(nocc * nvir**2)
t2_block = t2.reshape(nocc, -1)[:, p0:p1]
ovov_as_block = ovov.reshape(nocc, -1)[:, p0:p1] * 2.0
ovov_as_block -= ovov.swapaxes(1, 3).reshape(nocc, -1)[:, p0:p1]
h1 = np.dot(t2_block, ovov_as_block.T) * 0.5
h1 += h1.T
mpi_helper.barrier()
mpi_helper.allreduce_inplace(h1)
h1 += np.diag(helper.eo)
return h1
def matvec(y):
y = np.asarray(y, order='C')
r = np.zeros_like(y)
yi = y[:nocc]
ri = r[:nocc]
yija = y[nocc:].reshape(nocc, -1)
rija = r[nocc:].reshape(nocc, -1)
for i in mpi_helper.distr_iter(range(nocc)):
kja = np.dot(Loo[:, i].T, Lov.reshape(-1, nocc * nvir))
kia = np.dot(Loo.reshape(-1, nocc * nocc).T, Lov[:, i])
kia = kia.reshape(nocc, -1)
v = 2.0 * kja - kia
ri += np.dot(v, yija[i])
rija[i] += np.dot(kja.T, yi)
rija[i] += eija[i].ravel() * yija[i]
yija = yija.reshape(nocc, nocc, nvir)
rija = rija.reshape(nocc, nocc, nvir)
yija_as = 2.0 * yija - yija.swapaxes(0, 1)
for l in mpi_helper.distr_iter(range(nocc)):
o_vv = np.dot(Loo[:, :, l].T,
Lvv.reshape(-1,
nvir * nvir)).reshape(nocc, nvir, nvir)
ov_v = np.dot(Lov.reshape(-1, nocc * nvir).T,
Lov[:, l]).reshape(nocc, nvir, nvir)
ooo_ = np.dot(Loo.reshape(-1, nocc * nocc).T,
Loo[:, :, l]).reshape(nocc, nocc, nocc)
rija -= utils.einsum('ikj,ka->ija', ooo_, yija[:, l]) * sign
rija += utils.einsum('iba,jb->ija', o_vv, yija[l]) * sign
rija -= utils.einsum('jab,ib->ija', ov_v, yija_as[:, l]) * sign
rija += utils.einsum('jba,ib->ija', o_vv, yija[:, l]) * sign
mpi_helper.barrier()
mpi_helper.allreduce_inplace(ri)
mpi_helper.allreduce_inplace(rija)
ri += np.dot(h1, yi)
return r
def get_1h(helper, ki):
t2, ovov, ooov, eija = helper.unpack()
nocc = max(helper.nocc)
h1 = np.zeros((nocc, nocc), dtype=helper.dtype)
for kj, kk in mpi_helper.distr_iter(helper.kpt_loop(2)):
kl = helper.kconserv[ki, kj, kk]
vk = 2.0 * ovov[ki, kj, kk] - ovov[ki, kl, kk].swapaxes(1, 3)
vk = vk.reshape(nocc, -1)
t2k = t2[ki, kj, kk].reshape(nocc, -1)
h1 += np.dot(t2k, vk.T.conj()) * 0.5
mpi_helper.barrier()
mpi_helper.allreduce_inplace(h1)
h1 += h1.T.conj()
h1 += np.diag(helper.eo[ki])
return h1
def build(self):
nmo_per_kpt = [x.size for x in self.o]
nocc_per_kpt = [np.sum(x) for x in self.o]
nmo, nocc = max(nmo_per_kpt), max(nocc_per_kpt)
nvir = nmo - nocc
self.opad, self.vpad = _padding_k_idx(nmo_per_kpt,
nocc_per_kpt,
kind='split')
self.kconserv = tools.get_kconserv(self.mf.cell, self.mf.kpts)
dtype = np.complex128 # TODO
self.eo = [e[o] for e, o in zip(self.e, self.o)]
self.ev = [e[v] for e, v in zip(self.e, self.v)]
self.co = [c[:, o] for c, o in zip(self.c, self.o)]
self.cv = [c[:, v] for c, v in zip(self.c, self.v)]
self.eo = np.array([e[x] for e, x in zip(self.eo, self.opad)])
self.ev = np.array([e[x] for e, x in zip(self.ev, self.vpad)])
self.co = np.array([c[:, x] for c, x in zip(self.co, self.opad)])
self.cv = np.array([c[:, x] for c, x in zip(self.cv, self.vpad)])
self.ovov = np.zeros((self.nkpts, ) * 3 + (nocc, nvir, nocc, nvir),
dtype=dtype)
self.ooov = np.zeros((self.nkpts, ) * 3 + (nocc, nocc, nocc, nvir),
dtype=dtype)
self.eija = np.zeros((self.nkpts, ) * 3 + (nocc, nocc, nvir))
self.t2 = self.ovov.copy()
for ki, kj, kk in mpi_helper.distr_iter(self.kpt_loop(3)):
kl = self.kconserv[ki, kj, kk]
kpts = [ki, kj, kk, kl]
eo, ev, co, cv = self.eo, self.ev, self.co, self.cv
self.ovov[ki, kj, kk] = self.ao2mo(kpts, co[ki], cv[kj], co[kk],
cv[kl])
self.ooov[ki, kj, kk] = self.ao2mo(kpts, co[ki], co[kj], co[kk],
cv[kl])
self.eija[ki, kj, kk] = lib.direct_sum('i,j,a->ija', eo[kj],
eo[kk], -ev[kl])
eiajb = lib.direct_sum('i,a,j,b->iajb', eo[ki], -ev[kj], eo[kk],
-ev[kl])
self.t2[ki, kj, kk] = self.ovov[ki, kj, kk] / eiajb
mpi_helper.barrier()
mpi_helper.allreduce_inplace(self.ovov)
mpi_helper.allreduce_inplace(self.ooov)
mpi_helper.allreduce_inplace(self.eija)
mpi_helper.allreduce_inplace(self.t2)
self._to_unpack = ['t2', 'ovov', 'ooov', 'eija']
def get_1h(helper):
Lov, Loo, Lvv, eia, eija = helper.unpack()
nocc, nvir = helper.nocc, helper.nvir
h1 = np.zeros((nocc, nocc))
ejab = lib.direct_sum('ja,b->jab', eia, -helper.ev)
for i in mpi_helper.distr_iter(range(nocc)):
eiajb = (helper.eo[i] + ejab).reshape(nocc, -1)
iba = np.dot(Lov.reshape(-1, nocc * nvir).T, Lov[:,
i]).reshape(nocc, -1)
iab = iba.reshape(nocc, nvir, nvir).swapaxes(1, 2).reshape(nocc, -1)
t2 = iab / eiajb
v = 2.0 * iab - iba
h1 += np.dot(v, t2.T) * 0.5
mpi_helper.barrier()
mpi_helper.allreduce_inplace(h1)
h1 += h1.T
h1 += np.diag(helper.eo)
return h1
def matvec(y):
y = np.asarray(y, order='C')
r = np.zeros_like(y)
yi = y[:nocc]
ri = r[:nocc]
yija = y[nocc:]
rija = r[nocc:]
yija_block = yija[q0:q1]
rija_block = rija[q0:q1]
ri += np.dot(ooov_block, yija_block)
rija_block += np.dot(yi, ooov_as_block)
rija_block += eija_block * yija_block
mpi_helper.barrier()
mpi_helper.allreduce_inplace(ri)
mpi_helper.allreduce_inplace(rija)
ri += np.dot(h1, yi)
return r
def matvec(y):
y = np.asarray(y, order='C', dtype=helper.dtype)
r = np.zeros_like(y)
yi = y[:nocc]
ri = r[:nocc]
yija = y[nocc:].reshape(nkpts, nkpts, -1)
rija = r[nocc:].reshape(nkpts, nkpts, -1)
for kj, kk in mpi_helper.distr_iter(helper.kpt_loop(2)):
kl = helper.kconserv[ki, kj, kk]
vk = 2.0 * ooov[ki, kj, kk] - ooov[ki, kk, kj].swapaxes(1, 2)
vk = vk.reshape(nocc, -1)
ri += np.dot(ooov[ki, kj, kk].reshape(nocc, -1), yija[kj, kk])
rija[kj, kk] += np.dot(yi, vk.conj())
rija[kj, kk] += eija[kj, kk, kl].ravel() * yija[kj, kk]
mpi_helper.barrier()
mpi_helper.allreduce_inplace(ri)
mpi_helper.allreduce_inplace(rija)
ri += np.dot(h1, yi)
return r
def get_1h(helper):
t1_2, t2, t2_2, ovov, ooov, oovv, ovvv, eija = helper.unpack()
nocc = helper.nocc
sign = helper.sign
t2a = as1(t2)
t2_2_a = as1(t2_2)
h1 = np.diag(helper.eo)
tmp1 = dot_along_tail(t2, as2(ovov)) * 0.5
h1 += tmp1
h1 += tmp1.T
tmp1 = mpi_helper.einsum('ijld,ld->ij', as2(ooov, (0, 2)), t1_2)
h1 += sign * tmp1
h1 += sign * tmp1.T
tmp1 = dot_along_tail(t2_2_a, as1(ovov)) * 0.5
tmp1 += dot_along_tail(t2_2, ovov.swapaxes(1, 3)) * 0.5
tmp2 = lib.direct_sum('ijb,a->iajb', eija, -helper.ev)
tmp1 -= dot_along_tail(as2(t2_2) * tmp2, t2) * 0.5
h1 += sign * tmp1
h1 += sign * tmp1.T
tmp1 = dot_along_tail2(as2(t2), as2(t2))
tmp2 = dot_along_tail(tmp1, ovov)
tmp1 = dot_along_tail2(t2a, t2a)
tmp1 += dot_along_tail2(t2, t2)
tmp1 += dot_along_tail2(t2.swapaxes(1, 3), t2.swapaxes(1, 3))
tmp2 -= dot_along_tail(tmp1.swapaxes(1, 2), oovv)
h1 += sign * tmp2
h1 += sign * tmp2.T
tmp1 = helper._t2_oooo.copy().swapaxes(1, 2)
tmp1 += helper._t2_vvvv.copy().swapaxes(1, 2) * 0.5
tmp2 = dot_along_tail2(t2, oovv.swapaxes(1, 2)) * -0.5
tmp3 = dot_along_tail(t2, tmp1 + tmp2)
tmp2 = dot_along_tail2(t2a, oovv.swapaxes(1, 2)) * -0.5
tmp3 += dot_along_tail(t2a, 0.5 * as1(tmp1) + tmp2)
h1 += sign * tmp3
h1 += sign * tmp3.T
tmp1 = dot_along_tail(t2a, t2a)
tmp1 += dot_along_tail(t2, t2) * 2.0
tmp2 = np.zeros_like(h1)
for i in mpi_helper.distr_iter(range(nocc)):
v = np.dot(helper.Loo[:, i].T,
helper.Loo.reshape(-1,
nocc * nocc)).reshape(nocc, nocc, nocc)
v = 2.0 * v - v.swapaxes(0, 2)
tmp2[i] -= np.dot(v.reshape(nocc, -1), tmp1.ravel()) * sign * 0.5
mpi_helper.barrier()
mpi_helper.allreduce_inplace(tmp2)
h1 += tmp2
tmp1 = dot_along_tail(t2a.swapaxes(0, 1), t2a.swapaxes(0, 1))
tmp1 += dot_along_tail(t2.swapaxes(0, 1), t2.swapaxes(0, 1)) * 2.0
h1 += mpi_helper.einsum('ijbc,bc->ij', oovv, tmp1) * sign
h1 -= mpi_helper.einsum('ibjc,bc->ij', ovov, tmp1) * sign * 0.5
tmp1 = dot_along_tail2(as2(t2), ovov)
h1 += dot_along_tail(as2(t2), tmp1) * sign
tmp1 = dot_along_tail2(t2.swapaxes(1, 3), oovv.swapaxes(1, 2))
h1 -= dot_along_tail(t2.swapaxes(1, 3), tmp1) * sign
return h1
def build(self):
self.eo, self.ev = self.e[self.o], self.e[self.v]
self.co, self.cv = self.c[:, self.o], self.c[:, self.v]
self.Loo = self.ao2mo(self.co, self.co)
self.Lvv = self.ao2mo(self.cv, self.cv)
self.Lov = self.ao2mo(self.co, self.cv)
nocc, nvir = self.nocc, self.nvir
self.ovov = mpi_helper.dot(
self.Lov.reshape(-1, nocc * nvir).T,
self.Lov.reshape(-1, nocc * nvir),
).reshape(nocc, nvir, nocc, nvir)
self.oovv = mpi_helper.dot(
self.Loo.reshape(-1, nocc * nocc).T,
self.Lvv.reshape(-1, nvir * nvir),
).reshape(nocc, nocc, nvir, nvir)
self.ooov = mpi_helper.dot(
self.Loo.reshape(-1, nocc * nocc).T,
self.Lov.reshape(-1, nocc * nvir),
).reshape(nocc, nocc, nocc, nvir)
self.ovvv = mpi_helper.dot(
self.Lov.reshape(-1, nocc * nvir).T,
self.Lvv.reshape(-1, nvir * nvir),
).reshape(nocc, nvir, nvir, nvir)
self.eija = lib.direct_sum('i,j,a->ija', self.eo, self.eo, -self.ev)
self.eajb = lib.direct_sum('a,j,b->ajb', -self.ev, self.eo, -self.ev)
eia = self.eia = lib.direct_sum('i,a->ia', self.eo, -self.ev)
eiajb = lib.direct_sum('ia,jb->iajb', eia, eia)
self.t2 = self.ovov / eiajb
t2a = self.t2 - self.t2.swapaxes(0, 2).copy()
self.t1_2 = mpi_helper.einsum('kdac,ickd->ia', self.ovvv,
self.t2 + t2a * 0.5)
self.t1_2 -= mpi_helper.einsum('kcad,ickd->ia', self.ovvv, t2a) * 0.5
self.t1_2 -= mpi_helper.einsum('kilc,kalc->ia', self.ooov,
self.t2 + t2a * 0.5)
self.t1_2 -= mpi_helper.einsum('likc,lakc->ia', self.ooov, t2a) * 0.5
self.t1_2 /= eia
self._t2_oooo = np.zeros((nocc, nocc, nvir, nvir))
self._t2_vvvv = np.zeros((nocc, nocc, nvir, nvir))
for i in mpi_helper.distr_iter(range(nocc)):
v = np.dot(self.Loo[:, i].T,
self.Loo.reshape(-1, nocc * nocc)).reshape(
nocc, nocc, nocc)
self._t2_oooo += np.tensordot(v, self.t2[i], axes=((1, ), (1, )))
for a in mpi_helper.distr_iter(range(nvir)):
v = np.dot(self.Lvv[:, a].T,
self.Lvv.reshape(-1, nvir * nvir)).reshape(
nvir, nvir, nvir)
self._t2_vvvv += np.tensordot(self.t2[:, a],
v,
axes=((2, ), (2, )))
mpi_helper.barrier()
mpi_helper.allreduce_inplace(self._t2_oooo)
mpi_helper.allreduce_inplace(self._t2_vvvv)
self.t2_2 = self._t2_oooo.copy().swapaxes(1, 2)
self.t2_2 += self._t2_vvvv.copy().swapaxes(1, 2)
self.t2_2 += mpi_helper.einsum('kcjb,iakc->iajb', self.ovov,
self.t2 + t2a)
self.t2_2 -= mpi_helper.einsum('kjbc,iakc->iajb', self.oovv, self.t2)
self.t2_2 -= mpi_helper.einsum('kibc,kajc->iajb', self.oovv, self.t2)
self.t2_2 -= mpi_helper.einsum('kjac,ickb->iajb', self.oovv, self.t2)
self.t2_2 += mpi_helper.einsum('kcia,kcjb->iajb', self.ovov,
self.t2 + t2a)
self.t2_2 -= mpi_helper.einsum('kiac,kcjb->iajb', self.oovv, self.t2)
self.t2_2 /= eiajb
self.sign = 1
self.guess_high_order = True
self._to_unpack = [
't1_2', 't2', 't2_2', 'ovov', 'ooov', 'oovv', 'ovvv', 'eija'
]