Fix passing atomic numbers from PySCF to D4 interface (dftd4#248)

python/dftd4/pyscf.py

@@ -26,11 +26,11 @@
 except ModuleNotFoundError:
     raise ModuleNotFoundError("This submodule requires pyscf installed")
 
-import numpy as np
 from typing import Tuple
 
-from .interface import DispersionModel, DampingParam
+import numpy as np
 
+from .interface import DampingParam, DispersionModel
 
 GradientsBase = getattr(rhf_grad, "GradientsBase", rhf_grad.Gradients)
 
@@ -101,7 +101,7 @@ def kernel(self) -> Tuple[float, np.ndarray]:
         mol = self.mol
 
         disp = DispersionModel(
-            mol.atom_charges(),
+            np.asarray([gto.charge(sym) for sym in mol.elements]),
             mol.atom_coords(),
             mol.charge,
         )
@@ -177,8 +177,8 @@ def energy(mf):
     -110.917424528592
     """
 
-    from pyscf.scf import hf
     from pyscf.mcscf import casci
+    from pyscf.scf import hf
 
     if not isinstance(mf, (hf.SCF, casci.CASCI)):
         raise TypeError("mf must be an instance of SCF or CASCI")

python/dftd4/test_pyscf.py

@@ -14,20 +14,23 @@
 # You should have received a copy of the Lesser GNU General Public License
 # along with dftd4.  If not, see <https://www.gnu.org/licenses/>.
 
+from typing import Union
+
 import numpy as np
 import pytest
 from pytest import approx
 
 try:
+    import dftd4.pyscf as disp
     import pyscf
     from pyscf import gto, scf
-    import dftd4.pyscf as disp
 except ModuleNotFoundError:
     pyscf = None
 
 
 @pytest.mark.skipif(pyscf is None, reason="requires pyscf")
-def test_energy_r2scan_d4():
+@pytest.mark.parametrize("ecp", [None, "ccecp"])
+def test_energy_r2scan_d4(ecp: Union[None, str]) -> None:
     mol = gto.M(
         atom="""
              C   -0.755422531  -0.796459123  -1.023590391
@@ -48,15 +51,16 @@ def test_energy_r2scan_d4():
              H   -1.684433029   7.115457372  -0.460265708
              H   -2.683867206   5.816530502  -1.115183775
              H   -3.365330613   7.451201412  -0.890098894
-             """
+             """,
+        ecp=ecp,
     )
 
     d4 = disp.DFTD4Dispersion(mol, xc="r2SCAN")
     assert d4.kernel()[0] == approx(-0.005001101058518388, abs=1.0e-7)
 
 
 @pytest.mark.skipif(pyscf is None, reason="requires pyscf")
-def test_gradient_b97m_d4():
+def test_gradient_b97m_d4() -> None:
     mol = gto.M(
         atom="""
              H    0.002144194   0.361043475   0.029799709
@@ -105,7 +109,7 @@ def test_gradient_b97m_d4():
 
 
 @pytest.mark.skipif(pyscf is None, reason="requires pyscf")
-def test_energy_hf():
+def test_energy_hf() -> None:
     mol = gto.M(
         atom="""
              N  -1.57871857  -0.04661102   0.00000000
@@ -124,7 +128,7 @@ def test_energy_hf():
 
 
 @pytest.mark.skipif(pyscf is None, reason="requires pyscf")
-def test_gradient_hf():
+def test_gradient_hf() -> None:
     mol = gto.M(
         atom="""
              O  -1.65542061  -0.12330038   0.00000000