1#!/usr/bin/env python
2# Copyright 2017-2021 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import itertools
17import numpy as np
18from pyscf import lib
19from pyscf.pbc.lib import kpts_helper
20
21einsum = lib.einsum
22
23#FIXME: the dtype of each intermediates. When the khf is at gamma point, the
24# dtype is inconsistent between intermediates and t amplitudes
25
26def make_tau(cc, t2, t1, t1p, fac=1.):
27    t2aa, t2ab, t2bb = t2
28    nkpts = len(t2aa)
29
30    tauaa = t2aa.copy()
31    tauab = t2ab.copy()
32    taubb = t2bb.copy()
33    for ki in range(nkpts):
34        for kj in range(nkpts):
35            tauaa[ki,kj,ki] += einsum('ia,jb->ijab', fac*.5*t1[0][ki], t1p[0][kj])
36            tauaa[ki,kj,kj] -= einsum('ib,ja->ijab', fac*.5*t1[0][ki], t1p[0][kj])
37            tauaa[ki,kj,kj] -= einsum('ja,ib->ijab', fac*.5*t1[0][kj], t1p[0][ki])
38            tauaa[ki,kj,ki] += einsum('jb,ia->ijab', fac*.5*t1[0][kj], t1p[0][ki])
39
40            taubb[ki,kj,ki] += einsum('ia,jb->ijab', fac*.5*t1[1][ki], t1p[1][kj])
41            taubb[ki,kj,kj] -= einsum('ib,ja->ijab', fac*.5*t1[1][ki], t1p[1][kj])
42            taubb[ki,kj,kj] -= einsum('ja,ib->ijab', fac*.5*t1[1][kj], t1p[1][ki])
43            taubb[ki,kj,ki] += einsum('jb,ia->ijab', fac*.5*t1[1][kj], t1p[1][ki])
44
45            tauab[ki,kj,ki] += einsum('ia,jb->ijab', fac*.5*t1[0][ki], t1p[1][kj])
46            tauab[ki,kj,ki] += einsum('jb,ia->ijab', fac*.5*t1[1][kj], t1p[0][ki])
47    return tauaa, tauab, taubb
48
49def make_tau2(cc, t2, t1, t1p, fac=1.):
50    t2aa, t2ab, t2bb = t2
51    nkpts = len(t2aa)
52
53    tauaa = t2aa.copy()
54    tauab = t2ab.copy()
55    taubb = t2bb.copy()
56    for ki in range(nkpts):
57        for kj in range(nkpts):
58            tauaa[ki,kj,ki] += einsum('ia,jb->ijab', fac*.5*t1[0][ki], t1p[0][kj])
59            tauaa[ki,kj,ki] += einsum('jb,ia->ijab', fac*.5*t1[0][kj], t1p[0][ki])
60
61            taubb[ki,kj,ki] += einsum('ia,jb->ijab', fac*.5*t1[1][ki], t1p[1][kj])
62            taubb[ki,kj,ki] += einsum('jb,ia->ijab', fac*.5*t1[1][kj], t1p[1][ki])
63
64            tauab[ki,kj,ki] += einsum('ia,jb->ijab', fac*.5*t1[0][ki], t1p[1][kj])
65            tauab[ki,kj,ki] += einsum('jb,ia->ijab', fac*.5*t1[1][kj], t1p[0][ki])
66    return tauaa, tauab, taubb
67
68def cc_Fvv(cc, t1, t2, eris):
69    t1a, t1b = t1
70    t2aa, t2ab, t2bb = t2
71    nkpts, nocc_a, nvir_a = t1a.shape
72    nocc_b, nvir_b = t1b.shape[1:]
73
74    kconserv = cc.khelper.kconserv
75
76    fa = np.zeros((nkpts,nvir_a,nvir_a), dtype=np.complex128)
77    fb = np.zeros((nkpts,nvir_b,nvir_b), dtype=np.complex128)
78
79    tau_tildeaa,tau_tildeab,tau_tildebb=make_tau(cc,t2,t1,t1,fac=0.5)
80
81    fov = eris.fock[0][:,:nocc_a,nocc_a:]
82    fOV = eris.fock[1][:,:nocc_b,nocc_b:]
83    fvv = eris.fock[0][:,nocc_a:,nocc_a:]
84    fVV = eris.fock[1][:,nocc_b:,nocc_b:]
85
86    for ka in range(nkpts):
87        fa[ka]+=fvv[ka]
88        fb[ka]+=fVV[ka]
89        fa[ka]-=0.5*einsum('me,ma->ae',fov[ka],t1a[ka])
90        fb[ka]-=0.5*einsum('me,ma->ae',fOV[ka],t1b[ka])
91        for km in range(nkpts):
92            fa[ka]+=einsum('mf,fmea->ae',t1a[km], eris.vovv[km,km,ka].conj())
93            fa[ka]-=einsum('mf,emfa->ae',t1a[km], eris.vovv[ka,km,km].conj())
94            fa[ka]+=einsum('mf,fmea->ae',t1b[km], eris.VOvv[km,km,ka].conj())
95
96            fb[ka]+=einsum('mf,fmea->ae',t1b[km], eris.VOVV[km,km,ka].conj())
97            fb[ka]-=einsum('mf,emfa->ae',t1b[km], eris.VOVV[ka,km,km].conj())
98            fb[ka]+=einsum('mf,fmea->ae',t1a[km], eris.voVV[km,km,ka].conj())
99
100            for kn in range(nkpts):
101                kf = kconserv[km,ka,kn]
102                tmp = eris.ovov[km,ka,kn] - eris.ovov[km,kf,kn].transpose(0,3,2,1)
103                fa[ka] -= einsum('mnaf,menf->ae', tau_tildeaa[km,kn,ka], tmp) * .5
104                fa[ka] -= einsum('mNaF,meNF->ae', tau_tildeab[km,kn,ka], eris.ovOV[km,ka,kn])
105
106                tmp = eris.OVOV[km,ka,kn] - eris.OVOV[km,kf,kn].transpose(0,3,2,1)
107                fb[ka] -= einsum('mnaf,menf->ae', tau_tildebb[km,kn,ka], tmp) * .5
108                fb[ka] -= einsum('MnFa,MFne->ae', tau_tildeab[km,kn,kf], eris.ovOV[km,kf,kn])
109
110    return fa,fb
111
112
113def cc_Foo(cc, t1, t2, eris):
114    t1a, t1b = t1
115    t2aa, t2ab, t2bb = t2
116    nkpts, nocc_a, nvir_a = t1a.shape
117    nocc_b, nvir_b = t1b.shape[1:]
118
119    kconserv = cc.khelper.kconserv
120
121    fa = np.zeros((nkpts,nocc_a,nocc_a), dtype=np.complex128)
122    fb = np.zeros((nkpts,nocc_b,nocc_b), dtype=np.complex128)
123
124    tau_tildeaa,tau_tildeab,tau_tildebb=make_tau(cc,t2,t1,t1,fac=0.5)
125
126    fov = eris.fock[0][:,:nocc_a,nocc_a:]
127    fOV = eris.fock[1][:,:nocc_b,nocc_b:]
128    foo = eris.fock[0][:,:nocc_a,:nocc_a]
129    fOO = eris.fock[1][:,:nocc_b,:nocc_b]
130
131    for ka in range(nkpts):
132        fa[ka]+=foo[ka]
133        fb[ka]+=fOO[ka]
134        fa[ka]+=0.5*einsum('me,ne->mn',fov[ka],t1a[ka])
135        fb[ka]+=0.5*einsum('me,ne->mn',fOV[ka],t1b[ka])
136        for km in range(nkpts):
137            fa[ka]+=einsum('oa,mnoa->mn',t1a[km],eris.ooov[ka,ka,km])
138            fa[ka]+=einsum('oa,mnoa->mn',t1b[km],eris.ooOV[ka,ka,km])
139            fa[ka]-=einsum('oa,onma->mn',t1a[km],eris.ooov[km,ka,ka])
140
141            fb[ka]+=einsum('oa,mnoa->mn',t1b[km],eris.OOOV[ka,ka,km])
142            fb[ka]+=einsum('oa,mnoa->mn',t1a[km],eris.OOov[ka,ka,km])
143            fb[ka]-=einsum('oa,onma->mn',t1b[km],eris.OOOV[km,ka,ka])
144
145    for km in range(nkpts):
146        for kn in range(nkpts):
147            for ke in range(nkpts):
148                kf = kconserv[km,ke,kn]
149                tmp = eris.ovov[km,ke,kn] - eris.ovov[km,kf,kn].transpose(0,3,2,1)
150                fa[km] += einsum('inef,menf->mi', tau_tildeaa[km,kn,ke], tmp) * .5
151                fa[km] += einsum('iNeF,meNF->mi',tau_tildeab[km,kn,ke],eris.ovOV[km,ke,kn])
152
153                tmp = eris.OVOV[km,ke,kn] - eris.OVOV[km,kf,kn].transpose(0,3,2,1)
154                fb[km] += einsum('INEF,MENF->MI',tau_tildebb[km,kn,ke], tmp) * .5
155                fb[km] += einsum('nIeF,neMF->MI',tau_tildeab[kn,km,ke],eris.ovOV[kn,ke,km])
156
157    return fa,fb
158
159
160def cc_Fov(cc, t1, t2, eris):
161    t1a, t1b = t1
162    t2aa, t2ab, t2bb = t2
163    nkpts, nocc_a, nvir_a = t1a.shape
164    nocc_b, nvir_b = t1b.shape[1:]
165
166    fov = eris.fock[0][:,:nocc_a,nocc_a:]
167    fOV = eris.fock[1][:,:nocc_b,nocc_b:]
168
169    fa = np.zeros((nkpts,nocc_a,nvir_a), dtype=np.complex128)
170    fb = np.zeros((nkpts,nocc_b,nvir_b), dtype=np.complex128)
171
172    for km in range(nkpts):
173        fa[km]+=fov[km]
174        fb[km]+=fOV[km]
175        for kn in range(nkpts):
176            fa[km]+=einsum('nf,menf->me',t1a[kn],eris.ovov[km,km,kn])
177            fa[km]+=einsum('nf,menf->me',t1b[kn],eris.ovOV[km,km,kn])
178            fa[km]-=einsum('nf,mfne->me',t1a[kn],eris.ovov[km,kn,kn])
179            fb[km]+=einsum('nf,menf->me',t1b[kn],eris.OVOV[km,km,kn])
180            fb[km]+=einsum('nf,nfme->me',t1a[kn],eris.ovOV[kn,kn,km])
181            fb[km]-=einsum('nf,mfne->me',t1b[kn],eris.OVOV[km,kn,kn])
182
183    return fa,fb
184
185def cc_Woooo(cc, t1, t2, eris):
186    t1a, t1b = t1
187    t2aa, t2ab, t2bb = t2
188    nkpts, nocca, nvira = t1a.shape
189    noccb, nvirb = t1b.shape[1:]
190    dtype = np.result_type(t1a, t1b, t2aa, t2ab, t2bb)
191
192    Woooo = np.zeros(eris.oooo.shape, dtype=dtype)
193    WooOO = np.zeros(eris.ooOO.shape, dtype=dtype)
194    WOOOO = np.zeros(eris.OOOO.shape, dtype=dtype)
195
196    kconserv = cc.khelper.kconserv
197    tau_aa, tau_ab, tau_bb = make_tau(cc, t2, t1, t1)
198    for km in range(nkpts):
199        for kn in range(nkpts):
200            tmp_aaaaJ = einsum('xje, ymine->yxminj', t1a, eris.ooov[km,:,kn])
201            tmp_aaaaJ -= tmp_aaaaJ.transpose((1,0,2,5,4,3))
202            tmp_bbbbJ = einsum('xje, ymine->yxminj', t1b, eris.OOOV[km,:,kn])
203            tmp_bbbbJ -= tmp_bbbbJ.transpose((1,0,2,5,4,3))
204            tmp_aabbJ = einsum('xje, ymine->yxminj', t1b, eris.ooOV[km,:,kn])
205            tmp_baabJ = -einsum('yie,xmjne->yxminj', t1a, eris.OOov[km,:,kn])
206
207            Woooo[km,:,kn] += eris.oooo[km,:,kn]
208            WooOO[km,:,kn] += eris.ooOO[km,:,kn]
209            WOOOO[km,:,kn] += eris.OOOO[km,:,kn]
210
211            ki = range(nkpts)
212            kj = kconserv[km,ki,kn]
213            Woooo[km,ki,kn] += tmp_aaaaJ[ki,kj]
214            WOOOO[km,ki,kn] += tmp_bbbbJ[ki,kj]
215            WooOO[km,ki,kn] += tmp_aabbJ[ki,kj]
216            WooOO[kn,ki,km] -= tmp_baabJ[ki,kj].transpose(0,3,2,1,4)
217            Woooo[km,ki,kn] += 0.25*einsum('yxijef,xmenf->yminj', tau_aa[ki,kj], eris.ovov[km,ki,kn])
218            WOOOO[km,ki,kn] += 0.25*einsum('yxijef,xmenf->yminj', tau_bb[ki,kj], eris.OVOV[km,ki,kn])
219            WooOO[km,ki,kn] += 0.5*einsum('yxijef,xmenf->yminj', tau_ab[ki,kj], eris.ovOV[km,ki,kn])
220
221    Woooo = Woooo - Woooo.transpose(2,1,0,5,4,3,6)
222    WOOOO = WOOOO - WOOOO.transpose(2,1,0,5,4,3,6)
223    return Woooo, WooOO, WOOOO
224
225
226def cc_Wvvvv(cc, t1, t2, eris):
227    t1a, t1b = t1
228    kconserv = cc.khelper.kconserv
229
230    #:wvvvv = eris.vvvv.copy()
231    #:Wvvvv += np.einsum('ymb,zyxemfa,zxyw->wzyaebf', t1a, eris.vovv.conj(), P)
232    #:Wvvvv -= np.einsum('ymb,xyzfmea,xzyw->wzyaebf', t1a, eris.vovv.conj(), P)
233    #:Wvvvv = Wvvvv - Wvvvv.transpose(2,1,0,5,4,3,6)
234    Wvvvv = np.zeros_like(eris.vvvv)
235    for ka, kb, ke in kpts_helper.loop_kkk(cc.nkpts):
236        kf = kconserv[ka,ke,kb]
237        aebf = eris.vvvv[ka,ke,kb].copy()
238        aebf += einsum('mb,emfa->aebf', t1a[kb], eris.vovv[ke,kb,kf].conj())
239        aebf -= einsum('mb,fmea->aebf', t1a[kb], eris.vovv[kf,kb,ke].conj())
240        Wvvvv[ka,ke,kb] += aebf
241        Wvvvv[kb,ke,ka] -= aebf.transpose(2,1,0,3)
242
243    #:WvvVV = eris.vvVV.copy()
244    #:WvvVV -= np.einsum('xma,zxwemFB,zwxy->xzyaeBF', t1a, eris.voVV.conj(), P)
245    #:WvvVV -= np.einsum('yMB,wyzFMea,wzyx->xzyaeBF', t1b, eris.VOvv.conj(), P)
246    WvvVV = np.empty_like(eris.vvVV)
247    for ka, kb, ke in kpts_helper.loop_kkk(cc.nkpts):
248        kf = kconserv[ka,ke,kb]
249        aebf = eris.vvVV[ka,ke,kb].copy()
250        aebf -= einsum('ma,emfb->aebf', t1a[ka], eris.voVV[ke,ka,kf].conj())
251        aebf -= einsum('mb,fmea->aebf', t1b[kb], eris.VOvv[kf,kb,ke].conj())
252        WvvVV[ka,ke,kb] = aebf
253
254    #:WVVVV = eris.VVVV.copy()
255    #:WVVVV += np.einsum('ymb,zyxemfa,zxyw->wzyaebf', t1b, eris.VOVV.conj(), P)
256    #:WVVVV -= np.einsum('ymb,xyzfmea,xzyw->wzyaebf', t1b, eris.VOVV.conj(), P)
257    #:WVVVV = WVVVV - WVVVV.transpose(2,1,0,5,4,3,6)
258    WVVVV = np.zeros_like(eris.VVVV)
259    for ka, kb, ke in kpts_helper.loop_kkk(cc.nkpts):
260        kf = kconserv[ka,ke,kb]
261        aebf = eris.VVVV[ka,ke,kb].copy()
262        aebf += einsum('mb,emfa->aebf', t1b[kb], eris.VOVV[ke,kb,kf].conj())
263        aebf -= einsum('mb,fmea->aebf', t1b[kb], eris.VOVV[kf,kb,ke].conj())
264        WVVVV[ka,ke,kb] += aebf
265        WVVVV[kb,ke,ka] -= aebf.transpose(2,1,0,3)
266    return Wvvvv, WvvVV, WVVVV
267
268
269#TODO: merge cc_Wvvvv_half and cc_Wvvvv
270def cc_Wvvvv_half(cc, t1, t2, eris):
271    '''Similar to cc_Wvvvv, without anti-symmetrization'''
272    t1a, t1b = t1
273    kconserv = cc.khelper.kconserv
274
275    #:wvvvv = eris.vvvv.copy()
276    #:Wvvvv += np.einsum('ymb,zyxemfa,zxyw->wzyaebf', t1a, eris.vovv.conj(), P)
277    #:Wvvvv -= np.einsum('ymb,xyzfmea,xzyw->wzyaebf', t1a, eris.vovv.conj(), P)
278    #:Wvvvv = Wvvvv - Wvvvv.transpose(2,1,0,5,4,3,6)
279    Wvvvv = np.zeros_like(eris.vvvv)
280    for ka, kb, ke in kpts_helper.loop_kkk(cc.nkpts):
281        kf = kconserv[ka,ke,kb]
282        aebf = eris.vvvv[ka,ke,kb].copy()
283        aebf += einsum('mb,emfa->aebf', t1a[kb], eris.vovv[ke,kb,kf].conj())
284        aebf -= einsum('mb,fmea->aebf', t1a[kb], eris.vovv[kf,kb,ke].conj())
285        Wvvvv[ka,ke,kb] += aebf
286
287    #:WvvVV = eris.vvVV.copy()
288    #:WvvVV -= np.einsum('xma,zxwemFB,zwxy->xzyaeBF', t1a, eris.voVV.conj(), P)
289    #:WvvVV -= np.einsum('yMB,wyzFMea,wzyx->xzyaeBF', t1b, eris.VOvv.conj(), P)
290    WvvVV = np.empty_like(eris.vvVV)
291    for ka, kb, ke in kpts_helper.loop_kkk(cc.nkpts):
292        kf = kconserv[ka,ke,kb]
293        aebf = eris.vvVV[ka,ke,kb].copy()
294        aebf -= einsum('ma,emfb->aebf', t1a[ka], eris.voVV[ke,ka,kf].conj())
295        aebf -= einsum('mb,fmea->aebf', t1b[kb], eris.VOvv[kf,kb,ke].conj())
296        WvvVV[ka,ke,kb] = aebf
297
298    #:WVVVV = eris.VVVV.copy()
299    #:WVVVV += np.einsum('ymb,zyxemfa,zxyw->wzyaebf', t1b, eris.VOVV.conj(), P)
300    #:WVVVV -= np.einsum('ymb,xyzfmea,xzyw->wzyaebf', t1b, eris.VOVV.conj(), P)
301    #:WVVVV = WVVVV - WVVVV.transpose(2,1,0,5,4,3,6)
302    WVVVV = np.zeros_like(eris.VVVV)
303    for ka, kb, ke in kpts_helper.loop_kkk(cc.nkpts):
304        kf = kconserv[ka,ke,kb]
305        aebf = eris.VVVV[ka,ke,kb].copy()
306        aebf += einsum('mb,emfa->aebf', t1b[kb], eris.VOVV[ke,kb,kf].conj())
307        aebf -= einsum('mb,fmea->aebf', t1b[kb], eris.VOVV[kf,kb,ke].conj())
308        WVVVV[ka,ke,kb] += aebf
309    return Wvvvv, WvvVV, WVVVV
310
311def Wvvvv(cc, t1, t2, eris):
312    nkpts = cc.nkpts
313    kconserv = cc.khelper.kconserv
314
315    tauaa, tauab, taubb = make_tau(cc, t2, t1, t1)
316    Wvvvv, WvvVV, WVVVV = cc_Wvvvv(cc, t1, t2, eris)
317    for ka, kb, ke in kpts_helper.loop_kkk(cc.nkpts):
318        for km in range(nkpts):
319            kn = kconserv[ka,km,kb]
320            Wvvvv[ka,ke,kb] += einsum('mnab,menf->aebf', tauaa[km,kn,ka], eris.ovov[km,ke,kn])
321            WvvVV[ka,ke,kb] += einsum('mNaB,meNF->aeBF', tauab[km,kn,ka], eris.ovOV[km,ke,kn])
322            WVVVV[ka,ke,kb] += einsum('mnab,menf->aebf', taubb[km,kn,ka], eris.OVOV[km,ke,kn])
323    return Wvvvv, WvvVV, WVVVV
324
325def get_Wvvvv(cc, t1, t2, eris, ka, kb, kc):
326    t1a, t1b = t1
327    t2aa, t2ab, t2bb = t2
328    nocca, noccb = cc.nocc
329    nkpts = cc.nkpts
330    kconserv = cc.khelper.kconserv
331    kd = kconserv[ka, kc, kb]
332
333    if getattr(eris, 'Lpv', None) is not None:
334        # Using GDF to generate Wvvvv on the fly
335        Lpv = eris.Lpv
336        LPV = eris.LPV
337        Lac = (Lpv[ka,kc][:,nocca:] -
338               einsum('Lkc,ka->Lac', Lpv[ka,kc][:,:nocca], t1a[ka]))
339        Lbd = (Lpv[kb,kd][:,nocca:] -
340               einsum('Lkd,kb->Lbd', Lpv[kb,kd][:,:nocca], t1a[kb]))
341        Lbc = (Lpv[kb,kc][:,nocca:] -
342               einsum('Lkc,ka->Lac', Lpv[kb,kc][:,:nocca], t1a[kb]))
343        Lad = (Lpv[ka,kd][:,nocca:] -
344               einsum('Lkd,kb->Lbd', Lpv[ka,kd][:,:nocca], t1a[ka]))
345        LAC = (LPV[ka,kc][:,noccb:] -
346               einsum('Lkd,kb->Lbd', LPV[ka,kc][:,:noccb], t1b[ka]))
347        LBD = (LPV[kb,kd][:,noccb:] -
348               einsum('Lkd,kb->Lbd', LPV[kb,kd][:,:noccb], t1b[kb]))
349        LBC = (LPV[kb,kc][:,noccb:] -
350               einsum('Lkc,ka->Lac', LPV[kb,kc][:,:noccb], t1b[kb]))
351        LAD = (LPV[ka,kd][:,noccb:] -
352               einsum('Lkd,kb->Lbd', LPV[ka,kd][:,:noccb], t1b[ka]))
353        vvvv = einsum('Lac,Lbd->acbd', Lac, Lbd)
354        vvvv-= einsum('Lbc,Lad->acbd', Lbc, Lad)
355        vvVV = einsum('Lac,Lbd->acbd', Lac, LBD)
356        VVVV = einsum('Lac,Lbd->acbd', LAC, LBD)
357        VVVV-= einsum('Lbc,Lad->acbd', LBC, LAD)
358        vvvv *= (1./nkpts)
359        vvVV *= (1./nkpts)
360        VVVV *= (1./nkpts)
361    else:
362        vvvv  = einsum('emfa,mb->aebf', eris.vovv[kc,kb,kd].conj(), t1a[kb])
363        vvvv -= einsum('fmea,mb->aebf', eris.vovv[kd,kb,kc].conj(), t1a[kb])
364        vvvv -= einsum('emfb,ma->aebf', eris.vovv[kc,ka,kd].conj(), t1a[ka])
365        vvvv += einsum('fmeb,ma->aebf', eris.vovv[kd,ka,kc].conj(), t1a[ka])
366        vvvv += eris.vvvv[ka,kc,kb]
367        vvvv -= eris.vvvv[kb,kc,ka].transpose(2,1,0,3)
368        vvvv += einsum('mcnf,ma,nb->acbf', eris.ovov[ka,kc,kb], t1a[ka], t1a[kb])
369        vvvv -= einsum('mcnf,mb,na->acbf', eris.ovov[kb,kc,ka], t1a[kb], t1a[ka])
370
371        vvVV  = einsum('emfb,ma->aebf', eris.voVV[kc,ka,kd].conj(),-t1a[ka])
372        vvVV += einsum('fmea,mb->aebf', eris.VOvv[kd,kb,kc].conj(),-t1b[kb])
373        vvVV += einsum('mcnf,ma,nb->acbf', eris.ovOV[ka,kc,kb], t1a[ka], t1b[kb])
374        vvVV += eris.vvVV[ka,kc,kb]
375
376        VVVV  = einsum('emfa,mb->aebf', eris.VOVV[kc,kb,kd].conj(), t1b[kb])
377        VVVV -= einsum('fmea,mb->aebf', eris.VOVV[kd,kb,kc].conj(), t1b[kb])
378        VVVV -= einsum('emfb,ma->aebf', eris.VOVV[kc,ka,kd].conj(), t1b[ka])
379        VVVV += einsum('fmeb,ma->aebf', eris.VOVV[kd,ka,kc].conj(), t1b[ka])
380        VVVV += eris.VVVV[ka,kc,kb]
381        VVVV -= eris.VVVV[kb,kc,ka].transpose(2,1,0,3)
382        VVVV += einsum('mcnf,ma,nb->acbf', eris.OVOV[ka,kc,kb], t1b[ka], t1b[kb])
383        VVVV -= einsum('mcnf,mb,na->acbf', eris.OVOV[kb,kc,ka], t1b[kb], t1b[ka])
384
385    for km in range(nkpts):
386        kn = kconserv[ka,km,kb]
387        vvvv += einsum('mnab,mcnf->acbf', t2aa[km,kn,ka], eris.ovov[km,kc,kn])
388        vvVV += einsum('mNaB,mcNF->acBF', t2ab[km,kn,ka], eris.ovOV[km,kc,kn])
389        VVVV += einsum('mnab,mcnf->acbf', t2bb[km,kn,ka], eris.OVOV[km,kc,kn])
390    return vvvv, vvVV, VVVV
391
392def cc_Wovvo(cc, t1, t2, eris):
393    t1a, t1b = t1
394    t2aa, t2ab, t2bb = t2
395    nkpts, nocca, nvira = t1a.shape
396    noccb, nvirb = t1b.shape[1:]
397    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
398
399    dtype = np.result_type(*t2)
400    Wovvo = np.zeros((nkpts,nkpts,nkpts,nocca,nvira,nvira,nocca), dtype)
401    WovVO = np.zeros((nkpts,nkpts,nkpts,nocca,nvira,nvirb,noccb), dtype)
402    WOVvo = np.zeros((nkpts,nkpts,nkpts,noccb,nvirb,nvira,nocca), dtype)
403    WOVVO = np.zeros((nkpts,nkpts,nkpts,noccb,nvirb,nvirb,noccb), dtype)
404    WoVVo = np.zeros((nkpts,nkpts,nkpts,nocca,nvirb,nvirb,nocca), dtype)
405    WOvvO = np.zeros((nkpts,nkpts,nkpts,noccb,nvira,nvira,noccb), dtype)
406
407    for ka, ki, kj in kpts_helper.loop_kkk(nkpts):
408        kb = kconserv[ka,ki,kj]
409        Wovvo[ki,ka,kb] += eris.voov[ka,ki,kj].conj().transpose(1,0,3,2)
410        WovVO[ki,ka,kb] += eris.voOV[ka,ki,kj].conj().transpose(1,0,3,2)
411        WOVvo[ki,ka,kb] += eris.voOV[kb,kj,ki].transpose(2,3,0,1)
412        WOVVO[ki,ka,kb] += eris.VOOV[ka,ki,kj].conj().transpose(1,0,3,2)
413
414        kb = kconserv[ki,kj,ka]
415        Wovvo[ki,kb,ka] -= eris.oovv[ki,kj,ka].transpose(0,3,2,1)
416        WOVVO[ki,kb,ka] -= eris.OOVV[ki,kj,ka].transpose(0,3,2,1)
417        WoVVo[ki,kb,ka] -= eris.ooVV[ki,kj,ka].transpose(0,3,2,1)
418        WOvvO[ki,kb,ka] -= eris.OOvv[ki,kj,ka].transpose(0,3,2,1)
419
420    tauaa, tauab, taubb = make_tau2(cc, t2, t1, t1,fac=2.0)
421    for km in range(nkpts):
422        for kb in range(nkpts):
423            for ke in range(nkpts):
424                kj = kconserv[km,ke,kb]
425                vovv = eris.vovv[ke,km,kj].conj()
426                VOVV = eris.VOVV[ke,km,kj].conj()
427                voVV = eris.voVV[ke,km,kj].conj()
428                VOvv = eris.VOvv[ke,km,kj].conj()
429
430                Wovvo[km,ke,kb] += einsum('jf, emfb->mebj', t1a[kj], vovv)
431                WOVVO[km,ke,kb] += einsum('jf, emfb->mebj', t1b[kj], VOVV)
432                WovVO[km,ke,kb] += einsum('jf, emfb->mebj', t1b[kj], voVV)
433                WOVvo[km,ke,kb] += einsum('jf, emfb->mebj', t1a[kj], VOvv)
434                ##### warnings for Ks
435                Wovvo[km,kj,kb] -= einsum('je, emfb->mfbj', t1a[ke], vovv)
436                WOVVO[km,kj,kb] -= einsum('je, emfb->mfbj', t1b[ke], VOVV)
437                WOvvO[km,kj,kb] -= einsum('je, emfb->mfbj', t1b[ke], VOvv)
438                WoVVo[km,kj,kb] -= einsum('je, emfb->mfbj', t1a[ke], voVV)
439
440                WOVvo[km,ke,kb] -= einsum('nb, njme->mebj', t1a[kb], eris.ooOV[kb,kj,km])
441                WovVO[km,ke,kb] -= einsum('nb, njme->mebj', t1b[kb], eris.OOov[kb,kj,km])
442
443                WOvvO[km,ke,kb] += einsum('nb, mjne->mebj', t1a[kb], eris.OOov[km,kj,kb])
444                WoVVo[km,ke,kb] += einsum('nb, mjne->mebj', t1b[kb], eris.ooOV[km,kj,kb])
445
446                ooov_temp = eris.ooov[kb,kj,km] - eris.ooov[km,kj,kb].transpose((2,1,0,3))
447                Wovvo[km,ke,kb] -= einsum('nb, njme->mebj', t1a[kb], ooov_temp)
448                ooov_temp = None
449                OOOV_temp = eris.OOOV[kb,kj,km] - eris.OOOV[km,kj,kb].transpose((2,1,0,3))
450                WOVVO[km,ke,kb] -= einsum('nb, njme->mebj', t1b[kb], OOOV_temp)
451                OOOV_temp = None
452
453                Wovvo[km,ke,kb] += 0.5*einsum('xjnbf,xmenf->mebj', t2ab[kj,:,kb], eris.ovOV[km,ke,:])
454                WOvvO[km,ke,kb] += 0.5*einsum('xnjbf,xnemf->mebj', tauab[:,kj,kb], eris.ovOV[:,ke,km])
455                WovVO[km,ke,kb] -= 0.5*einsum('xnjbf,xmenf->mebj', taubb[:,kj,kb], eris.ovOV[km,ke,:])
456
457                temp_ovOV_1 = np.zeros([nkpts, nocca, nvira, noccb, nvirb], dtype=dtype)
458                temp_ovOV_2 = np.zeros([nkpts, nocca, nvira, noccb, nvirb], dtype=dtype)
459                for kn in range(nkpts):
460                    kf = kconserv[km,ke,kn]
461                    temp_ovOV_1[kn] += eris.ovOV[kn,kf,km].copy()
462                    temp_ovOV_2[kn] += eris.ovOV[km,kf,kn].copy()
463
464                kn = range(nkpts)
465                kf = kconserv[km,ke][kn]
466                WOVVO[km,ke,kb] += 0.5*einsum('xnjfb,xnfme->mebj', t2ab[kn,kj,kf], temp_ovOV_1)
467                WOVvo[km,ke,kb] -= 0.5*einsum('xnjbf,xnfme->mebj', tauaa[:,kj,kb], temp_ovOV_1)
468                WoVVo[km,ke,kb] += 0.5*einsum('xjnfb,xmfne->mebj', tauab[kj,kn,kf], temp_ovOV_2)
469
470                temp_OVOV = eris.OVOV[km,ke,:] - eris.OVOV[:,ke,km].transpose((0,3,2,1,4))
471                WOVVO[km,ke,kb] -= 0.5*einsum('xnjbf,xmenf->mebj', taubb[:,kj,kb], temp_OVOV)
472                WOVvo[km,ke,kb] += 0.5*einsum('xjnbf,xmenf->mebj', t2ab[kj,:,kb], temp_OVOV)
473                temp_OVOV = None
474
475                temp_ovov = eris.ovov[:,ke,km] - eris.ovov[km,ke,:].transpose((0,3,2,1,4))
476                Wovvo[km,ke,kb] += 0.5*einsum('xnjbf,xnemf->mebj', tauaa[:,kj,kb], temp_ovov)
477                WovVO[km,ke,kb] -= 0.5*einsum('xnjfb,xnemf->mebj', t2ab[kn,kj,kf], temp_ovov)
478                temp_ovov = None
479
480    return Wovvo, WovVO, WOVvo, WOVVO, WoVVo, WOvvO
481
482def _cc_Wovvo_k0k2(cc, t1, t2, eris, k0, k2):
483    t1a, t1b = t1
484    t2aa, t2ab, t2bb = t2
485    nkpts, nocca, nvira = t1a.shape
486    noccb, nvirb = t1b.shape[1:]
487    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
488
489    dtype = np.result_type(*t2)
490    Wovvo = np.zeros((nkpts,nocca,nvira,nvira,nocca), dtype)
491    WovVO = np.zeros((nkpts,nocca,nvira,nvirb,noccb), dtype)
492    WOVvo = np.zeros((nkpts,noccb,nvirb,nvira,nocca), dtype)
493    WOVVO = np.zeros((nkpts,noccb,nvirb,nvirb,noccb), dtype)
494    WoVVo = np.zeros((nkpts,nocca,nvirb,nvirb,nocca), dtype)
495    WOvvO = np.zeros((nkpts,noccb,nvira,nvira,noccb), dtype)
496
497    #:P = kconserv_mat(cc.nkpts, kconserv)
498    #:Wovvo = np.einsum('xyzaijb,xzyw->yxwiabj', eris.voov, P).conj()
499    #:WovVO = np.einsum('xyzaijb,xzyw->yxwiabj', eris.voOV, P).conj()
500    #:WOVvo = np.einsum('wzybjia,xzyw->yxwiabj', eris.voOV, P)
501    #:WOVVO = np.einsum('xyzaijb,xzyw->yxwiabj', eris.VOOV, P).conj()
502    #:Wovvo-= np.einsum('xyzijab,xzyw->xwzibaj', eris.oovv, P)
503    #:WOVVO-= np.einsum('xyzijab,xzyw->xwzibaj', eris.OOVV, P)
504    #:WoVVo = np.einsum('xyzijab,xzyw->xwzibaj', eris.ooVV, -P)
505    #:WOvvO = np.einsum('xyzijab,xzyw->xwzibaj', eris.OOvv, -P)
506    for kj in range(nkpts):
507        ka = kconserv[k2,kj,k0]
508        Wovvo[ka] += eris.voov[ka,k0,kj].conj().transpose(1,0,3,2)
509        WovVO[ka] += eris.voOV[ka,k0,kj].conj().transpose(1,0,3,2)
510        WOVvo[ka] += eris.voOV[k2,kj,k0].transpose(2,3,0,1)
511        WOVVO[ka] += eris.VOOV[ka,k0,kj].conj().transpose(1,0,3,2)
512
513    for kj in range(nkpts):
514        kb = kconserv[k0,kj,k2]
515        Wovvo[kb] -= eris.oovv[k0,kj,k2].transpose(0,3,2,1)
516        WOVVO[kb] -= eris.OOVV[k0,kj,k2].transpose(0,3,2,1)
517        WoVVo[kb] -= eris.ooVV[k0,kj,k2].transpose(0,3,2,1)
518        WOvvO[kb] -= eris.OOvv[k0,kj,k2].transpose(0,3,2,1)
519
520    for ke in range(nkpts):
521        kj = kconserv[k0,ke,k2]
522        vovv = eris.vovv[ke,k0,kj].conj()
523        VOVV = eris.VOVV[ke,k0,kj].conj()
524        voVV = eris.voVV[ke,k0,kj].conj()
525        VOvv = eris.VOvv[ke,k0,kj].conj()
526
527        Wovvo[ke] += einsum('jf, emfb->mebj', t1a[kj], vovv)
528        WOVVO[ke] += einsum('jf, emfb->mebj', t1b[kj], VOVV)
529        WovVO[ke] += einsum('jf, emfb->mebj', t1b[kj], voVV)
530        WOVvo[ke] += einsum('jf, emfb->mebj', t1a[kj], VOvv)
531
532        Wovvo[kj] -= einsum('je, emfb->mfbj', t1a[ke], vovv)
533        WOVVO[kj] -= einsum('je, emfb->mfbj', t1b[ke], VOVV)
534        WOvvO[kj] -= einsum('je, emfb->mfbj', t1b[ke], VOvv)
535        WoVVo[kj] -= einsum('je, emfb->mfbj', t1a[ke], voVV)
536
537        Wovvo[ke] -= einsum('nb, njme->mebj', t1a[k2], eris.ooov[k2,kj,k0])
538        WOVvo[ke] -= einsum('nb, njme->mebj', t1a[k2], eris.ooOV[k2,kj,k0])
539        WOVVO[ke] -= einsum('nb, njme->mebj', t1b[k2], eris.OOOV[k2,kj,k0])
540        WovVO[ke] -= einsum('nb, njme->mebj', t1b[k2], eris.OOov[k2,kj,k0])
541
542        Wovvo[ke] += einsum('nb, mjne->mebj', t1a[k2], eris.ooov[k0,kj,k2])
543        WOVVO[ke] += einsum('nb, mjne->mebj', t1b[k2], eris.OOOV[k0,kj,k2])
544        WoVVo[ke] += einsum('nb, mjne->mebj', t1b[k2], eris.ooOV[k0,kj,k2])
545        WOvvO[ke] += einsum('nb, mjne->mebj', t1a[k2], eris.OOov[k0,kj,k2])
546
547        for kn in range(nkpts):
548            kf = kconserv[k0,ke,kn]
549
550            tmp = eris.ovov[k0,ke,kn] - eris.ovov[kn,ke,k0].transpose(2,1,0,3)
551            Wovvo[ke] -= 0.5*einsum('jnfb,menf->mebj', t2aa[kj,kn,kf], tmp)
552            Wovvo[ke] += 0.5*einsum('jnbf,menf->mebj', t2ab[kj,kn,k2], eris.ovOV[k0,ke,kn])
553            tmp = eris.OVOV[k0,ke,kn] - eris.OVOV[kn,ke,k0].transpose(2,1,0,3)
554            WOVVO[ke] -= 0.5*einsum('jnfb,menf->mebj', t2bb[kj,kn,kf], tmp)
555            WOVVO[ke] += 0.5*einsum('njfb,nfme->mebj', t2ab[kn,kj,kf], eris.ovOV[kn,kf,k0])
556            tmp = eris.ovov[k0,ke,kn] - eris.ovov[kn,ke,k0].transpose(2,1,0,3)
557            WovVO[ke] += 0.5*einsum('njfb,menf->mebj', t2ab[kn,kj,kf], tmp)
558            WovVO[ke] -= 0.5*einsum('jnfb,menf->mebj', t2bb[kj,kn,kf], eris.ovOV[k0,ke,kn])
559            tmp = eris.OVOV[k0,ke,kn] - eris.OVOV[kn,ke,k0].transpose(2,1,0,3)
560            WOVvo[ke] += 0.5*einsum('jnbf,menf->mebj', t2ab[kj,kn,k2], tmp)
561            WOVvo[ke] -= 0.5*einsum('jnfb,nfme->mebj', t2aa[kj,kn,kf], eris.ovOV[kn,kf,k0])
562            WoVVo[ke] += 0.5*einsum('jnfb,mfne->mebj', t2ab[kj,kn,kf], eris.ovOV[k0,kf,kn])
563            WOvvO[ke] += 0.5*einsum('njbf,nemf->mebj', t2ab[kn,kj,k2], eris.ovOV[kn,ke,k0])
564
565            if kn == k2 and kf == kj:
566                tmp = einsum('menf,jf->menj', eris.ovov[k0,ke,kn], t1a[kj])
567                tmp-= einsum('nemf,jf->menj', eris.ovov[kn,ke,k0], t1a[kj])
568                Wovvo[ke] -= einsum('nb,menj->mebj', t1a[kn], tmp)
569                tmp = einsum('menf,jf->menj', eris.OVOV[k0,ke,kn], t1b[kj])
570                tmp-= einsum('nemf,jf->menj', eris.OVOV[kn,ke,k0], t1b[kj])
571                WOVVO[ke] -= einsum('nb,menj->mebj', t1b[kn], tmp)
572
573                WovVO[ke] -= einsum('jf,nb,menf->mebj',t1b[kj],t1b[kn], eris.ovOV[k0,ke,kn])
574                WOVvo[ke] -= einsum('jf,nb,nfme->mebj',t1a[kj],t1a[kn], eris.ovOV[kn,kf,k0])
575                WoVVo[ke] += einsum('jf,nb,mfne->mebj',t1a[kj],t1b[kn], eris.ovOV[k0,kf,kn])
576                WOvvO[ke] += einsum('jf,nb,nemf->mebj',t1b[kj],t1a[kn], eris.ovOV[kn,ke,k0])
577
578    return Wovvo, WovVO, WOVvo, WOVVO, WoVVo, WOvvO
579
580
581def kconserv_mat(nkpts, kconserv):
582    P = np.zeros((nkpts,nkpts,nkpts,nkpts))
583    for ki in range(nkpts):
584        for kj in range(nkpts):
585            for ka in range(nkpts):
586                kb = kconserv[ki,ka,kj]
587                P[ki,kj,ka,kb] = 1
588    return P
589
590def Foo(cc,t1,t2,eris):
591    t1a, t1b = t1
592    t2aa, t2ab, t2bb = t2
593    nkpts, nocca, noccb, nvira, nvirb = t2ab.shape[2:]
594
595    Fova, Fovb = cc_Fov(cc,t1,t2,eris)
596    Fooa, Foob = cc_Foo(cc,t1,t2,eris)
597    for ki in range(nkpts):
598        Fooa[ki] += 0.5*einsum('ie,me->mi',t1a[ki],Fova[ki])
599        Foob[ki] += 0.5*einsum('ie,me->mi',t1b[ki],Fovb[ki])
600    return Fooa, Foob
601
602def Fvv(cc,t1,t2,eris):
603    t1a, t1b = t1
604    t2aa, t2ab, t2bb = t2
605    nkpts, nocca, noccb, nvira, nvirb = t2ab.shape[2:]
606
607    Fova, Fovb = cc_Fov(cc,t1,t2,eris)
608    Fvva, Fvvb = cc_Fvv(cc,t1,t2,eris)
609    for ka in range(nkpts):
610        Fvva[ka] -= 0.5*lib.einsum('me,ma->ae', Fova[ka], t1a[ka])
611        Fvvb[ka] -= 0.5*lib.einsum('me,ma->ae', Fovb[ka], t1b[ka])
612    return Fvva, Fvvb
613
614def Fov(cc,t1,t2,eris):
615    Fme = cc_Fov(cc,t1,t2,eris)
616    return Fme
617
618def Wvvov(cc,t1,t2,eris):
619    kconserv = cc.khelper.kconserv
620    t1a, t1b = t1
621    t2aa, t2ab, t2bb = t2
622    nkpts, nocca, noccb, nvira, nvirb = t2ab.shape[2:]
623
624    Wvvov = np.zeros((nkpts,nkpts,nkpts,nvira,nvira,nocca,nvira),dtype=t1a.dtype)
625    WvvOV = np.zeros((nkpts,nkpts,nkpts,nvira,nvira,noccb,nvirb),dtype=t1a.dtype)
626    WVVov = np.zeros((nkpts,nkpts,nkpts,nvirb,nvirb,nocca,nvira),dtype=t1a.dtype)
627    WVVOV = np.zeros((nkpts,nkpts,nkpts,nvirb,nvirb,noccb,nvirb),dtype=t1a.dtype)
628
629    for kn, km, ke in itertools.product(range(nkpts),repeat=3):
630        kf = kconserv[kn, ke, km]
631        ka = kn
632        Wvvov[ka,ke,km] += eris.vovv[kf,km,ke].transpose(3,2,1,0).conj() - eris.vovv[ke,km,kf].transpose(3,0,1,2).conj()
633        WVVov[ka,ke,km] += eris.voVV[kf,km,ke].transpose(3,2,1,0).conj()
634        WvvOV[ka,ke,km] += eris.VOvv[kf,km,ke].transpose(3,2,1,0).conj()
635        WVVOV[ka,ke,km] += eris.VOVV[kf,km,ke].transpose(3,2,1,0).conj() - eris.VOVV[ke,km,kf].transpose(3,0,1,2).conj()
636
637        ovov = eris.ovov[kn, ke, km] - eris.ovov[kn, kf, km].transpose(0,3,2,1)
638        OVOV = eris.OVOV[kn, ke, km] - eris.OVOV[kn, kf, km].transpose(0,3,2,1)
639
640        Wvvov[ka,ke,km] += -lib.einsum('na,nemf->aemf',t1a[kn],ovov)
641        WvvOV[ka,ke,km] += -lib.einsum('na,neMF->aeMF',t1a[kn],eris.ovOV[kn,ke,km])
642        WVVov[ka,ke,km] += -lib.einsum('NA,NEmf->AEmf',t1b[kn],eris.OVov[kn,ke,km])
643        WVVOV[ka,ke,km] += -lib.einsum('NA,NEMF->AEMF',t1b[kn],OVOV)
644
645    return Wvvov, WvvOV, WVVov, WVVOV
646
647def Wvvvo(cc,t1,t2,eris):
648    kconserv = cc.khelper.kconserv
649    t1a, t1b = t1
650    t2aa, t2ab, t2bb = t2
651    nkpts, nocca, noccb, nvira, nvirb = t2ab.shape[2:]
652
653    fova, fovb = cc_Fov(cc, t1, t2, eris)
654    tauaa, tauab, taubb = make_tau(cc, t2, t1, t1)
655
656    Wvvvo = np.zeros((nkpts,nkpts,nkpts,nvira,nvira,nvira,nocca),dtype=t1a.dtype)
657    WvvVO = np.zeros((nkpts,nkpts,nkpts,nvira,nvira,nvirb,noccb),dtype=t1a.dtype)
658    WVVvo = np.zeros((nkpts,nkpts,nkpts,nvirb,nvirb,nvira,nocca),dtype=t1a.dtype)
659    WVVVO = np.zeros((nkpts,nkpts,nkpts,nvirb,nvirb,nvirb,noccb),dtype=t1a.dtype)
660
661    for ka, ke, kb in itertools.product(range(nkpts),repeat=3):
662        ki = kconserv[ka, ke, kb]
663        # - <mb||ef> t2(miaf)
664        for km in range(nkpts):
665            kf = kconserv[km,ke,kb]
666            ovvv = eris.vovv[ke,km,kf].transpose(1,0,3,2).conj() - eris.vovv[kf,km,ke].transpose(1,2,3,0).conj()
667            OVvv = eris.VOvv[ke,km,kf].transpose(1,0,3,2).conj()
668            ovVV = eris.voVV[ke,km,kf].transpose(1,0,3,2).conj()
669            OVVV = eris.VOVV[ke,km,kf].transpose(1,0,3,2).conj() - eris.VOVV[kf,km,ke].transpose(1,2,3,0).conj()
670
671            aebi = lib.einsum('mebf,miaf->aebi',ovvv,t2aa[km,ki,ka])
672            aebi += lib.einsum('MFbe,iMaF->aebi',eris.VOvv[kf,km,ke].transpose(1,0,3,2).conj(),t2ab[ki,km,ka])
673            Wvvvo[ka,ke,kb] -= aebi
674            # P(ab) for all alpha spin
675            Wvvvo[kb,ke,ka] += aebi.transpose(2,1,0,3)
676
677            WVVvo[ka,ke,kb] -= lib.einsum('MEbf,iMfA->AEbi',OVvv,t2ab[ki,km,kf])
678            WvvVO[ka,ke,kb] -= lib.einsum('meBF,mIaF->aeBI',ovVV,t2ab[km,ki,ka])
679
680            AEBI = lib.einsum('MEBF,MIAF->AEBI',OVVV,t2bb[km,ki,ka])
681            AEBI += lib.einsum('mfBE,mIfA->AEBI',eris.voVV[kf,km,ke].transpose(1,0,3,2).conj(),t2ab[km,ki,kf])
682            WVVVO[ka,ke,kb] -= AEBI
683            # P(ab) for all beta spin
684            WVVVO[kb,ke,ka] += AEBI.transpose(2,1,0,3)
685
686        # - t1(ma) (<mb||ei> - t2(nibf) <mn||ef>)
687        km = ka
688        ovvo = eris.voov[ke,km,ki].transpose(1,0,3,2).conj() - eris.oovv[km,ki,kb].transpose(0,3,2,1)
689        OVvo = eris.VOov[ke,km,ki].transpose(1,0,3,2).conj()
690        ovVO = eris.voOV[ke,km,ki].transpose(1,0,3,2).conj()
691        OVVO = eris.VOOV[ke,km,ki].transpose(1,0,3,2).conj() - eris.OOVV[km,ki,kb].transpose(0,3,2,1)
692
693        tmp1aa = np.zeros((nocca, nvira, nvira, nocca),dtype=t1a.dtype)
694        tmp1ab = np.zeros((nocca, nvira, nvirb, noccb),dtype=t1a.dtype)
695        tmp1ba = np.zeros((noccb, nvirb, nvira, nocca),dtype=t1a.dtype)
696        tmp1bb = np.zeros((noccb, nvirb, nvirb, noccb),dtype=t1a.dtype)
697
698        for kn in range(nkpts):
699            kf = kconserv[km,ke,kn]
700            ovov = eris.ovov[km,ke,kn] - eris.ovov[km,kf,kn].transpose(0,3,2,1)
701            OVov = eris.OVov[km,ke,kn]
702            ovOV = eris.ovOV[km,ke,kn]
703            OVOV = eris.OVOV[km,ke,kn] - eris.OVOV[km,kf,kn].transpose(0,3,2,1)
704
705            tmp1aa -= einsum('nibf,menf->mebi',t2aa[kn,ki,kb], ovov)
706            tmp1aa += einsum('iNbF,meNF->mebi',t2ab[ki,kn,kb], ovOV)
707
708            tmp1ab += einsum('nIfB,menf->meBI',t2ab[kn,ki,kf], ovov)
709            tmp1ab -= einsum('NIBF,meNF->meBI',t2bb[kn,ki,kb], ovOV)
710
711            tmp1ba += einsum('iNbF,MENF->MEbi',t2ab[ki,kn,kb], OVOV)
712            tmp1ba -= einsum('nibf,MEnf->MEbi',t2aa[kn,ki,kb], OVov)
713
714            tmp1bb -= einsum('NIBF,MENF->MEBI',t2bb[kn,ki,kb], OVOV)
715            tmp1bb += einsum('nIfB,MEnf->MEBI',t2ab[kn,ki,kf], OVov)
716
717        aebi = einsum('ma,mebi->aebi',t1a[km],(ovvo+tmp1aa))
718        Wvvvo[ka,ke,kb] -= aebi
719
720        WVVvo[ka,ke,kb] -= einsum('MA,MEbi->AEbi',t1b[km],OVvo+tmp1ba)
721        WvvVO[ka,ke,kb] -= einsum('ma,meBI->aeBI',t1a[km],ovVO+tmp1ab)
722
723        AEBI = einsum('MA,MEBI->AEBI',t1b[km],(OVVO+tmp1bb))
724        WVVVO[ka,ke,kb] -= AEBI
725
726
727    for ka, ke, kb in itertools.product(range(nkpts),repeat=3):
728        ki = kconserv[ka, ke, kb]
729        # P(ab) <mb||ef> t2(miaf) (alpha alpha beta beta) and (beta beta alpha alpha)
730        for km in range(nkpts):
731            kf = kconserv[km,ke,ka]
732            OVVV = (eris.VOVV[ke,km,kf].transpose(1,0,3,2).conj() -
733                    eris.VOVV[kf,km,ke].transpose(1,2,3,0).conj())
734            ovvv = (eris.vovv[ke,km,kf].transpose(1,0,3,2).conj() -
735                    eris.vovv[kf,km,ke].transpose(1,2,3,0).conj())
736
737            WVVvo[ka,ke,kb] -= lib.einsum('mfAE,mibf->AEbi',
738                                          eris.voVV[kf,km,ke].transpose(1,0,3,2).conj(), t2aa[km,ki,kb])
739            WVVvo[ka,ke,kb] -= lib.einsum('MEAF,iMbF->AEbi', OVVV, t2ab[ki,km,kb])
740
741            WvvVO[ka,ke,kb] -= lib.einsum('MFae,MIBF->aeBI',
742                                          eris.VOvv[kf,km,ke].transpose(1,0,3,2).conj(), t2bb[km,ki,kb])
743            WvvVO[ka,ke,kb] -= lib.einsum('meaf,mIfB->aeBI', ovvv, t2ab[km,ki,kf])
744
745        # P(ab) -t1(ma) (<mb||ei> - t2(nibf) <mn||ef>) for all spin configurations
746        km = kb
747        ovvo = (eris.voov[ke,km,ki].transpose(1,0,3,2).conj() -
748                eris.oovv[km,ki,ka].transpose(0,3,2,1))
749        OVVO = (eris.VOOV[ke,km,ki].transpose(1,0,3,2).conj() -
750                eris.OOVV[km,ki,ka].transpose(0,3,2,1))
751
752        tmp1aa = np.zeros((nocca, nvira, nvira, nocca),dtype=t1a.dtype)
753        tmp1ab = np.zeros((noccb, nvira, nvira, noccb),dtype=t1a.dtype)
754        tmp1ba = np.zeros((nocca, nvirb, nvirb, nocca),dtype=t1a.dtype)
755        tmp1bb = np.zeros((noccb, nvirb, nvirb, noccb),dtype=t1a.dtype)
756
757        for kn in range(nkpts):
758            kf = kconserv[km,ke,kn]
759            ovov = eris.ovov[km,ke,kn] - eris.ovov[km,kf,kn].transpose(0,3,2,1)
760            OVov = eris.OVov[km,ke,kn]
761            ovOV = eris.ovOV[km,ke,kn]
762            OVOV = eris.OVOV[km,ke,kn] - eris.OVOV[km,kf,kn].transpose(0,3,2,1)
763
764            tmp1aa -= einsum('niaf,menf->meai',t2aa[kn,ki,ka], ovov)
765            tmp1aa += einsum('iNaF,meNF->meai',t2ab[ki,kn,ka], ovOV)
766
767            tmp1ab += einsum('nIaF,MFne->MeaI',t2ab[kn,ki,ka], eris.OVov[km,kf,kn])
768
769            tmp1ba += einsum('iNfA,mfNE->mEAi',t2ab[ki,kn,kf], eris.ovOV[km,kf,kn])
770
771            tmp1bb -= einsum('NIAF,MENF->MEAI',t2bb[kn,ki,ka], OVOV)
772            tmp1bb += einsum('nIfA,MEnf->MEAI',t2ab[kn,ki,kf], OVov)
773
774        aebi = einsum('mb,meai->aebi',t1a[km],(ovvo+tmp1aa))
775        Wvvvo[ka,ke,kb] += aebi
776
777        WVVvo[ka,ke,kb] += einsum('mb,mEAi->AEbi',t1a[km], -eris.ooVV[km,ki,ka].transpose(0,3,2,1)+tmp1ba)
778        WvvVO[ka,ke,kb] += einsum('MB,MeaI->aeBI',t1b[km], -eris.OOvv[km,ki,ka].transpose(0,3,2,1)+tmp1ab)
779
780        AEBI = einsum('MB,MEAI->AEBI',t1b[km],(OVVO+tmp1bb))
781        WVVVO[ka,ke,kb] += AEBI
782
783    # Remaining terms
784    for ka, ke, kb in itertools.product(range(nkpts),repeat=3):
785        ki = kf = kconserv[ka, ke, kb]
786        Wvvvo[ka,ke,kb] += eris.vovv[kb,ki,ka].transpose(2,3,0,1) - eris.vovv[ka,ki,kb].transpose(0,3,2,1)
787        WVVvo[ka,ke,kb] += eris.voVV[kb,ki,ka].transpose(2,3,0,1)
788        WvvVO[ka,ke,kb] += eris.VOvv[kb,ki,ka].transpose(2,3,0,1)
789        WVVVO[ka,ke,kb] += eris.VOVV[kb,ki,ka].transpose(2,3,0,1) - eris.VOVV[ka,ki,kb].transpose(0,3,2,1)
790
791        Wvvvo[ka,ke,kb] -= lib.einsum('me,miab->aebi',fova[ke],t2aa[ke,ki,ka])
792        WVVvo[ka,ke,kb] -= lib.einsum('ME,iMbA->AEbi',fovb[ke],t2ab[ki,ke,kb])
793        WvvVO[ka,ke,kb] -= lib.einsum('me,mIaB->aeBI',fova[ke],t2ab[ke,ki,ka])
794        WVVVO[ka,ke,kb] -= lib.einsum('ME,MIAB->AEBI',fovb[ke],t2bb[ke,ki,ka])
795
796        Wvvvv, WvvVV, WVVVV = get_Wvvvv(cc, t1, t2, eris, ka, kb, ke)
797        Wvvvo[ka,ke,kb] += lib.einsum('if,aebf->aebi', t1a[ki], Wvvvv)
798        WVVvo[kb,kf,ka] += lib.einsum('ie,aeBF->BFai', t1a[ke], WvvVV)
799        WvvVO[ka,ke,kb] += lib.einsum('IF,aeBF->aeBI', t1b[ki], WvvVV)
800        WVVVO[ka,ke,kb] += lib.einsum('IF,AEBF->AEBI', t1b[ki], WVVVV)
801
802        for km in range(nkpts):
803            kn = kconserv[ka,km,kb]
804            ovoo = eris.ooov[kn,ki,km].transpose(2,3,0,1) - eris.ooov[km,ki,kn].transpose(0,3,2,1)
805            ovOO = eris.OOov[kn,ki,km].transpose(2,3,0,1)
806
807            ooOV = eris.ooOV[km,ki,kn]
808            OOov = eris.OOov[km,ki,kn]
809
810            OVoo = eris.ooOV[kn,ki,km].transpose(2,3,0,1)
811            OVOO = eris.OOOV[kn,ki,km].transpose(2,3,0,1) - eris.OOOV[km,ki,kn].transpose(0,3,2,1)
812
813            Wvvvo[ka,ke,kb] += 0.5*lib.einsum('meni,mnab->aebi',ovoo,tauaa[km,kn,ka])
814
815            WVVvo[ka,ke,kb] += 0.5*lib.einsum('MEni,nMbA->AEbi',OVoo,tauab[kn,km,kb])
816            WVVvo[ka,ke,kb] += 0.5*lib.einsum('miNE,mNbA->AEbi',ooOV,tauab[km,kn,kb])
817
818            WvvVO[ka,ke,kb] += 0.5*lib.einsum('meNI,mNaB->aeBI',ovOO,tauab[km,kn,ka])
819            WvvVO[ka,ke,kb] += 0.5*lib.einsum('MIne,nMaB->aeBI',OOov,tauab[kn,km,ka])
820
821            WVVVO[ka,ke,kb] += 0.5*lib.einsum('MENI,MNAB->AEBI',OVOO,taubb[km,kn,ka])
822
823    return Wvvvo, WvvVO, WVVvo, WVVVO
824
825
826def Woooo(cc,t1,t2,eris):
827    kconserv = cc.khelper.kconserv
828    t1a, t1b = t1
829    t2aa, t2ab, t2bb = t2
830    _, _, nkpts, nocca, noccb, nvira, nvirb = t2ab.shape
831
832    dtype = np.result_type(*t2)
833    Woooo = np.zeros(eris.oooo.shape, dtype=dtype)
834    WooOO = np.zeros(eris.ooOO.shape, dtype=dtype)
835    WOOOO = np.zeros(eris.OOOO.shape, dtype=dtype)
836
837    tau_aa, tau_ab, tau_bb = make_tau(cc, t2, t1, t1)
838    for km in range(nkpts):
839        for kn in range(nkpts):
840            tmp_aaaaJ = einsum('xje, ymine->yxminj', t1a, eris.ooov[km,:,kn])
841            tmp_aaaaJ-= einsum('yie, xmjne->yxminj', t1a, eris.ooov[km,:,kn])
842            tmp_bbbbJ = einsum('xje, ymine->yxminj', t1b, eris.OOOV[km,:,kn])
843            tmp_bbbbJ-= einsum('yie, xmjne->yxminj', t1b, eris.OOOV[km,:,kn])
844            #tmp_aabbJ = einsum('xje, ymine->yxminj', t1b, eris.ooOV[km,:,kn])
845            #tmp_bbaaJ = einsum('xje, ymine->yxminj', t1a, eris.OOov[km,:,kn])
846            #tmp_abbaJ = -einsum('yie,xmjne->yxminj', t1b, eris.ooOV[km,:,kn])
847            tmp_baabJ = -einsum('yie,xmjne->yxminj', t1a, eris.OOov[km,:,kn])
848            tmp_aabbJ = einsum('xje, ymine->yxminj', t1b, eris.ooOV[km,:,kn])
849
850            for ki in range(nkpts):
851                kj = kconserv[km,ki,kn]
852                Woooo[km,ki,kn] += tmp_aaaaJ[ki,kj]
853                WOOOO[km,ki,kn] += tmp_bbbbJ[ki,kj]
854                WooOO[km,ki,kn] += tmp_aabbJ[ki,kj]
855                WooOO[kn,ki,km] -= tmp_baabJ[ki,kj].transpose(2,1,0,3)
856                Woooo[km,ki,kn] += eris.oooo[km,ki,kn]
857                WooOO[km,ki,kn] += eris.ooOO[km,ki,kn]
858                WOOOO[km,ki,kn] += eris.OOOO[km,ki,kn]
859
860    Woooo = Woooo - Woooo.transpose(2,1,0,5,4,3,6)
861    WOOOO = WOOOO - WOOOO.transpose(2,1,0,5,4,3,6)
862
863    for km, ki, kn in itertools.product(range(nkpts), repeat=3):
864        kj = kconserv[km, ki, kn]
865
866        for ke in range(nkpts):
867            kf = kconserv[km, ke, kn]
868
869            ovov = eris.ovov[km, ke, kn] - eris.ovov[km, kf, kn].transpose(0,3,2,1)
870            OVOV = eris.OVOV[km, ke, kn] - eris.OVOV[km, kf, kn].transpose(0,3,2,1)
871
872            Woooo[km, ki, kn] += 0.5*lib.einsum('ijef,menf->minj', tau_aa[ki, kj, ke],      ovov)
873            WOOOO[km, ki, kn] += 0.5*lib.einsum('IJEF,MENF->MINJ', tau_bb[ki, kj, ke],      OVOV)
874            WooOO[km, ki, kn] +=     lib.einsum('iJeF,meNF->miNJ', tau_ab[ki, kj, ke], eris.ovOV[km, ke, kn])
875
876    WOOoo = None
877    return Woooo, WooOO, WOOoo, WOOOO
878
879def Woovo(cc,t1,t2,eris):
880    kconserv = cc.khelper.kconserv
881    t1a, t1b = t1
882    t2aa, t2ab, t2bb = t2
883    nkpts, nocca, noccb, nvira, nvirb = t2ab.shape[2:]
884
885    dtype = np.result_type(*t2)
886    Woovo = np.zeros((nkpts, nkpts, nkpts, nocca, nocca, nvira, nocca), dtype=dtype)
887    WooVO = np.zeros((nkpts, nkpts, nkpts, nocca, nocca, nvirb, noccb), dtype=dtype)
888    WOOvo = np.zeros((nkpts, nkpts, nkpts, noccb, noccb, nvira, nocca), dtype=dtype)
889    WOOVO = np.zeros((nkpts, nkpts, nkpts, noccb, noccb, nvirb, noccb), dtype=dtype)
890
891    for km, kb, ki in kpts_helper.loop_kkk(nkpts):
892        kj = kconserv[km, ki, kb]
893        Woovo[km,ki,kb] += eris.ooov[ki,km,kj].transpose(1,0,3,2).conj()
894        Woovo[km,ki,kb] -= eris.ooov[kj,km,ki].transpose(1,2,3,0).conj()
895        WooVO[km,ki,kb] += eris.ooOV[ki,km,kj].transpose(1,0,3,2).conj()
896        WOOvo[km,ki,kb] += eris.OOov[ki,km,kj].transpose(1,0,3,2).conj()
897        WOOVO[km,ki,kb] += eris.OOOV[ki,km,kj].transpose(1,0,3,2).conj()
898        WOOVO[km,ki,kb] -= eris.OOOV[kj,km,ki].transpose(1,2,3,0).conj()
899        for kn in range(nkpts):
900            ke = kconserv[km,ki,kn]
901            ooov = eris.ooov[km,ki,kn] - eris.ooov[kn,ki,km].transpose(2,1,0,3)
902            OOOV = eris.OOOV[km,ki,kn] - eris.OOOV[kn,ki,km].transpose(2,1,0,3)
903
904            Woovo[km,ki,kb] += einsum('mine,jnbe->mibj', ooov, t2aa[kj,kn,kb])
905            Woovo[km,ki,kb] += einsum('miNE,jNbE->mibj', eris.ooOV[km,ki,kn], t2ab[kj,kn,kb])
906            WooVO[km,ki,kb] += einsum('mine,nJeB->miBJ', ooov, t2ab[kn,kj,ke])
907            WooVO[km,ki,kb] += einsum('miNE,JNBE->miBJ', eris.ooOV[km,ki,kn], t2bb[kj,kn,kb])
908            WOOvo[km,ki,kb] += einsum('MINE,jNbE->MIbj', OOOV, t2ab[kj,kn,kb])
909            WOOvo[km,ki,kb] += einsum('MIne,jnbe->MIbj', eris.OOov[km,ki,kn], t2aa[kj,kn,kb])
910            WOOVO[km,ki,kb] += einsum('MINE,JNBE->MIBJ', OOOV, t2bb[kj,kn,kb])
911            WOOVO[km,ki,kb] += einsum('MIne,nJeB->MIBJ', eris.OOov[km,ki,kn], t2ab[kn,kj,ke])
912            # P(ij)
913            ke = kconserv[km,kj,kn]
914            ooov = eris.ooov[km,kj,kn] - eris.ooov[kn,kj,km].transpose(2,1,0,3)
915            OOOV = eris.OOOV[km,kj,kn] - eris.OOOV[kn,kj,km].transpose(2,1,0,3)
916
917            Woovo[km,ki,kb] -= einsum('mjne,inbe->mibj', ooov, t2aa[ki,kn,kb])
918            Woovo[km,ki,kb] -= einsum('mjNE,iNbE->mibj', eris.ooOV[km,kj,kn], t2ab[ki,kn,kb])
919            WooVO[km,ki,kb] -= einsum('NJme,iNeB->miBJ', eris.OOov[kn,kj,km], t2ab[ki,kn,ke])
920            WOOvo[km,ki,kb] -= einsum('njME,nIbE->MIbj', eris.ooOV[kn,kj,km], t2ab[kn,ki,kb])
921            WOOVO[km,ki,kb] -= einsum('MJNE,INBE->MIBJ', OOOV, t2bb[ki,kn,kb])
922            WOOVO[km,ki,kb] -= einsum('MJne,nIeB->MIBJ', eris.OOov[km,kj,kn], t2ab[kn,ki,ke])
923
924        ovvo = eris.voov[ki,km,kj].transpose(1,0,3,2).conj() - eris.oovv[km,kj,kb].transpose(0,3,2,1)
925        OVVO = eris.VOOV[ki,km,kj].transpose(1,0,3,2).conj() - eris.OOVV[km,kj,kb].transpose(0,3,2,1)
926        ovVO = eris.voOV[ki,km,kj].transpose(1,0,3,2).conj()
927        OVvo = eris.VOov[ki,km,kj].transpose(1,0,3,2).conj()
928        Woovo[km,ki,kb] += einsum('ie,mebj->mibj', t1a[ki], ovvo)
929        WooVO[km,ki,kb] += einsum('ie,meBJ->miBJ', t1a[ki], ovVO)
930        WOOvo[km,ki,kb] += einsum('IE,MEbj->MIbj', t1b[ki], OVvo)
931        WOOVO[km,ki,kb] += einsum('IE,MEBJ->MIBJ', t1b[ki], OVVO)
932        #P(ij)
933        ovvo = eris.voov[kj,km,ki].transpose(1,0,3,2).conj() - eris.oovv[km,ki,kb].transpose(0,3,2,1)
934        OVVO = eris.VOOV[kj,km,ki].transpose(1,0,3,2).conj() - eris.OOVV[km,ki,kb].transpose(0,3,2,1)
935        Woovo[km,ki,kb] -= einsum('je,mebi->mibj', t1a[kj], ovvo)
936        WooVO[km,ki,kb] -= -einsum('JE,miBE->miBJ', t1b[kj], eris.ooVV[km,ki,kb])
937        WOOvo[km,ki,kb] -= -einsum('je,MIbe->MIbj', t1a[kj], eris.OOvv[km,ki,kb])
938        WOOVO[km,ki,kb] -= einsum('JE,MEBI->MIBJ', t1b[kj], OVVO)
939
940
941        for kf in range(nkpts):
942            kn = kconserv[kb, kj, kf]
943            ovov = eris.ovov[km,ki,kn] - eris.ovov[km,kf,kn].transpose(0,3,2,1)
944            OVOV = eris.OVOV[km,ki,kn] - eris.OVOV[km,kf,kn].transpose(0,3,2,1)
945            Woovo[km,ki,kb] -= (einsum('ie,njbf,menf->mibj', t1a[ki], t2aa[kn,kj,kb], ovov) -
946                                einsum('ie,jNbF,meNF->mibj', t1a[ki], t2ab[kj,kn,kb], eris.ovOV[km,ki,kn]))
947            WooVO[km,ki,kb] -= (-einsum('ie,nJfB,menf->miBJ', t1a[ki], t2ab[kn,kj,kf], ovov) +
948                                einsum('ie,NJBF,meNF->miBJ', t1a[ki], t2bb[kn,kj,kb], eris.ovOV[km,ki,kn]))
949            WOOvo[km,ki,kb] -= (-einsum('IE,jNbF,MENF->MIbj', t1b[ki], t2ab[kj,kn,kb], OVOV) +
950                                einsum('IE,njbf,MEnf->MIbj', t1b[ki], t2aa[kn,kj,kb], eris.OVov[km,ki,kn]))
951            WOOVO[km,ki,kb] -= (einsum('IE,NJBF,MENF->MIBJ', t1b[ki], t2bb[kn,kj,kb], OVOV) -
952                                einsum('IE,nJfB,MEnf->MIBJ', t1b[ki], t2ab[kn,kj,kf], eris.OVov[km,ki,kn]))
953            #P(ij)
954            kn = kconserv[kb, ki, kf]
955            ovov = eris.ovov[km,kj,kn] - eris.ovov[km,kf,kn].transpose(0,3,2,1)
956            OVOV = eris.OVOV[km,kj,kn] - eris.OVOV[km,kf,kn].transpose(0,3,2,1)
957            Woovo[km,ki,kb] += (einsum('je,nibf,menf->mibj', t1a[kj], t2aa[kn,ki,kb], ovov) -
958                                einsum('je,iNbF,meNF->mibj', t1a[kj], t2ab[ki,kn,kb], eris.ovOV[km,kj,kn]))
959            WooVO[km,ki,kb] += -einsum('JE,iNfB,mfNE->miBJ', t1b[kj], t2ab[ki,kn,kf], eris.ovOV[km, kf, kn])
960            WOOvo[km,ki,kb] += -einsum('je,nIbF,MFne->MIbj', t1a[kj], t2ab[kn,ki,kb], eris.OVov[km, kf, kn])
961            WOOVO[km,ki,kb] += (einsum('JE,NIBF,MENF->MIBJ', t1b[kj], t2bb[kn,ki,kb], OVOV) -
962                                einsum('JE,nIfB,MEnf->MIBJ', t1b[kj], t2ab[kn,ki,kf], eris.OVov[km,kj,kn]))
963
964    Fme, FME = Fov(cc, t1, t2, eris)
965    Wminj, WmiNJ, WMInj, WMINJ = Woooo(cc,t1,t2,eris)
966    tauaa, tauab, taubb = make_tau(cc, t2, t1, t1, fac=1.)
967    for km, kb, ki in kpts_helper.loop_kkk(nkpts):
968        kj = kconserv[km, ki, kb]
969
970        Woovo[km,ki,kb] -= einsum('me,ijbe->mibj', Fme[km], t2aa[ki,kj,kb])
971        WooVO[km,ki,kb] -= -einsum('me,iJeB->miBJ', Fme[km], t2ab[ki,kj,km])
972        WOOvo[km,ki,kb] -= -einsum('ME,jIbE->MIbj', FME[km], t2ab[kj,ki,kb])
973        WOOVO[km,ki,kb] -= einsum('ME,IJBE->MIBJ', FME[km], t2bb[ki,kj,kb])
974
975        Woovo[km,ki,kb] -= einsum('nb, minj->mibj', t1a[kb], Wminj[km, ki, kb])
976        WooVO[km,ki,kb] -= einsum('NB, miNJ->miBJ', t1b[kb], WmiNJ[km, ki, kb])
977        WOOvo[km,ki,kb] -= einsum('nb, njMI->MIbj', t1a[kb], WmiNJ[kb, kj, km])
978        WOOVO[km,ki,kb] -= einsum('NB, MINJ->MIBJ', t1b[kb], WMINJ[km, ki, kb])
979
980    for km, kb, ki in kpts_helper.loop_kkk(nkpts):
981        kj = kconserv[km, ki, kb]
982        for ke in range(nkpts):
983            kf = kconserv[km,ke,kb]
984            ovvv = eris.vovv[ke,km,kf].transpose(1,0,3,2).conj() - eris.vovv[kf,km,ke].transpose(1,2,3,0).conj()
985            OVVV = eris.VOVV[ke,km,kf].transpose(1,0,3,2).conj() - eris.VOVV[kf,km,ke].transpose(1,2,3,0).conj()
986            ovVV = eris.voVV[ke,km,kf].transpose(1,0,3,2).conj()
987            OVvv = eris.VOvv[ke,km,kf].transpose(1,0,3,2).conj()
988            Woovo[km,ki,kb] += 0.5 * einsum('mebf,ijef->mibj', ovvv, tauaa[ki,kj,ke])
989            WooVO[km,ki,kb] += einsum('meBF,iJeF->miBJ', ovVV, tauab[ki,kj,ke])
990            WOOvo[km,ki,kb] += einsum('MEbf,jIfE->MIbj', OVvv, tauab[kj,ki,kf])
991            WOOVO[km,ki,kb] += 0.5 * einsum('MEBF,IJEF->MIBJ', OVVV, taubb[ki,kj,ke])
992
993    return Woovo, WooVO, WOOvo, WOOVO
994
995
996def Wooov(cc, t1, t2, eris, kconserv):
997    t1a, t1b = t1
998
999    Wooov = eris.ooov - np.asarray(eris.ooov).transpose(2,1,0,5,4,3,6)
1000    WooOV = np.array(eris.ooOV, copy=True)
1001    WOOov = np.array(eris.OOov, copy=True)
1002    WOOOV = eris.OOOV - np.asarray(eris.OOOV).transpose(2,1,0,5,4,3,6)
1003
1004    Wooov += einsum('yif,xyzmfne->xyzmine', t1a, eris.ovov) - einsum('yif, zyxnfme->xyzmine', t1a, eris.ovov)
1005    WooOV += einsum('yif,xyzmfNE->xyzmiNE', t1a, eris.ovOV)
1006    WOOov += einsum('yIF,xyzMFne->xyzMIne', t1b, eris.OVov)
1007    WOOOV += einsum('yIF,xyzMFNE->xyzMINE', t1b, eris.OVOV) - einsum('yIF, zyxNFME->xyzMINE', t1b, eris.OVOV)
1008
1009    return Wooov, WooOV, WOOov, WOOOV
1010
1011def Wovvo(cc, t1, t2, eris):
1012    kconserv = cc.khelper.kconserv
1013    t1a, t1b = t1
1014    t2aa, t2ab, t2bb = t2
1015    nkpts, nocca, noccb, nvira, nvirb = t2ab.shape[2:]
1016
1017    Wovvo, WovVO, WOVvo, WOVVO, WoVVo, WOvvO = cc_Wovvo(cc,t1,t2,eris)
1018    for km, kb, ke in kpts_helper.loop_kkk(nkpts):
1019        kj = kconserv[km, ke, kb]
1020        for kn in range(nkpts):
1021            kf = kconserv[km, ke, kn]
1022            Wovvo[km,ke,kb] += 0.5 * einsum('jnbf,menf->mebj', t2aa[kj,kn,kb], eris.ovov[km,ke,kn])
1023            Wovvo[km,ke,kb] -= 0.5 * einsum('jnbf,mfne->mebj', t2aa[kj,kn,kb], eris.ovov[km,kf,kn])
1024            Wovvo[km,ke,kb] += 0.5 * einsum('jNbF,meNF->mebj', t2ab[kj,kn,kb], eris.ovOV[km,ke,kn])
1025
1026            WOVvo[km,ke,kb] += 0.5 * einsum('jNbF,MENF->MEbj', t2ab[kj,kn,kb], eris.OVOV[km,ke,kn])
1027            WOVvo[km,ke,kb] -= 0.5 * einsum('jNbF,MFNE->MEbj', t2ab[kj,kn,kb], eris.OVOV[km,kf,kn])
1028            WOVvo[km,ke,kb] += 0.5 * einsum('jnbf,MEnf->MEbj', t2aa[kj,kn,kb], eris.OVov[km,ke,kn])
1029
1030            WovVO[km,ke,kb] += 0.5 * einsum('nJfB,menf->meBJ', t2ab[kn,kj,kf], eris.ovov[km,ke,kn])
1031            WovVO[km,ke,kb] -= 0.5 * einsum('nJfB,mfne->meBJ', t2ab[kn,kj,kf], eris.ovov[km,kf,kn])
1032            WovVO[km,ke,kb] += 0.5 * einsum('JNBF,meNF->meBJ', t2bb[kj,kn,kb], eris.ovOV[km,ke,kn])
1033
1034            WOVVO[km,ke,kb] += 0.5 * einsum('JNBF,MENF->MEBJ', t2bb[kj,kn,kb], eris.OVOV[km,ke,kn])
1035            WOVVO[km,ke,kb] -= 0.5 * einsum('JNBF,MFNE->MEBJ', t2bb[kj,kn,kb], eris.OVOV[km,kf,kn])
1036            WOVVO[km,ke,kb] += 0.5 * einsum('nJfB,MEnf->MEBJ', t2ab[kn,kj,kf], eris.OVov[km,ke,kn])
1037
1038    return Wovvo, WovVO, WOVvo, WOVVO
1039
1040def W1oovv(cc, t1, t2, eris):
1041    kconserv = cc.khelper.kconserv
1042    dtype = np.result_type(*t1)
1043    t1a, t1b = t1
1044    t2aa, t2ab, t2bb = t2
1045    nkpts, nocc, nvir = t1a.shape
1046    Woovv = np.zeros(eris.oovv.shape, dtype=dtype)
1047    WooVV = np.zeros(eris.ooVV.shape, dtype=dtype)
1048    WOOvv = np.zeros(eris.OOvv.shape, dtype=dtype)
1049    WOOVV = np.zeros(eris.OOVV.shape, dtype=dtype)
1050    for kk in range(nkpts):
1051        for ki in range(nkpts):
1052            for kb in range(nkpts):
1053                kd = kconserv[kk,ki,kb]
1054                Woovv[kk,ki,kb] += eris.oovv[kk,ki,kb]
1055                Woovv[kk,ki,kb] -= eris.voov[kb,ki,kk].transpose(2,1,0,3)
1056                WooVV[kk,ki,kb] += eris.ooVV[kk,ki,kb]
1057                WOOvv[kk,ki,kb] += eris.OOvv[kk,ki,kb]
1058                WOOVV[kk,ki,kb] += eris.OOVV[kk,ki,kb]
1059                WOOVV[kk,ki,kb] -= eris.VOOV[kb,ki,kk].transpose(2,1,0,3)
1060
1061                for kl in range(nkpts):
1062                    kc = kconserv[ki,kb,kl]
1063                    Woovv[kk,ki,kb] -= einsum('lckd,ilbc->kibd', eris.ovov[kl,kc,kk], t2aa[ki,kl,kb])
1064                    Woovv[kk,ki,kb] += einsum('ldkc,ilbc->kibd', eris.ovov[kl,kd,kk], t2aa[ki,kl,kb])
1065                    Woovv[kk,ki,kb] -= einsum('LCkd,iLbC->kibd', eris.OVov[kl,kc,kk], t2ab[ki,kl,kb])
1066
1067                    WooVV[kk,ki,kb] -= einsum('kcLD,iLcB->kiBD', eris.ovOV[kk,kc,kl], t2ab[ki,kl,kc])
1068                    WOOvv[kk,ki,kb] -= einsum('KCld,lIbC->KIbd', eris.OVov[kk,kc,kl], t2ab[kl,ki,kb])
1069
1070                    WOOVV[kk,ki,kb] -= einsum('LCKD,ILBC->KIBD', eris.OVOV[kl,kc,kk], t2bb[ki,kl,kb])
1071                    WOOVV[kk,ki,kb] += einsum('LDKC,ILBC->KIBD', eris.OVOV[kl,kd,kk], t2bb[ki,kl,kb])
1072                    WOOVV[kk,ki,kb] -= einsum('lcKD,lIcB->KIBD', eris.ovOV[kl,kc,kk], t2ab[kl,ki,kc])
1073
1074    return Woovv, WooVV, WOOvv, WOOVV
1075
1076def W2oovv(cc, t1, t2, eris):
1077    kconserv = cc.khelper.kconserv
1078    dtype = np.result_type(*t1)
1079    t1a, t1b = t1
1080    t2aa, t2ab, t2bb = t2
1081    nkpts, nocc, nvir = t1a.shape
1082    Woovv = np.zeros(eris.oovv.shape, dtype=dtype)
1083    WooVV = np.zeros(eris.ooVV.shape, dtype=dtype)
1084    WOOvv = np.zeros(eris.OOvv.shape, dtype=dtype)
1085    WOOVV = np.zeros(eris.OOVV.shape, dtype=dtype)
1086    WWooov, WWooOV, WWOOov, WWOOOV = Wooov(cc, t1, t2, eris, kconserv)
1087    for kk in range(nkpts):
1088        for ki in range(nkpts):
1089            for kb in range(nkpts):
1090                kd = kconserv[kk,ki,kb]
1091                Woovv[kk,ki,kb] += einsum('kild,lb->kibd',WWooov[kk,ki,kb],-t1a[kb])
1092                WooVV[kk,ki,kb] += einsum('kiLD,LB->kiBD',WWooOV[kk,ki,kb],-t1b[kb])
1093                WOOvv[kk,ki,kb] += einsum('KIld,lb->KIbd',WWOOov[kk,ki,kb],-t1a[kb])
1094                WOOVV[kk,ki,kb] += einsum('KILD,LB->KIBD',WWOOOV[kk,ki,kb],-t1b[kb])
1095
1096                Woovv[kk,ki,kb] += einsum('ckdb,ic->kibd', eris.vovv[ki,kk,kd].conj(), t1a[ki])
1097                Woovv[kk,ki,kb] -= einsum('dkcb,ic->kibd', eris.vovv[kd,kk,ki].conj(), t1a[ki])
1098
1099                WooVV[kk,ki,kb] += einsum('ckDB,ic->kiBD', eris.voVV[ki,kk,kd].conj(), t1a[ki])
1100                WOOvv[kk,ki,kb] += einsum('CKdb,IC->KIbd', eris.VOvv[ki,kk,kd].conj(), t1b[ki])
1101
1102                WOOVV[kk,ki,kb] += einsum('CKDB,IC->KIBD', eris.VOVV[ki,kk,kd].conj(), t1b[ki])
1103                WOOVV[kk,ki,kb] -= einsum('DKCB,IC->KIBD', eris.VOVV[kd,kk,ki].conj(), t1b[ki])
1104
1105    return Woovv, WooVV, WOOvv, WOOVV
1106
1107def Woovv(cc, t1, t2, eris):
1108    t1a, t1b = t1
1109    nkpts, nocc, nvir = t1a.shape
1110    Woovv, WooVV, WOOvv, WOOVV = W1oovv(cc, t1, t2, eris)
1111    WWoovv, WWooVV, WWOOvv, WWOOVV = W2oovv(cc, t1, t2, eris)
1112    for kk,ki,kb in itertools.product(range(nkpts), repeat=3):
1113        Woovv[kk,ki,kb] = Woovv[kk,ki,kb] + WWoovv[kk,ki,kb]
1114        WooVV[kk,ki,kb] = WooVV[kk,ki,kb] + WWooVV[kk,ki,kb]
1115        WOOvv[kk,ki,kb] = WOOvv[kk,ki,kb] + WWOOvv[kk,ki,kb]
1116        WOOVV[kk,ki,kb] = WOOVV[kk,ki,kb] + WWOOVV[kk,ki,kb]
1117    return Woovv, WooVV, WOOvv, WOOVV
1118
1119# vvvv is a string, ('oooo', 'ooov', ..., 'vvvv')
1120# orbspin can be accessed through general spin-orbital kintermediates eris
1121# orbspin = eris.mo_coeff.orbspin
1122def _eri_spin2spatial(chemist_eri_spin, vvvv, eris, nocc, orbspin, cross_ab=False):
1123    nocc_a, nocc_b = nocc
1124    nocc = nocc_a + nocc_b
1125    nkpts = len(orbspin)
1126    idxoa = [np.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)]
1127    idxob = [np.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)]
1128    idxva = [np.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)]
1129    idxvb = [np.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)]
1130
1131    def select_idx(s):
1132        if s.lower() == 'o':
1133            return idxoa, idxob
1134        else:
1135            return idxva, idxvb
1136
1137    if len(vvvv) == 2:
1138        idx1a, idx1b = select_idx(vvvv[0])
1139        idx2a, idx2b = select_idx(vvvv[1])
1140
1141        fa = np.zeros((nkpts,len(idx1a[0]),len(idx2a[0])), dtype=np.complex128)
1142        fb = np.zeros((nkpts,len(idx1b[0]),len(idx2b[0])), dtype=np.complex128)
1143        for k in range(nkpts):
1144            fa[k] = chemist_eri_spin[k, idx1a[k][:,None],idx2a[k]]
1145            fb[k] = chemist_eri_spin[k, idx1b[k][:,None],idx2b[k]]
1146        return fa, fb
1147
1148    idx1a, idx1b = select_idx(vvvv[0])
1149    idx2a, idx2b = select_idx(vvvv[1])
1150    idx3a, idx3b = select_idx(vvvv[2])
1151    idx4a, idx4b = select_idx(vvvv[3])
1152
1153    eri_aaaa = np.zeros((nkpts,nkpts,nkpts,len(idx1a[0]),len(idx2a[0]),len(idx3a[0]),len(idx4a[0])), dtype=np.complex128)  # noqa: E501
1154    eri_aabb = np.zeros((nkpts,nkpts,nkpts,len(idx1a[0]),len(idx2a[0]),len(idx3b[0]),len(idx4b[0])), dtype=np.complex128)  # noqa: E501
1155    eri_bbaa = np.zeros((nkpts,nkpts,nkpts,len(idx1b[0]),len(idx2b[0]),len(idx3a[0]),len(idx4a[0])), dtype=np.complex128)  # noqa: E501
1156    eri_bbbb = np.zeros((nkpts,nkpts,nkpts,len(idx1b[0]),len(idx2b[0]),len(idx3b[0]),len(idx4b[0])), dtype=np.complex128)  # noqa: E501
1157    if cross_ab:
1158        eri_abba = np.zeros((nkpts,nkpts,nkpts,len(idx1a[0]),len(idx2b[0]),len(idx3b[0]),len(idx4a[0])), dtype=np.complex128)  # noqa: E501
1159        eri_baab = np.zeros((nkpts,nkpts,nkpts,len(idx1b[0]),len(idx2a[0]),len(idx3a[0]),len(idx4b[0])), dtype=np.complex128)  # noqa: E501
1160    kconserv = kpts_helper.get_kconserv(eris.cell, eris.kpts)
1161    for ki, kj, kk in kpts_helper.loop_kkk(nkpts):
1162        kl = kconserv[ki, kj, kk]
1163        eri_aaaa[ki,kj,kk] = chemist_eri_spin[ki,kj,kk, idx1a[ki][:,None,None,None],idx2a[kj][:,None,None],idx3a[kk][:,None],idx4a[kl]]  # noqa: E501
1164        eri_aabb[ki,kj,kk] = chemist_eri_spin[ki,kj,kk, idx1a[ki][:,None,None,None],idx2a[kj][:,None,None],idx3b[kk][:,None],idx4b[kl]]  # noqa: E501
1165        eri_bbaa[ki,kj,kk] = chemist_eri_spin[ki,kj,kk, idx1b[ki][:,None,None,None],idx2b[kj][:,None,None],idx3a[kk][:,None],idx4a[kl]]  # noqa: E501
1166        eri_bbbb[ki,kj,kk] = chemist_eri_spin[ki,kj,kk, idx1b[ki][:,None,None,None],idx2b[kj][:,None,None],idx3b[kk][:,None],idx4b[kl]]  # noqa: E501
1167        if cross_ab:
1168            eri_abba[ki,kj,kk] = chemist_eri_spin[ki,kj,kk, idx1a[ki][:,None,None,None],idx2b[kj][:,None,None],idx3b[kk][:,None],idx4a[kl]]  # noqa: E501
1169            eri_baab[ki,kj,kk] = chemist_eri_spin[ki,kj,kk, idx1b[ki][:,None,None,None],idx2a[kj][:,None,None],idx3a[kk][:,None],idx4b[kl]]  # noqa: E501
1170    if cross_ab:
1171        return eri_aaaa, eri_aabb, eri_bbaa, eri_bbbb, eri_abba, eri_baab
1172    else:
1173        return eri_aaaa, eri_aabb, eri_bbaa, eri_bbbb
1174
1175def _eri_spatial2spin(eri_aa_ab_ba_bb, vvvv, eris, orbspin, cross_ab=False):
1176    nocc_a, nocc_b = eris.nocc
1177    nocc = nocc_a + nocc_b
1178    nkpts = len(orbspin)
1179    idxoa = [np.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)]
1180    idxob = [np.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)]
1181    idxva = [np.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)]
1182    idxvb = [np.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)]
1183
1184    def select_idx(s):
1185        if s.lower() == 'o':
1186            return idxoa, idxob
1187        else:
1188            return idxva, idxvb
1189
1190    if len(vvvv) == 2:
1191        idx1a, idx1b = select_idx(vvvv[0])
1192        idx2a, idx2b = select_idx(vvvv[1])
1193
1194        fa, fb = eri_aa_ab_ba_bb
1195        f = np.zeros((nkpts, len(idx1a[0])+len(idx1b[0]),
1196                      len(idx2a[0])+len(idx2b[0])), dtype=np.complex128)
1197        for k in range(nkpts):
1198            f[k, idx1a[k][:,None],idx2a[k]] = fa[k]
1199            f[k, idx1b[k][:,None],idx2b[k]] = fb[k]
1200        return f
1201
1202    idx1a, idx1b = select_idx(vvvv[0])
1203    idx2a, idx2b = select_idx(vvvv[1])
1204    idx3a, idx3b = select_idx(vvvv[2])
1205    idx4a, idx4b = select_idx(vvvv[3])
1206
1207    if cross_ab:
1208        eri_aaaa, eri_aabb, eri_bbaa, eri_bbbb, eri_abba, eri_baab = eri_aa_ab_ba_bb
1209    else:
1210        eri_aaaa, eri_aabb, eri_bbaa, eri_bbbb = eri_aa_ab_ba_bb
1211    eri = np.zeros((nkpts,nkpts,nkpts, len(idx1a[0])+len(idx1b[0]),
1212                    len(idx2a[0])+len(idx2b[0]),
1213                    len(idx3a[0])+len(idx3b[0]),
1214                    len(idx4a[0])+len(idx4b[0])), dtype=np.complex128)
1215    kconserv = kpts_helper.get_kconserv(eris.cell, eris.kpts)
1216    for ki, kj, kk in kpts_helper.loop_kkk(nkpts):
1217        kl = kconserv[ki, kj, kk]
1218        eri[ki,kj,kk, idx1a[ki][:,None,None,None],idx2a[kj][:,None,None],idx3a[kk][:,None],idx4a[kl]] = eri_aaaa[ki,kj,kk]  # noqa: E501
1219        eri[ki,kj,kk, idx1a[ki][:,None,None,None],idx2a[kj][:,None,None],idx3b[kk][:,None],idx4b[kl]] = eri_aabb[ki,kj,kk]  # noqa: E501
1220        eri[ki,kj,kk, idx1b[ki][:,None,None,None],idx2b[kj][:,None,None],idx3a[kk][:,None],idx4a[kl]] = eri_bbaa[ki,kj,kk]  # noqa: E501
1221        eri[ki,kj,kk, idx1b[ki][:,None,None,None],idx2b[kj][:,None,None],idx3b[kk][:,None],idx4b[kl]] = eri_bbbb[ki,kj,kk]  # noqa: E501
1222        if cross_ab:
1223            eri[ki,kj,kk, idx1a[ki][:,None,None,None],idx2b[kj][:,None,None],idx3b[kk][:,None],idx4a[kl]] = eri_abba[ki,kj,kk]  # noqa: E501
1224            eri[ki,kj,kk, idx1b[ki][:,None,None,None],idx2a[kj][:,None,None],idx3a[kk][:,None],idx4b[kl]] = eri_baab[ki,kj,kk]  # noqa: E501
1225    return eri
1226