Add pass-through arguments and allow setting parameters in PySCF wrapper (#44)

Author: awvwgk

python/dftd3/pyscf.py

@@ -21,12 +21,13 @@
 """
 
 try:
-    from pyscf import lib, gto
+    from pyscf import gto, lib, mcscf, scf
+    from pyscf.grad import rhf as rhf_grad
 except ModuleNotFoundError:
     raise ModuleNotFoundError("This submodule requires pyscf installed")
 
 import numpy as np
-from typing import Tuple
+from typing import Dict, Optional, Tuple
 
 from .interface import (
     DispersionModel,
@@ -66,6 +67,52 @@ class DFTD3Dispersion(lib.StreamObject):
     ``"d3op"``
         Optimized power damping function
 
+    Custom parameters can be provided with the `param` dictionary.
+    The `param` dict contains the damping parameters, at least s8, a1 and a2
+    must be provided for rational damping, while s8 and rs6 are required in case
+    of zero damping.
+
+    Parameters for (modified) rational damping are:
+
+    ======================== =========== ============================================
+    Tweakable parameter      Default     Description
+    ======================== =========== ============================================
+    s6                       1.0         Scaling of the dipole-dipole dispersion
+    s8                       None        Scaling of the dipole-quadrupole dispersion
+    s9                       1.0         Scaling of the three-body dispersion energy
+    a1                       None        Scaling of the critical radii
+    a2                       None        Offset of the critical radii
+    alp                      14.0        Exponent of the zero damping (ATM only)
+    ======================== =========== ============================================
+
+    Parameters for (modified) zero damping are:
+
+    ======================== =========== ===================================================
+    Tweakable parameter      Default     Description
+    ======================== =========== ===================================================
+    s6                       1.0         Scaling of the dipole-dipole dispersion
+    s8                       None        Scaling of the dipole-quadrupole dispersion
+    s9                       1.0         Scaling of the three-body dispersion energy
+    rs6                      None        Scaling of the dipole-dipole damping
+    rs8                      1.0         Scaling of the dipole-quadrupole damping
+    alp                      14.0        Exponent of the zero damping
+    bet                      None        Offset for damping radius (modified zero damping)
+    ======================== =========== ===================================================
+
+    Parameters for optimized power damping are:
+
+    ======================== =========== ============================================
+    Tweakable parameter      Default     Description
+    ======================== =========== ============================================
+    s6                       1.0         Scaling of the dipole-dipole dispersion
+    s8                       None        Scaling of the dipole-quadrupole dispersion
+    s9                       1.0         Scaling of the three-body dispersion energy
+    a1                       None        Scaling of the critical radii
+    a2                       None        Offset of the critical radii
+    alp                      14.0        Exponent of the zero damping (ATM only)
+    bet                      None        Power for the zero-damping component
+    ======================== =========== ============================================
+
     The version of the damping can be changed after constructing the dispersion correction.
     With the `atm` boolean the three-body dispersion energy can be enabled, which is
     generally recommended.
@@ -107,14 +154,22 @@ class DFTD3Dispersion(lib.StreamObject):
     array(-0.00574289)
     """
 
-    def __init__(self, mol, xc="hf", version="d3bj", atm=False):
+    def __init__(
+        self,
+        mol: gto.Mole,
+        xc: str = "hf",
+        version: str = "d3bj",
+        atm: bool = False,
+        param: Optional[Dict[str, float]] = None,
+    ):
         self.mol = mol
         self.verbose = mol.verbose
         self.xc = xc
+        self.param = param
         self.atm = atm
         self.version = version
 
-    def dump_flags(self, verbose=None):
+    def dump_flags(self, verbose: Optional[bool] = None):
         """
         Show options used for the DFT-D3 dispersion correction.
         """
@@ -168,16 +223,19 @@ def kernel(self) -> Tuple[float, np.ndarray]:
             mol.atom_coords(),
         )
 
-        param = _damping_param[self.version](
-            method=self.xc,
-            atm=self.atm,
-        )
+        if self.param is not None:
+            param = _damping_param[self.version](**self.param)
+        else:
+            param = _damping_param[self.version](
+                method=self.xc,
+                atm=self.atm,
+            )
 
         res = disp.get_dispersion(param=param, grad=True)
 
         return res.get("energy"), res.get("gradient")
 
-    def reset(self, mol):
+    def reset(self, mol: gto.Mole):
         """Reset mol and clean up relevant attributes for scanner mode"""
         self.mol = mol
         return self
@@ -199,7 +257,7 @@ class _DFTD3Grad:
     pass
 
 
-def energy(mf):
+def energy(mf: scf.hf.SCF, **kwargs) -> scf.hf.SCF:
     """
     Apply DFT-D3 corrections to SCF or MCSCF methods by returning an
     instance of a new class built from the original instances class.
@@ -208,8 +266,10 @@ def energy(mf):
 
     Parameters
     ----------
-    mf
+    mf: scf.hf.SCF
         The method to which DFT-D3 corrections will be applied.
+    **kwargs
+        Keyword arguments passed to the `DFTD3Dispersion` class.
 
     Returns
     -------
@@ -237,17 +297,15 @@ def energy(mf):
     -110.93260361702605
     """
 
