Add calculator option for dispersion correction (#573)

* Add option to add dispersion to all calculators

* Add d3 optional extra

* Test D3 dispersion

* Add D3 to all extra

* Add m3gnet D3 test

* Add D3 tests to CI

* Add test skip for missing torch_dftd

* Add d3 to extras list

* Add MLIP parameters to sum calculator

* Move D3 import into function

* Test mace_mp dispersion is consistent

* Improve dispersion defaults

* Modify arch saved when adding D3

* Test adding D3 calculator

* Test mace_mp with dispersion kwargs

* Add docs for dispersion

* Change dispersion label

* Update arch for dispersion in tests

* Remove unnecessary dispersion check

* Update torch-dftd dependency

Author: ElliottKasoar

.github/workflows/ci.yml

@@ -42,7 +42,7 @@ jobs:
 
       - name: Install updated e3nn dependencies
         run: |
-          uv sync --extra mattersim --extra fairchem
+          uv sync --extra mattersim --extra fairchem --extra d3
           uv pip install --reinstall pynvml
           uv pip install fairchem-core[torch-extras] --no-build-isolation
 
@@ -55,7 +55,7 @@ jobs:
 
       - name: Install dgl dependencies
         run: |
-          uv sync --extra mace --extra m3gnet --extra alignn
+          uv sync --extra mace --extra m3gnet --extra alignn --extra d3
           uv pip install --reinstall pynvml
 
       - name: Run test suite for dgl dependencies

.github/workflows/mac.yml

@@ -42,7 +42,7 @@ jobs:
 
       - name: Install updated e3nn dependencies
         run: |
-          uv sync --extra mattersim --extra fairchem
+          uv sync --extra mattersim --extra fairchem --extra d3
           uv pip install --reinstall pynvml
           uv pip install "fairchem-core[torch-extras]" --no-build-isolation
 
@@ -55,7 +55,7 @@ jobs:
 
       - name: Install dgl dependencies
         run: |
-          uv sync --extra mace --extra m3gnet --extra alignn
+          uv sync --extra mace --extra m3gnet --extra alignn --extra d3
           uv pip install --reinstall pynvml
 
       - name: Run test suite for dgl dependencies

docs/source/conf.py

@@ -243,4 +243,5 @@
     ("py:class", "Progress"),
     ("py:class", "ProgressBar"),
     ("py:class", "typer.models.CallbackParam"),
+    ("py:class", "SumCalculator"),
 ]

docs/source/user_guide/get_started.rst

@@ -71,6 +71,7 @@ Currently supported MLIP ``extras`` are:
 
 Additional features can also be enabled as ``extras``:
 
+- ``d3``: `DFTD3 <https://github.com/pfnet-research/torch-dftd>`_
 - ``visualise``: `WEAS Widget <https://github.com/superstar54/weas-widget>`_
 - ``plumed``: `PLUMED <https://www.plumed.org>`_
 

docs/source/user_guide/python.rst

@@ -98,6 +98,60 @@ will return
     unless the same ``arch`` is chosen, in which case these values will also be overwritten.
 
 
+D3 Dispersion
+=============
+
+A PyTorch implementation of DFTD2 and DFTD3, using the `TorchDFTD3Calculator <https://github.com/pfnet-research/torch-dftd>`_,
+can be used to add dispersion corrections to MLIP predictions.
+
+The required Python pacakge is included with ``mace_mp``, but can also be installed as its own extra:
+
+.. code-block:: bash
+
+    pip install janus-core[d3]
+
+
+Once installed, dispersion can be added through ``calc_kwargs`` through the ``dispersion`` keyword,
+with ``dispersion_kwargs`` used to pass any further keywords to the ``TorchDFTD3Calculator``:
+
+.. code-block:: python
+
+    from ase import units
+
+    from janus_core.calculations.single_point import SinglePoint
+
+    single_point = SinglePoint(
+        struct="tests/data/NaCl.cif",
+        arch="mace_mp",
+        model="tests/models/mace_mp_small.model",
+        calc_kwargs={"dispersion": True, "dispersion_kwargs": {"cutoff":  95.0 * units.Bohr}}
+    )
+
+.. note::
+    In most cases, defaults for ``dispersion_kwargs`` are those set within ``TorchDFTD3Calculator``,
+    but in the case of ``mace_mp``, we mirror the corresponding defaults from the
+    ``mace.calculators.mace_mp`` function.
+
+
+The ``TorchDFTD3Calculator`` can also be added to any existing calculator if required:
+
+.. note::
+    Keyword arguments for ``TorchDFTD3Calculator`` should be passed directly here,
+    as shown with ``cutoff``. This will not have access to ``mace_mp`` default values,
+    so will always use defaults from ``TorchDFTD3Calculator``.
+
+
+.. code-block:: python
+
+    from ase import units
+
+    from janus_core.helpers.mlip_calculators import add_dispersion, choose_calculator
+
+    mace_calc = choose_calculator("mace_mp")
+    calc = add_dispersion(mace_calc, device="cpu", cutoff=95 * units.Bohr)
+
+
+
 Additional Calculators
 ======================
 

