1#!/usr/bin/env python
2# Copyright 2014-2020 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 numpy
17
18from pyscf import lib
19from pyscf.cc import rccsd
20from pyscf.cc import uccsd
21from pyscf.cc import gccsd
22from pyscf.pbc import mp
23
24class RCCSD(rccsd.RCCSD):
25    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
26        from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
27        warn_pbc2d_eri(self._scf)
28        if mbpt2:
29            pt = mp.RMP2(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
30            self.e_corr, self.t2 = pt.kernel(eris=eris)
31            nocc, nvir = self.t2.shape[1:3]
32            self.t1 = numpy.zeros((nocc,nvir))
33            return self.e_corr, self.t1, self.t2
34        return rccsd.RCCSD.ccsd(self, t1, t2, eris)
35
36    def ao2mo(self, mo_coeff=None):
37        from pyscf.pbc import tools
38        ao2mofn = mp.mp2._gen_ao2mofn(self._scf)
39        # _scf.exxdiv affects eris.fock. HF exchange correction should be
40        # excluded from the Fock matrix.
41        with lib.temporary_env(self._scf, exxdiv=None):
42            eris = rccsd._make_eris_incore(self, mo_coeff, ao2mofn=ao2mofn)
43
44        # eris.mo_energy so far is just the diagonal part of the Fock matrix
45        # without the exxdiv treatment. Here to add the exchange correction to
46        # get better orbital energies. It is important for the low-dimension
47        # systems since their occupied and the virtual orbital energies may
48        # overlap which may lead to numerical issue in the CCSD iterations.
49        #if mo_coeff is self._scf.mo_coeff:
50        #    eris.mo_energy = self._scf.mo_energy[self.get_frozen_mask()]
51        #else:
52
53        # Add the HFX correction of Ewald probe charge method.
54        # FIXME: Whether to add this correction for other exxdiv treatments?
55        # Without the correction, MP2 energy may be largely off the
56        # correct value.
57        madelung = tools.madelung(self._scf.cell, self._scf.kpt)
58        eris.mo_energy = _adjust_occ(eris.mo_energy, eris.nocc, -madelung)
59        return eris
60
61class UCCSD(uccsd.UCCSD):
62    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
63        from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
64        warn_pbc2d_eri(self._scf)
65        if mbpt2:
66            pt = mp.UMP2(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
67            self.e_corr, self.t2 = pt.kernel(eris=eris)
68            nocca, noccb = self.nocc
69            nmoa, nmob = self.nmo
70            nvira, nvirb = nmoa-nocca, nmob-noccb
71            self.t1 = (numpy.zeros((nocca,nvira)), numpy.zeros((noccb,nvirb)))
72            return self.e_corr, self.t1, self.t2
73        return uccsd.UCCSD.ccsd(self, t1, t2, eris)
74
75    def ao2mo(self, mo_coeff=None):
76        from pyscf.pbc import tools
77        ao2mofn = mp.mp2._gen_ao2mofn(self._scf)
78        # _scf.exxdiv affects eris.fock. HF exchange correction should be
79        # excluded from the Fock matrix.
80        with lib.temporary_env(self._scf, exxdiv=None):
81            eris = uccsd._make_eris_incore(self, mo_coeff, ao2mofn=ao2mofn)
82
83        #if mo_coeff is self._scf.mo_coeff:
84        #    idxa, idxb = self.get_frozen_mask()
85        #    mo_e_a, mo_e_b = self._scf.mo_energy
86        #    eris.mo_energy = (mo_e_a[idxa], mo_e_b[idxb])
87        #else:
88        nocca, noccb = eris.nocc
89        madelung = tools.madelung(self._scf.cell, self._scf.kpt)
90        eris.mo_energy = (_adjust_occ(eris.mo_energy[0], nocca, -madelung),
91                          _adjust_occ(eris.mo_energy[1], noccb, -madelung))
92        return eris
93
94class GCCSD(gccsd.GCCSD):
95    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
96        from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
97        warn_pbc2d_eri(self._scf)
98        if mbpt2:
99            from pyscf.pbc.mp import mp2
100            pt = mp2.GMP2(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
101            self.e_corr, self.t2 = pt.kernel(eris=eris)
102            nocc, nvir = self.t2.shape[1:3]
103            self.t1 = numpy.zeros((nocc,nvir))
104            return self.e_corr, self.t1, self.t2
105        return gccsd.GCCSD.ccsd(self, t1, t2, eris)
106
107    def ao2mo(self, mo_coeff=None):
108        from pyscf.pbc import tools
109        with_df = self._scf.with_df
110        kpt = self._scf.kpt
111        def ao2mofn(mo_coeff):
112            nao, nmo = mo_coeff.shape
113            mo_a = mo_coeff[:nao//2]
114            mo_b = mo_coeff[nao//2:]
115            orbspin = getattr(mo_coeff, 'orbspin', None)
116            if orbspin is None:
117                eri  = with_df.ao2mo(mo_a, kpt, compact=False)
118                eri += with_df.ao2mo(mo_b, kpt, compact=False)
119                eri1 = with_df.ao2mo((mo_a,mo_a,mo_b,mo_b), kpt, compact=False)
120                eri += eri1
121                eri += eri1.T
122                eri = eri.reshape([nmo]*4)
123            else:
124                # If GHF orbitals have orbspin labels, alpha and beta orbitals
125                # occupy different columns. Here merging them into one set of
126                # orbitals then zero out spin forbidden MO integrals
127                mo = mo_a + mo_b
128                eri  = with_df.ao2mo(mo, kpt, compact=False).reshape([nmo]*4)
129                sym_forbid = (orbspin[:,None] != orbspin)
130                eri[sym_forbid,:,:] = 0
131                eri[:,:,sym_forbid] = 0
132            return eri
133
134        # _scf.exxdiv affects eris.fock. HF exchange correction should be
135        # excluded from the Fock matrix.
136        with lib.temporary_env(self._scf, exxdiv=None):
137            eris = gccsd._make_eris_incore(self, mo_coeff, ao2mofn=ao2mofn)
138
139        #if mo_coeff is self._scf.mo_coeff:
140        #    eris.mo_energy = self._scf.mo_energy[self.get_frozen_mask()]
141        #else:
142        madelung = tools.madelung(self._scf.cell, self._scf.kpt)
143        eris.mo_energy = _adjust_occ(eris.mo_energy, eris.nocc, -madelung)
144        return eris
145
146def _adjust_occ(mo_energy, nocc, shift):
147    '''Modify occupied orbital energy'''
148    mo_energy = mo_energy.copy()
149    mo_energy[:nocc] += shift
150    return mo_energy
151
152
153from pyscf.pbc import scf
154scf.hf.RHF.CCSD = lib.class_as_method(RCCSD)
155scf.uhf.UHF.CCSD = lib.class_as_method(UCCSD)
156scf.ghf.GHF.CCSD = lib.class_as_method(GCCSD)
157scf.rohf.ROHF.CCSD = None
158
159