-    from pyscf.scf import hf
-    from pyscf.mcscf import casci
-
-    if not isinstance(mf, (hf.SCF, casci.CASCI)):
+    if not isinstance(mf, (scf.hf.SCF, mcscf.casci.CASCI)):
         raise TypeError("mf must be an instance of SCF or CASCI")
 
     with_dftd3 = DFTD3Dispersion(
         mf.mol,
         xc="hf"
-        if isinstance(mf, casci.CASCI)
+        if isinstance(mf, mcscf.casci.CASCI)
         else getattr(mf, "xc", "HF").upper().replace(" ", ""),
+        **kwargs,
     )
 
     if isinstance(mf, _DFTD3):
@@ -287,7 +345,7 @@ def nuc_grad_method(self):
     return DFTD3(mf, with_dftd3)
 
 
-def grad(scf_grad):
+def grad(scf_grad: rhf_grad.Gradients, **kwargs):
     """
     Apply DFT-D3 corrections to SCF or MCSCF nuclear gradients methods
     by returning an instance of a new class built from the original class.
@@ -296,8 +354,10 @@ def grad(scf_grad):
 
     Parameters
     ----------
-    mfgrad
+    scf_grad: rhf_grad.Gradients
         The method to which DFT-D3 corrections will be applied.
+    **kwargs
+        Keyword arguments passed to the `DFTD3Dispersion` class.
 
     Returns
     -------
@@ -330,14 +390,13 @@ def grad(scf_grad):
     5 H    -0.0154527822     0.0229409425    -0.0215141991
     ----------------------------------------------
     """
-    from pyscf.grad import rhf as rhf_grad
 
     if not isinstance(scf_grad, rhf_grad.Gradients):
         raise TypeError("scf_grad must be an instance of Gradients")
 
     # Ensure that the zeroth order results include DFTD3 corrections
     if not getattr(scf_grad.base, "with_dftd3", None):
-        scf_grad.base = dftd3(scf_grad.base)
+        scf_grad.base = energy(scf_grad.base, **kwargs)
 
     class DFTD3Grad(_DFTD3Grad, scf_grad.__class__):
         def grad_nuc(self, mol=None, atmlst=None):

python/dftd3/test_pyscf.py

@@ -56,6 +56,52 @@ def test_energy_r2scan_d3():
     assert d3.kernel()[0] == approx(-0.00578401192369041, abs=1.0e-7)
 
 
+@pytest.mark.skipif(pyscf is None, reason="requires pyscf")
+@pytest.mark.parametrize("atm", [True, False])
+def test_energy_bp_d3zero(atm):
+    thr = 1e-9
+
+    mol = gto.M(
+        atom=[
+            ("C", [-1.42754169820131, -1.50508961850828, -1.93430551124333]),
+            ("C", [+1.19860572924150, -1.66299114873979, -2.03189643761298]),
+            ("C", [+2.65876001301880, +0.37736955363609, -1.23426391650599]),
+            ("C", [+1.50963368042358, +2.57230374419743, -0.34128058818180]),
+            ("C", [-1.12092277855371, +2.71045691257517, -0.25246348639234]),
+            ("C", [-2.60071517756218, +0.67879949508239, -1.04550707592673]),
+            ("I", [-2.86169588073340, +5.99660765711210, +1.08394899986031]),
+            ("H", [+2.09930989272956, -3.36144811062374, -2.72237695164263]),
+            ("H", [+2.64405246349916, +4.15317840474646, +0.27856972788526]),
+            ("H", [+4.69864865613751, +0.26922271535391, -1.30274048619151]),
+            ("H", [-4.63786461351839, +0.79856258572808, -0.96906659938432]),
+            ("H", [-2.57447518692275, -3.08132039046931, -2.54875517521577]),
+            ("S", [-5.88211879210329, 11.88491819358157, +2.31866455902233]),
+            ("H", [-8.18022701418703, 10.95619984550779, +1.83940856333092]),
+            ("C", [-5.08172874482867, 12.66714386256482, -0.92419491629867]),
+            ("H", [-3.18311711399702, 13.44626574330220, -0.86977613647871]),
+            ("H", [-5.07177399637298, 10.99164969235585, -2.10739192258756]),
+            ("H", [-6.35955320518616, 14.08073002965080, -1.68204314084441]),
+        ],
+        unit="bohr",
+    )
+
+    d3 = disp.DFTD3Dispersion(
+        mol,
+        param={
+            "s6": 1.0,
+            "s8": 1.683,
+            "rs6": 1.139,
+            "rs8": 1.0,
+            "alp": 14.0,
+            "s9": 1.0 if atm else 0.0,
+        },
+        version="d3zero",
+    )
+    ref = -0.01410721853585842 if atm else -0.014100267345314462
+
+    assert approx(d3.kernel()[0], abs=thr) == ref
+
+
 @pytest.mark.skipif(pyscf is None, reason="requires pyscf")
 @pytest.mark.parametrize("xc", ["b3lyp", "b3lypg", "b3lyp5", "b3lyp3"])
 def test_energy_b3lyp_d3(xc: str):