Improve pbc.x2c get_pnucp integrals

pyscf/pbc/df/mdf.py

@@ -90,6 +90,18 @@ def __init__(self, cell, kpts=np.zeros((1,3))):
         self._rsh_df = {}  # Range separated Coulomb DF objects
         self._keys = set(self.__dict__.keys())
 
+    def build(self, j_only=None, with_j3c=True, kpts_band=None):
+        df.GDF.build(self, j_only, with_j3c, kpts_band)
+        cell = self.cell
+        if any(x % 2 == 0 for x in self.mesh[:cell.dimension]):
+            # Even number in mesh can produce planewaves without couterparts
+            # (see np.fft.fftfreq). MDF mesh is typically not enough to capture
+            # all basis. The singular planewaves can break the symmetry in
+            # potential (leads to non-real density) and thereby break the
+            # hermitian of J and K matrices
+            logger.warn(self, 'MDF with even number in mesh may have significant errors')
+        return self
+
     def _make_j3c(self, cell=None, auxcell=None, kptij_lst=None, cderi_file=None):
         if cell is None: cell = self.cell
         if auxcell is None: auxcell = self.auxcell

pyscf/pbc/df/mdf_ao2mo.py

@@ -78,6 +78,9 @@ def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None):
     else:
         assert (out.shape == eri_shape)
 
+    if mydf._cderi is None:
+        mydf.build()
+
     kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts])
     kptis_lst = kptij_lst[:,0]
     kptjs_lst = kptij_lst[:,1]

pyscf/pbc/mp/test/test_kpoint.py

@@ -185,13 +185,12 @@ def test_kmp2_with_cderi(self):
         kmf2.conv_tol = 1e-12
         kmf2.with_df._cderi = kmf.with_df._cderi
         ekpt2 = kmf2.scf()
-        mp = pyscf.pbc.mp.kmp2.KMP2(kmf2).run()        
+        mp = pyscf.pbc.mp.kmp2.KMP2(kmf2).run()
 
-        self.assertAlmostEqual(ekpt2, -1.2053666821021261, 9)
+        self.assertAlmostEqual(ekpt2, -1.2053666821021261, 7)
         self.assertAlmostEqual(mp.e_corr, -6.9881475423322723e-06, 9)
 
 
 if __name__ == '__main__':
     print("Full kpoint test")
     unittest.main()
-

pyscf/pbc/x2c/sfx2c1e.py

@@ -33,9 +33,10 @@
 from pyscf.pbc import gto as pbcgto
 from pyscf.pbc import tools
 from pyscf.pbc.df import aft
-from pyscf.pbc.df import incore
 from pyscf.pbc.df import aft_jk
 from pyscf.pbc.df import ft_ao
+from pyscf.pbc.df import incore
+from pyscf.pbc.df import gdf_builder
 from pyscf.pbc.scf import ghf
 from pyscf import __config__
 
@@ -250,42 +251,57 @@ def get_pnucp(mydf, kpts=None):
     nao = cell.nao_nr()
     nao_pair = nao * (nao+1) // 2
 
-    Gv, Gvbase, kws = cell.get_Gv_weights(mydf.mesh)
-    charge = -cell.atom_charges()
+    eta, mesh, ke_cutoff = gdf_builder._guess_eta(cell, kpts_lst)
+    log.debug1('get_pnucp eta = %s mesh = %s', eta, mesh)
+
+    nuccell = incore._compensate_nuccell(cell, eta)
+    wj = mydf._int_nuc_vloc(nuccell, kpts_lst, 'int3c2e_pvp1')
+    t1 = log.timer_debug1('pnucp pass1: analytic int', *t1)
+
+    charge = -cell.atom_charges() # Apply Koseki effective charge?
+    if cell.dimension == 3:
+        nucbar = sum([z/nuccell.bas_exp(i)[0] for i,z in enumerate(charge)])
+        nucbar *= numpy.pi/cell.vol
+
+        ovlp = cell.pbc_intor('int1e_kin', 1, lib.HERMITIAN, kpts_lst)
+        for k in range(nkpts):
+            s = lib.pack_tril(ovlp[k])
+            # *2 due to the factor 1/2 in T
+            wj[k] -= nucbar*2 * s
+
+    dfbuilder = incore._IntNucBuilder(cell, kpts_lst).build()
+    supmol_ft = dfbuilder.supmol.strip_basis()
+    ft_kern = supmol_ft.gen_ft_kernel('s2', intor='GTO_ft_pdotp',
+                                      return_complex=False, verbose=log)
+
+    Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
+    gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
+    ngrids = Gv.shape[0]
     kpt_allow = numpy.zeros(3)
