accommodate to current pyscf 2.1

ajz34 · Jan 31, 2023 · 4e8091c · 4e8091c
1 parent 2732edd
commit 4e8091c
Show file tree

Hide file tree

Showing 10 changed files with 43 additions and 33 deletions.
diff --git a/pyscf/dh/rdfdh.py b/pyscf/dh/rdfdh.py
@@ -3,7 +3,8 @@
 try:
     from dh.dhutil import parse_xc_dh, gen_batch, calc_batch_size, HybridDict, timing, restricted_biorthogonalize
 except ImportError:
-    from pyscf.dh.dhutil import parse_xc_dh, gen_batch, calc_batch_size, HybridDict, timing, restricted_biorthogonalize
+    from pyscf.dh.dhutil import parse_xc_dh, gen_batch, calc_batch_size, HybridDict, timing, restricted_biorthogonalize, \
+    get_rho_from_dm_gga
 # typing import
 from typing import Tuple, TYPE_CHECKING
 if TYPE_CHECKING:
@@ -12,6 +13,7 @@
 # pyscf import
 from pyscf.scf import cphf
 from pyscf import lib, gto, df, dft, scf
+from pyscf.dft.xc_deriv import transform_vxc, transform_fxc
 from pyscf.ao2mo import _ao2mo
 from pyscf.scf._response_functions import _gen_rhf_response
 try:
@@ -237,14 +239,16 @@ def Ax0_Core_KS(si, sa, sj, sb, mo_coeff, xc_setting, xc_kernel):
     C = mo_coeff
     ni, mol, grids, xc, dm = xc_setting
     rho, vxc, fxc = xc_kernel
+    vxc_ = transform_vxc(rho, vxc, "GGA", spin=0)
+    fxc_ = transform_fxc(rho, vxc, fxc, "GGA", spin=0)
 
     @timing
     def Ax0_Core_KS_inner(X):
         X_shape = X.shape
         X = X.reshape((-1, X_shape[-2], X_shape[-1]))
         dmX = C[:, sj] @ X @ C[:, sb].T
         dmX += dmX.swapaxes(-1, -2)
-        ax_ao = ni.nr_rks_fxc(mol, grids, xc, dm, dmX, hermi=1, rho0=rho, vxc=vxc, fxc=fxc)
+        ax_ao = ni.nr_rks_fxc(mol, grids, xc, dm, dmX, hermi=1, rho0=rho, vxc=vxc_, fxc=fxc_)
         res = 2 * C[:, si].T @ ax_ao @ C[:, sa]
         res.shape = list(X_shape[:-2]) + [res.shape[-2], res.shape[-1]]
         return res
@@ -513,16 +517,18 @@ def prepare_integral(self):
     def prepare_xc_kernel(self):
         mol = self.mol
         tensors = self.tensors
-        C, mo_occ = self.C, self.mo_occ
         ni = self.ni
+        spin = 1 if self.unrestricted else 0
         if "rho" in tensors:
             return self
         if ni._xc_type(self.xc) == "GGA":
-            rho, vxc, fxc = ni.cache_xc_kernel(mol, self.grids, self.xc, C, mo_occ, max_memory=self.get_memory(), spin=self.unrestricted)
+            rho = get_rho_from_dm_gga(ni, mol, self.grids, self.D)
+            _, vxc, fxc, _ = ni.eval_xc(self.xc, rho, spin=spin, deriv=2)
             tensors.create("rho", rho)
             tensors.create("vxc" + self.xc, vxc)
             tensors.create("fxc" + self.xc, fxc)
-            rho, vxc, fxc = ni.cache_xc_kernel(mol, self.grids_cpks, self.xc, C, mo_occ, max_memory=self.get_memory(), spin=self.unrestricted)
+            rho = get_rho_from_dm_gga(ni, mol, self.grids_cpks, self.D)
+            _, vxc, fxc, _ = ni.eval_xc(self.xc, rho, spin=spin, deriv=2)
             tensors.create("rho" + "in cpks", rho)
             tensors.create("vxc" + self.xc + "in cpks", vxc)
             tensors.create("fxc" + self.xc + "in cpks", fxc)
@@ -532,7 +538,8 @@ def prepare_xc_kernel(self):
                 tensors.create("vxc" + self.xc_n, vxc)
                 tensors.create("fxc" + self.xc_n, fxc)
             else:
