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#
16# Authors: James D. McClain
17#          Timothy Berkelbach <tim.berkelbach@gmail.com>
18#
19
20
21import numpy
22from functools import reduce
23
24from pyscf import lib
25from pyscf.lib import logger
26from pyscf.pbc import scf
27from pyscf.cc import gccsd
28from pyscf.cc import ccsd
29from pyscf.pbc.mp.kmp2 import (get_frozen_mask, get_nmo, get_nocc,
30                               padded_mo_coeff, padding_k_idx)  # noqa
31from pyscf.pbc.cc import kintermediates as imdk
32from pyscf.lib.parameters import LOOSE_ZERO_TOL, LARGE_DENOM  # noqa
33from pyscf.pbc.lib import kpts_helper
34
35DEBUG = False
36
37#
38# FIXME: When linear dependence is found in KHF and handled by function
39# pyscf.scf.addons.remove_linear_dep_, different k-point may have different
40# number of orbitals.
41#
42
43#einsum = numpy.einsum
44einsum = lib.einsum
45
46
47def energy(cc, t1, t2, eris):
48    nkpts, nocc, nvir = t1.shape
49    fock = eris.fock
50    eris_oovv = eris.oovv.copy()
51    e = 0.0 + 0j
52    for ki in range(nkpts):
53        e += einsum('ia,ia', fock[ki, :nocc, nocc:], t1[ki, :, :])
54    t1t1 = numpy.zeros(shape=t2.shape, dtype=t2.dtype)
55    for ki in range(nkpts):
56        ka = ki
57        for kj in range(nkpts):
58            #kb = kj
59            t1t1[ki, kj, ka, :, :, :, :] = einsum('ia,jb->ijab', t1[ki, :, :], t1[kj, :, :])
60    tau = t2 + 2 * t1t1
61    e += 0.25 * numpy.dot(tau.flatten(), eris_oovv.flatten())
62    e /= nkpts
63    if abs(e.imag) > 1e-4:
64        logger.warn(cc, 'Non-zero imaginary part found in KCCSD energy %s', e)
65    return e.real
66
67
68def update_amps(cc, t1, t2, eris):
69    time0 = logger.process_clock(), logger.perf_counter()
70    log = logger.Logger(cc.stdout, cc.verbose)
71    nkpts, nocc, nvir = t1.shape
72    fock = eris.fock
73    mo_e_o = [e[:nocc] for e in eris.mo_energy]
74    mo_e_v = [e[nocc:] + cc.level_shift for e in eris.mo_energy]
75
76    # Get location of padded elements in occupied and virtual space
77    nonzero_opadding, nonzero_vpadding = padding_k_idx(cc, kind="split")
78
79    fov = fock[:, :nocc, nocc:].copy()
80
81    # Get the momentum conservation array
82    # Note: chemist's notation for momentum conserving t2(ki,kj,ka,kb), even though
83    # integrals are in physics notation
84    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
85
86    tau = imdk.make_tau(cc, t2, t1, t1, kconserv)
87
88    Fvv = imdk.cc_Fvv(cc, t1, t2, eris, kconserv)
89    Foo = imdk.cc_Foo(cc, t1, t2, eris, kconserv)
90    Fov = imdk.cc_Fov(cc, t1, t2, eris, kconserv)
91    Woooo = imdk.cc_Woooo(cc, t1, t2, eris, kconserv)
92    Wvvvv = imdk.cc_Wvvvv(cc, t1, t2, eris, kconserv)
93    Wovvo = imdk.cc_Wovvo(cc, t1, t2, eris, kconserv)
94
95    # Move energy terms to the other side
96    for k in range(nkpts):
97        Foo[k][numpy.diag_indices(nocc)] -= mo_e_o[k]
98        Fvv[k][numpy.diag_indices(nvir)] -= mo_e_v[k]
99
100    eris_ovvo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nocc, nvir, nvir, nocc), dtype=t2.dtype)
101    eris_oovo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nocc, nocc, nvir, nocc), dtype=t2.dtype)
102    eris_vvvo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nvir, nvir, nvir, nocc), dtype=t2.dtype)
103    for km, kb, ke in kpts_helper.loop_kkk(nkpts):
104        kj = kconserv[km, ke, kb]
105        # <mb||je> -> -<mb||ej>
106        eris_ovvo[km, kb, ke] = -eris.ovov[km, kb, kj].transpose(0, 1, 3, 2)
107        # <mn||je> -> -<mn||ej>
108        # let kb = kn as a dummy variable
109        eris_oovo[km, kb, ke] = -eris.ooov[km, kb, kj].transpose(0, 1, 3, 2)
110        # <ma||be> -> - <be||am>*
111        # let kj = ka as a dummy variable
112        kj = kconserv[km, ke, kb]
113        eris_vvvo[ke, kj, kb] = -eris.ovvv[km, kb, ke].transpose(2, 3, 1, 0).conj()
114
115    # T1 equation
116    t1new = numpy.zeros(shape=t1.shape, dtype=t1.dtype)
117    for ka in range(nkpts):
118        ki = ka
119        t1new[ka] += numpy.array(fov[ka, :, :]).conj()
120        t1new[ka] += einsum('ie,ae->ia', t1[ka], Fvv[ka])
121        t1new[ka] += -einsum('ma,mi->ia', t1[ka], Foo[ka])
122        for km in range(nkpts):
123            t1new[ka] += einsum('imae,me->ia', t2[ka, km, ka], Fov[km])
124            t1new[ka] += -einsum('nf,naif->ia', t1[km], eris.ovov[km, ka, ki])
125            for kn in range(nkpts):
126                ke = kconserv[km, ki, kn]
127                t1new[ka] += -0.5 * einsum('imef,maef->ia', t2[ki, km, ke], eris.ovvv[km, ka, ke])
128                t1new[ka] += -0.5 * einsum('mnae,nmei->ia', t2[km, kn, ka], eris_oovo[kn, km, ke])
129
130    # T2 equation
131    t2new = numpy.array(eris.oovv).conj()
132    for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
133        # Chemist's notation for momentum conserving t2(ki,kj,ka,kb)
134        kb = kconserv[ki, ka, kj]
135
136        Ftmp = Fvv[kb] - 0.5 * einsum('mb,me->be', t1[kb], Fov[kb])
137        tmp = einsum('ijae,be->ijab', t2[ki, kj, ka], Ftmp)
138        t2new[ki, kj, ka] += tmp
139
140        #t2new[ki,kj,kb] -= tmp.transpose(0,1,3,2)
141        Ftmp = Fvv[ka] - 0.5 * einsum('ma,me->ae', t1[ka], Fov[ka])
142        tmp = einsum('ijbe,ae->ijab', t2[ki, kj, kb], Ftmp)
143        t2new[ki, kj, ka] -= tmp
144
145        Ftmp = Foo[kj] + 0.5 * einsum('je,me->mj', t1[kj], Fov[kj])
146        tmp = einsum('imab,mj->ijab', t2[ki, kj, ka], Ftmp)
147        t2new[ki, kj, ka] -= tmp
148
149        #t2new[kj,ki,ka] += tmp.transpose(1,0,2,3)
150        Ftmp = Foo[ki] + 0.5 * einsum('ie,me->mi', t1[ki], Fov[ki])
151        tmp = einsum('jmab,mi->ijab', t2[kj, ki, ka], Ftmp)
152        t2new[ki, kj, ka] += tmp
153
154        for km in range(nkpts):
155            # Wminj
156            #   - km - kn + ka + kb = 0
157            # =>  kn = ka - km + kb
158            kn = kconserv[ka, km, kb]
159            t2new[ki, kj, ka] += 0.5 * einsum('mnab,mnij->ijab', tau[km, kn, ka], Woooo[km, kn, ki])
160            ke = km
161            t2new[ki, kj, ka] += 0.5 * einsum('ijef,abef->ijab', tau[ki, kj, ke], Wvvvv[ka, kb, ke])
162
163            # Wmbej
164            #     - km - kb + ke + kj = 0
165            #  => ke = km - kj + kb
166            ke = kconserv[km, kj, kb]
167            tmp = einsum('imae,mbej->ijab', t2[ki, km, ka], Wovvo[km, kb, ke])
168            #     - km - kb + ke + kj = 0
169            # =>  ke = km - kj + kb
170            #
171            # t[i,e] => ki = ke
172            # t[m,a] => km = ka
173            if km == ka and ke == ki:
174                tmp -= einsum('ie,ma,mbej->ijab', t1[ki], t1[km], eris_ovvo[km, kb, ke])
175            t2new[ki, kj, ka] += tmp
176            t2new[ki, kj, kb] -= tmp.transpose(0, 1, 3, 2)
177            t2new[kj, ki, ka] -= tmp.transpose(1, 0, 2, 3)
178            t2new[kj, ki, kb] += tmp.transpose(1, 0, 3, 2)
179
180        ke = ki
181        tmp = einsum('ie,abej->ijab', t1[ki], eris_vvvo[ka, kb, ke])
182        t2new[ki, kj, ka] += tmp
183        # P(ij) term
184        ke = kj
185        tmp = einsum('je,abei->ijab', t1[kj], eris_vvvo[ka, kb, ke])
186        t2new[ki, kj, ka] -= tmp
187
188        km = ka
189        tmp = einsum('ma,mbij->ijab', t1[ka], eris.ovoo[km, kb, ki])
190        t2new[ki, kj, ka] -= tmp
191        # P(ab) term
192        km = kb
193        tmp = einsum('mb,maij->ijab', t1[kb], eris.ovoo[km, ka, ki])
194        t2new[ki, kj, ka] += tmp
195
196    for ki in range(nkpts):
197        ka = ki
198        # Remove zero/padded elements from denominator
199        eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
200        n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
201        eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]
202        t1new[ki] /= eia
203
204    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
205    for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
206        kb = kconserv[ki, ka, kj]
207        # For LARGE_DENOM, see t1new update above
208        eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
209        n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
210        eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]
211
212        ejb = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
213        n0_ovp_jb = numpy.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
214        ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
215        eijab = eia[:, None, :, None] + ejb[:, None, :]
216
217        t2new[ki, kj, ka] /= eijab
218
219    time0 = log.timer_debug1('update t1 t2', *time0)
220
221    return t1new, t2new
222
223def spatial2spin(tx, orbspin, kconserv):
224    '''Convert T1/T2 of spatial orbital representation to T1/T2 of
225    spin-orbital representation
226    '''
227    if isinstance(tx, numpy.ndarray) and tx.ndim == 3:
228        # KRCCSD t1 amplitudes
229        return spatial2spin((tx,tx), orbspin, kconserv)
230    elif isinstance(tx, numpy.ndarray) and tx.ndim == 7:
231        # KRCCSD t2 amplitudes
232        t2aa = numpy.zeros_like(tx)
233        nkpts = t2aa.shape[2]
234        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
235            kb = kconserv[ki,ka,kj]
236            t2aa[ki,kj,ka] = tx[ki,kj,ka] - tx[ki,kj,kb].transpose(0,1,3,2)
237        return spatial2spin((t2aa,tx,t2aa), orbspin, kconserv)
238    elif len(tx) == 2:  # KUCCSD t1
239        t1a, t1b = tx
240        nocc_a, nvir_a = t1a.shape[1:]
241        nocc_b, nvir_b = t1b.shape[1:]
242    else:  # KUCCSD t2
243        t2aa, t2ab, t2bb = tx
244        nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape[3:]
245
246    nkpts = len(orbspin)
247    nocc = nocc_a + nocc_b
248    nvir = nvir_a + nvir_b
249    idxoa = [numpy.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)]
250    idxob = [numpy.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)]
251    idxva = [numpy.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)]
252    idxvb = [numpy.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)]
253
254    if len(tx) == 2:  # t1
255        t1 = numpy.zeros((nkpts,nocc,nvir), dtype=t1a.dtype)
256        for k in range(nkpts):
257            lib.takebak_2d(t1[k], t1a[k], idxoa[k], idxva[k])
258            lib.takebak_2d(t1[k], t1b[k], idxob[k], idxvb[k])
259        t1 = lib.tag_array(t1, orbspin=orbspin)
260        return t1
261
262    else:
263        t2 = numpy.zeros((nkpts,nkpts,nkpts,nocc**2,nvir**2), dtype=t2aa.dtype)
264        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
265            kb = kconserv[ki,ka,kj]
266            idxoaa = idxoa[ki][:,None] * nocc + idxoa[kj]
267            idxoab = idxoa[ki][:,None] * nocc + idxob[kj]
268            idxoba = idxob[kj][:,None] * nocc + idxoa[ki]
269            idxobb = idxob[ki][:,None] * nocc + idxob[kj]
270            idxvaa = idxva[ka][:,None] * nvir + idxva[kb]
271            idxvab = idxva[ka][:,None] * nvir + idxvb[kb]
272            idxvba = idxvb[kb][:,None] * nvir + idxva[ka]
273            idxvbb = idxvb[ka][:,None] * nvir + idxvb[kb]
274            tmp2aa = t2aa[ki,kj,ka].reshape(nocc_a*nocc_a,nvir_a*nvir_a)
275            tmp2bb = t2bb[ki,kj,ka].reshape(nocc_b*nocc_b,nvir_b*nvir_b)
276            tmp2ab = t2ab[ki,kj,ka].reshape(nocc_a*nocc_b,nvir_a*nvir_b)
277            lib.takebak_2d(t2[ki,kj,ka], tmp2aa, idxoaa.ravel()  , idxvaa.ravel()  )
278            lib.takebak_2d(t2[ki,kj,ka], tmp2bb, idxobb.ravel()  , idxvbb.ravel()  )
279            lib.takebak_2d(t2[ki,kj,ka], tmp2ab, idxoab.ravel()  , idxvab.ravel()  )
280            lib.takebak_2d(t2[kj,ki,kb], tmp2ab, idxoba.T.ravel(), idxvba.T.ravel())
281
282            abba = -tmp2ab
283            lib.takebak_2d(t2[ki,kj,kb], abba, idxoab.ravel()  , idxvba.T.ravel())
284            lib.takebak_2d(t2[kj,ki,ka], abba, idxoba.T.ravel(), idxvab.ravel()  )
285        t2 = t2.reshape(nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir)
286        t2 = lib.tag_array(t2, orbspin=orbspin)
287        return t2
288
289def spin2spatial(tx, orbspin, kconserv):
290    if tx.ndim == 3:  # t1
291        nocc, nvir = tx.shape[1:]
292    else:
293        nocc, nvir = tx.shape[4:6]
294    nkpts = len(tx)
295
296    idxoa = [numpy.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)]
297    idxob = [numpy.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)]
298    idxva = [numpy.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)]
299    idxvb = [numpy.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)]
300    nocc_a = len(idxoa[0])
301    nocc_b = len(idxob[0])
302    nvir_a = len(idxva[0])
303    nvir_b = len(idxvb[0])
304
305    if tx.ndim == 3:  # t1
306        t1a = numpy.zeros((nkpts,nocc_a,nvir_a), dtype=tx.dtype)
307        t1b = numpy.zeros((nkpts,nocc_b,nvir_b), dtype=tx.dtype)
308        for k in range(nkpts):
309            lib.take_2d(tx[k], idxoa[k], idxva[k], out=t1a[k])
310            lib.take_2d(tx[k], idxob[k], idxvb[k], out=t1b[k])
311        return t1a, t1b
312
313    else:
314        t2aa = numpy.zeros((nkpts,nkpts,nkpts,nocc_a,nocc_a,nvir_a,nvir_a), dtype=tx.dtype)
315        t2ab = numpy.zeros((nkpts,nkpts,nkpts,nocc_a,nocc_b,nvir_a,nvir_b), dtype=tx.dtype)
316        t2bb = numpy.zeros((nkpts,nkpts,nkpts,nocc_b,nocc_b,nvir_b,nvir_b), dtype=tx.dtype)
317        t2 = tx.reshape(nkpts,nkpts,nkpts,nocc**2,nvir**2)
318        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
319            kb = kconserv[ki,ka,kj]
320            idxoaa = idxoa[ki][:,None] * nocc + idxoa[kj]
321            idxoab = idxoa[ki][:,None] * nocc + idxob[kj]
322            idxobb = idxob[ki][:,None] * nocc + idxob[kj]
323            idxvaa = idxva[ka][:,None] * nvir + idxva[kb]
324            idxvab = idxva[ka][:,None] * nvir + idxvb[kb]
325            idxvbb = idxvb[ka][:,None] * nvir + idxvb[kb]
326            lib.take_2d(t2[ki,kj,ka], idxoaa.ravel(), idxvaa.ravel(), out=t2aa[ki,kj,ka])
327            lib.take_2d(t2[ki,kj,ka], idxobb.ravel(), idxvbb.ravel(), out=t2bb[ki,kj,ka])
328            lib.take_2d(t2[ki,kj,ka], idxoab.ravel(), idxvab.ravel(), out=t2ab[ki,kj,ka])
329        return t2aa, t2ab, t2bb
330
331
332class GCCSD(gccsd.GCCSD):
333    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
334        assert (isinstance(mf, scf.khf.KSCF))
335        if not isinstance(mf, scf.kghf.KGHF):
336            mf = scf.addons.convert_to_ghf(mf)
337        self.kpts = mf.kpts
338        self.khelper = kpts_helper.KptsHelper(mf.cell, mf.kpts)
339        gccsd.GCCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)
340
341    @property
342    def nkpts(self):
343        return len(self.kpts)
344
345    get_nocc = get_nocc
346    get_nmo = get_nmo
347    get_frozen_mask = get_frozen_mask
348
349    def dump_flags(self, verbose=None):
350        logger.info(self, '\n')
351        logger.info(self, '******** PBC CC flags ********')
352        gccsd.GCCSD.dump_flags(self, verbose)
353        return self
354
355    def init_amps(self, eris):
356        time0 = logger.process_clock(), logger.perf_counter()
357        nocc = self.nocc
358        nvir = self.nmo - nocc
359        nkpts = self.nkpts
360        mo_e_o = [eris.mo_energy[k][:nocc] for k in range(nkpts)]
361        mo_e_v = [eris.mo_energy[k][nocc:] for k in range(nkpts)]
362        t1 = numpy.zeros((nkpts, nocc, nvir), dtype=numpy.complex128)
363        t2 = numpy.zeros((nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir), dtype=numpy.complex128)
364        self.emp2 = 0
365        eris_oovv = eris.oovv.copy()
366
367        # Get location of padded elements in occupied and virtual space
368        nonzero_opadding, nonzero_vpadding = padding_k_idx(self, kind="split")
369
370        kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts)
371        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
372            kb = kconserv[ki, ka, kj]
373            # For LARGE_DENOM, see t1new update above
374            eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
375            n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
376            eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]
377
378            ejb = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
379            n0_ovp_jb = numpy.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
380            ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
381            eijab = eia[:, None, :, None] + ejb[:, None, :]
382
383            t2[ki, kj, ka] = eris_oovv[ki, kj, ka] / eijab
384
385        t2 = numpy.conj(t2)
386        self.emp2 = 0.25 * numpy.einsum('pqrijab,pqrijab', t2, eris_oovv).real
387        self.emp2 /= nkpts
388
389        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2.real)
390        logger.timer(self, 'init mp2', *time0)
391        return self.emp2, t1, t2
392
393    def ccsd(self, t1=None, t2=None, eris=None, **kwargs):
394        if eris is None: eris = self.ao2mo(self.mo_coeff)
395        e_corr, self.t1, self.t2 = ccsd.CCSD.ccsd(self, t1, t2, eris)
396        if getattr(eris, 'orbspin', None) is not None:
397            self.t1 = lib.tag_array(self.t1, orbspin=eris.orbspin)
398            self.t2 = lib.tag_array(self.t2, orbspin=eris.orbspin)
399        return e_corr, self.t1, self.t2
400
401    update_amps = update_amps
402
403    energy = energy
404
405    def ao2mo(self, mo_coeff=None):
406        nkpts = self.nkpts
407        nmo = self.nmo
408        mem_incore = nkpts**3 * nmo**4 * 8 / 1e6
409        mem_now = lib.current_memory()[0]
410
411        if (mem_incore + mem_now < self.max_memory) or self.mol.incore_anyway:
412            return _make_eris_incore(self, mo_coeff)
413        else:
414            raise NotImplementedError
415
416    def ccsd_t(self, t1=None, t2=None, eris=None):
417        from pyscf.pbc.cc import kccsd_t
418        if t1 is None: t1 = self.t1
419        if t2 is None: t2 = self.t2
420        if eris is None: eris = self.ao2mo(self.mo_coeff)
421        return kccsd_t.kernel(self, eris, t1, t2, self.verbose)
422
423    def amplitudes_to_vector(self, t1, t2):
424        return numpy.hstack((t1.ravel(), t2.ravel()))
425
426    def vector_to_amplitudes(self, vec, nmo=None, nocc=None):
427        if nocc is None: nocc = self.nocc
428        if nmo is None: nmo = self.nmo
429        nvir = nmo - nocc
430        nkpts = self.nkpts
431        nov = nkpts * nocc * nvir
432        t1 = vec[:nov].reshape(nkpts, nocc, nvir)
433        t2 = vec[nov:].reshape(nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir)
434        return t1, t2
435
436    def spatial2spin(self, tx, orbspin=None, kconserv=None):
437        if orbspin is None:
438            if getattr(self.mo_coeff[0], 'orbspin', None) is not None:
439                orbspin = [self.mo_coeff[k].orbspin[idx]
440                           for k, idx in enumerate(self.get_frozen_mask())]
441            else:
442                orbspin = numpy.zeros((self.nkpts,self.nmo), dtype=int)
443                orbspin[:,1::2] = 1
444        if kconserv is None:
445            kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts)
446        return spatial2spin(tx, orbspin, kconserv)
447
448    def spin2spatial(self, tx, orbspin=None, kconserv=None):
449        if orbspin is None:
450            if getattr(self.mo_coeff[0], 'orbspin', None) is not None:
451                orbspin = [self.mo_coeff[k].orbspin[idx]
452                           for k, idx in enumerate(self.get_frozen_mask())]
453            else:
454                orbspin = numpy.zeros((self.nkpts,self.nmo), dtype=int)
455                orbspin[:,1::2] = 1
456        if kconserv is None:
457            kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts)
458        return spin2spatial(tx, orbspin, kconserv)
459
460    def from_uccsd(self, t1, t2, orbspin=None):
461        return self.spatial2spin(t1, orbspin), self.spatial2spin(t2, orbspin)
462
463    def to_uccsd(self, t1, t2, orbspin=None):
464        return spin2spatial(t1, orbspin), spin2spatial(t2, orbspin)
465
466CCSD = KCCSD = KGCCSD = GCCSD
467
468
469def _make_eris_incore(cc, mo_coeff=None):
470    from pyscf.pbc import tools
471    from pyscf.pbc.cc.ccsd import _adjust_occ
472
473    log = logger.Logger(cc.stdout, cc.verbose)
474    cput0 = (logger.process_clock(), logger.perf_counter())
475    eris = gccsd._PhysicistsERIs()
476    cell = cc._scf.cell
477    kpts = cc.kpts
478    nkpts = cc.nkpts
479    nocc = cc.nocc
480    nmo = cc.nmo
481    eris.nocc = nocc
482
483    #if any(nocc != numpy.count_nonzero(cc._scf.mo_occ[k] > 0) for k in range(nkpts)):
484    #    raise NotImplementedError('Different occupancies found for different k-points')
485
486    if mo_coeff is None:
487        mo_coeff = cc.mo_coeff
488
489    nao = mo_coeff[0].shape[0]
490    dtype = mo_coeff[0].dtype
491
492    moidx = get_frozen_mask(cc)
493    nocc_per_kpt = numpy.asarray(get_nocc(cc, per_kpoint=True))
494    nmo_per_kpt  = numpy.asarray(get_nmo(cc, per_kpoint=True))
495
496    padded_moidx = []
497    for k in range(nkpts):
498        kpt_nocc = nocc_per_kpt[k]
499        kpt_nvir = nmo_per_kpt[k] - kpt_nocc
500        kpt_padded_moidx = numpy.concatenate((numpy.ones(kpt_nocc, dtype=numpy.bool),
501                                              numpy.zeros(nmo - kpt_nocc - kpt_nvir, dtype=numpy.bool),
502                                              numpy.ones(kpt_nvir, dtype=numpy.bool)))
503        padded_moidx.append(kpt_padded_moidx)
504
505    eris.mo_coeff = []
506    eris.orbspin = []
507    # Generate the molecular orbital coefficients with the frozen orbitals masked.
508    # Each MO is tagged with orbspin, a list of 0's and 1's that give the overall
509    # spin of each MO.
510    #
511    # Here we will work with two index arrays; one is for our original (small) moidx
512    # array while the next is for our new (large) padded array.
513    for k in range(nkpts):
514        kpt_moidx = moidx[k]
515        kpt_padded_moidx = padded_moidx[k]
516
517        mo = numpy.zeros((nao, nmo), dtype=dtype)
518        mo[:, kpt_padded_moidx] = mo_coeff[k][:, kpt_moidx]
519        if getattr(mo_coeff[k], 'orbspin', None) is not None:
520            orbspin_dtype = mo_coeff[k].orbspin[kpt_moidx].dtype
521            orbspin = numpy.zeros(nmo, dtype=orbspin_dtype)
522            orbspin[kpt_padded_moidx] = mo_coeff[k].orbspin[kpt_moidx]
523            mo = lib.tag_array(mo, orbspin=orbspin)
524            eris.orbspin.append(orbspin)
525        # FIXME: What if the user freezes all up spin orbitals in
526        # an RHF calculation?  The number of electrons will still be
527        # even.
528        else:  # guess orbital spin - assumes an RHF calculation
529            assert (numpy.count_nonzero(kpt_moidx) % 2 == 0)
530            orbspin = numpy.zeros(mo.shape[1], dtype=int)
531            orbspin[1::2] = 1
532            mo = lib.tag_array(mo, orbspin=orbspin)
533            eris.orbspin.append(orbspin)
534        eris.mo_coeff.append(mo)
535
536    # Re-make our fock MO matrix elements from density and fock AO
537    dm = cc._scf.make_rdm1(cc.mo_coeff, cc.mo_occ)
538    with lib.temporary_env(cc._scf, exxdiv=None):
539        # _scf.exxdiv affects eris.fock. HF exchange correction should be
540        # excluded from the Fock matrix.
541        vhf = cc._scf.get_veff(cell, dm)
542    fockao = cc._scf.get_hcore() + vhf
543    eris.fock = numpy.asarray([reduce(numpy.dot, (mo.T.conj(), fockao[k], mo))
544                               for k, mo in enumerate(eris.mo_coeff)])
545    eris.e_hf = cc._scf.energy_tot(dm=dm, vhf=vhf)
546
547    eris.mo_energy = [eris.fock[k].diagonal().real for k in range(nkpts)]
548    # Add HFX correction in the eris.mo_energy to improve convergence in
549    # CCSD iteration. It is useful for the 2D systems since their occupied and
550    # the virtual orbital energies may overlap which may lead to numerical
551    # issue in the CCSD iterations.
552    # FIXME: Whether to add this correction for other exxdiv treatments?
553    # Without the correction, MP2 energy may be largely off the correct value.
554    madelung = tools.madelung(cell, kpts)
555    eris.mo_energy = [_adjust_occ(mo_e, nocc, -madelung)
556                      for k, mo_e in enumerate(eris.mo_energy)]
557
558    # Get location of padded elements in occupied and virtual space.
559    nocc_per_kpt = get_nocc(cc, per_kpoint=True)
560    nonzero_padding = padding_k_idx(cc, kind="joint")
561
562    # Check direct and indirect gaps for possible issues with CCSD convergence.
563    mo_e = [eris.mo_energy[kp][nonzero_padding[kp]] for kp in range(nkpts)]
564    mo_e = numpy.sort([y for x in mo_e for y in x])  # Sort de-nested array
565    gap = mo_e[numpy.sum(nocc_per_kpt)] - mo_e[numpy.sum(nocc_per_kpt)-1]
566    if gap < 1e-5:
567        logger.warn(cc, 'HOMO-LUMO gap %s too small for KCCSD. '
568                        'May cause issues in convergence.', gap)
569
570    kconserv = kpts_helper.get_kconserv(cell, kpts)
571    if getattr(mo_coeff[0], 'orbspin', None) is None:
572        # The bottom nao//2 coefficients are down (up) spin while the top are up (down).
573        mo_a_coeff = [mo[:nao // 2] for mo in eris.mo_coeff]
574        mo_b_coeff = [mo[nao // 2:] for mo in eris.mo_coeff]
575
576        eri = numpy.empty((nkpts, nkpts, nkpts, nmo, nmo, nmo, nmo), dtype=numpy.complex128)
577        fao2mo = cc._scf.with_df.ao2mo
578        for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
579            ks = kconserv[kp, kq, kr]
580            eri_kpt = fao2mo(
581                (mo_a_coeff[kp], mo_a_coeff[kq], mo_a_coeff[kr], mo_a_coeff[ks]),
582                (kpts[kp], kpts[kq], kpts[kr], kpts[ks]),
583                compact=False)
584            eri_kpt += fao2mo(
585                (mo_b_coeff[kp], mo_b_coeff[kq], mo_b_coeff[kr], mo_b_coeff[ks]),
586                (kpts[kp], kpts[kq], kpts[kr], kpts[ks]),
587                compact=False)
588            eri_kpt += fao2mo(
589                (mo_a_coeff[kp], mo_a_coeff[kq], mo_b_coeff[kr], mo_b_coeff[ks]),
590                (kpts[kp], kpts[kq], kpts[kr], kpts[ks]),
591                compact=False)
592            eri_kpt += fao2mo(
593                (mo_b_coeff[kp], mo_b_coeff[kq], mo_a_coeff[kr], mo_a_coeff[ks]),
594                (kpts[kp], kpts[kq], kpts[kr], kpts[ks]),
595                compact=False)
596
597            eri_kpt = eri_kpt.reshape(nmo, nmo, nmo, nmo)
598            eri[kp, kq, kr] = eri_kpt
599    else:
600        mo_a_coeff = [mo[:nao // 2] + mo[nao // 2:] for mo in eris.mo_coeff]
601
602        eri = numpy.empty((nkpts, nkpts, nkpts, nmo, nmo, nmo, nmo), dtype=numpy.complex128)
603        fao2mo = cc._scf.with_df.ao2mo
604        for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
605            ks = kconserv[kp, kq, kr]
606            eri_kpt = fao2mo(
607                (mo_a_coeff[kp], mo_a_coeff[kq], mo_a_coeff[kr], mo_a_coeff[ks]),
608                (kpts[kp], kpts[kq], kpts[kr], kpts[ks]),
609                compact=False)
610
611            eri_kpt[(eris.orbspin[kp][:, None] != eris.orbspin[kq]).ravel()] = 0
612            eri_kpt[:, (eris.orbspin[kr][:, None] != eris.orbspin[ks]).ravel()] = 0
613            eri_kpt = eri_kpt.reshape(nmo, nmo, nmo, nmo)
614            eri[kp, kq, kr] = eri_kpt
615
616    # Check some antisymmetrized properties of the integrals
617    if DEBUG:
618        check_antisymm_3412(cc, cc.kpts, eri)
619
620    # Antisymmetrizing (pq|rs)-(ps|rq), where the latter integral is equal to
621    # (rq|ps); done since we aren't tracking the kpoint of orbital 's'
622    eri = eri - eri.transpose(2, 1, 0, 5, 4, 3, 6)
623    # Chemist -> physics notation
624    eri = eri.transpose(0, 2, 1, 3, 5, 4, 6)
625
626    # Set the various integrals
627    eris.dtype = eri.dtype
628    eris.oooo = eri[:, :, :, :nocc, :nocc, :nocc, :nocc].copy() / nkpts
629    eris.ooov = eri[:, :, :, :nocc, :nocc, :nocc, nocc:].copy() / nkpts
630    eris.ovoo = eri[:, :, :, :nocc, nocc:, :nocc, :nocc].copy() / nkpts
631    eris.oovv = eri[:, :, :, :nocc, :nocc, nocc:, nocc:].copy() / nkpts
632    eris.ovov = eri[:, :, :, :nocc, nocc:, :nocc, nocc:].copy() / nkpts
633    eris.ovvv = eri[:, :, :, :nocc, nocc:, nocc:, nocc:].copy() / nkpts
634    eris.vvvv = eri[:, :, :, nocc:, nocc:, nocc:, nocc:].copy() / nkpts
635
636    log.timer('CCSD integral transformation', *cput0)
637    return eris
638
639
640def check_antisymm_3412(cc, kpts, integrals):
641    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
642    nkpts = len(kpts)
643    diff = 0.0
644    for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
645        ks = kconserv[kp, kr, kq]
646        for p in range(integrals.shape[3]):
647            for q in range(integrals.shape[4]):
648                for r in range(integrals.shape[5]):
649                    for s in range(integrals.shape[6]):
650                        pqrs = integrals[kp, kq, kr, p, q, r, s]
651                        rspq = integrals[kq, kp, kr, q, p, r, s]
652                        cdiff = numpy.linalg.norm(pqrs - rspq).real
653                        if diff > 1e-5:
654                            print("AS diff = %.15g" % cdiff, pqrs, rspq, kp, kq, kr, ks, p, q, r, s)
655                        diff = max(diff, cdiff)
656    print("antisymmetrization : max diff = %.15g" % diff)
657    if diff > 1e-5:
658        print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.")
659    return diff
660
661
662def check_antisymm_12(cc, kpts, integrals):
663    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
664    nkpts = len(kpts)
665    diff = 0.0
666    for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
667        ks = kconserv[kp, kr, kq]
668        for p in range(integrals.shape[3]):
669            for q in range(integrals.shape[4]):
670                for r in range(integrals.shape[5]):
671                    for s in range(integrals.shape[6]):
672                        pqrs = integrals[kp, kq, kr, p, q, r, s]
673                        qprs = integrals[kq, kp, kr, q, p, r, s]
674                        cdiff = numpy.linalg.norm(pqrs + qprs).real
675                        if diff > 1e-5:
676                            print("AS diff = %.15g" % cdiff, pqrs, qprs, kp, kq, kr, ks, p, q, r, s)
677                        diff = max(diff, cdiff)
678    print("antisymmetrization : max diff = %.15g" % diff)
679    if diff > 1e-5:
680        print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.")
681
682
683def check_antisymm_34(cc, kpts, integrals):
684    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
685    nkpts = len(kpts)
686    diff = 0.0
687    for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
688        ks = kconserv[kp, kr, kq]
689        for p in range(integrals.shape[3]):
690            for q in range(integrals.shape[4]):
691                for r in range(integrals.shape[5]):
692                    for s in range(integrals.shape[6]):
693                        pqrs = integrals[kp, kq, kr, p, q, r, s]
694                        pqsr = integrals[kp, kq, ks, p, q, s, r]
695                        cdiff = numpy.linalg.norm(pqrs + pqsr).real
696                        if diff > 1e-5:
697                            print("AS diff = %.15g" % cdiff, pqrs, pqsr, kp, kq, kr, ks, p, q, r, s)
698                        diff = max(diff, cdiff)
699    print("antisymmetrization : max diff = %.15g" % diff)
700    if diff > 1e-5:
701        print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.")
702
703imd = imdk
704class _IMDS:
705    # Identical to molecular rccsd_slow
706    def __init__(self, cc):
707        self.verbose = cc.verbose
708        self.stdout = cc.stdout
709        self.t1 = cc.t1
710        self.t2 = cc.t2
711        self.eris = cc.eris
712        self.kconserv = cc.khelper.kconserv
713        self.made_ip_imds = False
714        self.made_ea_imds = False
715        self._made_shared_2e = False
716        self._fimd = None
717
718    def _make_shared_1e(self):
719        cput0 = (logger.process_clock(), logger.perf_counter())
720        log = logger.Logger(self.stdout, self.verbose)
721
722        t1,t2,eris = self.t1, self.t2, self.eris
723        kconserv = self.kconserv
724        self.Loo = imd.Loo(t1,t2,eris,kconserv)
725        self.Lvv = imd.Lvv(t1,t2,eris,kconserv)
726        self.Fov = imd.cc_Fov(t1,t2,eris,kconserv)
727
728        log.timer('EOM-CCSD shared one-electron intermediates', *cput0)
729
730    def _make_shared_2e(self):
731        cput0 = (logger.process_clock(), logger.perf_counter())
732        log = logger.Logger(self.stdout, self.verbose)
733
734        t1,t2,eris = self.t1, self.t2, self.eris
735        kconserv = self.kconserv
736
737        # TODO: check whether to hold Wovov Wovvo in memory
738        if self._fimd is None:
739            self._fimd = lib.H5TmpFile()
740        nkpts, nocc, nvir = t1.shape
741        self._fimd.create_dataset('ovov', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), t1.dtype.char)
742        self._fimd.create_dataset('ovvo', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), t1.dtype.char)
743
744        # 2 virtuals
745        self.Wovov = imd.Wovov(t1,t2,eris,kconserv, self._fimd['ovov'])
746        self.Wovvo = imd.Wovvo(t1,t2,eris,kconserv, self._fimd['ovvo'])
747        self.Woovv = eris.oovv
748
749        log.timer('EOM-CCSD shared two-electron intermediates', *cput0)
750
751    def make_ip(self, ip_partition=None):
752        self._make_shared_1e()
753        if self._made_shared_2e is False and ip_partition != 'mp':
754            self._make_shared_2e()
755            self._made_shared_2e = True
756
757        cput0 = (logger.process_clock(), logger.perf_counter())
758        log = logger.Logger(self.stdout, self.verbose)
759
760        t1,t2,eris = self.t1, self.t2, self.eris
761        kconserv = self.kconserv
762
763        nkpts, nocc, nvir = t1.shape
764        self._fimd.create_dataset('oooo', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nocc), t1.dtype.char)
765        self._fimd.create_dataset('ooov', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), t1.dtype.char)
766        self._fimd.create_dataset('ovoo', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nocc), t1.dtype.char)
767
768        # 0 or 1 virtuals
769        if ip_partition != 'mp':
770            self.Woooo = imd.Woooo(t1,t2,eris,kconserv, self._fimd['oooo'])
771        self.Wooov = imd.Wooov(t1,t2,eris,kconserv, self._fimd['ooov'])
772        self.Wovoo = imd.Wovoo(t1,t2,eris,kconserv, self._fimd['ovoo'])
773        self.made_ip_imds = True
774        log.timer('EOM-CCSD IP intermediates', *cput0)
775
776    def make_ea(self, ea_partition=None):
777        self._make_shared_1e()
778        if self._made_shared_2e is False and ea_partition != 'mp':
779            self._make_shared_2e()
780            self._made_shared_2e = True
781
782        cput0 = (logger.process_clock(), logger.perf_counter())
783        log = logger.Logger(self.stdout, self.verbose)
784
785        t1,t2,eris = self.t1, self.t2, self.eris
786        kconserv = self.kconserv
787
788        nkpts, nocc, nvir = t1.shape
789        self._fimd.create_dataset('vovv', (nkpts,nkpts,nkpts,nvir,nocc,nvir,nvir), t1.dtype.char)
790        self._fimd.create_dataset('vvvo', (nkpts,nkpts,nkpts,nvir,nvir,nvir,nocc), t1.dtype.char)
791        self._fimd.create_dataset('vvvv', (nkpts,nkpts,nkpts,nvir,nvir,nvir,nvir), t1.dtype.char)
792
793        # 3 or 4 virtuals
794        self.Wvovv = imd.Wvovv(t1,t2,eris,kconserv, self._fimd['vovv'])
795        if ea_partition == 'mp' and numpy.all(t1 == 0):
796            self.Wvvvo = imd.Wvvvo(t1,t2,eris,kconserv, self._fimd['vvvo'])
797        else:
798            self.Wvvvv = imd.Wvvvv(t1,t2,eris,kconserv, self._fimd['vvvv'])
799            self.Wvvvo = imd.Wvvvo(t1,t2,eris,kconserv,self.Wvvvv, self._fimd['vvvo'])
800        self.made_ea_imds = True
801        log.timer('EOM-CCSD EA intermediates', *cput0)
802
803
804scf.kghf.KGHF.CCSD = lib.class_as_method(KGCCSD)
805
806
807if __name__ == '__main__':
808    from pyscf.pbc import gto
809
810    cell = gto.Cell()
811    cell.atom='''
812    C 0.000000000000   0.000000000000   0.000000000000
813    C 1.685068664391   1.685068664391   1.685068664391
814    '''
815    cell.basis = 'gth-szv'
816    cell.pseudo = 'gth-pade'
817    cell.a = '''
818    0.000000000, 3.370137329, 3.370137329
819    3.370137329, 0.000000000, 3.370137329
820    3.370137329, 3.370137329, 0.000000000'''
821    cell.unit = 'B'
822    cell.verbose = 5
823    cell.build()
824
825    # Running HF and CCSD with 1x1x2 Monkhorst-Pack k-point mesh
826    kmf = scf.KRHF(cell, kpts=cell.make_kpts([1,1,2]), exxdiv=None)
827    ehf = kmf.kernel()
828    kmf = scf.addons.convert_to_ghf(kmf)
829
830    mycc = KGCCSD(kmf)
831    ecc, t1, t2 = mycc.kernel()
832    print(ecc - -0.155298393321855)
833