-    coulG = tools.get_coulG(cell, kpt_allow, mesh=mydf.mesh, Gv=Gv)
+    coulG = tools.get_coulG(cell, kpt_allow, mesh=mesh, Gv=Gv)
     coulG *= kws
-    if mydf.eta == 0:
-        wj = numpy.zeros((nkpts,nao_pair), dtype=numpy.complex128)
-        SI = cell.get_SI(Gv)
-        vG = numpy.einsum('i,ix->x', charge, SI) * coulG
-    else:
-        nuccell = incore._compensate_nuccell(mydf.cell, mydf.eta)
-        wj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst, 'int3c2e_pvp1'))
-        t1 = log.timer_debug1('pnucp pass1: analytic int', *t1)
-
-        aoaux = ft_ao.ft_ao(nuccell, Gv)
-        vG = numpy.einsum('i,xi->x', charge, aoaux) * coulG
-        if cell.dimension == 3:
-            nucbar = sum([z/nuccell.bas_exp(i)[0] for i,z in enumerate(charge)])
-            nucbar *= numpy.pi/cell.vol
-
-            ovlp = cell.pbc_intor('int1e_kin', 1, lib.HERMITIAN, kpts_lst)
-            for k in range(nkpts):
-                s = lib.pack_tril(ovlp[k])
-                # *2 due to the factor 1/2 in T
-                wj[k] -= nucbar*2 * s
-
+    aoaux = ft_ao.ft_ao(nuccell, Gv)
+    vG = numpy.einsum('i,xi->x', charge, aoaux) * coulG
+    vGR = vG.real
+    vGI = vG.imag
     max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
-    for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, kpt_allow, kpts_lst,
-                                       max_memory=max_memory, aosym='s2',
-                                       intor='GTO_ft_pdotp'):
-        for k, aoao in enumerate(aoaoks):
-            if aft_jk.gamma_point(kpts_lst[k]):
-                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real)
-                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag)
-            else:
-                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao)
+    Gblksize = max(16, int(max_memory*1e6/16/nao_pair/nkpts))
+    Gblksize = min(Gblksize, ngrids, 200000)
+    log.debug1('max_memory = %s  Gblksize = %s  ngrids = %s',
+               max_memory, Gblksize, ngrids)
+
+    buf = numpy.empty((2, nkpts, Gblksize, nao_pair))
+    for p0, p1 in lib.prange(0, ngrids, Gblksize):
+        # shape of Gpq (nkpts, nGv, nao_pair)
+        Gpq = ft_kern(Gv[p0:p1], gxyz[p0:p1], Gvbase, kpt_allow, kpts_lst, out=buf)
+        for k, (GpqR, GpqI) in enumerate(zip(*Gpq)):
+            vR  = numpy.einsum('k,kx->x', vGR[p0:p1], GpqR)
+            vR += numpy.einsum('k,kx->x', vGI[p0:p1], GpqI)
+            wj[k] += vR
+            if not aft_jk.gamma_point(kpts_lst[k]):
+                vI  = numpy.einsum('k,kx->x', vGR[p0:p1], GpqI)
+                vI -= numpy.einsum('k,kx->x', vGI[p0:p1], GpqR)
+                wj[k] += vI * 1j
     t1 = log.timer_debug1('contracting pnucp', *t1)
 
     wj_kpts = []

