Keep PySCF dispersion state consistent across molecule resets

python/dftd3/pyscf.py

@@ -455,7 +455,8 @@ def energy_nuc(self) -> float:
             return enuc
 
         def reset(self, mol: Optional[Union[gto.Mole, pbc.gto.Cell]] = None) -> "DFTD3":
-            self.with_dftd3.reset(mol)
+            if mol is not None and self.with_dftd3:
+                self.with_dftd3.reset(mol)
             return mf.__class__.reset(self, mol)
 
         def nuc_grad_method(self) -> GradientsBase:
@@ -601,6 +602,14 @@ def d3_hess(scf_hess, **kwargs):
         scf_hess.base = d3_energy(scf_hess.base, **kwargs)
 
     class DFTD3Hess(_DFTD3Hess, scf_hess.__class__):
+        def reset(
+            self, mol: Optional[Union[gto.Mole, pbc.gto.Cell]] = None
+        ) -> "DFTD3Hess":
+            if mol is not None:
+                self.mol = mol
+            self.base.reset(mol)
+            return self
+
         def hess_nuc(
             self,
             mol: Optional[Union[gto.Mole, pbc.gto.Cell]] = None,

python/dftd3/test_pyscf.py

@@ -28,6 +28,57 @@
     pyscf = None
 
 
+@pytest.mark.skipif(pyscf is None, reason="requires pyscf")
+@pytest.mark.parametrize("wrapper", ["energy", "grad", "hess"])
+def test_d3_reset(wrapper: str) -> None:
+    mol = gto.M(atom="He 0 0 0; He 0 0 2", verbose=0)
+    mf = scf.RHF(mol)
+    if wrapper == "energy":
+        obj = disp.d3_energy(mf)
+    elif wrapper == "grad":
+        obj = disp.d3_grad(mf.nuc_grad_method())
+    else:
+        obj = disp.d3_hess(mf.Hessian())
+    base = obj if wrapper == "energy" else obj.base
+
+    assert obj.reset() is obj
+    assert base.with_dftd3.mol is mol
+
+    new_mol = mol.set_geom_("He 0 0 0; He 0 0 3", inplace=False)
+    assert obj.reset(new_mol) is obj
+    assert obj.mol is new_mol
+    assert base.mol is new_mol
+    assert base.with_dftd3.mol is new_mol
+    assert base.with_dftd3.kernel()[0] == approx(
+        disp.DFTD3Dispersion(new_mol).kernel()[0], abs=1.0e-14
+    )
+
+    assert obj.reset() is obj
+    assert base.with_dftd3.mol is new_mol
+
+
+@pytest.mark.skipif(pyscf is None, reason="requires pyscf")
+@pytest.mark.parametrize("gradient", [False, True])
+def test_d3_scanner_reset(gradient: bool) -> None:
+    mol = gto.M(atom="He 0 0 0; He 0 0 2", verbose=0)
+    mf = disp.d3_energy(scf.RHF(mol))
+    scanner = (mf.nuc_grad_method() if gradient else mf).as_scanner()
+
+    for distance in [3, 4]:
+        new_mol = mol.set_geom_(f"He 0 0 0; He 0 0 {distance}", inplace=False)
+        result = scanner(new_mol)
+        base = scanner.base if gradient else scanner
+        assert base.with_dftd3.mol is new_mol
+
+        reference = disp.d3_energy(scf.RHF(new_mol)).run()
+        energy = result[0] if gradient else result
+        assert energy == approx(reference.e_tot, abs=1.0e-10)
+        if gradient:
+            assert result[1] == approx(
+                reference.nuc_grad_method().kernel(), abs=1.0e-10
+            )
+
+
 @pytest.mark.skipif(pyscf is None, reason="requires pyscf")
 def test_hessian_layout() -> None:
     """The dispersion hessian must use the pyscf (natm, natm, 3, 3) layout"""