1# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14#
15# Author: Oliver J. Backhouse <olbackhouse@gmail.com>
16#         George H. Booth <george.booth@kcl.ac.uk>
17#
18
19'''
20Auxiliary second-order Green's function perturbation theory for
21arbitrary moment consistency
22'''
23
24import numpy as np
25from pyscf import lib
26from pyscf.lib import logger
27from pyscf import __config__
28from pyscf.agf2 import aux, ragf2
29
30
31def build_se_part(agf2, eri, gf_occ, gf_vir, os_factor=1.0, ss_factor=1.0):
32    ''' Builds either the auxiliaries of the occupied self-energy,
33        or virtual if :attr:`gf_occ` and :attr:`gf_vir` are swapped.
34
35    Args:
36        eri : _ChemistsERIs
37            Electronic repulsion integrals
38        gf_occ : GreensFunction
39            Occupied Green's function
40        gf_vir : GreensFunction
41            Virtual Green's function
42
43    Kwargs:
44        os_factor : float
45            Opposite-spin factor for spin-component-scaled (SCS)
46            calculations. Default 1.0
47        ss_factor : float
48            Same-spin factor for spin-component-scaled (SCS)
49            calculations. Default 1.0
50
51    Returns:
52        :class:`SelfEnergy`
53    '''
54
55    cput0 = (logger.process_clock(), logger.perf_counter())
56    log = logger.Logger(agf2.stdout, agf2.verbose)
57
58    assert type(gf_occ) is aux.GreensFunction
59    assert type(gf_vir) is aux.GreensFunction
60
61    nmo = agf2.nmo
62    nocc = gf_occ.naux
63    nvir = gf_vir.naux
64    naux = nocc * nocc * nvir
65    tol = agf2.weight_tol
66
67    if not (agf2.frozen is None or agf2.frozen == 0):
68        mask = ragf2.get_frozen_mask(agf2)
69        nmo -= np.sum(~mask)
70
71    e = np.zeros((naux))
72    v = np.zeros((nmo, naux))
73
74    fpos = np.sqrt(0.5 * os_factor)
75    fneg = np.sqrt(0.5 * os_factor + ss_factor)
76    fdia = np.sqrt(os_factor)
77
78    eja = lib.direct_sum('j,a->ja', gf_occ.energy, -gf_vir.energy)
79
80    coeffs = (gf_occ.coupling, gf_occ.coupling, gf_vir.coupling)
81    qeri = _make_qmo_eris_incore(agf2, eri, coeffs)
82
83    p1 = 0
84    for i in range(nocc):
85        xija = qeri[:,i,:i].reshape(nmo, -1)
86        xjia = qeri[:,:i,i].reshape(nmo, -1)
87        xiia = qeri[:,i,i].reshape(nmo, -1)
88        eija = gf_occ.energy[i] + eja[:i+1]
89
90        p0, p1 = p1, p1 + i*nvir
91        e[p0:p1] = eija[:i].ravel()
92        v[:,p0:p1] = fneg * (xija - xjia)
93
94        p0, p1 = p1, p1 + i*nvir
95        e[p0:p1] = eija[:i].ravel()
96        v[:,p0:p1] = fpos * (xija + xjia)
97
98        p0, p1 = p1, p1 + nvir
99        e[p0:p1] = eija[i].ravel()
100        v[:,p0:p1] = fdia * xiia
101
102    se = aux.SelfEnergy(e, v, chempot=gf_occ.chempot)
103    se.remove_uncoupled(tol=tol)
104
105    if not (agf2.frozen is None or agf2.frozen == 0):
106        coupling = np.zeros((agf2.nmo, se.naux))
107        coupling[mask] = se.coupling
108        se = aux.SelfEnergy(se.energy, coupling, chempot=se.chempot)
109
110    log.timer('se part', *cput0)
111
112    return se
113
114
115class RAGF2(ragf2.RAGF2):
116    ''' Restricted AGF2 with canonical HF reference for arbitrary
117        moment consistency
118
119    Attributes:
120        nmom : tuple of int
121            Compression level of the Green's function and
122            self-energy, respectively
123        verbose : int
124            Print level. Default value equals to :class:`Mole.verbose`
125        max_memory : float or int
126            Allowed memory in MB. Default value equals to :class:`Mole.max_memory`
127        conv_tol : float
128            Convergence threshold for AGF2 energy. Default value is 1e-7
129        conv_tol_rdm1 : float
130            Convergence threshold for first-order reduced density matrix.
131            Default value is 1e-8.
132        conv_tol_nelec : float
133            Convergence threshold for the number of electrons. Default
134            value is 1e-6.
135        max_cycle : int
136            Maximum number of AGF2 iterations. Default value is 50.
137        max_cycle_outer : int
138            Maximum number of outer Fock loop iterations. Default
139            value is 20.
140        max_cycle_inner : int
141            Maximum number of inner Fock loop iterations. Default
142            value is 50.
143        weight_tol : float
144            Threshold in spectral weight of auxiliaries to be considered
145            zero. Default 1e-11.
146        fock_diis_space : int
147            DIIS space size for Fock loop iterations. Default value is 6.
148        fock_diis_min_space :
149            Minimum space of DIIS. Default value is 1.
150        os_factor : float
151            Opposite-spin factor for spin-component-scaled (SCS)
152            calculations. Default 1.0
153        ss_factor : float
154            Same-spin factor for spin-component-scaled (SCS)
155            calculations. Default 1.0
156        damping : float
157            Damping factor for the self-energy. Default value is 0.0
158
159    Saved results
160
161        e_corr : float
162            AGF2 correlation energy
163        e_tot : float
164            Total energy (HF + correlation)
165        e_1b : float
166            One-body part of :attr:`e_tot`
167        e_2b : float
168            Two-body part of :attr:`e_tot`
169        e_init : float
170            Initial correlation energy (truncated MP2)
171        converged : bool
172            Whether convergence was successful
173        se : SelfEnergy
174            Auxiliaries of the self-energy
175        gf : GreensFunction
176            Auxiliaries of the Green's function
177    '''
178
179    def __init__(self, mf, nmom=(None,0), frozen=None, mo_energy=None, mo_coeff=None, mo_occ=None):
180
181        ragf2.RAGF2.__init__(self, mf, frozen=frozen, mo_energy=mo_energy,
182                             mo_coeff=mo_coeff, mo_occ=mo_occ)
183
184        self.nmom = nmom
185
186        self._keys.update(['nmom'])
187
188    build_se_part = build_se_part
189
190    def build_se(self, eri=None, gf=None, os_factor=None, ss_factor=None, se_prev=None):
191        ''' Builds the auxiliaries of the self-energy.
192
193        Args:
194            eri : _ChemistsERIs
195                Electronic repulsion integrals
196            gf : GreensFunction
197                Auxiliaries of the Green's function
198
199        Kwargs:
200            os_factor : float
201                Opposite-spin factor for spin-component-scaled (SCS)
202                calculations. Default 1.0
203            ss_factor : float
204                Same-spin factor for spin-component-scaled (SCS)
205                calculations. Default 1.0
206            se_prev : SelfEnergy
207                Previous self-energy for damping. Default value is None
208
209        Returns:
210            :class:`SelfEnergy`
211        '''
212
213        if eri is None: eri = self.ao2mo()
214        if gf is None: gf = self.gf
215        if gf is None: gf = self.init_gf()
216
217        fock = None
218        if self.nmom[0] is not None:
219            fock = self.get_fock(eri=eri, gf=gf)
220
221        if os_factor is None: os_factor = self.os_factor
222        if ss_factor is None: ss_factor = self.ss_factor
223
224        facs = dict(os_factor=os_factor, ss_factor=ss_factor)
225        gf_occ = gf.get_occupied()
226        gf_vir = gf.get_virtual()
227
228        se_occ = self.build_se_part(eri, gf_occ, gf_vir, **facs)
229        se_occ = se_occ.compress(n=(None, self.nmom[1]))
230
231        se_vir = self.build_se_part(eri, gf_vir, gf_occ, **facs)
232        se_vir = se_vir.compress(n=(None, self.nmom[1]))
233
234        se = aux.combine(se_vir, se_occ)
235        se = se.compress(phys=fock, n=(self.nmom[0], None))
236
237        if se_prev is not None and self.damping != 0.0:
238            se.coupling *= np.sqrt(1.0-self.damping)
239            se_prev.coupling *= np.sqrt(self.damping)
240            se = aux.combine(se, se_prev)
241            se = se.compress(n=self.nmom)
242
243        return se
244
245    def dump_flags(self, verbose=None):
246        ragf2.RAGF2.dump_flags(self, verbose=verbose)
247        logger.info(self, 'nmom = %s', repr(self.nmom))
248        return self
249
250    def run_diis(self, se, diis=None):
251        return se
252
253
254class _ChemistsERIs(ragf2._ChemistsERIs):
255    pass
256
257_make_qmo_eris_incore = ragf2._make_qmo_eris_incore
258
259
260
261if __name__ == '__main__':
262    from pyscf import gto, scf, mp
263
264    mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='cc-pvdz', verbose=3)
265    rhf = scf.RHF(mol)
266    rhf.conv_tol = 1e-11
267    rhf.run()
268
269    agf2 = RAGF2(rhf, nmom=(None,0))
270    agf2.run()
271
272    agf2 = ragf2.RAGF2(rhf)
273    agf2.run()
274