pyscf/pbc/x2c/test/test_x2c.py

@@ -173,11 +173,7 @@ def test_pnucp(self):
 
         mydf = df.AFTDF(cell1)
         dat = sfx2c1e.get_pnucp(mydf)
-        self.assertAlmostEqual(abs(dat-vne_ref).max(), 0, 7)
-
-        mydf.eta = 0
-        dat = sfx2c1e.get_pnucp(mydf)
-        self.assertAlmostEqual(abs(dat-vne_ref).max(), 0, 7)
+        self.assertAlmostEqual(abs(dat-vne_ref).max(), 0, 6)
 
     def test_pvxp(self):
         # w_soc should not depend on eta and the type of DF class
@@ -189,10 +185,6 @@ def test_pvxp(self):
         self.assertAlmostEqual(dat[1].sum(), 0.0 + -0.19650430913542424j, 7)
         self.assertAlmostEqual(dat[2].sum(), 0.0 + 0.25706456053958415j, 7)
 
-        mydf.eta = 0.0
-        dat = x2c1e.get_pbc_pvxp(mydf, kpts[1])
-        self.assertAlmostEqual(abs(dat - ref).max(), 0, 6)
-
         # GDF
         mydf = df.GDF(cell1)
         mydf.kpts = kpts

pyscf/pbc/x2c/x2c1e.py

@@ -38,6 +38,8 @@
 from pyscf.pbc.df import aft
 from pyscf.pbc.df import aft_jk
 from pyscf.pbc.df import ft_ao
+from pyscf.pbc.df import incore
+from pyscf.pbc.df import gdf_builder
 from pyscf.pbc.scf import ghf
 from pyscf.pbc.x2c import sfx2c1e
 from pyscf import __config__
@@ -209,52 +211,62 @@ def get_pbc_pvxp(mydf, kpts=None):
 
     log = logger.Logger(cell.stdout, cell.verbose)
     t1 = (logger.process_clock(), logger.perf_counter())
-    #t1 = (time.clock(), time.time())
 
     nkpts = len(kpts_lst)
     nao = cell.nao_nr()
 
-    Gv, Gvbase, kws = cell.get_Gv_weights(mydf.mesh)
-    charge = -cell.atom_charges() # Apply Koseki effective charge?
+    eta, mesh, ke_cutoff = gdf_builder._guess_eta(cell, kpts_lst)
+    log.debug1('get_pnucp eta = %s mesh = %s', eta, mesh)
+    nuccell = copy.copy(cell)
+    half_sph_norm = .5/numpy.sqrt(numpy.pi)
+    norm = half_sph_norm/mole.gaussian_int(2, eta)
+    chg_env = [eta, norm]
+    ptr_eta = cell._env.size
+    ptr_norm = ptr_eta + 1
+    chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0] for ia in range(cell.natm)]
+    nuccell._atm = cell._atm
+    nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32)
+    nuccell._env = numpy.hstack((cell._env, chg_env))
+
+    soc_mat = mydf._int_nuc_vloc(nuccell, kpts_lst, 'int3c2e_pvxp1_sph',
+                                 aosym='s1', comp=3)
+    soc_mat = soc_mat.reshape(nkpts,3,nao,nao)
+    t1 = log.timer_debug1('pnucp pass1: analytic int', *t1)
+
+    dfbuilder = incore._IntNucBuilder(cell, kpts_lst).build()
+    supmol_ft = dfbuilder.supmol.strip_basis()
+    ft_kern = supmol_ft.gen_ft_kernel('s1', intor='GTO_ft_pxp_sph', comp=3,
+                                      return_complex=False, verbose=log)
+
+    Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
+    gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
+    ngrids = Gv.shape[0]
     kpt_allow = numpy.zeros(3)
