fix pbc df with KPoints input

Author: fishjojo

pyscf/pbc/df/df.py

@@ -53,6 +53,7 @@
 from pyscf.pbc.df import df_ao2mo
 from pyscf.pbc.df.aft import estimate_eta, _check_kpts
 from pyscf.pbc.df.df_jk import zdotCN
+from pyscf.pbc.lib.kpts import KPoints
 from pyscf.pbc.lib.kpts_helper import (is_zero, gamma_point, member, unique,
                                        KPT_DIFF_TOL)
 from pyscf.pbc.df.gdf_builder import libpbc, _CCGDFBuilder, _CCNucBuilder
@@ -146,6 +147,8 @@ def __init__(self, cell, kpts=numpy.zeros((1,3))):
         self.verbose = cell.verbose
         self.max_memory = cell.max_memory
 
+        if isinstance(kpts, KPoints):
+            kpts = kpts.kpts
         self.kpts = kpts  # default is gamma point
         self.kpts_band = None
         self._auxbasis = None

pyscf/pbc/df/df_jk.py

@@ -28,7 +28,6 @@
 from pyscf import lib
 from pyscf.lib import logger, zdotNN, zdotCN, zdotNC
 from pyscf.pbc import tools
-from pyscf.pbc.lib.kpts import KPoints
 from pyscf.pbc.lib.kpts_helper import is_zero, gamma_point, member, get_kconserv_ria
 from pyscf import __config__
 
@@ -53,8 +52,6 @@ def density_fit(mf, auxbasis=None, mesh=None, with_df=None):
         else:
             kpts = numpy.reshape(mf.kpt, (1,3))
 
-        if isinstance(kpts, KPoints):
-            kpts = kpts.kpts
         with_df = df.DF(mf.cell, kpts)
         with_df.max_memory = mf.max_memory
         with_df.stdout = mf.stdout

pyscf/pbc/df/fft.py

@@ -30,6 +30,7 @@
 from pyscf.pbc.df import fft_jk
 from pyscf.pbc.df import aft
 from pyscf.pbc.df.aft import _check_kpts
+from pyscf.pbc.lib.kpts import KPoints
 from pyscf.pbc.lib.kpts_helper import is_zero
 from pyscf import __config__
 
@@ -168,6 +169,8 @@ def __init__(self, cell, kpts=numpy.zeros((1,3))):
         self.verbose = cell.verbose
         self.max_memory = cell.max_memory
 
+        if isinstance(kpts, KPoints):
+            kpts = kpts.kpts
         self.kpts = kpts
 
         self.grids = gen_grid.UniformGrids(cell)

pyscf/pbc/df/mdf.py

@@ -38,6 +38,7 @@
 from pyscf.pbc.df.gdf_builder import _CCGDFBuilder
 from pyscf.pbc.df.rsdf_builder import _RSGDFBuilder
 from pyscf.pbc.df.incore import libpbc, make_auxcell
+from pyscf.pbc.lib.kpts import KPoints
 from pyscf.pbc.lib.kpts_helper import is_zero, member, unique
 from pyscf.pbc.df import mdf_jk
 from pyscf.pbc.df import mdf_ao2mo
@@ -55,6 +56,8 @@ def __init__(self, cell, kpts=np.zeros((1,3))):
         self.verbose = cell.verbose
         self.max_memory = cell.max_memory
 
+        if isinstance(kpts, KPoints):
+            kpts = kpts.kpts
         self.kpts = kpts  # default is gamma point
         self.kpts_band = None
         self._auxbasis = None

pyscf/pbc/df/mdf_jk.py

@@ -26,7 +26,6 @@
 from pyscf.lib import logger
 from pyscf.pbc.df import df_jk
 from pyscf.pbc.df import aft_jk
-from pyscf.pbc.lib.kpts import KPoints
 
 #
 # Divide the Coulomb potential to two parts.  Computing short range part in
@@ -52,8 +51,6 @@ def density_fit(mf, auxbasis=None, mesh=None, with_df=None):
         else:
             kpts = numpy.reshape(mf.kpt, (1,3))
 
-        if isinstance(kpts, KPoints):
-            kpts = kpts.kpts
         with_df = mdf.MDF(mf.cell, kpts)
         with_df.max_memory = mf.max_memory
         with_df.stdout = mf.stdout

pyscf/pbc/df/test/test_df.py

@@ -20,7 +20,7 @@
 from pyscf import ao2mo, gto
 from pyscf.pbc import gto as pgto
 from pyscf.pbc import scf as pscf
-from pyscf.pbc.df import df, aug_etb, FFTDF
+from pyscf.pbc.df import df, aug_etb, FFTDF, mdf
 from pyscf.pbc.df import gdf_builder
 #from mpi4pyscf.pbc.df import df
 pyscf.pbc.DEBUG = False
@@ -250,6 +250,21 @@ def test_cell_with_cart(self):
         eri1 = df.GDF(cell).set(auxbasis=aug_etb(cell)).get_eri()
         self.assertAlmostEqual(abs(eri1-eri0).max(), 0, 2)
 
+    def test_kpoints_input(sef):
+        cell.space_group_symmetry = True
+        cell.build()
+        kpts = cell.make_kpts([2,2,2],
+                              space_group_symmetry=True,
+                              time_reversal_symmetry=True)
+
+        mydf = df.GDF(cell, kpts=kpts)
+        assert mydf.kpts.shape == (8,3)
+
+        mydf = FFTDF(cell, kpts=kpts)
+        assert mydf.kpts.shape == (8,3)
+
+        mydf = mdf.MDF(cell, kpts=kpts)
+        assert mydf.kpts.shape == (8,3)
 
 if __name__ == '__main__':
     print("Full Tests for df")