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#
16# Author: Xing Zhang <zhangxing.nju@gmail.com>
17#
18
19'''
20density fitting GMP2,  3-center integrals incore.
21'''
22
23import numpy as np
24from pyscf import __config__
25from pyscf import lib
26from pyscf.lib import logger
27from pyscf.ao2mo import _ao2mo
28from pyscf.mp import gmp2, dfmp2
29from pyscf.mp.gmp2 import make_rdm1, make_rdm2
30
31WITH_T2 = getattr(__config__, 'mp_dfgmp2_with_t2', True)
32
33def kernel(mp, mo_energy=None, mo_coeff=None, eris=None, with_t2=WITH_T2,
34           verbose=logger.NOTE):
35    if mo_energy is not None or mo_coeff is not None:
36        # For backward compatibility.  In pyscf-1.4 or earlier, mp.frozen is
37        # not supported when mo_energy or mo_coeff is given.
38        assert(mp.frozen == 0 or mp.frozen is None)
39
40    if eris is None:      eris = mp.ao2mo(mo_coeff)
41    if mo_energy is None: mo_energy = eris.mo_energy
42    if mo_coeff is None:  mo_coeff = eris.mo_coeff
43
44    nocc = mp.nocc
45    nvir = mp.nmo - nocc
46    naux = mp.with_df.get_naoaux()
47    eia = mo_energy[:nocc,None] - mo_energy[None,nocc:]
48
49    if with_t2:
50        t2 = np.empty((nocc,nocc,nvir,nvir), dtype=mo_coeff.dtype)
51    else:
52        t2 = None
53
54    orbspin = eris.orbspin
55    Lova = np.empty((naux, nocc*nvir), dtype=mo_coeff.dtype)
56    if orbspin is None:
57        Lovb = np.empty((naux, nocc*nvir), dtype=mo_coeff.dtype)
58
59    p1 = 0
60    for qova, qovb in mp.loop_ao2mo(mo_coeff, nocc, orbspin):
61        p0, p1 = p1, p1 + qova.shape[0]
62        Lova[p0:p1] = qova
63        if orbspin is None:
64            Lovb[p0:p1] = qovb
65
66    if orbspin is not None:
67        sym_forbid = (orbspin[:nocc,None] != orbspin[nocc:]).flatten()
68        Lova[:,sym_forbid] = 0
69
70    emp2 = 0
71    for i in range(nocc):
72        gi = np.dot(Lova[:,i*nvir:(i+1)*nvir].T, Lova)
73        if orbspin is None:
74            gi += np.dot(Lovb[:,i*nvir:(i+1)*nvir].T, Lovb)
75            gi += np.dot(Lova[:,i*nvir:(i+1)*nvir].T, Lovb)
76            gi += np.dot(Lovb[:,i*nvir:(i+1)*nvir].T, Lova)
77        gi = gi.reshape(nvir,nocc,nvir)
78        gi = gi.transpose(1,0,2) - gi.transpose(1,2,0)
79        t2i = np.conj(gi/lib.direct_sum('jb+a->jba', eia, eia[i]))
80        emp2 += np.einsum('jab,jab', t2i, gi) * .25
81        if with_t2:
82            t2[i] = t2i
83    return emp2.real, t2
84
85
86class DFGMP2(dfmp2.DFMP2):
87    def loop_ao2mo(self, mo_coeff, nocc, orbspin):
88        nao, nmo = mo_coeff.shape
89        complex_orb = mo_coeff.dtype == np.complex128
90        if orbspin is None:
91            moa = mo_coeff[:nao//2]
92            mob = mo_coeff[nao//2:]
93        else:
94            moa = mo_coeff[:nao//2] + mo_coeff[nao//2:]
95
96        if complex_orb:
97            moa_RR = moa.real
98            moa_II = moa.imag
99            moa_RI = np.concatenate((moa.real[:,:nocc], moa.imag[:,nocc:]), axis=1)
100            moa_IR = np.concatenate((moa.imag[:,:nocc], moa.real[:,nocc:]), axis=1)
101            if orbspin is None:
102                mob_RR = mob.real
103                mob_II = mob.imag
104                mob_RI = np.concatenate((mob.real[:,:nocc], mob.imag[:,nocc:]), axis=1)
105                mob_IR = np.concatenate((mob.imag[:,:nocc], mob.real[:,nocc:]), axis=1)
106
107        ijslice = (0, nocc, nocc, nmo)
108        Lova = Lovb = None
109        if complex_orb:
110            Lova_RR = Lova_II = Lova_RI = Lova_IR = None
111            Lovb_RR = Lovb_II = Lovb_RI = Lovb_IR = None
112        with_df = self.with_df
113
114        nvir = nmo - nocc
115        naux = with_df.get_naoaux()
116        mem_now = lib.current_memory()[0]
117        max_memory = max(2000, self.max_memory*.9-mem_now)
118        fac = 1
119        if complex_orb:
120            fac = 6
121        if orbspin is None:
122            fac *= 2
123        blksize = int(min(naux, max(with_df.blockdim,
124                                    (max_memory*1e6/8-nocc*nvir**2*2)/(fac*nocc*nvir))))
125
126        for eri1 in with_df.loop(blksize=blksize):
127            if complex_orb:
128                Lova_RR = _ao2mo.nr_e2(eri1, moa_RR, ijslice, aosym='s2', out=Lova_RR)
129                Lova_II = _ao2mo.nr_e2(eri1, moa_II, ijslice, aosym='s2', out=Lova_II)
130                Lova_RI = _ao2mo.nr_e2(eri1, moa_RI, ijslice, aosym='s2', out=Lova_RI)
131                Lova_IR = _ao2mo.nr_e2(eri1, moa_IR, ijslice, aosym='s2', out=Lova_IR)
132                Lova = Lova_RR + Lova_II + 1j * (Lova_RI - Lova_IR)
133            else:
134                Lova = _ao2mo.nr_e2(eri1, moa, ijslice, aosym='s2', out=Lova)
135            if orbspin is None:
136                if complex_orb:
137                    Lovb_RR = _ao2mo.nr_e2(eri1, mob_RR, ijslice, aosym='s2', out=Lovb_RR)
138                    Lovb_II = _ao2mo.nr_e2(eri1, mob_II, ijslice, aosym='s2', out=Lovb_II)
139                    Lovb_RI = _ao2mo.nr_e2(eri1, mob_RI, ijslice, aosym='s2', out=Lovb_RI)
140                    Lovb_IR = _ao2mo.nr_e2(eri1, mob_IR, ijslice, aosym='s2', out=Lovb_IR)
141                    Lovb = Lovb_RR + Lovb_II + 1j * (Lovb_RI - Lovb_IR)
142                else:
143                    Lovb = _ao2mo.nr_e2(eri1, mob, ijslice, aosym='s2', out=Lovb)
144            yield Lova, Lovb
145
146    def ao2mo(self, mo_coeff=None):
147        eris = gmp2._PhysicistsERIs()
148        # Initialize only the mo_coeff and
149        eris._common_init_(self, mo_coeff)
150        return eris
151
152    def make_rdm1(self, t2=None, ao_repr=False):
153        if t2 is None:
154            t2 = self.t2
155        assert t2 is not None
156        return make_rdm1(self, t2, ao_repr=ao_repr)
157
158    def make_rdm2(self, t2=None, ao_repr=False):
159        if t2 is None:
160            t2 = self.t2
161        assert t2 is not None
162        return make_rdm2(self, t2, ao_repr=ao_repr)
163
164    def init_amps(self, mo_energy=None, mo_coeff=None, eris=None, with_t2=WITH_T2):
165        return kernel(self, mo_energy, mo_coeff, eris, with_t2)
166
167
168if __name__ == "__main__":
169    from pyscf import gto, scf, mp
170    mol = gto.Mole()
171    mol.atom = [
172        ['Li', (0., 0., 0.)],
173        ['H', (1., 0., 0.)]]
174    mol.basis = 'cc-pvdz'
175    mol.build()
176
177    mf = scf.GHF(mol).run()
178    mymp = DFGMP2(mf)
179    mymp.kernel()
180
181    mf = scf.RHF(mol).run()
182    mf = mf.to_ghf()
183    mymp = DFGMP2(mf)
184    mymp.kernel()
185
186    mymp = mp.GMP2(mf).density_fit()
187    mymp.kernel()
188
189    mf = scf.RHF(mol).density_fit().run()
190    mymp = mp.GMP2(mf)
191    mymp.kernel()
192
193    mf = scf.GHF(mol)
194    dm = mf.get_init_guess() + 0j
195    dm[0,:] += .1j
196    dm[:,0] -= .1j
197    mf.kernel(dm0=dm)
198    mymp = DFGMP2(mf)
199    mymp.kernel()
200