-    coulG = tools.get_coulG(cell, kpt_allow, mesh=mydf.mesh, Gv=Gv)
+    coulG = tools.get_coulG(cell, kpt_allow, mesh=mesh, Gv=Gv)
     coulG *= kws
-    if mydf.eta == 0:
-        soc_mat = numpy.zeros((nkpts,3,nao*nao), dtype=numpy.complex128)
-        SI = cell.get_SI(Gv)
-        vG = numpy.einsum('i,ix->x', charge, SI) * coulG
-    else:
-        nuccell = copy.copy(cell)
-        half_sph_norm = .5/numpy.sqrt(numpy.pi)
-        norm = half_sph_norm/mole.gaussian_int(2, mydf.eta)
-        chg_env = [mydf.eta, norm]
-        ptr_eta = cell._env.size
-        ptr_norm = ptr_eta + 1
-        chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0] for ia in range(cell.natm)]
-        nuccell._atm = cell._atm
-        nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32)
-        nuccell._env = numpy.hstack((cell._env, chg_env))
-
-        soc_mat = mydf._int_nuc_vloc(nuccell, kpts_lst, 'int3c2e_pvxp1_sph',
-                                     aosym='s1', comp=3)
-        # Some corrections are needed.
-        soc_mat = numpy.asarray(soc_mat).reshape(nkpts,3,nao**2)
-        t1 = log.timer_debug1('pnucp pass1: analytic int', *t1)
-
-        aoaux = ft_ao.ft_ao(nuccell, Gv)
-        vG = numpy.einsum('i,xi->x', charge, aoaux) * coulG
+    aoaux = ft_ao.ft_ao(nuccell, Gv)
+    charge = -cell.atom_charges() # Apply Koseki effective charge?
+    vG = numpy.einsum('i,xi->x', charge, aoaux) * coulG
+    vGR = vG.real
+    vGI = vG.imag
 
     max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
-    for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, kpt_allow, kpts_lst,
-                                       max_memory=max_memory, aosym='s1',
-                                       intor='GTO_ft_pxp_sph', comp=3):
-        for k, aoao in enumerate(aoaoks):
-            aoao = aoao.reshape(3,-1,nao**2)
-            if aft_jk.gamma_point(kpts_lst[k]):
-                soc_mat[k] += numpy.einsum('k,ckx->cx', vG[p0:p1].real, aoao.real)
-                soc_mat[k] += numpy.einsum('k,ckx->cx', vG[p0:p1].imag, aoao.imag)
-            else:
-                soc_mat[k] += numpy.einsum('k,ckx->cx', vG[p0:p1].conj(), aoao)
+    Gblksize = max(16, int(max_memory*1e6/16/3/nao**2/nkpts))
+    Gblksize = min(Gblksize, ngrids, 200000)
+    log.debug1('max_memory = %s  Gblksize = %s  ngrids = %s',
+               max_memory, Gblksize, ngrids)
+
+    for p0, p1 in lib.prange(0, ngrids, Gblksize):
+        # shape of Gpq (nkpts, nGv, nao_pair)
+        Gpq = ft_kern(Gv[p0:p1], gxyz[p0:p1], Gvbase, kpt_allow, kpts_lst)
+        for k, (GpqR, GpqI) in enumerate(zip(*Gpq)):
+            vR  = numpy.einsum('k,ckpq->cpq', vGR[p0:p1], GpqR)
+            vR += numpy.einsum('k,ckpq->cpq', vGI[p0:p1], GpqI)
+            soc_mat[k] += vR
+            if not aft_jk.gamma_point(kpts_lst[k]):
+                vI  = numpy.einsum('k,ckpq->cpq', vGR[p0:p1], GpqI)
+                vI -= numpy.einsum('k,ckpq->cpq', vGI[p0:p1], GpqR)
+                soc_mat[k] += vI * 1j
     t1 = log.timer_debug1('contracting pnucp', *t1)
 
     soc_mat_kpts = []