API update to allow pyscf.dh as pyscf extension

examples/00-dh-energy.py

@@ -1,5 +1,8 @@
+# for development, import dh
+# for general user, import pyscf.dh
+
 from pyscf import gto, dft, mp
-from dh import RDFDH
+from dh import DFDH
 
 
 def get_energy_xDH(mol: gto.Mole, xc_scf: str, c_pt2: float, xc_dh: str or None=None):
@@ -22,19 +25,19 @@ def get_energy_xDH(mol: gto.Mole, xc_scf: str, c_pt2: float, xc_dh: str or None=
     #          PySCF   -76.2910535257
     #          RDFDH   -76.2910470614
     print(get_energy_xDH(mol, "B3LYPg", 0.3211, "0.8033*HF - 0.0140*LDA + 0.2107*B88, 0.6789*LYP"))
-    print(RDFDH(mol, xc="XYG3").run().e_tot)
+    print(DFDH(mol, xc="XYG3").run().e_tot)
     # B2PLYP   QChem   -76.29075997
     #          PySCF   -76.29075969
     #          RDFDH   -76.29076324
     print(get_energy_xDH(mol, "0.53*HF + 0.47*B88, 0.73*LYP", 0.27))
-    print(RDFDH(mol, xc="B2PLYP").run().e_tot)
+    print(DFDH(mol, xc="B2PLYP").run().e_tot)
     # XYGJ-OS  QChem   -76.1460262831
     #          RDFDH   -76.1460317123
-    print(RDFDH(mol, xc="XYGJ-OS").run().e_tot)
+    print(DFDH(mol, xc="XYGJ-OS").run().e_tot)
     # MP2      PySCF   -76.1108060780191
     #          RDFDH   -76.1107838767765
     print(mp.MP2(mol).run().e_tot)
-    print(RDFDH(mol, xc="MP2").run().e_tot)
+    print(DFDH(mol, xc="MP2").run().e_tot)
     # Details:  QChem uses 99, 590 grid; i.e. XC_GRID 000099000590
     #                 no density fitting used in QChem for XYG3, B2PLYP, MP2
     #                 LT-SOS-RI-MP2 for XYGJ-OS

examples/01-dh-grad.py

@@ -1,5 +1,8 @@
+# for development, import dh
+# for general user, import pyscf.dh
+
 from pyscf import gto
-from dh import RDFDH
+from dh import DFDH
 import numpy as np
 
 np.set_printoptions(5, suppress=True, linewidth=180)
@@ -11,25 +14,25 @@
     #  [-0.16752  0.02649  0.02154]
     #  [ 0.02384 -0.03202  0.01631]
     #  [ 0.01762  0.01483  0.04038]]
-    print(RDFDH(mol, xc="XYG3").run().e_tot)
-    mf = RDFDH(mol, xc="XYG3").nuc_grad_method().run()
+    print(DFDH(mol, xc="XYG3").run().e_tot)
+    mf = DFDH(mol, xc="XYG3").nuc_grad_method().run()
     print(mf.e_tot)
     print(mf.de)
     # B2PLYP
     # [[ 0.13023 -0.00424 -0.07277]
     #  [-0.17088  0.02598  0.0211 ]
     #  [ 0.02338 -0.03631  0.01602]
     #  [ 0.01726  0.01457  0.03565]]
-    print(RDFDH(mol, xc="B2PLYP").run().e_tot)
-    mf = RDFDH(mol, xc="B2PLYP").nuc_grad_method().run()
+    print(DFDH(mol, xc="B2PLYP").run().e_tot)
+    mf = DFDH(mol, xc="B2PLYP").nuc_grad_method().run()
     print(mf.e_tot)
     print(mf.de)
     # MP2
     # [[ 0.1345  -0.00321 -0.07428]
     #  [-0.17439  0.02557  0.02063]
     #  [ 0.02301 -0.03636  0.01541]
     #  [ 0.01688  0.01401  0.03824]]
-    print(RDFDH(mol, xc="MP2").run().e_tot)
-    mf = RDFDH(mol, xc="MP2").nuc_grad_method().run()
+    print(DFDH(mol, xc="MP2").run().e_tot)
+    mf = DFDH(mol, xc="MP2").nuc_grad_method().run()
     print(mf.e_tot)
     print(mf.de)

examples/02-dh-polar.py

@@ -1,5 +1,8 @@
+# for development, import dh
+# for general user, import pyscf.dh
+
 from pyscf import gto
-from dh import RDFDH
+from dh import DFDH
 import numpy as np
 
 np.set_printoptions(5, suppress=True, linewidth=180)
@@ -12,22 +15,22 @@
     # [[ 7.26795 -0.15045 -0.21885]
     #  [-0.15045  8.75195 -0.367  ]
     #  [-0.21885 -0.367   10.47337]]
-    mf = RDFDH(mol, xc="XYG3").polar_method().run()
+    mf = DFDH(mol, xc="XYG3").polar_method().run()
     print(mf.dipole())
     print(mf.pol_tot)
     # B2PLYP
     # [0.50747 0.48579 0.45491]
     # [[ 7.38219 -0.15075 -0.21973]
     #  [-0.15075  8.86653 -0.36795]
     #  [-0.21973 -0.36795 10.5606 ]]
-    mf = RDFDH(mol, xc="B2PLYP").polar_method().run()
+    mf = DFDH(mol, xc="B2PLYP").polar_method().run()
     print(mf.dipole())
     print(mf.pol_tot)
     # MP2
     # [0.51397 0.49493 0.46905]
     # [[ 7.28733 -0.14367 -0.2035 ]
     #  [-0.14367  8.72405 -0.33421]
     #  [-0.2035  -0.33421 10.34546]]
-    mf = RDFDH(mol, xc="MP2").polar_method().run()
+    mf = DFDH(mol, xc="MP2").polar_method().run()
     print(mf.dipole())
     print(mf.pol_tot)

examples/04-opt.py

@@ -1,3 +1,6 @@
+# for development, import dh
+# for general user, import pyscf.dh
+
 from pyscf import gto, lib
 from dh import DFDH
 from pyscf.geomopt.geometric_solver import optimize

pyscf/dh/__init__.py

@@ -1,12 +1,11 @@
-from dh.rdfdh import RDFDH
-from dh.udfdh import UDFDH
+from . import rdfdh, udfdh, dhutil
 from pyscf import gto
 
 
 def DFDH(mol: gto.Mole, *args, **kwargs):
     if mol.spin != 0:
-        return UDFDH(mol, *args, **kwargs)
+        return udfdh.UDFDH(mol, *args, **kwargs)
     else:
-        return RDFDH(mol, *args, **kwargs)
+        return rdfdh.RDFDH(mol, *args, **kwargs)
 
 

pyscf/dh/dhutil.py

@@ -226,7 +226,7 @@ def restricted_biorthogonalize(t_ijab, cc, c_os, c_ss):
 def hermi_sum_last2dim(tsr, inplace=True, hermi=HERMITIAN):
     # shameless call lib.hermi_sum, just for a tensor wrapper
     tsr_shape = tsr.shape
-    tsr.shape = (-1, tsr.shape[-1], tsr.shape[-2])
+    tsr = tsr.reshape(-1, tsr.shape[-2], tsr.shape[-1])
     res = lib.hermi_sum(tsr, axes=(0, 2, 1), hermi=hermi, inplace=inplace)
     tsr.shape = tsr_shape
     res.shape = tsr_shape

pyscf/dh/grad/rdfdh.py

@@ -1,11 +1,16 @@
 from __future__ import annotations
-
-from pyscf.dft.numint import _dot_ao_dm, _contract_rho
-
-from dh import RDFDH
-from dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, timing, as_scanner_grad
+# dh import
+try:
+    from dh.rdfdh import RDFDH
+    from dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, timing, as_scanner_grad
+except ImportError:
+    from pyscf.dh.rdfdh import RDFDH
+    from pyscf.dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, timing, as_scanner_grad
+# pyscf import
 from pyscf import gto, lib, df
+from pyscf.dft.numint import _dot_ao_dm, _contract_rho
 from pyscf.df.grad.rhf import _int3c_wrapper as int3c_wrapper
+# other import
 import numpy as np
 
 einsum = lib.einsum

pyscf/dh/grad/udfdh.py

@@ -1,11 +1,19 @@
 from __future__ import annotations
-
-from dh import UDFDH
-from dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, tot_size, timing
-from dh.grad.rdfdh import contract_multiple_rho, get_H_1_ao, get_S_1_ao, generator_L_1
-import dh.grad.rdfdh
+# dh import
+try:
+    from dh.udfdh import UDFDH
+    from dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, tot_size, timing
+    from dh.grad.rdfdh import get_H_1_ao, get_S_1_ao, generator_L_1
+    from dh.grad.rdfdh import Gradients as RGradients
+except ImportError:
+    from pyscf.dh.udfdh import UDFDH
+    from pyscf.dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, tot_size, timing
+    from pyscf.dh.grad.rdfdh import get_H_1_ao, get_S_1_ao, generator_L_1
+    from pyscf.dh.grad.rdfdh import Gradients as RGradients
+# pyscf import
 from pyscf import gto, lib, df
 from pyscf.df.grad.rhf import _int3c_wrapper as int3c_wrapper
+# other import
 import numpy as np
 import itertools
 
@@ -93,7 +101,7 @@ def get_gradient_jk(dfobj: df.DF, C, D, D_r, Y_mo, cx, cx_n, max_memory=2000):
     return grad_contrib
 
 
-class Gradients(UDFDH, dh.grad.rdfdh.Gradients):
+class Gradients(UDFDH, RGradients):
 
     def __init__(self, mol: gto.Mole, *args, skip_construct=False, **kwargs):
         if not skip_construct:

pyscf/dh/polar/rdfdh.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
-
-from dh import RDFDH
-from dh.dhutil import gen_batch, get_rho_from_dm_gga, restricted_biorthogonalize, hermi_sum_last2dim
+# dh import
+try:
+    from dh.rdfdh import RDFDH
+    from dh.dhutil import gen_batch, get_rho_from_dm_gga, restricted_biorthogonalize, hermi_sum_last2dim
+except ImportError:
+    from pyscf.dh.rdfdh import RDFDH
+    from pyscf.dh.dhutil import gen_batch, get_rho_from_dm_gga, restricted_biorthogonalize, hermi_sum_last2dim
 from pyscf import gto, lib, dft
 import numpy as np
 

pyscf/dh/polar/udfdh.py

@@ -1,12 +1,18 @@
 from __future__ import annotations
-
+# dh import
+try:
+    from dh.udfdh import UDFDH
+    from dh.polar.rdfdh import Polar as RPolar
+    from dh.dhutil import gen_batch, get_rho_from_dm_gga, tot_size, hermi_sum_last2dim
+except ImportError:
+    from pyscf.dh.udfdh import UDFDH
+    from pyscf.dh.polar.rdfdh import Polar as RPolar
+    from pyscf.dh.dhutil import gen_batch, get_rho_from_dm_gga, tot_size, hermi_sum_last2dim
+# pyscf import
+from pyscf import gto, lib, dft
 from pyscf.dft.numint import _scale_ao
 from pyscf.lib.numpy_helper import ANTIHERMI
-
-from dh import UDFDH
-import dh.polar.rdfdh
-from dh.dhutil import gen_batch, get_rho_from_dm_gga, tot_size, hermi_sum_last2dim
-from pyscf import gto, lib, dft
+# other import
 import numpy as np
 import itertools
 
@@ -107,7 +113,7 @@ def _uks_gga_wv2_inner(rho0, rho1, rho2):
     return _uks_gga_wv2_inner
 
 
-class Polar(UDFDH, dh.polar.rdfdh.Polar):
+class Polar(UDFDH, RPolar):
 
     def __init__(self, mol: gto.Mole, *args, skip_construct=False, **kwargs):
         if not skip_construct:

pyscf/dh/rdfdh.py

@@ -1,18 +1,22 @@
 from __future__ import annotations
-
-import os
-import pickle
+# dh import
+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
+# typing import
 from typing import Tuple, TYPE_CHECKING
 if TYPE_CHECKING:
     from dh.grad.rdfdh import Gradients
     from dh.polar.rdfdh import Polar
-
+# pyscf import
 from pyscf.scf import cphf
-
-from dh.dhutil import parse_xc_dh, gen_batch, calc_batch_size, HybridDict, timing, restricted_biorthogonalize
 from pyscf import lib, gto, df, dft, scf
 from pyscf.ao2mo import _ao2mo
 from pyscf.scf._response_functions import _gen_rhf_response
+# other import
+import os
+import pickle
 import numpy as np
 
 einsum = lib.einsum
@@ -630,18 +634,23 @@ def load_intermediates(self, dir_path="scratch", rerun_scf=False):
     # to avoid cyclic imports in typing https://stackoverflow.com/questions/39740632/
 
     def nuc_grad_method(self) -> Gradients:
-        from dh.grad.rdfdh import Gradients
+        try:
+            from dh.grad.rdfdh import Gradients
+        except ImportError:
+            from pyscf.dh.grad.rdfdh import Gradients
         self.__class__ = Gradients
         Gradients.__init__(self, self.mol, skip_construct=True)
         return self
 
     def polar_method(self) -> Polar:
-        from dh.polar.rdfdh import Polar
+        try:
+            from dh.polar.rdfdh import Polar
+        except ImportError:
+            from pyscf.dh.polar.rdfdh import Polar
         self.__class__ = Polar
         Polar.__init__(self, self.mol, skip_construct=True)
         return self
 
-
     # endregion first derivative related in class
 
     energy_elec_nc = energy_elec_nc

pyscf/dh/udfdh.py

@@ -1,26 +1,27 @@
 from __future__ import annotations
-
+# dh import
+try:
+    from dh.dhutil import gen_batch, calc_batch_size, timing, tot_size, hermi_sum_last2dim
+    from dh.rdfdh import get_cderi_mo, kernel, RDFDH
+except ImportError:
+    from pyscf.dh.dhutil import gen_batch, calc_batch_size, timing, tot_size, hermi_sum_last2dim
+    from pyscf.dh.rdfdh import get_cderi_mo, kernel, RDFDH
+# typing import
 from typing import Tuple, TYPE_CHECKING
-
-from pyscf.lib.numpy_helper import ANTIHERMI
-
 if TYPE_CHECKING:
     from dh.grad.udfdh import Gradients
     from dh.polar.udfdh import Polar
-
+# pyscf import
+from pyscf import lib, gto, df, dft
 from pyscf.scf import ucphf
+from pyscf.lib.numpy_helper import ANTIHERMI
+# other import
 import h5py
-from dh import RDFDH
-from dh.dhutil import gen_batch, calc_batch_size, timing, tot_size, hermi_sum_last2dim
-from pyscf import lib, gto, df, dft
 import numpy as np
 
-from dh.rdfdh import get_cderi_mo, kernel
-
 einsum = lib.einsum
 α, β = 0, 1
 αα, αβ, ββ = 0, 1, 2
-
 ndarray = np.ndarray or h5py.Dataset
 
 
@@ -485,13 +486,19 @@ def dipole(self):
     # to avoid cyclic imports in typing https://stackoverflow.com/questions/39740632/
 
     def nuc_grad_method(self) -> Gradients:
-        from dh.grad.udfdh import Gradients
+        try:
+            from dh.grad.udfdh import Gradients
+        except ImportError:
+            from pyscf.dh.grad.udfdh import Gradients
         self.__class__ = Gradients
         Gradients.__init__(self, self.mol, skip_construct=True)
         return self
 
     def polar_method(self) -> Polar:
-        from dh.polar.udfdh import Polar
+        try:
+            from dh.polar.udfdh import Polar
+        except ImportError:
+            from pyscf.dh.polar.udfdh import Polar
         self.__class__ = Polar
         Polar.__init__(self, self.mol, skip_construct=True)
         return self