GFN2: Non-self-consistent D4 ATM term (#200)

.github/workflows/release.yaml

@@ -20,18 +20,52 @@ on:
   push:
     branches:
       - main
+      - master
     tags:
       - "v*"
+    paths-ignore:
+      - "doc*/**"
+      - "./*.ya?ml"
+      - "**/*.md"
+      - "**/*.rst"
 
   pull_request:
-    branches:
-      - main
+    paths-ignore:
+      - "doc*/**"
+      - "./*.ya?ml"
+      - "**/*.md"
+      - "**/*.rst"
 
   workflow_dispatch:
 
 jobs:
+  sdist:
+    permissions:
+      contents: read
+
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - name: Build source distribution (sdist)
+        run: pipx run build --sdist
+
+      - name: Upload source distribution as artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ github.event.repository.name }}-sdist
+          path: dist
+
   wheel:
+    permissions:
+      contents: read
+
     runs-on: ubuntu-latest
+
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -44,38 +78,68 @@ jobs:
       - name: Upload wheel as artifact
         uses: actions/upload-artifact@v4
         with:
-          name: dxtb-wheel
-          path: dist/*.whl
+          name: ${{ github.event.repository.name }}-wheel
+          path: dist
+
+  install_wheel:
+    needs: [wheel]
+
+    permissions:
+      contents: read
 
-  sdist:
     runs-on: ubuntu-latest
+
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
         with:
           persist-credentials: false
 
-      - name: Build source distribution (sdist)
-        run: pipx run build --sdist
-
-      - name: Upload source distribution as artifact
-        uses: actions/upload-artifact@v4
+      - name: Download build artifacts
+        uses: actions/download-artifact@v4
         with:
-          name: dxtb-sdist
-          path: dist/*.tar.gz
+          path: dist
+          merge-multiple: true
+
+      - name: Show downloaded artifacts
+        run: ls -lcahFR --color=auto dist
+
+      - name: Install wheel
+        run: |
+          pip install torch --index-url https://download.pytorch.org/whl/cpu
+          pip install dist/*.whl
+
+      - name: Determine package name
+        run: |
+          name=$(echo "${REPO_NAME}" | tr '-' '_')
+          echo "PKG_NAME=$name" >> "$GITHUB_ENV"
+          echo "PKG_NAME is set to '${name}'."
+        env:
+          REPO_NAME: ${{ github.event.repository.name }}
+
+      - name: Test import
+        run: python -c "import ${PKG_NAME}; print(${PKG_NAME}.__version__)"
+        env:
+          PKG_NAME: ${{ env.PKG_NAME }}
 
   upload_test_pypi:
     needs: [sdist, wheel]
+
     runs-on: ubuntu-latest
+
     environment: release
+
     permissions:
+      contents: read
       id-token: write
+
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
     steps:
       - name: Download build artifacts
         uses: actions/download-artifact@v4
         with:
           path: dist
+          merge-multiple: true
 
       - name: Publish to Test PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
@@ -84,17 +148,22 @@ jobs:
 
   upload_pypi:
     needs: [sdist, wheel, upload_test_pypi]
+
     runs-on: ubuntu-latest
+
     environment: release
+
     permissions:
+      contents: read
       id-token: write
+
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
     steps:
       - name: Download build artifacts
         uses: actions/download-artifact@v4
         with:
-          name: artifact
           path: dist
+          merge-multiple: true
 
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1

setup.cfg

@@ -46,7 +46,7 @@ install_requires =
     pydantic>=2.0.0
     scipy
     tad-dftd3>=0.3.0
-    tad-dftd4>=0.5.1
+    tad-dftd4>=0.5.2
     tad-mctc>=0.3.2
     tad-multicharge
     tomli

src/dxtb/_src/components/classicals/dispersion/base.py

@@ -28,7 +28,7 @@
 import torch
 
 from dxtb import IndexHelper
-from dxtb._src.typing import Any, Tensor
+from dxtb._src.typing import Any, Literal, Tensor
 
 from ...classicals import Classical, ClassicalCache
 
@@ -49,20 +49,30 @@ class Dispersion(Classical):
     charge: Tensor | None
     """Total charge of the system."""
 
-    __slots__ = ["numbers", "param", "charge"]
+    ref_charges: Literal["eeq", "gfn2"]
+    """
+    Reference charges for the dispersion model.
+    This is only required for charge-dependent models.
+
+    :default: ``"eeq"``
+    """
+
+    __slots__ = ["numbers", "param", "charge", "ref_charges"]
 
     def __init__(
         self,
         numbers: Tensor,
         param: dict[str, Tensor],
         charge: Tensor | None = None,
+        ref_charges: Literal["eeq", "gfn2"] = "eeq",
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__(device, dtype)
         self.numbers = numbers
         self.param = param
         self.charge = charge
+        self.ref_charges = ref_charges
 
     @abstractmethod
     def get_cache(

src/dxtb/_src/components/classicals/dispersion/d4.py

@@ -118,36 +118,41 @@ def get_cache(
 
         self._cachevars = cachvars
 
-        model: d4.model.D4Model = (
-            kwargs.pop(
-                "model",
-                d4.model.D4Model(numbers, **self.dd),
+        model = kwargs.pop("model", None)
+        if model is not None and not isinstance(model, d4.model.D4Model):
+            raise TypeError("D4: Model is not of type 'd4.model.D4Model'.")
+        if model is None:
+            model = d4.model.D4Model(
+                numbers, ref_charges=self.ref_charges, **self.dd
             )
-            .type(self.dtype)
-            .to(self.device)
-        )
-        rcov: Tensor = (
-            kwargs.pop(
-                "rcov",
-                d4.data.COV_D3.to(**self.dd)[numbers],
-            )
-            .type(self.dtype)
-            .to(self.device)
-        )
+        else:
+            model = model.type(self.dtype).to(self.device)
+
+        rcov = kwargs.pop("rcov", None)
+        if rcov is not None and not isinstance(rcov, Tensor):
+            raise TypeError("D4: 'rcov' is not of type 'Tensor'.")
+        if rcov is None:
+            rcov = d4.data.COV_D3.to(**self.dd)[numbers]
+        else:
+            rcov = rcov.to(**self.dd)
+
+        r4r2 = kwargs.pop("r4r2", None)
+        if r4r2 is not None and not isinstance(r4r2, Tensor):
+            raise TypeError("D4: 'r4r2' is not of type 'Tensor'.")
+        if r4r2 is None:
+            r4r2 = d4.data.R4R2.to(**self.dd)[numbers]
+        else:
+            r4r2 = r4r2.to(**self.dd)
+
+        cutoff = kwargs.pop("cutoff", None)
+        if cutoff is not None and not isinstance(cutoff, d4.cutoff.Cutoff):
+            raise TypeError("D4: 'cutoff' is not of type 'd4.cutoff.Cutoff'.")
+        if cutoff is None:
+            cutoff = d4.cutoff.Cutoff(**self.dd)
+        else:
+            cutoff = cutoff.type(self.dtype).to(self.device)
+
         q = kwargs.pop("q", None)
-        r4r2: Tensor = (
-            kwargs.pop(
-                "r4r2",
-                d4.data.R4R2.to(**self.dd)[numbers],
-            )
-            .type(self.dtype)
-            .to(self.device)
-        )
-        cutoff: d4.cutoff.Cutoff = (
-            (kwargs.pop("cutoff", d4.cutoff.Cutoff(**self.dd)))
-            .type(self.dtype)
-            .to(self.device)
-        )
 
         cf = kwargs.pop("counting_function", d4.ncoord.erf_count)
         df = kwargs.pop("damping_function", d4.damping.rational_damping)

src/dxtb/_src/components/classicals/dispersion/factory.py

@@ -28,7 +28,7 @@
 import torch
 
 from dxtb._src.param import Param
-from dxtb._src.typing import DD, Tensor, get_default_dtype
+from dxtb._src.typing import DD, Literal, Tensor, get_default_dtype
 from dxtb._src.typing.exceptions import ParameterWarning
 from dxtb._src.utils import convert_float_tensor
 
@@ -43,6 +43,7 @@ def new_dispersion(
     numbers: Tensor,
     par: Param,
     charge: Tensor | None = None,
+    ref_charges: Literal["eeq", "gfn2"] = "eeq",
     device: torch.device | None = None,
     dtype: torch.dtype | None = None,
 ) -> Dispersion | None:
@@ -55,6 +56,9 @@ def new_dispersion(
         Atomic numbers for all atoms in the system (shape: ``(..., nat)``).
     par : Param
         Representation of an extended tight-binding model.
+    ref_charges : Literal["eeq", "gfn2"], optional
+        Reference charges for the dispersion model. This is only required for
+        charge-dependent models. Default is ``"eeq"``.
     device : torch.device | None, optional
         Device to store the tensor on. If ``None`` (default), the default
         device is used.
@@ -109,15 +113,30 @@ def new_dispersion(
             **dd,
         )
 
+        if isinstance(par.dispersion.d4.sc, bool) is False:
+            raise ValueError("D4 self-consistency flag is not a boolean.")
+
         # only non-self-consistent D4 is a classical component
         if par.dispersion.d4.sc is False:
             if charge is None:
                 raise ValueError("The total charge is required for DFT-D4.")
 
             return DispersionD4(
-                numbers, param, charge=charge, device=device, dtype=dtype
+                numbers,
+                param,
+                ref_charges=ref_charges,
+                charge=charge,
+                device=device,
+                dtype=dtype,
             )
 
+        # Classical part of self-consistent D4 is only ATM term
+        param["s6"] = torch.tensor(0.0, **dd)
+        param["s8"] = torch.tensor(0.0, **dd)
+        return DispersionD4(
+            numbers, param, ref_charges="gfn2", device=device, dtype=dtype
+        )
+
     if par.dispersion.d3 is not None and par.dispersion.d4 is not None:
         raise ValueError("Parameters for both D3 and D4 found. Please decide.")
 

src/dxtb/_src/components/interactions/dispersion/d4sc.py

@@ -71,10 +71,7 @@ class constructor.
     except for multiplication with C6 and C8.
     """
 
-    __slots__ = [
-        "__store",
-        "cn",
-    ]
+    __slots__ = ["__store", "cn", "dispmat"]
 
     def __init__(
         self,