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
21unrestricted references
22'''
23
24import time
25import numpy as np
26from pyscf import lib
27from pyscf.lib import logger
28from pyscf import __config__
29from pyscf import ao2mo
30from pyscf.scf import _vhf
31from pyscf.agf2 import aux, ragf2, _agf2, mpi_helper
32from pyscf.agf2.chempot import binsearch_chempot, minimize_chempot
33from pyscf.mp.ump2 import get_frozen_mask as _get_frozen_mask
34
35BLKMIN = getattr(__config__, 'agf2_blkmin', 1)
36
37
38def build_se_part(agf2, eri, gf_occ, gf_vir, os_factor=1.0, ss_factor=1.0):
39    ''' Builds either the auxiliaries of the occupied self-energy,
40        or virtual if :attr:`gf_occ` and :attr:`gf_vir` are swapped,
41        for a single spin.
42
43    Args:
44        eri : _ChemistsERIs
45            Electronic repulsion integrals
46        gf_occ : tuple of GreensFunction
47            Occupied Green's function for each spin
48        gf_vir : tuple of GreensFunction
49            Virtual Green's function for each spin
50
51    Kwargs:
52        os_factor : float
53            Opposite-spin factor for spin-component-scaled (SCS)
54            calculations. Default 1.0
55        ss_factor : float
56            Same-spin factor for spin-component-scaled (SCS)
57            calculations. Default 1.0
58
59    Returns:
60        :class:`SelfEnergy`
61    '''
62
63    cput0 = (logger.process_clock(), logger.perf_counter())
64    log = logger.Logger(agf2.stdout, agf2.verbose)
65
66    assert type(gf_occ[0]) is aux.GreensFunction
67    assert type(gf_occ[1]) is aux.GreensFunction
68    assert type(gf_vir[0]) is aux.GreensFunction
69    assert type(gf_vir[1]) is aux.GreensFunction
70
71    nmo = eri.nmo
72    noa, nob = gf_occ[0].naux, gf_occ[1].naux
73    nva, nvb = gf_vir[0].naux, gf_vir[1].naux
74    tol = agf2.weight_tol
75    facs = dict(os_factor=os_factor, ss_factor=ss_factor)
76
77    ci_a, ei_a = gf_occ[0].coupling, gf_occ[0].energy
78    ci_b, ei_b = gf_occ[1].coupling, gf_occ[1].energy
79    ca_a, ea_a = gf_vir[0].coupling, gf_vir[0].energy
80    ca_b, ea_b = gf_vir[1].coupling, gf_vir[1].energy
81
82    mem_incore = (nmo[0]*noa*(noa*nva+nob*nvb)) * 8/1e6
83    mem_now = lib.current_memory()[0]
84    if (mem_incore+mem_now < agf2.max_memory) or agf2.incore_complete:
85        qeri = _make_qmo_eris_incore(agf2, eri, (ci_a, ci_a, ca_a), (ci_b, ci_b, ca_b), spin=0)
86    else:
87        qeri = _make_qmo_eris_outcore(agf2, eri, (ci_a, ci_a, ca_a), (ci_b, ci_b, ca_b), spin=0)
88
89    if isinstance(qeri[0], np.ndarray):
90        vv, vev = _agf2.build_mats_uagf2_incore(qeri, (ei_a, ei_b), (ea_a, ea_b), **facs)
91    else:
92        vv, vev = _agf2.build_mats_uagf2_outcore(qeri, (ei_a, ei_b), (ea_a, ea_b), **facs)
93
94    e, c = _agf2.cholesky_build(vv, vev)
95    se_a = aux.SelfEnergy(e, c, chempot=gf_occ[0].chempot)
96    se_a.remove_uncoupled(tol=tol)
97
98    if not (agf2.frozen is None or agf2.frozen == 0):
99        mask = get_frozen_mask(agf2)
100        coupling = np.zeros((nmo[0], se_a.naux))
101        coupling[mask[0]] = se_a.coupling
102        se_a = aux.SelfEnergy(se_a.energy, coupling, chempot=se_a.chempot)
103
104    cput0 = log.timer('se part (alpha)', *cput0)
105
106    mem_incore = (nmo[1]*nob*(nob*nvb+noa*nva)) * 8/1e6
107    mem_now = lib.current_memory()[0]
108    if (mem_incore+mem_now < agf2.max_memory) or agf2.incore_complete:
109        qeri = _make_qmo_eris_incore(agf2, eri, (ci_a, ci_a, ca_a), (ci_b, ci_b, ca_b), spin=1)
110    else:
111        qeri = _make_qmo_eris_outcore(agf2, eri, (ci_a, ci_a, ca_a), (ci_b, ci_b, ca_b), spin=1)
112
113    if isinstance(qeri[0], np.ndarray):
114        vv, vev = _agf2.build_mats_uagf2_incore(qeri, (ei_b, ei_a), (ea_b, ea_a), **facs)
115    else:
116        vv, vev = _agf2.build_mats_uagf2_outcore(qeri, (ei_b, ei_a), (ea_b, ea_a), **facs)
117
118    e, c = _agf2.cholesky_build(vv, vev)
119    se_b = aux.SelfEnergy(e, c, chempot=gf_occ[1].chempot)
120    se_b.remove_uncoupled(tol=tol)
121
122    if not (agf2.frozen is None or agf2.frozen == 0):
123        mask = get_frozen_mask(agf2)
124        coupling = np.zeros((nmo[1], se_b.naux))
125        coupling[mask[1]] = se_b.coupling
126        se_b = aux.SelfEnergy(se_b.energy, coupling, chempot=se_b.chempot)
127
128    cput0 = log.timer('se part (beta)', *cput0)
129
130    return (se_a, se_b)
131
132
133def get_fock(agf2, eri, gf=None, rdm1=None):
134    ''' Computes the physical space Fock matrix in MO basis. If :attr:`rdm1`
135        is not supplied, it is built from :attr:`gf`, which defaults to
136        the mean-field Green's function.
137
138    Args:
139        eri : _ChemistsERIs
140            Electronic repulsion integrals
141
142    Kwargs:
143        gf : GreensFunction
144            Auxiliaries of the Green's function
145        rdm1 : 2D array
146            Reduced density matrix
147
148    Returns:
149        ndarray of physical space Fock matrix
150    '''
151
152    if rdm1 is None:
153        rdm1 = agf2.make_rdm1(gf)
154
155    vj_aa, vk_aa = agf2.get_jk(eri.eri_aa, rdm1=rdm1[0])
156    vj_bb, vk_bb = agf2.get_jk(eri.eri_bb, rdm1=rdm1[1])
157    vj_ab = agf2.get_jk(eri.eri_ab, rdm1=rdm1[1], with_k=False)[0]
158    vj_ba = agf2.get_jk(eri.eri_ba, rdm1=rdm1[0], with_k=False)[0]
159
160    fock_a = eri.h1e[0] + vj_aa + vj_ab - vk_aa
161    fock_b = eri.h1e[1] + vj_bb + vj_ba - vk_bb
162
163    fock = (fock_a, fock_b)
164
165    return fock
166
167
168def fock_loop(agf2, eri, gf, se):
169    ''' Self-consistent loop for the density matrix via the HF self-
170        consistent field.
171
172    Args:
173        eri : _ChemistsERIs
174            Electronic repulsion integrals
175        gf : tuple of GreensFunction
176            Auxiliaries of the Green's function for each spin
177        se : tuple of SelfEnergy
178            Auxiliaries of the self-energy for each spin
179
180    Returns:
181        :class:`SelfEnergy`, :class:`GreensFunction` and a boolean
182        indicating whether convergence was successful.
183    '''
184
185    assert type(gf[0]) is aux.GreensFunction
186    assert type(gf[1]) is aux.GreensFunction
187    assert type(se[0]) is aux.SelfEnergy
188    assert type(se[1]) is aux.SelfEnergy
189
190    cput0 = cput1 = (logger.process_clock(), logger.perf_counter())
191    log = logger.Logger(agf2.stdout, agf2.verbose)
192
193    diis = lib.diis.DIIS(agf2)
194    diis.space = agf2.fock_diis_space
195    diis.min_space = agf2.fock_diis_min_space
196    focka, fockb = agf2.get_fock(eri, gf)
197    sea, seb = se
198    gfa, gfb = gf
199
200    nalph, nbeta = agf2.nocc
201    nmoa, nmob = eri.nmo
202    nauxa, nauxb = sea.naux, seb.naux
203    nqmoa, nqmob = nauxa+nmoa, nauxb+nmob
204    bufa, bufb = np.zeros((nqmoa, nqmoa)), np.zeros((nqmob, nqmob))
205    rdm1a_prev = 0
206    rdm1b_prev = 0
207    converged = False
208    opts = dict(tol=agf2.conv_tol_nelec, maxiter=agf2.max_cycle_inner)
209
210    for niter1 in range(1, agf2.max_cycle_outer+1):
211        sea, opt = minimize_chempot(sea, focka, nalph, x0=sea.chempot,
212                                    occupancy=1, **opts)
213        seb, opt = minimize_chempot(seb, fockb, nbeta, x0=seb.chempot,
214                                    occupancy=1, **opts)
215
216        for niter2 in range(1, agf2.max_cycle_inner+1):
217            wa, va = sea.eig(focka, chempot=0.0, out=bufa)
218            wb, vb = seb.eig(fockb, chempot=0.0, out=bufb)
219            sea.chempot, nerra = \
220                    binsearch_chempot((wa, va), nmoa, nalph, occupancy=1)
221            seb.chempot, nerrb = \
222                    binsearch_chempot((wb, vb), nmob, nbeta, occupancy=1)
223            nerr = max(nerra, nerrb)
224
225            wa, va = sea.eig(focka, out=bufa)
226            wb, vb = seb.eig(fockb, out=bufb)
227            gfa = aux.GreensFunction(wa, va[:nmoa], chempot=sea.chempot)
228            gfb = aux.GreensFunction(wb, vb[:nmob], chempot=seb.chempot)
229
230            gf = (gfa, gfb)
231            focka, fockb = agf2.get_fock(eri, gf)
232            rdm1a, rdm1b = agf2.make_rdm1(gf)
233            focka, fockb = diis.update(np.array((focka, fockb)), xerr=None)
234
235            if niter2 > 1:
236                derra = np.max(np.absolute(rdm1a - rdm1a_prev))
237                derrb = np.max(np.absolute(rdm1b - rdm1b_prev))
238                derr = max(derra, derrb)
239
240                if derr < agf2.conv_tol_rdm1:
241                    break
242
243            rdm1a_prev = rdm1a.copy()
244            rdm1b_prev = rdm1b.copy()
245
246        log.debug1('fock loop %d  cycles = %d  dN = %.3g  |ddm| = %.3g',
247                   niter1, niter2, nerr, derr)
248        cput1 = log.timer_debug1('fock loop %d'%niter1, *cput1)
249
250        if derr < agf2.conv_tol_rdm1 and abs(nerr) < agf2.conv_tol_nelec:
251            converged = True
252            break
253
254    se = (sea, seb)
255
256    log.info('fock converged = %s' % converged)
257    log.info(' alpha: chempot = %.9g  dN = %.3g  |ddm| = %.3g',
258             sea.chempot, nerra, derra)
259    log.info(' beta:  chempot = %.9g  dN = %.3g  |ddm| = %.3g',
260             seb.chempot, nerrb, derrb)
261    log.timer('fock loop', *cput0)
262
263    return gf, se, converged
264
265
266def energy_1body(agf2, eri, gf):
267    ''' Calculates the one-body energy according to the UHF form.
268
269    Args:
270        eri : _ChemistsERIs
271            Electronic repulsion integrals
272        gf : tuple of GreensFunction
273            Auxiliaries of the Green's function for each spin
274
275    Returns:
276        One-body energy
277    '''
278
279    assert type(gf[0]) is aux.GreensFunction
280    assert type(gf[1]) is aux.GreensFunction
281
282    rdm1 = agf2.make_rdm1(gf)
283    fock = agf2.get_fock(eri, gf)
284
285    e1b_a = 0.5 * np.sum(rdm1[0] * (eri.h1e[0] + fock[0]))
286    e1b_b = 0.5 * np.sum(rdm1[1] * (eri.h1e[1] + fock[1]))
287
288    e1b = e1b_a + e1b_b
289    e1b += agf2.energy_nuc()
290
291    return e1b
292
293
294def energy_2body(agf2, gf, se):
295    ''' Calculates the two-body energy using analytically integrated
296        Galitskii-Migdal formula. The formula is symmetric and only
297        one side needs to be calculated.
298
299    Args:
300        gf : tuple of GreensFunction
301            Auxiliaries of the Green's function for each spin
302        se : tuple of SelfEnergy
303            Auxiliaries of the self-energy for each spin
304
305    Returns:
306        Two-body energy
307    '''
308
309    e2b_a = ragf2.energy_2body(agf2, gf[0], se[0])
310    e2b_b = ragf2.energy_2body(agf2, gf[1], se[1])
311
312    e2b = (e2b_a + e2b_b) * 0.5
313
314    return e2b
315
316
317def energy_mp2(agf2, gf, se):
318    ''' Calculates the two-bdoy energy using analytically integrated
319        Galitskii-Migdal formula for an MP2 self-energy. Per the
320        definition of one- and two-body partitioning in the Dyson
321        equation, this reuslt is half of :func:`energy_2body`.
322
323    Args:
324        gf : tuple of GreensFunction
325            Auxiliaries of the Green's function for each spin
326        se : tuple of SelfEnergy
327            Auxiliaries of the self-energy for each spin
328
329    Returns:
330        MP2 energy
331    '''
332
333    emp2_a = ragf2.energy_mp2(agf2, gf[0], se[0])
334    emp2_b = ragf2.energy_mp2(agf2, gf[1], se[1])
335
336    emp2 = (emp2_a + emp2_b) * 0.5
337
338    return emp2
339
340
341class UAGF2(ragf2.RAGF2):
342    ''' Unrestricted AGF2 with canonical HF reference
343
344    Attributes:
345        verbose : int
346            Print level. Default value equals to :class:`Mole.verbose`
347        max_memory : float or int
348            Allowed memory in MB. Default value equals to :class:`Mole.max_memory`
349        incore_complete : bool
350            Avoid all I/O. Default is False.
351        conv_tol : float
352            Convergence threshold for AGF2 energy. Default value is 1e-7
353        conv_tol_rdm1 : float
354            Convergence threshold for first-order reduced density matrix.
355            Default value is 1e-8.
356        conv_tol_nelec : float
357            Convergence threshold for the number of electrons. Default
358            value is 1e-6.
359        max_cycle : int
360            Maximum number of AGF2 iterations. Default value is 50.
361        max_cycle_outer : int
362            Maximum number of outer Fock loop iterations. Default
363            value is 20.
364        max_cycle_inner : int
365            Maximum number of inner Fock loop iterations. Default
366            value is 50.
367        weight_tol : float
368            Threshold in spectral weight of auxiliaries to be considered
369            zero. Default 1e-11.
370        diis : bool or lib.diis.DIIS
371            Whether to use DIIS, can also be a lib.diis.DIIS object. Default
372            value is True.
373        diis_space : int
374            DIIS space size. Default value is 8.
375        diis_min_space : int
376            Minimum space of DIIS. Default value is 1.
377        fock_diis_space : int
378            DIIS space size for Fock loop iterations. Default value is 6.
379        fock_diis_min_space :
380            Minimum space of DIIS. Default value is 1.
381        os_factor : float
382            Opposite-spin factor for spin-component-scaled (SCS)
383            calculations. Default 1.0
384        ss_factor : float
385            Same-spin factor for spin-component-scaled (SCS)
386            calculations. Default 1.0
387        damping : float
388            Damping factor for the self-energy. Default value is 0.0
389
390    Saved results
391
392        e_corr : float
393            AGF2 correlation energy
394        e_tot : float
395            Total energy (HF + correlation)
396        e_1b : float
397            One-body part of :attr:`e_tot`
398        e_2b : float
399            Two-body part of :attr:`e_tot`
400        e_init : float
401            Initial correlation energy (truncated MP2)
402        converged : bool
403            Whether convergence was successful
404        se : tuple of SelfEnergy
405            Auxiliaries of the self-energy for each spin
406        gf : tuple of GreensFunction
407            Auxiliaries of the Green's function for each spin
408    '''
409
410    energy_1body = energy_1body
411    energy_2body = energy_2body
412    fock_loop = fock_loop
413    build_se_part = build_se_part
414
415    def ao2mo(self, mo_coeff=None):
416        ''' Get the electronic repulsion integrals in MO basis.
417        '''
418
419        nmo = max(self.nmo)
420        mem_incore = ((nmo*(nmo+1)//2)**2) * 8/1e6
421        mem_now = lib.current_memory()[0]
422
423        if (self._scf._eri is not None and
424                (mem_incore+mem_now < self.max_memory or self.incore_complete)):
425            eri = _make_mo_eris_incore(self, mo_coeff)
426        else:
427            logger.warn(self, 'MO eris are outcore - this may be very '
428                              'slow for agf2. increasing max_memory or '
429                              'using density fitting is recommended.')
430            eri = _make_mo_eris_outcore(self, mo_coeff)
431
432        return eri
433
434    def make_rdm1(self, gf=None):
435        ''' Compute the one-body reduced density matrix in MO basis.
436
437        Kwargs:
438            gf : tuple of GreensFunction
439                Auxiliaries of the Green's functions for each spin
440
441        Returns:
442            tuple of ndarray of density matrices
443        '''
444
445        if gf is None: gf = self.gf
446        if gf is None: gf = self.init_gf()
447
448        rdm1_a = gf[0].make_rdm1(occupancy=1)
449        rdm1_b = gf[1].make_rdm1(occupancy=1)
450
451        return (rdm1_a, rdm1_b)
452
453    def get_fock(self, eri=None, gf=None, rdm1=None):
454        ''' Computes the physical space Fock matrix in MO basis.
455        '''
456
457        if eri is None: eri = self.ao2mo()
458        if gf is None: gf = self.gf
459
460        return get_fock(self, eri, gf=gf, rdm1=rdm1)
461
462    def energy_mp2(self, mo_energy=None, se=None):
463        if mo_energy is None: mo_energy = self.mo_energy
464        if se is None: se = self.build_se(gf=self.gf)
465
466        self.e_init = energy_mp2(self, mo_energy, se)
467
468        return self.e_init
469
470    def init_gf(self, frozen=False):
471        ''' Builds the Hartree-Fock Green's function.
472
473        Returns:
474            tuple of :class:`GreensFunction`, tuple of :class:`SelfEnergy`
475        '''
476
477        nmoa, nmob = self.nmo
478        nocca, noccb = self.nocc
479
480        energy = self.mo_energy
481        coupling = (np.eye(nmoa), np.eye(nmob))
482
483        focka = np.diag(energy[0])
484        fockb = np.diag(energy[1])
485
486        cpt_a = binsearch_chempot(focka, nmoa, nocca, occupancy=1)[0]
487        cpt_b = binsearch_chempot(fockb, nmob, noccb, occupancy=1)[1]
488
489        if frozen:
490            mask = get_frozen_mask(self)
491            energy = (energy[0][mask[0]], energy[1][mask[1]])
492            coupling = (coupling[0][:,mask[0]], coupling[1][:,mask[1]])
493
494        gf_a = aux.GreensFunction(energy[0], coupling[0], chempot=cpt_a)
495        gf_b = aux.GreensFunction(energy[1], coupling[1], chempot=cpt_b)
496
497        gf = (gf_a, gf_b)
498
499        return gf
500
501    def build_gf(self, eri=None, gf=None, se=None):
502        ''' Builds the auxiliaries of the Green's functions by solving
503            the Dyson equation for each spin.
504
505        Kwargs:
506            eri : _ChemistsERIs
507                Electronic repulsion integrals
508            gf : tuple of GreensFunction
509                Auxiliaries of the Green's function for each spin
510            se : tuple of SelfEnergy
511                Auxiliaries of the self-energy for each spin
512
513        Returns:
514            tuple of :class:`GreensFunction`
515        '''
516
517        if eri is None: eri = self.ao2mo()
518        if gf is None: gf = self.gf
519        if gf is None: gf = self.init_gf()
520        if se is None: se = self.build_se(eri, gf)
521
522        focka, fockb = self.get_fock(eri, gf)
523
524        gf_a = se[0].get_greens_function(focka)
525        gf_b = se[1].get_greens_function(fockb)
526
527        return (gf_a, gf_b)
528
529    def build_se(self, eri=None, gf=None, os_factor=None, ss_factor=None, se_prev=None):
530        ''' Builds the auxiliaries of the self-energy.
531
532        Args:
533            eri : _ChemistsERIs
534                Electronic repulsion integrals
535            gf : tuple of GreensFunction
536                Auxiliaries of the Green's function
537
538        Kwargs:
539            os_factor : float
540                Opposite-spin factor for spin-component-scaled (SCS)
541                calculations. Default 1.0
542            ss_factor : float
543                Same-spin factor for spin-component-scaled (SCS)
544                calculations. Default 1.0
545            se_prev : SelfEnergy
546                Previous self-energy for damping. Default value is None
547
548        Returns
549            tuple of :class:`SelfEnergy`
550        '''
551
552        if eri is None: eri = self.ao2mo()
553        if gf is None: gf = self.gf
554        if gf is None: gf = self.init_gf()
555
556        if os_factor is None: os_factor = self.os_factor
557        if ss_factor is None: ss_factor = self.ss_factor
558
559        facs = dict(os_factor=os_factor, ss_factor=ss_factor)
560        gf_occ = (gf[0].get_occupied(), gf[1].get_occupied())
561        gf_vir = (gf[0].get_virtual(), gf[1].get_virtual())
562
563        se_occ = self.build_se_part(eri, gf_occ, gf_vir, **facs)
564        se_vir = self.build_se_part(eri, gf_vir, gf_occ, **facs)
565
566        se_a = aux.combine(se_occ[0], se_vir[0])
567        se_b = aux.combine(se_occ[1], se_vir[1])
568
569        if se_prev is not None and self.damping != 0.0:
570            se_a_prev, se_b_prev = se_prev
571            se_a.coupling *= np.sqrt(1.0-self.damping)
572            se_b.coupling *= np.sqrt(1.0-self.damping)
573            se_a_prev.coupling *= np.sqrt(self.damping)
574            se_b_prev.coupling *= np.sqrt(self.damping)
575            se_a = aux.combine(se_a, se_a_prev)
576            se_b = aux.combine(se_b, se_b_prev)
577            se_a = se_a.compress(n=(None,0))
578            se_b = se_b.compress(n=(None,0))
579
580        return (se_a, se_b)
581
582    def run_diis(self, se, diis=None):
583        ''' Runs the direct inversion of the iterative subspace for the
584            self-energy.
585
586        Args:
587            se : SelfEnergy
588                Auxiliaries of the self-energy
589            diis : lib.diis.DIIS
590                DIIS object
591
592        Returns:
593            tuple of :class:`SelfEnergy`
594        '''
595
596        if diis is None:
597            return se
598
599        se_occ_a, se_occ_b = (se[0].get_occupied(), se[1].get_occupied())
600        se_vir_a, se_vir_b = (se[0].get_virtual(), se[1].get_virtual())
601
602        vv_occ_a = np.dot(se_occ_a.coupling, se_occ_a.coupling.T)
603        vv_occ_b = np.dot(se_occ_b.coupling, se_occ_b.coupling.T)
604        vv_vir_a = np.dot(se_vir_a.coupling, se_vir_a.coupling.T)
605        vv_vir_b = np.dot(se_vir_b.coupling, se_vir_b.coupling.T)
606
607        vev_occ_a = np.dot(se_occ_a.coupling * se_occ_a.energy[None], se_occ_a.coupling.T)
608        vev_occ_b = np.dot(se_occ_b.coupling * se_occ_b.energy[None], se_occ_b.coupling.T)
609        vev_vir_a = np.dot(se_vir_a.coupling * se_vir_a.energy[None], se_vir_a.coupling.T)
610        vev_vir_b = np.dot(se_vir_b.coupling * se_vir_b.energy[None], se_vir_b.coupling.T)
611
612        dat = np.array([vv_occ_a, vv_vir_a, vev_occ_a, vev_vir_a,
613                        vv_occ_b, vv_vir_b, vev_occ_b, vev_vir_b])
614        dat = diis.update(dat)
615        vv_occ_a, vv_vir_a, vev_occ_a, vev_vir_a, \
616                vv_occ_b, vv_vir_b, vev_occ_b, vev_vir_b = dat
617
618        se_occ_a = aux.SelfEnergy(*_agf2.cholesky_build(vv_occ_a, vev_occ_a), chempot=se[0].chempot)
619        se_vir_a = aux.SelfEnergy(*_agf2.cholesky_build(vv_vir_a, vev_vir_a), chempot=se[0].chempot)
620        se_occ_b = aux.SelfEnergy(*_agf2.cholesky_build(vv_occ_b, vev_occ_b), chempot=se[1].chempot)
621        se_vir_b = aux.SelfEnergy(*_agf2.cholesky_build(vv_vir_b, vev_vir_b), chempot=se[1].chempot)
622        se = (aux.combine(se_occ_a, se_vir_a), aux.combine(se_occ_b, se_vir_b))
623
624        return se
625
626    def density_fit(self, auxbasis=None, with_df=None):
627        from pyscf.agf2 import dfuagf2
628        myagf2 = dfuagf2.DFUAGF2(self._scf)
629        myagf2.__dict__.update(self.__dict__)
630
631        if with_df is not None:
632            myagf2.with_df = with_df
633
634        if auxbasis is not None and myagf2.with_df.auxbasis != auxbasis:
635            import copy
636            myagf2.with_df = copy.copy(myagf2.with_df)
637            myagf2.with_df.auxbasis = auxbasis
638
639        return myagf2
640
641    def get_ip(self, gf, nroots=5):
642        gf_occ = (gf[0].get_occupied(), gf[1].get_occupied())
643        spin = np.array([0,]*gf_occ[0].naux + [1,]*gf_occ[1].naux)
644        e_ip = np.concatenate([gf_occ[0].energy, gf_occ[1].energy], axis=0)
645        v_ip = np.concatenate([gf_occ[0].coupling, gf_occ[1].coupling], axis=1)
646
647        mask = np.argsort(e_ip)
648        spin = list(spin[mask][-nroots:])[::-1]
649        e_ip = list(-e_ip[mask][-nroots:])[::-1]
650        v_ip = list(v_ip[:,mask][:,-nroots:].T)[::-1]
651
652        return e_ip, v_ip, spin
653
654    def ipagf2(self, nroots=5):
655        e_ip, v_ip, spin = self.get_ip(self.gf, nroots=nroots)
656
657        for n, en, vn, sn in zip(range(nroots), e_ip, v_ip, spin):
658            qpwt = np.linalg.norm(vn)**2
659            tag = ['alpha', 'beta'][sn]
660            logger.note(self, 'IP energy level %d E = %.16g  QP weight = %0.6g  (%s)', n, en, qpwt, tag)
661
662        if nroots == 1:
663            return e_ip[0], v_ip[0]
664        else:
665            return e_ip, v_ip
666
667    def get_ea(self, gf, nroots=5):
668        gf_vir = (gf[0].get_virtual(), gf[1].get_virtual())
669        spin = np.array([0,]*gf_vir[0].naux + [1,]*gf_vir[1].naux)
670        e_ea = np.concatenate([gf_vir[0].energy, gf_vir[1].energy], axis=0)
671        v_ea = np.concatenate([gf_vir[0].coupling, gf_vir[1].coupling], axis=1)
672
673        mask = np.argsort(e_ea)
674        spin = list(spin[mask][:nroots])
675        e_ea = list(e_ea[mask][:nroots])
676        v_ea = list(v_ea[:,mask][:,:nroots].T)
677
678        return e_ea, v_ea, spin
679
680    def eaagf2(self, nroots=5):
681        e_ea, v_ea, spin = self.get_ea(self.gf, nroots=nroots)
682
683        for n, en, vn, sn in zip(range(nroots), e_ea, v_ea, spin):
684            qpwt = np.linalg.norm(vn)**2
685            tag = ['alpha', 'beta'][sn]
686            logger.note(self, 'EA energy level %d E = %.16g  QP weight = %0.6g  (%s)', n, en, qpwt, tag)
687
688        if nroots == 1:
689            return e_ea[0], v_ea[0]
690        else:
691            return e_ea, v_ea
692
693
694    @property
695    def nocc(self):
696        if self._nocc is None:
697            self._nocc = (np.sum(self.mo_occ[0] > 0), np.sum(self.mo_occ[1] > 0))
698        return self._nocc
699    @nocc.setter
700    def nocc(self, val):
701        self._nocc = val
702
703    @property
704    def nmo(self):
705        if self._nmo is None:
706            self._nmo = (self.mo_occ[0].size, self.mo_occ[1].size)
707        return self._nmo
708    @nmo.setter
709    def nmo(self, val):
710        self._nmo = val
711
712    @property
713    def qmo_energy(self):
714        return (self.gf[0].energy, self.gf[1].energy)
715
716    @property
717    def qmo_coeff(self):
718        ''' Gives the couplings in AO basis '''
719        return (np.dot(self.mo_coeff[0], self.gf[0].coupling),
720                np.dot(self.mo_coeff[1], self.gf[1].coupling))
721
722    @property
723    def qmo_occ(self):
724        coeff_a = self.gf[0].get_occupied().coupling
725        coeff_b = self.gf[1].get_occupied().coupling
726        occ_a = np.linalg.norm(coeff_a, axis=0) ** 2
727        occ_b = np.linalg.norm(coeff_b, axis=0) ** 2
728        vir_a = np.zeros_like(self.gf[0].get_virtual().energy)
729        vir_b = np.zeros_like(self.gf[1].get_virtual().energy)
730        qmo_occ_a = np.concatenate([occ_a, vir_a])
731        qmo_occ_b = np.concatenate([occ_b, vir_b])
732        return qmo_occ_a, qmo_occ_b
733
734
735def get_frozen_mask(agf2):
736    with lib.temporary_env(agf2, _nocc=None, _nmo=None):
737        return _get_frozen_mask(agf2)
738
739
740class _ChemistsERIs:
741    ''' (pq|rs)
742
743    MO integrals stored in s4 symmetry, we only need QMO integrals
744    in low-symmetry tensors and s4 is highest supported by _vhf
745    '''
746
747    def __init__(self, mol=None):
748        self.mol = mol
749        self.mo_coeff = None
750        self.nocc = None
751        self.nmo = None
752
753        self.fock = None
754        self.h1e = None
755        self.eri = None
756        self.e_hf = None
757
758    def _common_init_(self, agf2, mo_coeff=None):
759        if mo_coeff is None:
760            mo_coeff = agf2.mo_coeff
761
762        self.mo_coeff = mo_coeff
763
764        dm = agf2._scf.make_rdm1(agf2.mo_coeff, agf2.mo_occ)
765        h1e_ao = agf2._scf.get_hcore()
766        vhf = agf2._scf.get_veff(agf2.mol, dm)
767        fock_ao = agf2._scf.get_fock(vhf=vhf, dm=dm)
768
769        self.h1e = (np.dot(np.dot(mo_coeff[0].conj().T, h1e_ao), mo_coeff[0]),
770                    np.dot(np.dot(mo_coeff[1].conj().T, h1e_ao), mo_coeff[1]))
771        self.fock = (np.dot(np.dot(mo_coeff[0].conj().T, fock_ao[0]), mo_coeff[0]),
772                     np.dot(np.dot(mo_coeff[1].conj().T, fock_ao[1]), mo_coeff[1]))
773
774        self.h1e = (mpi_helper.bcast(self.h1e[0]), mpi_helper.bcast(self.h1e[1]))
775        self.fock = (mpi_helper.bcast(self.fock[0]), mpi_helper.bcast(self.fock[1]))
776
777        self.e_hf = mpi_helper.bcast(agf2._scf.e_tot)
778
779        self.nmo = agf2.nmo
780        nocca, noccb = self.nocc = agf2.nocc
781        self.mol = agf2.mol
782
783        mo_e = (self.fock[0].diagonal(), self.fock[1].diagonal())
784        gap_a = abs(mo_e[0][:nocca,None] - mo_e[0][None,nocca:]).min()
785        gap_b = abs(mo_e[1][:noccb,None] - mo_e[1][None,noccb:]).min()
786        gap = min(gap_a, gap_b)
787        if gap < 1e-5:
788            logger.warn(agf2, 'HOMO-LUMO gap %s may be too small for AGF2', gap)
789
790        return self
791
792def _make_mo_eris_incore(agf2, mo_coeff=None):
793    ''' Returns _ChemistsERIs
794    '''
795
796    cput0 = (logger.process_clock(), logger.perf_counter())
797    log = logger.Logger(agf2.stdout, agf2.verbose)
798
799    eris = _ChemistsERIs()
800    eris._common_init_(agf2, mo_coeff)
801    moa, mob = eris.mo_coeff
802    nmoa, nmob = eris.nmo
803
804    eri_aa = ao2mo.incore.full(agf2._scf._eri, moa, verbose=log)
805    eri_bb = ao2mo.incore.full(agf2._scf._eri, mob, verbose=log)
806
807    eri_aa = ao2mo.addons.restore('s4', eri_aa, nmoa)
808    eri_bb = ao2mo.addons.restore('s4', eri_bb, nmob)
809
810    eri_ab = ao2mo.incore.general(agf2._scf._eri, (moa,moa,mob,mob), verbose=log)
811    assert eri_ab.shape == (nmoa*(nmob+1)//2, nmob*(nmob+1)//2)
812    eri_ba = np.transpose(eri_ab)
813
814    eris.eri_aa = eri_aa
815    eris.eri_ab = eri_ab
816    eris.eri_ba = eri_ba
817    eris.eri_bb = eri_bb
818    eris.eri = ((eri_aa, eri_ab), (eri_ba, eri_bb))
819
820    log.timer('MO integral transformation', *cput0)
821
822    return eris
823
824def _make_mo_eris_outcore(agf2, mo_coeff=None):
825    ''' Returns _ChemistsERIs
826    '''
827
828    log = logger.Logger(agf2.stdout, agf2.verbose)
829
830    eris = _ChemistsERIs()
831    eris._common_init_(agf2, mo_coeff)
832
833    mol = agf2.mol
834    moa = np.asarray(eris.mo_coeff[0], order='F')
835    mob = np.asarray(eris.mo_coeff[1], order='F')
836    nmoa, nmob = eris.nmo
837
838    eris.feri = lib.H5TmpFile()
839
840    ao2mo.outcore.full(mol, moa, eris.feri, dataname='mo/aa')
841    ao2mo.outcore.full(mol, mob, eris.feri, dataname='mo/bb')
842    ao2mo.outcore.general(mol, (moa,moa,mob,mob), eris.feri, dataname='mo/ab', verbose=log)
843    ao2mo.outcore.general(mol, (mob,mob,moa,moa), eris.feri, dataname='mo/ba', verbose=log)
844
845    eris.eri_aa = eris.feri['mo/aa']
846    eris.eri_ab = eris.feri['mo/ab']
847    eris.eri_ba = eris.feri['mo/ba']
848    eris.eri_bb = eris.feri['mo/bb']
849
850    eris.eri = ((eris.eri_aa, eris.eri_ab), (eris.eri_ba, eris.eri_bb))
851
852    return eris
853
854def _make_qmo_eris_incore(agf2, eri, coeffs_a, coeffs_b, spin=None):
855    ''' Returns nested tuple of ndarray
856
857    spin = None: ((aaaa, aabb), (bbaa, bbbb))
858    spin = 0: (aaaa, aabb)
859    spin = 1: (bbbb, bbaa)
860    '''
861
862    cput0 = (logger.process_clock(), logger.perf_counter())
863    log = logger.Logger(agf2.stdout, agf2.verbose)
864
865    nmo = eri.nmo
866    nmoa, nmob = nmo
867
868    cxa = np.eye(nmoa)
869    cxb = np.eye(nmob)
870    if not (agf2.frozen is None or agf2.frozen == 0):
871        mask = get_frozen_mask(agf2)
872        cxa = cxa[:,mask[0]]
873        cxb = cxb[:,mask[1]]
874
875    # npaira, npairb = nmoa*(nmoa+1)//2, nmob*(nmob+1)//2
876    cia, cja, caa = coeffs_a
877    cib, cjb, cab = coeffs_b
878    nia, nja, naa = [x.shape[1] for x in coeffs_a]
879    nib, njb, nab = [x.shape[1] for x in coeffs_b]
880
881    if spin is None or spin == 0:
882        c_aa = (cxa, cia, cja, caa)
883        c_ab = (cxa, cia, cjb, cab)
884
885        qeri_aa = ao2mo.incore.general(eri.eri_aa, c_aa, compact=False, verbose=log)
886        qeri_ab = ao2mo.incore.general(eri.eri_ab, c_ab, compact=False, verbose=log)
887
888        qeri_aa = qeri_aa.reshape(cxa.shape[1], nia, nja, naa)
889        qeri_ab = qeri_ab.reshape(cxa.shape[1], nia, njb, nab)
890
891    if spin is None or spin == 1:
892        c_bb = (cxb, cib, cjb, cab)
893        c_ba = (cxb, cib, cja, caa)
894
895        qeri_bb = ao2mo.incore.general(eri.eri_bb, c_bb, compact=False, verbose=log)
896        qeri_ba = ao2mo.incore.general(eri.eri_ba, c_ba, compact=False, verbose=log)
897
898        qeri_bb = qeri_bb.reshape(cxb.shape[1], nib, njb, nab)
899        qeri_ba = qeri_ba.reshape(cxb.shape[1], nib, nja, naa)
900
901    if spin is None:
902        qeri = ((qeri_aa, qeri_ab), (qeri_ba, qeri_bb))
903    elif spin == 0:
904        qeri = (qeri_aa, qeri_ab)
905    elif spin == 1:
906        qeri = (qeri_bb, qeri_ba)
907
908    log.timer('QMO integral transformation', *cput0)
909
910    return qeri
911
912def _make_qmo_eris_outcore(agf2, eri, coeffs_a, coeffs_b, spin=None):
913    ''' Returns nested tuple of H5 dataset
914
915    spin = None: ((aaaa, aabb), (bbaa, bbbb))
916    spin = 0: (aaaa, aabb)
917    spin = 1: (bbbb, bbaa)
918    '''
919
920    cput0 = (logger.process_clock(), logger.perf_counter())
921    log = logger.Logger(agf2.stdout, agf2.verbose)
922
923    nmo = eri.nmo
924    nmoa, nmob = nmo
925
926    mask = get_frozen_mask(agf2)
927    frozena = np.sum(~mask[0])
928    frozenb = np.sum(~mask[1])
929
930    # npaira, npairb = nmoa*(nmoa+1)//2, nmob*(nmob+1)//2
931    cia, cja, caa = coeffs_a
932    cib, cjb, cab = coeffs_b
933    nia, nja, naa = [x.shape[1] for x in coeffs_a]
934    nib, njb, nab = [x.shape[1] for x in coeffs_b]
935
936    # possible to have incore MO, outcore QMO
937    if getattr(eri, 'feri', None) is None:
938        eri.feri = lib.H5TmpFile()
939    else:
940        for key in ['aa', 'ab', 'ba', 'bb']:
941            if 'qmo/%s'%key in eri.feri:
942                del eri.feri['qmo/%s'%key]
943
944    if spin is None or spin == 0:
945        eri.feri.create_dataset('qmo/aa', (nmoa-frozena, nia, nja, naa), 'f8')
946        eri.feri.create_dataset('qmo/ab', (nmoa-frozena, nia, njb, nab), 'f8')
947
948        blksize = _agf2.get_blksize(agf2.max_memory, (nmoa**3, nmoa*nja*naa),
949                                                (nmoa*nmob**2, nmoa*njb*nab))
950        blksize = min(nmoa, max(BLKMIN, blksize))
951        log.debug1('blksize (uagf2._make_qmo_eris_outcore) = %d', blksize)
952
953        tril2sq = lib.square_mat_in_trilu_indices(nmoa)
954        q1 = 0
955        for p0, p1 in lib.prange(0, nmoa, blksize):
956            if not np.any(mask[0][p0:p1]):
957                # block is fully frozen
958                continue
959
960            inds = np.arange(p0, p1)[mask[0][p0:p1]]
961            q0, q1 = q1, q1 + len(inds)
962            idx = list(np.concatenate(tril2sq[inds]))
963
964            # aa
965            buf = eri.eri_aa[idx] # (blk, nmoa, npaira)
966            buf = buf.reshape((q1-q0)*nmoa, -1) # (blk*nmoa, npaira)
967
968            jasym_aa, nja_aa, cja_aa, sja_aa = ao2mo.incore._conc_mos(cja, caa)
969            buf = ao2mo._ao2mo.nr_e2(buf, cja_aa, sja_aa, 's2kl', 's1')
970            buf = buf.reshape(q1-q0, nmoa, nja, naa)
971
972            buf = lib.einsum('xpja,pi->xija', buf, cia)
973            eri.feri['qmo/aa'][q0:q1] = np.asarray(buf, order='C')
974
975            # ab
976            buf = eri.eri_ab[idx] # (blk, nmoa, npairb)
977            buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmoa, npairb)
978
979            jasym_ab, nja_ab, cja_ab, sja_ab = ao2mo.incore._conc_mos(cjb, cab)
980            buf = ao2mo._ao2mo.nr_e2(buf, cja_ab, sja_ab, 's2kl', 's1')
981            buf = buf.reshape(q1-q0, nmoa, njb, nab)
982
983            buf = lib.einsum('xpja,pi->xija', buf, cia)
984            eri.feri['qmo/ab'][q0:q1] = np.asarray(buf, order='C')
985
986    if spin is None or spin == 1:
987        eri.feri.create_dataset('qmo/ba', (nmob-frozenb, nib, nja, naa), 'f8')
988        eri.feri.create_dataset('qmo/bb', (nmob-frozenb, nib, njb, nab), 'f8')
989
990        max_memory = agf2.max_memory - lib.current_memory()[0]
991        blksize = int((max_memory/8e-6) / max(nmob**3+nmob*njb*nab,
992                                              nmob*nmoa**2*nja*naa))
993        blksize = min(nmob, max(BLKMIN, blksize))
994        log.debug1('blksize (uagf2._make_qmo_eris_outcore) = %d', blksize)
995
996        tril2sq = lib.square_mat_in_trilu_indices(nmob)
997        q1 = 0
998        for p0, p1 in lib.prange(0, nmob, blksize):
999            if not np.any(mask[1][p0:p1]):
1000                # block is fully frozen
1001                continue
1002
1003            inds = np.arange(p0, p1)[mask[1][p0:p1]]
1004            q0, q1 = q1, q1 + len(inds)
1005            idx = list(np.concatenate(tril2sq[inds]))
1006
1007            # ba
1008            buf = eri.eri_ba[idx] # (blk, nmob, npaira)
1009            buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmob, npaira)
1010
1011            jasym_ba, nja_ba, cja_ba, sja_ba = ao2mo.incore._conc_mos(cja, caa)
1012            buf = ao2mo._ao2mo.nr_e2(buf, cja_ba, sja_ba, 's2kl', 's1')
1013            buf = buf.reshape(q1-q0, nmob, nja, naa)
1014
1015            buf = lib.einsum('xpja,pi->xija', buf, cib)
1016            eri.feri['qmo/ba'][q0:q1] = np.asarray(buf, order='C')
1017
1018            # bb
1019            buf = eri.eri_bb[idx] # (blk, nmob, npairb)
1020            buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmob, npairb)
1021
1022            jasym_bb, nja_bb, cja_bb, sja_bb = ao2mo.incore._conc_mos(cjb, cab)
1023            buf = ao2mo._ao2mo.nr_e2(buf, cja_bb, sja_bb, 's2kl', 's1')
1024            buf = buf.reshape(q1-q0, nmob, njb, nab)
1025
1026            buf = lib.einsum('xpja,pi->xija', buf, cib)
1027            eri.feri['qmo/bb'][q0:q1] = np.asarray(buf, order='C')
1028
1029    if spin is None:
1030        qeri = ((eri.feri['qmo/aa'], eri.feri['qmo/ab']),
1031                (eri.feri['qmo/ba'], eri.feri['qmo/bb']))
1032    elif spin == 0:
1033        qeri = (eri.feri['qmo/aa'], eri.feri['qmo/ab'])
1034    elif spin == 1:
1035        qeri = (eri.feri['qmo/bb'], eri.feri['qmo/ba'])
1036
1037    log.timer('QMO integral transformation', *cput0)
1038
1039    return qeri
1040
1041
1042
1043if __name__ == '__main__':
1044    from pyscf import gto, scf, mp
1045
1046    mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='cc-pvdz', charge=-1, spin=1, verbose=3)
1047    uhf = scf.UHF(mol)
1048    uhf.conv_tol = 1e-11
1049    uhf.run()
1050
1051    uagf2 = UAGF2(uhf, frozen=0)
1052
1053    uagf2.run()
1054    uagf2.ipagf2(nroots=5)
1055    uagf2.eaagf2(nroots=5)
1056
1057
1058    uagf2 = uagf2.density_fit()
1059    uagf2.run()
1060