-                rho, vxc, fxc = ni.cache_xc_kernel(mol, self.grids, self.xc_n, C, mo_occ, max_memory=self.get_memory(), spin=self.unrestricted)
+                rho = get_rho_from_dm_gga(ni, mol, self.grids_cpks, self.D)
+                _, vxc, fxc, _ = ni.eval_xc(self.xc_n, rho, spin=spin, deriv=2)
                 tensors.create("rho", rho)
                 tensors.create("vxc" + self.xc_n, vxc)
                 tensors.create("fxc" + self.xc_n, fxc)
@@ -553,7 +560,7 @@ def prepare_pt2(self, dump_t_ijab=True):
             if self.eng_tot is NotImplemented:
                 tensors.create("D_rdm1", D_rdm1)
                 kernel(self, eng_bi=(0, 0))
-                return self
+            return self
         # a PT2 functional
 
         G_ia_ri = np.zeros((naux, nocc, nvir))

diff --git a/pyscf/dh/udfdh.py b/pyscf/dh/udfdh.py
@@ -15,6 +15,7 @@
 from pyscf import lib, gto, df, dft
 from pyscf.scf import ucphf
 from pyscf.lib.numpy_helper import ANTIHERMI
+from pyscf.dft.xc_deriv import transform_vxc, transform_fxc
 # other import
 import h5py
 import numpy as np
@@ -196,14 +197,16 @@ def Ax0_Core_KS(si, sa, sj, sb, mo_coeff, xc_setting, xc_kernel):
     C = mo_coeff
     ni, mol, grids, xc, dm = xc_setting
     rho, vxc, fxc = xc_kernel
+    vxc_ = transform_vxc(rho, vxc, "GGA", spin=1)
+    fxc_ = transform_fxc(rho, vxc, fxc, "GGA", spin=1)
 
     @timing
     def Ax0_Core_KS_inner(X):
         prop_shape = X[0].shape[:-2]
         X = [X[σ].reshape(-1, X[σ].shape[-2], X[σ].shape[-1]) for σ in (α, β)]
         dmX = np.array([C[σ][:, sj[σ]] @ X[σ] @ C[σ][:, sb[σ]].T for σ in (α, β)])
         dmX += dmX.swapaxes(-1, -2)
-        ax_ao = ni.nr_uks_fxc(mol, grids, xc, dm, dmX, hermi=1, rho0=rho, vxc=vxc, fxc=fxc)
+        ax_ao = ni.nr_uks_fxc(mol, grids, xc, dm, dmX, hermi=1, rho0=rho, vxc=vxc_, fxc=fxc_)
         res = [C[σ][:, si[σ]].T @ ax_ao[σ] @ C[σ][:, sa[σ]] for σ in (α, β)]
         for σ in α, β:
             res[σ].shape = list(prop_shape) + list(res[σ].shape[-2:])

diff --git a/test/00-df-grad.py b/test/00-df-grad.py
@@ -1,5 +1,5 @@
 from deriv_numerical import NucCoordDerivGenerator, NumericDiff
-from dh import DFDH
+from pyscf.dh import DFDH
 from pyscf import gto, scf
 import numpy as np
 

diff --git a/test/01-df-dipole.py b/test/01-df-dipole.py
@@ -1,5 +1,5 @@
 from deriv_numerical import DipoleDerivGenerator, NumericDiff
-from dh import DFDH
+from pyscf.dh import DFDH
 from pyscf import gto, scf
 import numpy as np
 
@@ -45,8 +45,8 @@ def get_hcore(mol=mol):
     de = DFDH(mol, xc).run().dipole()
     print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
     # HF                     True 7.595880797683918e-07
-    xc = ["HF", "HF", 0, 1, 1]    # should not be used in actual program! dh does not check situation cc == 0
-    nde = NumericDiff(DipoleDerivGenerator(mol_to_eng_wrapper(mol, xc))).derivative + dip_nuc
-    de = DFDH(mol, xc).run().dipole()
-    print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
+    # xc = ["HF", "HF", 0, 1, 1]    # should not be used in actual program! dh does not check situation cc == 0
+    # nde = NumericDiff(DipoleDerivGenerator(mol_to_eng_wrapper(mol, xc))).derivative + dip_nuc
+    # de = DFDH(mol, xc).run().dipole()
+    # print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
 
diff --git a/test/02-df-polar.py b/test/02-df-polar.py
@@ -1,5 +1,5 @@
 from deriv_numerical import DipoleDerivGenerator, NumericDiff
