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, -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 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.ovov = self.ao2mo(self.co, self.cv, self.co, self.cv)
self.ooov = self.ao2mo(self.co, self.co, self.co, self.cv)
self.oooo = self.ao2mo(self.co, self.co, self.co, self.co)
self.oovv = self.ao2mo(self.co, self.co, self.cv, self.cv)
self.ovvv = self.ao2mo(self.co, self.cv, self.cv, self.cv)
self.vvvv = self.ao2mo(self.cv, self.cv, self.cv, self.cv)
self.eija = lib.direct_sum('i,j,a->ija', self.eo, self.eo, -self.ev)
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
self._t2_oooo = np.tensordot(self.oooo, self.t2, axes=((0, 2), (0, 2)))
self._t2_vvvv = np.tensordot(self.t2, self.vvvv, axes=((1, 3), (0, 2)))
t2a = self.t2 - self.t2.swapaxes(0, 2).copy()
self.t1_2 = utils.einsum('kdac,ickd->ia', self.ovvv, self.t2 + t2a)
self.t1_2 -= utils.einsum('kilc,kalc->ia', self.ooov, self.t2 + t2a)
self.t1_2 /= eia
self.t2_2 = utils.einsum('ijab->iajb', self._t2_oooo.copy())
self.t2_2 += utils.einsum('ijab->iajb', self._t2_vvvv.copy())
self.t2_2 += utils.einsum('kcjb,iakc->iajb', self.ovov, self.t2 + t2a)
self.t2_2 -= utils.einsum('kjcb,iakc->iajb', self.oovv, self.t2)
self.t2_2 -= utils.einsum('kicb,kajc->iajb', self.oovv, self.t2)
self.t2_2 -= utils.einsum('kjac,ickb->iajb', self.oovv, self.t2)
self.t2_2 += utils.einsum('kcia,kcjb->iajb', self.ovov, self.t2 + t2a)
self.t2_2 -= utils.einsum('kica,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', 'oooo', 'oovv', 'ovvv',
'vvvv', 'eija'
]
def get_moments(helper, nmax):
t2, ovov, ooov, oooo, oovv, eija = helper.unpack()
nocc = helper.nocc
sign = helper.sign
vl = 2.0 * ooov - ooov.swapaxes(1, 2)
vr = ooov.copy()
t = np.zeros((nmax + 1, nocc, nocc), dtype=ovov.dtype)
t[0] = np.dot(vl.reshape(nocc, -1), vr.reshape(nocc, -1).T.conj())
for n in range(1, nmax + 1):
vr = (-utils.einsum('ikjl,xkla->xija', oooo, vr) * sign +
utils.einsum('ilba,xljb->xija', oovv, vr) * sign -
utils.einsum('jalb,xilb->xija', ovov, vr) * sign * 2 +
utils.einsum('jalb,xlib->xija', ovov, vr) * sign +
utils.einsum('jlba,xilb->xija', oovv, vr) * sign +
utils.einsum('ija,xija->xija', eija, vr))
t[n] += utils.einsum('xija,yija->xy', vl, vr)
return t
def get_matvec(helper):
t2, ovov, ooov, oooo, oovv, eija = helper.unpack()
nocc, nvir = helper.nocc, helper.nvir
sign = helper.sign
p0, p1 = mpi_helper.distr_blocks(nocc * nvir**2)
ovov_as_block = ovov.reshape(nocc, -1)[:, p0:p1] * 2.0
ovov_as_block -= ovov.swapaxes(1, 3).reshape(nocc, -1)[:, p0:p1]
q0, q1 = mpi_helper.distr_blocks(nocc**2 * nvir)
ooov_block = ooov.reshape(nocc, -1)[:, q0:q1]
ooov_as_block = ooov_block * 2.0
ooov_as_block -= ooov.swapaxes(1, 2).reshape(nocc, -1)[:, q0:q1]
eija_block = eija.ravel()[q0:q1]
h1 = get_1h(helper)
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
diag = np.concatenate([np.diag(h1), eija.ravel()])
if helper.guess_high_order:
# According to A. Sokolov these might actually make things worse
# See https://github.com/pyscf/pyscf/commit/994e325159866bc74319418033db270a6b6a9d57#r45037621
diag_ija = diag[nocc:].reshape(nocc, nocc, nvir)
diag_ija -= utils.einsum('iijj->ij', oooo)[:, :, None] * sign
diag_ija += utils.einsum('jjaa->ja', oovv)[None, :, :] * sign
diag_ija += utils.einsum('iiaa->ia', oovv)[:, None, :] * sign
return matvec, diag
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 += utils.einsum('kija,ija->k', as2(ooov, (1, 2)), yija)
rija += utils.einsum('k,kija->ija', yi, ooov)
rija += eija * yija
rija -= utils.einsum('ikjl,kla->ija', oooo, yija) * sign
rija += utils.einsum('ilba,ljb->ija', oovv, yija) * sign
rija -= utils.einsum('jalb,ilb->ija', ovov, as2(yija, (0, 1))) * sign
rija += utils.einsum('jlba,ilb->ija', oovv, yija) * sign
tmp1 = utils.einsum('ibjc,jia->cab', t2a, as1(yija, (0, 1))) * 0.25
ri += utils.einsum('cab,icab->i', tmp1, as1(ovvv)) * sign
tmp1 = utils.einsum('kbjc,jka->cab', t2, yija)
ri += utils.einsum('cab,icab->i', tmp1, ovvv) * sign
tmp1 = utils.einsum('i,ibac->cab', yi, ovvv)
rija += utils.einsum('cab,jcib->ija', tmp1, t2) * sign
tmp1 = utils.einsum('jalb,kja->blk', t2a, as1(yija, (0, 1)))
tmp1 += utils.einsum('jalb,kja->blk', t2, yija)
ri += utils.einsum('blk,iklb->i', tmp1, as1(ooov, (0, 2))) * sign
tmp1 = utils.einsum('jalb,kja->blk', t2, as1(yija, (0, 1)))
tmp1 += utils.einsum('jalb,kja->blk', t2a, yija)
ri += utils.einsum('blk,iklb->i', tmp1, ooov) * sign
tmp1 = utils.einsum('jbla,jka->blk', t2, yija)
ri -= utils.einsum('blk,lkib->i', tmp1, ooov) * sign
tmp1 = utils.einsum('i,iklb->kbl', yi, as1(ooov, (0, 2)))
rija += utils.einsum('kbl,jalb->kja', tmp1, t2) * sign
tmp1 = utils.einsum('i,iklb->kbl', yi, ooov)
rija += utils.einsum('kbl,jalb->kja', tmp1, t2a) * sign
rija -= utils.einsum('i,ljib,kbla->kja', yi, ooov, t2) * sign
return r
def get_matvec(helper):
t1_2, t2, t2_2, ovov, ooov, oooo, oovv, ovvv, vvvv, eija = helper.unpack()
nocc, nvir = helper.nocc, helper.nvir
sign = helper.sign
t2a = as1(t2)
h1 = get_1h(helper)
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 += utils.einsum('kija,ija->k', as2(ooov, (1, 2)), yija)
rija += utils.einsum('k,kija->ija', yi, ooov)
rija += eija * yija
rija -= utils.einsum('ikjl,kla->ija', oooo, yija) * sign
rija += utils.einsum('ilba,ljb->ija', oovv, yija) * sign
rija -= utils.einsum('jalb,ilb->ija', ovov, as2(yija, (0, 1))) * sign
rija += utils.einsum('jlba,ilb->ija', oovv, yija) * sign
tmp1 = utils.einsum('ibjc,jia->cab', t2a, as1(yija, (0, 1))) * 0.25
ri += utils.einsum('cab,icab->i', tmp1, as1(ovvv)) * sign
tmp1 = utils.einsum('kbjc,jka->cab', t2, yija)
ri += utils.einsum('cab,icab->i', tmp1, ovvv) * sign
tmp1 = utils.einsum('i,ibac->cab', yi, ovvv)
rija += utils.einsum('cab,jcib->ija', tmp1, t2) * sign
tmp1 = utils.einsum('jalb,kja->blk', t2a, as1(yija, (0, 1)))
tmp1 += utils.einsum('jalb,kja->blk', t2, yija)
ri += utils.einsum('blk,iklb->i', tmp1, as1(ooov, (0, 2))) * sign
tmp1 = utils.einsum('jalb,kja->blk', t2, as1(yija, (0, 1)))
tmp1 += utils.einsum('jalb,kja->blk', t2a, yija)
ri += utils.einsum('blk,iklb->i', tmp1, ooov) * sign
tmp1 = utils.einsum('jbla,jka->blk', t2, yija)
ri -= utils.einsum('blk,lkib->i', tmp1, ooov) * sign
tmp1 = utils.einsum('i,iklb->kbl', yi, as1(ooov, (0, 2)))
rija += utils.einsum('kbl,jalb->kja', tmp1, t2) * sign
tmp1 = utils.einsum('i,iklb->kbl', yi, ooov)
rija += utils.einsum('kbl,jalb->kja', tmp1, t2a) * sign
rija -= utils.einsum('i,ljib,kbla->kja', yi, ooov, t2) * sign
return r
diag = np.concatenate([np.diag(h1), eija.ravel()])
#FIXME: is this broken? not sure it is correct
if helper.guess_high_order:
# According to A. Sokolov these might actually make things worse
# See https://github.com/pyscf/pyscf/commit/994e325159866bc74319418033db270a6b6a9d57#r45037621
diag_ija = diag[nocc:].reshape(nocc, nocc, nvir)
diag_ija -= utils.einsum('iijj->ij', oooo)[:, :, None] * sign
diag_ija += utils.einsum('jjaa->ja', oovv)[None, :, :] * sign
diag_ija += utils.einsum('iiaa->ia', oovv)[:, None, :] * sign
return matvec, diag
def get_1h(helper):
t1_2, t2, t2_2, ovov, ooov, oooo, oovv, ovvv, vvvv, eija = helper.unpack()
sign = helper.sign
t2a = as1(t2)
t2_2_a = as1(t2_2)
h1 = np.diag(helper.eo)
tmp1 = utils.einsum('iakb,jakb->ij', t2, as2(ovov)) * 0.5
h1 += tmp1
h1 += tmp1.T
tmp1 = utils.einsum('ld,jild->ij', t1_2, as2(ooov, (0, 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 -= utils.einsum('iakb,jakb->ij', as2(t2_2) * tmp2, t2) * 0.5
h1 += sign * tmp1
h1 += sign * tmp1.T
tmp1 = utils.einsum('lckb,jalc->jakb', as2(t2), as2(t2))
tmp2 = utils.einsum('jakb,iakb->ij', tmp1, ovov)
tmp1 = utils.einsum('lckb,jalc->jkab', t2a, t2a)
tmp1 += utils.einsum('lckb,jalc->jkab', t2, t2)
tmp1 += utils.einsum('lakc,jclb->jkab', t2, t2)
tmp2 -= utils.einsum('ikab,jkab->ij', tmp1, oovv)
h1 += sign * tmp2
h1 += sign * tmp2.T
tmp1 = utils.einsum('jkab->jakb', helper._t2_oooo.copy())
tmp1 += utils.einsum('jkab->jakb', helper._t2_vvvv.copy()) * 0.5
tmp2 = utils.einsum('jalc,klbc->jakb', t2, oovv) * -0.5
tmp3 = utils.einsum('iakb,jakb->ij', t2, tmp1 + tmp2)
tmp2 = utils.einsum('jalc,klbc->jakb', t2a, oovv) * -0.5
tmp3 += utils.einsum('iakb,jakb->ij', t2a, 0.5 * as1(tmp1) + tmp2)
h1 += sign * tmp3
h1 += sign * tmp3.T
tmp1 = utils.einsum('iakb,ialb->kl', t2a, t2a)
tmp1 += utils.einsum('iakb,ialb->kl', t2, t2) * 2.0
h1 -= utils.einsum('ijkl,kl->ij', as2(oooo), tmp1) * sign * 0.5
tmp1 = utils.einsum('iakc,iakb->bc', t2a, t2a)
tmp1 += utils.einsum('iakc,iakb->bc', t2, t2) * 2.0
h1 += utils.einsum('jibc,bc->ij', oovv, tmp1) * sign
h1 -= utils.einsum('jcib,bc->ij', ovov, tmp1) * sign * 0.5
tmp1 = utils.einsum('jalc,kblc->jakb', as2(t2), ovov)
h1 += utils.einsum('iakb,jakb->ij', as2(t2), tmp1) * sign
tmp1 = utils.einsum('jcla,klbc->jakb', t2, oovv)
h1 -= utils.einsum('ibka,jakb->ij', t2, tmp1) * sign
return h1
def get_moments(helper, nmax):
t1_2, t2, t2_2, ovov, ooov, oooo, oovv, ovvv, vvvv, eija = helper.unpack()
nocc = helper.nocc
sign = helper.sign
t2a = as2(t2)
vl = as2(ooov, axis=(1, 2))
vl += utils.einsum('jbic,kcba->kija', t2a, ovvv) * sign
vl += utils.einsum('jalb,kilb->kija', t2a, ooov) * sign * 2
vl -= utils.einsum('jalb,likb->kija', t2a, ooov) * sign
vl -= utils.einsum('ialb,kjlb->kija', t2a, ooov) * sign
vl -= utils.einsum('ibla,jlkb->kija', t2a, ooov) * sign
vr = ooov.copy()
vr += utils.einsum('ibjc,kbca->kija', t2, ovvv) * sign
vr -= utils.einsum('ibla,jlkb->kija', t2, ooov) * sign
vr += utils.einsum('jalb,kilb->kija', t2a, ooov) * sign
vr -= utils.einsum('jalb,ilkb->kija', t2, ooov) * sign
t = np.zeros((nmax + 1, nocc, nocc), dtype=ovov.dtype)
t[0] = np.dot(vl.reshape(nocc, -1), vr.reshape(nocc, -1).T.conj())
for n in range(1, nmax + 1):
vr = (-utils.einsum('ikjl,xkla->xija', oooo, vr) * sign +
utils.einsum('ilba,xljb->xija', oovv, vr) * sign -
utils.einsum('jalb,xilb->xija', ovov, as2(vr, (1, 2))) * sign +
utils.einsum('jlba,xilb->xija', oovv, vr) * sign +
utils.einsum('ija,xija->xija', eija, vr))
t[n] += utils.einsum('xija,yija->xy', vl, vr)
return t
def pick(w, v, nroots, envs):
x0 = lib.linalg_helper._gen_x0(envs['v'], envs['xs'])
s = np.dot(np.asarray(guess).conj(), np.asarray(x0).T)
snorm = utils.einsum('pi,pi->i', s.conj(), s)
idx = np.argsort(-snorm)[:nroots]
return lib.linalg_helper._eigs_cmplx2real(w, v, idx, real_system)
def get_matvec(helper):
t1_2, t2, t2_2, ovov, ooov, oovv, ovvv, eija = helper.unpack()
nocc, nvir = helper.nocc, helper.nvir
sign = helper.sign
t2a = as1(t2)
h1 = get_1h(helper)
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?
diag = np.concatenate([np.diag(h1), eija.ravel()])
if helper.guess_high_order:
# According to A. Sokolov these might actually make things worse
# See https://github.com/pyscf/pyscf/commit/994e325159866bc74319418033db270a6b6a9d57#r45037621
diag_ija = diag[nocc:].reshape(nocc, nocc, nvir)
diag_ija -= utils.einsum('Lii,Ljj->ij', helper.Loo,
helper.Loo)[:, :, None] * sign
diag_ija += utils.einsum('jjaa->ja', oovv)[None, :, :] * sign
diag_ija += utils.einsum('iiaa->ia', oovv)[:, None, :] * sign
return matvec, diag
def get_matvec(helper):
Lov, Loo, Lvv, eia, eija = helper.unpack()
nocc, nvir = helper.nocc, helper.nvir
sign = helper.sign
h1 = get_1h(helper)
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
diag = np.concatenate([np.diag(h1), eija.ravel()])
if helper.guess_high_order:
# According to A. Sokolov these might actually make things worse
# See https://github.com/pyscf/pyscf/commit/994e325159866bc74319418033db270a6b6a9d57#r45037621
diag_ija = diag[nocc:].reshape(nocc, nocc, nvir)
diag_ija -= utils.einsum('Lii,Ljj->ij', Loo, Loo)[:, :, None] * sign
diag_ija += utils.einsum('Ljj,Laa->ja', Loo, Lvv)[None, :, :] * sign
diag_ija += utils.einsum('Lii,Laa->ia', Loo, Lvv)[:, None, :] * sign
return matvec, diag
