1#!/usr/bin/env python
2# Copyright 2014-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# Author: Timothy Berkelbach <tim.berkelbach@gmail.com>
17#
18
19'''
20Restricted CCSD implementation which supports both real and complex integrals.
21The 4-index integrals are saved on disk entirely (without using any symmetry).
22This code is slower than the pyscf.cc.ccsd implementation.
23
24Note MO integrals are treated in chemist's notation
25
26Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004)
27'''
28
29
30import numpy as np
31
32from pyscf import lib
33from pyscf import ao2mo
34from pyscf.lib import logger
35from pyscf.cc import ccsd
36from pyscf.cc import rintermediates as imd
37from pyscf.mp import mp2
38from pyscf import __config__
39
40BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
41MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)
42
43def update_amps(cc, t1, t2, eris):
44    # Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004) Eqs.(35)-(36)
45    assert(isinstance(eris, ccsd._ChemistsERIs))
46    nocc, nvir = t1.shape
47    fock = eris.fock
48    mo_e_o = eris.mo_energy[:nocc]
49    mo_e_v = eris.mo_energy[nocc:] + cc.level_shift
50
51    fov = fock[:nocc,nocc:].copy()
52    foo = fock[:nocc,:nocc].copy()
53    fvv = fock[nocc:,nocc:].copy()
54
55    Foo = imd.cc_Foo(t1,t2,eris)
56    Fvv = imd.cc_Fvv(t1,t2,eris)
57    Fov = imd.cc_Fov(t1,t2,eris)
58
59    # Move energy terms to the other side
60    Foo[np.diag_indices(nocc)] -= mo_e_o
61    Fvv[np.diag_indices(nvir)] -= mo_e_v
62
63    # T1 equation
64    t1new  =-2*np.einsum('kc,ka,ic->ia', fov, t1, t1)
65    t1new +=   np.einsum('ac,ic->ia', Fvv, t1)
66    t1new +=  -np.einsum('ki,ka->ia', Foo, t1)
67    t1new += 2*np.einsum('kc,kica->ia', Fov, t2)
68    t1new +=  -np.einsum('kc,ikca->ia', Fov, t2)
69    t1new +=   np.einsum('kc,ic,ka->ia', Fov, t1, t1)
70    t1new += fov.conj()
71    t1new += 2*np.einsum('kcai,kc->ia', eris.ovvo, t1)
72    t1new +=  -np.einsum('kiac,kc->ia', eris.oovv, t1)
73    eris_ovvv = np.asarray(eris.get_ovvv())
74    t1new += 2*lib.einsum('kdac,ikcd->ia', eris_ovvv, t2)
75    t1new +=  -lib.einsum('kcad,ikcd->ia', eris_ovvv, t2)
76    t1new += 2*lib.einsum('kdac,kd,ic->ia', eris_ovvv, t1, t1)
77    t1new +=  -lib.einsum('kcad,kd,ic->ia', eris_ovvv, t1, t1)
78    eris_ovoo = np.asarray(eris.ovoo, order='C')
79    t1new +=-2*lib.einsum('lcki,klac->ia', eris_ovoo, t2)
80    t1new +=   lib.einsum('kcli,klac->ia', eris_ovoo, t2)
81    t1new +=-2*lib.einsum('lcki,lc,ka->ia', eris_ovoo, t1, t1)
82    t1new +=   lib.einsum('kcli,lc,ka->ia', eris_ovoo, t1, t1)
83
84    # T2 equation
85    tmp2  = lib.einsum('kibc,ka->abic', eris.oovv, -t1)
86    tmp2 += np.asarray(eris_ovvv).conj().transpose(1,3,0,2)
87    tmp = lib.einsum('abic,jc->ijab', tmp2, t1)
88    t2new = tmp + tmp.transpose(1,0,3,2)
89    tmp2  = lib.einsum('kcai,jc->akij', eris.ovvo, t1)
90    tmp2 += eris_ovoo.transpose(1,3,0,2).conj()
91    tmp = lib.einsum('akij,kb->ijab', tmp2, t1)
92    t2new -= tmp + tmp.transpose(1,0,3,2)
93    t2new += np.asarray(eris.ovov).conj().transpose(0,2,1,3)
94    if cc.cc2:
95        Woooo2 = np.asarray(eris.oooo).transpose(0,2,1,3).copy()
96        Woooo2 += lib.einsum('lcki,jc->klij', eris_ovoo, t1)
97        Woooo2 += lib.einsum('kclj,ic->klij', eris_ovoo, t1)
98        Woooo2 += lib.einsum('kcld,ic,jd->klij', eris.ovov, t1, t1)
99        t2new += lib.einsum('klij,ka,lb->ijab', Woooo2, t1, t1)
100        Wvvvv = lib.einsum('kcbd,ka->abcd', eris_ovvv, -t1)
101        Wvvvv = Wvvvv + Wvvvv.transpose(1,0,3,2)
102        Wvvvv += np.asarray(eris.vvvv).transpose(0,2,1,3)
103        t2new += lib.einsum('abcd,ic,jd->ijab', Wvvvv, t1, t1)
104        Lvv2 = fvv - np.einsum('kc,ka->ac', fov, t1)
105        Lvv2 -= np.diag(np.diag(fvv))
106        tmp = lib.einsum('ac,ijcb->ijab', Lvv2, t2)
107        t2new += (tmp + tmp.transpose(1,0,3,2))
108        Loo2 = foo + np.einsum('kc,ic->ki', fov, t1)
109        Loo2 -= np.diag(np.diag(foo))
110        tmp = lib.einsum('ki,kjab->ijab', Loo2, t2)
111        t2new -= (tmp + tmp.transpose(1,0,3,2))
112    else:
113        Loo = imd.Loo(t1, t2, eris)
114        Lvv = imd.Lvv(t1, t2, eris)
115        Loo[np.diag_indices(nocc)] -= mo_e_o
116        Lvv[np.diag_indices(nvir)] -= mo_e_v
117
118        Woooo = imd.cc_Woooo(t1, t2, eris)
119        Wvoov = imd.cc_Wvoov(t1, t2, eris)
120        Wvovo = imd.cc_Wvovo(t1, t2, eris)
121        Wvvvv = imd.cc_Wvvvv(t1, t2, eris)
122
123        tau = t2 + np.einsum('ia,jb->ijab', t1, t1)
124        t2new += lib.einsum('klij,klab->ijab', Woooo, tau)
125        t2new += lib.einsum('abcd,ijcd->ijab', Wvvvv, tau)
126        tmp = lib.einsum('ac,ijcb->ijab', Lvv, t2)
127        t2new += (tmp + tmp.transpose(1,0,3,2))
128        tmp = lib.einsum('ki,kjab->ijab', Loo, t2)
129        t2new -= (tmp + tmp.transpose(1,0,3,2))
130        tmp  = 2*lib.einsum('akic,kjcb->ijab', Wvoov, t2)
131        tmp -=   lib.einsum('akci,kjcb->ijab', Wvovo, t2)
132        t2new += (tmp + tmp.transpose(1,0,3,2))
133        tmp = lib.einsum('akic,kjbc->ijab', Wvoov, t2)
134        t2new -= (tmp + tmp.transpose(1,0,3,2))
135        tmp = lib.einsum('bkci,kjac->ijab', Wvovo, t2)
136        t2new -= (tmp + tmp.transpose(1,0,3,2))
137
138    eia = mo_e_o[:,None] - mo_e_v
139    eijab = lib.direct_sum('ia,jb->ijab',eia,eia)
140    t1new /= eia
141    t2new /= eijab
142
143    return t1new, t2new
144
145
146def energy(cc, t1=None, t2=None, eris=None):
147    '''RCCSD correlation energy'''
148    if t1 is None: t1 = cc.t1
149    if t2 is None: t2 = cc.t2
150    if eris is None: eris = cc.ao2mo()
151
152    nocc, nvir = t1.shape
153    fock = eris.fock
154    e = 2*np.einsum('ia,ia', fock[:nocc,nocc:], t1)
155    tau = np.einsum('ia,jb->ijab',t1,t1)
156    tau += t2
157    eris_ovov = np.asarray(eris.ovov)
158    e += 2*np.einsum('ijab,iajb', tau, eris_ovov)
159    e +=  -np.einsum('ijab,ibja', tau, eris_ovov)
160    if abs(e.imag) > 1e-4:
161        logger.warn(cc, 'Non-zero imaginary part found in RCCSD energy %s', e)
162    return e.real
163
164
165class RCCSD(ccsd.CCSD):
166    '''restricted CCSD with IP-EOM, EA-EOM, EE-EOM, and SF-EOM capabilities
167
168    Ground-state CCSD is performed in optimized ccsd.CCSD and EOM is performed here.
169    '''
170
171    def kernel(self, t1=None, t2=None, eris=None, mbpt2=False):
172        return self.ccsd(t1, t2, eris, mbpt2)
173    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
174        '''Ground-state CCSD.
175
176        Kwargs:
177            mbpt2 : bool
178                Use one-shot MBPT2 approximation to CCSD.
179        '''
180        if mbpt2:
181            pt = mp2.MP2(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
182            self.e_corr, self.t2 = pt.kernel(eris=eris)
183            nocc, nvir = self.t2.shape[1:3]
184            self.t1 = np.zeros((nocc,nvir))
185            return self.e_corr, self.t1, self.t2
186
187        if eris is None:
188            eris = self.ao2mo(self.mo_coeff)
189        return ccsd.CCSD.ccsd(self, t1, t2, eris)
190
191    def ao2mo(self, mo_coeff=None):
192        nmo = self.nmo
193        nao = self.mo_coeff.shape[0]
194        nmo_pair = nmo * (nmo+1) // 2
195        nao_pair = nao * (nao+1) // 2
196        mem_incore = (max(nao_pair**2, nmo**4) + nmo_pair**2) * 8/1e6
197        mem_now = lib.current_memory()[0]
198        if (self._scf._eri is not None and
199            (mem_incore+mem_now < self.max_memory) or self.mol.incore_anyway):
200            return _make_eris_incore(self, mo_coeff)
201
202        elif getattr(self._scf, 'with_df', None):
203            logger.warn(self, 'CCSD detected DF being used in the HF object. '
204                        'MO integrals are computed based on the DF 3-index tensors.\n'
205                        'It\'s recommended to use dfccsd.CCSD for the '
206                        'DF-CCSD calculations')
207            raise NotImplementedError
208            #return _make_df_eris_outcore(self, mo_coeff)
209
210        else:
211            return _make_eris_outcore(self, mo_coeff)
212
213    energy = energy
214    update_amps = update_amps
215
216    def solve_lambda(self, t1=None, t2=None, l1=None, l2=None,
217                     eris=None):
218        from pyscf.cc import rccsd_lambda
219        if t1 is None: t1 = self.t1
220        if t2 is None: t2 = self.t2
221        if eris is None: eris = self.ao2mo(self.mo_coeff)
222        self.converged_lambda, self.l1, self.l2 = \
223                rccsd_lambda.kernel(self, eris, t1, t2, l1, l2,
224                                    max_cycle=self.max_cycle,
225                                    tol=self.conv_tol_normt,
226                                    verbose=self.verbose)
227        return self.l1, self.l2
228
229    def ccsd_t(self, t1=None, t2=None, eris=None):
230        return ccsd.CCSD.ccsd_t(self, t1, t2, eris)
231
232    def density_fit(self, auxbasis=None, with_df=None):
233        raise NotImplementedError
234
235
236class _ChemistsERIs(ccsd._ChemistsERIs):
237
238    def get_ovvv(self, *slices):
239        '''To access a subblock of ovvv tensor'''
240        if slices:
241            return self.ovvv[slices]
242        else:
243            return self.ovvv
244
245def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
246    cput0 = (logger.process_clock(), logger.perf_counter())
247    eris = _ChemistsERIs()
248    eris._common_init_(mycc, mo_coeff)
249    nocc = eris.nocc
250    nmo = eris.fock.shape[0]
251
252    if callable(ao2mofn):
253        eri1 = ao2mofn(eris.mo_coeff).reshape([nmo]*4)
254    else:
255        eri1 = ao2mo.incore.full(mycc._scf._eri, eris.mo_coeff)
256        eri1 = ao2mo.restore(1, eri1, nmo)
257    eris.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
258    eris.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
259    eris.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
260    eris.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
261    eris.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
262    eris.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
263    eris.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
264    logger.timer(mycc, 'CCSD integral transformation', *cput0)
265    return eris
266
267def _make_eris_outcore(mycc, mo_coeff=None):
268    cput0 = (logger.process_clock(), logger.perf_counter())
269    log = logger.Logger(mycc.stdout, mycc.verbose)
270    eris = _ChemistsERIs()
271    eris._common_init_(mycc, mo_coeff)
272
273    mol = mycc.mol
274    mo_coeff = eris.mo_coeff
275    nocc = eris.nocc
276    nao, nmo = mo_coeff.shape
277    nvir = nmo - nocc
278    eris.feri1 = lib.H5TmpFile()
279    eris.oooo = eris.feri1.create_dataset('oooo', (nocc,nocc,nocc,nocc), 'f8')
280    eris.ovoo = eris.feri1.create_dataset('ovoo', (nocc,nvir,nocc,nocc), 'f8', chunks=(nocc,1,nocc,nocc))
281    eris.ovov = eris.feri1.create_dataset('ovov', (nocc,nvir,nocc,nvir), 'f8', chunks=(nocc,1,nocc,nvir))
282    eris.ovvo = eris.feri1.create_dataset('ovvo', (nocc,nvir,nvir,nocc), 'f8', chunks=(nocc,1,nvir,nocc))
283    eris.ovvv = eris.feri1.create_dataset('ovvv', (nocc,nvir,nvir,nvir), 'f8')
284    eris.oovv = eris.feri1.create_dataset('oovv', (nocc,nocc,nvir,nvir), 'f8', chunks=(nocc,nocc,1,nvir))
285    eris.vvvv = eris.feri1.create_dataset('vvvv', (nvir,nvir,nvir,nvir), 'f8')
286    max_memory = max(MEMORYMIN, mycc.max_memory-lib.current_memory()[0])
287
288    ftmp = lib.H5TmpFile()
289    ao2mo.full(mol, mo_coeff, ftmp, max_memory=max_memory, verbose=log)
290    eri = ftmp['eri_mo']
291
292    nocc_pair = nocc*(nocc+1)//2
293    tril2sq = lib.square_mat_in_trilu_indices(nmo)
294    oo = eri[:nocc_pair]
295    eris.oooo[:] = ao2mo.restore(1, oo[:,:nocc_pair], nocc)
296    oovv = lib.take_2d(oo, tril2sq[:nocc,:nocc].ravel(), tril2sq[nocc:,nocc:].ravel())
297    eris.oovv[:] = oovv.reshape(nocc,nocc,nvir,nvir)
298    oo = oovv = None
299
300    tril2sq = lib.square_mat_in_trilu_indices(nmo)
301    blksize = min(nvir, max(BLKMIN, int(max_memory*1e6/8/nmo**3/2)))
302    for p0, p1 in lib.prange(0, nvir, blksize):
303        q0, q1 = p0+nocc, p1+nocc
304        off0 = q0*(q0+1)//2
305        off1 = q1*(q1+1)//2
306        buf = lib.unpack_tril(eri[off0:off1])
307
308        tmp = buf[ tril2sq[q0:q1,:nocc] - off0 ]
309        eris.ovoo[:,p0:p1] = tmp[:,:,:nocc,:nocc].transpose(1,0,2,3)
310        eris.ovvo[:,p0:p1] = tmp[:,:,nocc:,:nocc].transpose(1,0,2,3)
311        eris.ovov[:,p0:p1] = tmp[:,:,:nocc,nocc:].transpose(1,0,2,3)
312        eris.ovvv[:,p0:p1] = tmp[:,:,nocc:,nocc:].transpose(1,0,2,3)
313
314        tmp = buf[ tril2sq[q0:q1,nocc:q1] - off0 ]
315        eris.vvvv[p0:p1,:p1] = tmp[:,:,nocc:,nocc:]
316        if p0 > 0:
317            eris.vvvv[:p0,p0:p1] = tmp[:,:p0,nocc:,nocc:].transpose(1,0,2,3)
318        buf = tmp = None
319    log.timer('CCSD integral transformation', *cput0)
320    return eris
321
322
323if __name__ == '__main__':
324    from pyscf import scf
325    from pyscf import gto
326    from pyscf.cc import gccsd
327
328    mol = gto.Mole()
329    mol.atom = [
330        [8 , (0. , 0.     , 0.)],
331        [1 , (0. , -0.757 , 0.587)],
332        [1 , (0. , 0.757  , 0.587)]]
333    mol.basis = 'cc-pvdz'
334    #mol.basis = '3-21G'
335    mol.verbose = 0
336    mol.spin = 0
337    mol.build()
338    mf = scf.RHF(mol).run(conv_tol=1e-14)
339
340    mycc = RCCSD(mf)
341    mycc.max_memory = 0
342    eris = mycc.ao2mo()
343    emp2, t1, t2 = mycc.init_amps(eris)
344    print(lib.finger(t2) - 0.044540097905897198)
345    np.random.seed(1)
346    t1 = np.random.random(t1.shape)*.1
347    t2 = np.random.random(t2.shape)*.1
348    t2 = t2 + t2.transpose(1,0,3,2)
349    t1, t2 = update_amps(mycc, t1, t2, eris)
350    print(lib.finger(t1) - 0.25118555558133576)
351    print(lib.finger(t2) - 0.02352137419932243)
352
353    ecc, t1, t2 = mycc.kernel()
354    print(ecc - -0.21334326214236796)
355
356    e, v = mycc.ipccsd(nroots=3)
357    print(e[0] - 0.43356041409195489)
358    print(e[1] - 0.51876598058509493)
359    print(e[2] - 0.6782879569941862 )
360
361    e, v = mycc.eeccsd(nroots=4)
362    print(e[0] - 0.2757159395886167)
363    print(e[1] - 0.2757159395886167)
364    print(e[2] - 0.2757159395886167)
365    print(e[3] - 0.3005716731825082)
366
367
368    mol = gto.Mole()
369    mol.verbose = 0
370    mol.atom = [
371        [8 , (0. , 0.     , 0.)],
372        [1 , (0. , -0.757 , 0.587)],
373        [1 , (0. , 0.757  , 0.587)]]
374    mol.basis = '631g'
375    mol.build()
376    mf = scf.RHF(mol)
377    mf.conv_tol = 1e-16
378    mf.scf()
379    mo_coeff = mf.mo_coeff + np.sin(mf.mo_coeff) * .01j
380    nao = mo_coeff.shape[0]
381    eri = ao2mo.restore(1, mf._eri, nao)
382    eri0 = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_coeff.conj(), mo_coeff,
383                      mo_coeff.conj(), mo_coeff)
384
385    nocc, nvir = 5, nao-5
386    eris = _ChemistsERIs(mol)
387    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
388    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
389    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
390    eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
391    eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
392    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:].copy()
393    eris.vvvv = eri0[nocc:,nocc:,nocc:,nocc:].copy()
394    eris.fock = np.diag(mf.mo_energy)
395    eris.mo_energy = mf.mo_energy
396
397    np.random.seed(1)
398    t1 = np.random.random((nocc,nvir)) + np.random.random((nocc,nvir))*.1j - .5
399    t2 = np.random.random((nocc,nocc,nvir,nvir)) - .5
400    t2 = t2 + np.sin(t2) * .1j
401    t2 = t2 + t2.transpose(1,0,3,2)
402
403    mycc = RCCSD(mf)
404    t1new_ref, t2new_ref = update_amps(mycc, t1, t2, eris)
405
406    orbspin = np.zeros(nao*2, dtype=int)
407    orbspin[1::2] = 1
408    eri1 = np.zeros([nao*2]*4, dtype=np.complex128)
409    eri1[0::2,0::2,0::2,0::2] = eri1[0::2,0::2,1::2,1::2] = eri0
410    eri1[1::2,1::2,0::2,0::2] = eri1[1::2,1::2,1::2,1::2] = eri0
411    eri1 = eri1.transpose(0,2,1,3) - eri1.transpose(0,2,3,1)
412    erig = gccsd._PhysicistsERIs(mol)
413    nocc *= 2
414    nvir *= 2
415    erig.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
416    erig.ooov = eri1[:nocc,:nocc,:nocc,nocc:].copy()
417    erig.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
418    erig.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
419    erig.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
420    erig.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
421    erig.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
422    mo_e = np.array([mf.mo_energy]*2)
423    erig.fock = np.diag(mo_e.T.ravel())
424
425    myccg = gccsd.GCCSD(scf.addons.convert_to_ghf(mf))
426    t1, t2 = myccg.amplitudes_from_ccsd(t1, t2)
427    t1new, t2new = gccsd.update_amps(myccg, t1, t2, erig)
428    print(abs(t1new[0::2,0::2]-t1new_ref).max())
429    t2aa = t2new[0::2,0::2,0::2,0::2]
430    t2ab = t2new[0::2,1::2,0::2,1::2]
431    print(abs(t2ab-t2new_ref).max())
432    print(abs(t2ab-t2ab.transpose(1,0,2,3) - t2aa).max())
433
434