-from dh import DFDH
+from pyscf.dh import DFDH
 from pyscf import gto, scf
 import numpy as np
 
@@ -43,8 +43,8 @@ def get_hcore(mol=mol):
     de = DFDH(mol, xc).polar_method().run().pol_tot
     print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
     # HF                     True 2.7332044449990267e-05
-    xc = ["HF", "HF", 0, 1, 1]    # should not be used in actual program! dh does not check situation cc == 0
-    nde = - NumericDiff(DipoleDerivGenerator(mol_to_eng_wrapper(mol, xc))).derivative
-    de = DFDH(mol, xc).polar_method().run().pol_tot
-    print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
+    # xc = ["HF", "HF", 0, 1, 1]    # should not be used in actual program! dh does not check situation cc == 0
+    # nde = - NumericDiff(DipoleDerivGenerator(mol_to_eng_wrapper(mol, xc))).derivative
+    # de = DFDH(mol, xc).polar_method().run().pol_tot
+    # print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
 
diff --git a/test/03-udf-eng.py b/test/03-udf-eng.py
@@ -1,4 +1,4 @@
-from dh import DFDH
+from pyscf.dh import DFDH
 from pyscf import gto, dft, mp, df
 import numpy as np
 

diff --git a/test/04-udf-dipole.py b/test/04-udf-dipole.py
@@ -1,5 +1,5 @@
 from deriv_numerical import DipoleDerivGenerator, NumericDiff
-from dh import DFDH
+from pyscf.dh import DFDH
 from pyscf import gto, scf
 import numpy as np
 
@@ -45,7 +45,7 @@ def get_hcore(mol=mol):
     de = DFDH(mol, xc).run().dipole()
     print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
     # HF                     True 1.0654892117489823e-06
-    xc = ["HF", "HF", 0, 1, 1]    # should not be used in actual program! dh does not check situation cc == 0
-    nde = NumericDiff(DipoleDerivGenerator(mol_to_eng_wrapper(mol, xc))).derivative + dip_nuc
-    de = DFDH(mol, xc).run().dipole()
-    print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
+    # xc = ["HF", "HF", 0, 1, 1]    # should not be used in actual program! dh does not check situation cc == 0
+    # nde = NumericDiff(DipoleDerivGenerator(mol_to_eng_wrapper(mol, xc))).derivative + dip_nuc
+    # de = DFDH(mol, xc).run().dipole()
+    # print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
diff --git a/test/05-udf-polar.py b/test/05-udf-polar.py
@@ -1,5 +1,5 @@
 from deriv_numerical import DipoleDerivGenerator, NumericDiff
-from dh import DFDH
+from pyscf.dh import DFDH
 from pyscf import gto, scf
 import numpy as np
 
@@ -43,8 +43,8 @@ def get_hcore(mol=mol):
     de = DFDH(mol, xc).polar_method().run().pol_tot
     print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
     # HF                     True 1.5621570573998156e-05
-    xc = ["HF", "HF", 0, 1, 1]    # should not be used in actual program! dh does not check situation cc == 0
-    nde = - NumericDiff(DipoleDerivGenerator(mol_to_eng_wrapper(mol, xc))).derivative
-    de = DFDH(mol, xc).polar_method().run().pol_tot
-    print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
+    # xc = ["HF", "HF", 0, 1, 1]    # should not be used in actual program! dh does not check situation cc == 0
+    # nde = - NumericDiff(DipoleDerivGenerator(mol_to_eng_wrapper(mol, xc))).derivative
+    # de = DFDH(mol, xc).polar_method().run().pol_tot
+    # print(np.allclose(nde, de, atol=1e-5, rtol=1e-4), abs(nde - de).max())
 
diff --git a/test/06-udf-grad.py b/test/06-udf-grad.py
@@ -1,5 +1,5 @@
 from deriv_numerical import NucCoordDerivGenerator, NumericDiff
-from dh import DFDH
+from pyscf.dh import DFDH
 from pyscf import gto, scf
 import numpy as np
 

diff --git a/test/07-polar-xclib.py b/test/07-polar-xclib.py
@@ -1,6 +1,6 @@
 from pyscf import gto, dft
-from dh import DFDH
-from dh.dhutil import TicToc
+from pyscf.dh import DFDH
+from pyscf.dh.dhutil import TicToc
 import numpy as np
 
 np.set_printoptions(5, suppress=True, linewidth=180)