janus_core/helpers/mlip_calculators.py

@@ -12,7 +12,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, get_args
 
+from ase import units
 from ase.calculators.mixing import SumCalculator
+from torch import get_default_dtype
 
 from janus_core.helpers.janus_types import Architectures, Devices, PathLike
 from janus_core.helpers.utils import none_to_dict
@@ -84,10 +86,59 @@ def _set_no_weights_only_load():
     environ.setdefault("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1")
 
 
+def add_dispersion(
+    calc: Calculator,
+    device: Devices = "cpu",
+    dtype: torch.dtype | None = None,
+    **kwargs,
+) -> SumCalculator:
+    """
+    Add D3 dispersion calculator to existing calculator.
+
+    Parameters
+    ----------
+    calc
+        Calculator to add D3 correction to.
+    device
+        Device to run calculator on. Default is "cpu".
+    dtype
+        Calculation precision. Default is current torch dtype.
+    **kwargs
+        Additional keyword arguments passed to `TorchDFTD3Calculator`.
+
+    Returns
+    -------
+    SumCalculator
+        Configured calculator with D3 dispersion correction added.
+    """
+    try:
+        from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
+    except ImportError as err:
+        raise ImportError("Please install the d3 extra.") from err
+
+    dtype = dtype if dtype else get_default_dtype()
+
+    d3_calc = TorchDFTD3Calculator(
+        device=device,
+        dtype=dtype,
+        **kwargs,
+    )
+    sum_calc = SumCalculator([calc, d3_calc])
+
+    # Copy calculator parameters to make more accessible
+    sum_calc.parameters = calc.parameters
+    if "arch" in sum_calc.parameters:
+        sum_calc.parameters["arch"] = sum_calc.parameters["arch"] + "_d3"
+
+    return sum_calc
+
+
 def choose_calculator(
     arch: Architectures,
     device: Devices = "cpu",
     model: PathLike | None = None,
+    dispersion: bool = False,
+    dispersion_kwargs: dict[str, Any] | None = None,
     **kwargs,
 ) -> Calculator:
     """
@@ -101,6 +152,11 @@ def choose_calculator(
         Device to run calculator on. Default is "cpu".
     model
         MLIP model label, path to model, or loaded model. Default is `None`.
+    dispersion
+        Whether to add D3 dispersion.
+    dispersion_kwargs
+        Additional keyword arguments for `TorchDFTD3Calculator`. Defaults for mace_mp
+        are taken from mace_mp's defaults.
     **kwargs
         Additional keyword arguments passed to the selected calculator.
 
@@ -116,6 +172,8 @@ def choose_calculator(
     ValueError
         Invalid architecture specified.
     """
+    dispersion_kwargs = dispersion_kwargs if dispersion_kwargs else {}
+
     model = _set_model(model, kwargs)
 
     if device not in get_args(Devices):
