NEO-MP2: support multiple quantum nuclei

pyscf/neo/mp2.py

@@ -2,68 +2,78 @@
 # Author: Kurt Brorsen (brorsenk@missouri.edu)
 
 import numpy
-from pyscf import lib, ao2mo, neo
+from pyscf import lib, neo
 from pyscf import mp
 from timeit import default_timer as timer
-from pyscf.neo.ao2mo import *
+from pyscf.neo.ao2mo import ep_ovov
 
 class MP2(lib.StreamObject):
-    def __init__(self, mf):
+    def __init__(self, mf, nuclei_only=False):
 
         self.mf  = mf
-        self.mp_ee = mp.MP2(self.mf.mf_elec)
+        self.mol = mf.mol
+        self.nuclei_only = nuclei_only
+        self.mp_e = mp.MP2(self.mf.mf_elec)
+        self.mp_n = []
+        for i in range(self.mol.nuc_num):
+            self.mp_n.append(mp.MP2(mf.mf_nuc[i]))
+
+
+        self.verbose = self.mol.verbose
+        self.stdout = self.mol.stdout
+        self.max_memory = mf.max_memory
 
         if(self.mf.unrestricted==True):
             raise NotImplementedError('NEO-MP2 is for RHF wave functions only')
 
-        if(len(self.mf.mol.nuc)>1):
-            raise NotImplementedError('NEO-MP2 is for single quantum nuclei')
 
     def kernel(self):
 
-        emp2_ee = self.mp_ee.kernel()[0]
+        emp2_ee = 0
+        if not self.nuclei_only:
+            emp2_ee = self.mp_e.kernel()[0]
 
-        e_nocc = self.mf.mf_elec.mo_coeff[:,self.mf.mf_elec.mo_occ>0].shape[1]
-        e_tot  = self.mf.mf_elec.mo_coeff[0,:].shape[0]
-        e_nvir = e_tot - e_nocc
+        e_nocc = self.mp_e.nocc
+        e_nvir = self.mp_e.nmo - e_nocc
 
-        p_nocc = self.mf.mf_nuc[0].mo_coeff[:,self.mf.mf_nuc[0].mo_occ>0].shape[1]
-        p_tot  = self.mf.mf_nuc[0].mo_coeff[0,:].shape[0]
-        p_nvir = p_tot - p_nocc
+        emp2_ep = 0.0
+        for i in range(self.mol.nuc_num):
+            n_nocc = self.mp_n[i].nocc
+            n_nvir = self.mp_n[i].nmo - n_nocc
 
-        eia = self.mf.mf_elec.mo_energy[:e_nocc,None] - self.mf.mf_elec.mo_energy[None,e_nocc:]
-        ejb = self.mf.mf_nuc[0].mo_energy[:p_nocc,None] - self.mf.mf_nuc[0].mo_energy[None,p_nocc:]
+            eia = self.mf.mf_elec.mo_energy[:e_nocc,None] - self.mf.mf_elec.mo_energy[None,e_nocc:]
+            ejb = self.mf.mf_nuc[i].mo_energy[:n_nocc,None] - self.mf.mf_nuc[i].mo_energy[None,n_nocc:]
 
-        start = timer()
-        eri_ep = neo.ao2mo.ep_ovov(self.mf)
-        finish = timer()
+            start = timer()
+            eri_ep = ep_ovov(self.mf, i)
+            finish = timer()
 
-        print('time for ep ao2mo transform = ',finish-start)
+            print('time for ep ao2mo transform = ', finish-start)
 
-        start = timer()
-        emp2_ep = 0.0
-        for i in range(e_nocc):
-            gi = numpy.asarray(eri_ep[i*e_nvir:(i+1)*e_nvir])
-            gi = gi.reshape(e_nvir,p_nocc,p_nvir)
-            t2i = gi/lib.direct_sum('a+jb->ajb', eia[i], ejb)
-            emp2_ep += numpy.einsum('ajb,ajb', t2i, gi)
+            start = timer()
+            
+            for i in range(e_nocc):
+                gi = numpy.asarray(eri_ep[i*e_nvir:(i+1)*e_nvir])
+                gi = gi.reshape(e_nvir, n_nocc, n_nvir)
+                t2i = gi/lib.direct_sum('a+jb->ajb', eia[i], ejb)
+                emp2_ep += numpy.einsum('ajb,ajb', t2i, gi)
 
-        emp2_ep = 2.0 * emp2_ep
-        end = timer()
-        print('time for python mp2',end-start)
+            emp2_ep += 2.0 * emp2_ep
+            end = timer()
+            print('time for python mp2', end-start)
 
         return emp2_ee, emp2_ep
 
 
 if __name__ == '__main__':
 
     mol = neo.Mole()
-    mol.build(atom='''H 0 0 0; F 0 0 1.15; F 0 0 -1.15''', basis='ccpvdz', quantum_nuc=[0], charge=-1)
+    mol.build(atom='''O 0 0 0; H 0 0.7 0.7; H 0 0.7 -0.7''', basis='ccpvdz')
     mf = neo.HF(mol)
     energy = mf.scf()
 
     emp2_ee, emp2_ep = MP2(mf).kernel()
 
-    print('emp2_ee = ',emp2_ee)
-    print('emp2_ep = ',emp2_ep)
-    print('total neo-mp2 = ',energy+emp2_ee+emp2_ep)
+    print('emp2_ee = ', emp2_ee)
+    print('emp2_ep = ', emp2_ep)
+    print('total neo-mp2 = ', energy + emp2_ee + emp2_ep)