@@ -147,6 +205,13 @@ def choose_calculator(
             model = model if model else "small"
             kwargs.setdefault("default_dtype", "float64")
 
+            # Set mace_mp dispersion defaults
+            dispersion_kwargs.setdefault("damping", kwargs.pop("damping", "bj"))
+            dispersion_kwargs.setdefault("xc", kwargs.pop("dispersion_xc", "pbe"))
+            dispersion_kwargs.setdefault(
+                "cutoff", kwargs.pop("dispersion_cutoff", 40.0 * units.Bohr)
+            )
+
             calculator = mace_mp(model=model, device=device, **kwargs)
 
         case "mace_off":
@@ -420,6 +485,9 @@ def choose_calculator(
     calculator.parameters["arch"] = arch
     calculator.parameters["model"] = str(model)
 
+    if dispersion:
+        return add_dispersion(calc=calculator, device=device, **dispersion_kwargs)
+
     return calculator
 
 

pyproject.toml

@@ -46,12 +46,15 @@ chgnet = [
 dpa3 = [
     "deepmd-kit == 3.1.0",
 ]
+d3 = [
+    "torch-dftd==0.5.1",
+]
 grace = [
     "tensorpotential == 0.5.1",
 ]
 mace = [
     "mace-torch==0.3.13",
-    "torch-dftd==0.4.0",
+    "janus-core[d3]",
 ]
 nequip = [
     "nequip == 0.6.1",
@@ -73,8 +76,8 @@ visualise = [
 ]
 all = [
     "janus-core[chgnet]",
-    "janus-core[dpa3]",
     "janus-core[grace]",
+    "janus-core[d3]",
     "janus-core[mace]",
     "janus-core[nequip]",
     "janus-core[orb]",

tests/test_descriptors.py

@@ -119,7 +119,7 @@ def test_dispersion():
     descriptors_disp.run()
 
     assert (
-        descriptors_disp.struct.info["mace_mp_descriptor"]
+        descriptors_disp.struct.info["mace_mp_d3_descriptor"]
         == descriptors.struct.info["mace_mp_descriptor"]
     )
 

tests/test_mlip_calculators.py

@@ -8,7 +8,7 @@
 
 import pytest
 
-from janus_core.helpers.mlip_calculators import choose_calculator
+from janus_core.helpers.mlip_calculators import add_dispersion, choose_calculator
 from tests.utils import skip_extras
 
 MODEL_PATH = Path(__file__).parent / "models"
@@ -279,3 +279,24 @@ def test_invalid_device(arch):
     """Test error raised if invalid device is specified."""
     with pytest.raises(ValueError):
         choose_calculator(arch=arch, device="invalid")
+
+
+def test_d3():
+    """Test adding D3 dispersion calculator automatically."""
+    skip_extras("mace_mp")
+
+    calculator = choose_calculator(arch="mace_mp", dispersion=True)
+    assert calculator.parameters["version"] is not None
+    assert calculator.parameters["model"] is not None
+    assert calculator.parameters["arch"] == "mace_mp_d3"
+
+
+def test_d3_manual():
+    """Test adding D3 dispersion calculator manually."""
+    skip_extras("mace_mp")
+
+    calculator = choose_calculator(arch="mace_mp")
+    calculator = add_dispersion(calculator)
+    assert calculator.parameters["version"] is not None
+    assert calculator.parameters["model"] is not None
+    assert calculator.parameters["arch"] == "mace_mp_d3"

tests/test_single_point.py

@@ -6,6 +6,8 @@
 import shutil
 from urllib.error import HTTPError, URLError
 
+from ase import units
+from ase.calculators.mixing import SumCalculator
 from ase.calculators.singlepoint import SinglePointCalculator
 from ase.io import read
 from numpy import isfinite
@@ -549,3 +551,81 @@ def test_missing_arch(struct):
 
     with pytest.raises(ValueError, match="A calculator must be attached"):
         SinglePoint(struct=struct)
+
+
+@pytest.mark.parametrize(
+    "arch, kwargs, pred",
+    [
+        ("m3gnet", {}, -0.08281749),
+        (
+            "mace_mp",
+            {"damping": "zero", "xc": "pbe", "cutoff": 95 * units.Bohr},
+            -0.08281749,
+        ),
+        ("mace_off", {}, -0.08281747),
+        ("mattersim", {}, -0.08281749),
+        ("sevennet", {}, -0.08281749),
+    ],
+)
+def test_dispersion(arch, kwargs, pred):
+    """Test dispersion correction."""
+    skip_extras(arch)
+    pytest.importorskip("torch_dftd")
+
+    data_path = DATA_PATH / "benzene.xyz"
+    sp_no_d3 = SinglePoint(
+        struct=data_path,
+        arch=arch,
+        properties="energy",
+        calc_kwargs={"dispersion": False},
+    )
+    assert not isinstance(sp_no_d3.struct.calc, SumCalculator)
+    no_d3_results = sp_no_d3.run()
+
+    sp_d3 = SinglePoint(
+        struct=data_path,
+        arch=arch,
+        properties="energy",
+        calc_kwargs={"dispersion": True, "dispersion_kwargs": {**kwargs}},
+    )
+    assert isinstance(sp_d3.struct.calc, SumCalculator)
+    d3_results = sp_d3.run()
+
+    assert (d3_results["energy"] - no_d3_results["energy"]) == pytest.approx(pred)
+
+
+def test_mace_mp_dispersion():
+    """Test mace_mp dispersion correction matches default."""
+    skip_extras("mace_mp")
+    pytest.importorskip("torch_dftd")
+
+    from mace.calculators import mace_mp
+
+    data_path = DATA_PATH / "benzene.xyz"
+
+    no_d3_energy = SinglePoint(
+        struct=data_path,
+        arch="mace_mp",
+        properties="energy",
+        calc_kwargs={"dispersion": False},
+    ).run()["energy"]
+
+    d3_energy = SinglePoint(
+        struct=data_path,
+        arch="mace_mp",
+        properties="energy",
+        calc_kwargs={"dispersion": True},
+    ).run()["energy"]
+
+    struct = read(data_path)
+    struct.calc = mace_mp(model="small", dispersion=True)
+
+    mace_d3_energy = SinglePoint(
+        struct=struct,
+        properties="energy",
+        calc_kwargs={"dispersion": False},
+    ).run()["energy"]
+
+    # Different default to other architectures
+    assert d3_energy - no_d3_energy == pytest.approx(-0.29815768)
+    assert d3_energy == pytest.approx(mace_d3_energy)