[pull] main from dftd3:main

.github/workflows/release.yaml

@@ -140,7 +140,7 @@ jobs:
           merge-multiple: true
 
       - name: Publish to Test PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
+        uses: pypa/gh-action-pypi-publish@76f52bc884231f62b9a034ebfe128415bbaabdfc  # v1.12
         with:
           repository-url: https://test.pypi.org/legacy/
 
@@ -164,4 +164,4 @@ jobs:
           merge-multiple: true
 
       - name: Publish to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
+        uses: pypa/gh-action-pypi-publish@76f52bc884231f62b9a034ebfe128415bbaabdfc  # v1.12

.pre-commit-config.yaml

@@ -42,7 +42,7 @@ repos:
           ]
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.19.1
+    rev: v3.20.0
     hooks:
       - id: pyupgrade
         args: [--py37-plus, --keep-runtime-typing]
@@ -52,21 +52,21 @@ repos:
     hooks:
       - id: isort
         name: isort (python)
-        args: ["--profile", "black", "--filter-files"]
+        args: ["--profile", "black", "--line-length", "80", "--filter-files"]
 
   - repo: https://github.com/psf/black
     rev: 25.1.0
     hooks:
       - id: black
-        stages: [pre-commit]
+        args: ["--line-length", "80"]
 
   - repo: https://github.com/woodruffw/zizmor-pre-commit
-    rev: v1.5.2
+    rev: v1.11.0
     hooks:
       - id: zizmor
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.15.0
+    rev: v1.17.1
     hooks:
       - id: mypy
         pass_filenames: false

docs/source/modules/data/index.rst

@@ -5,4 +5,3 @@
 .. toctree::
 
    r4r2
-   radii

docs/source/modules/typing/index.rst

@@ -4,6 +4,5 @@
 
 .. toctree::
 
-   builtin
    d3
    pytorch

examples/batch.py

@@ -60,30 +60,18 @@
     )
 )
 ref = d3.reference.Reference()
-rcov = d3.data.COV_D3[numbers]
-rvdw = d3.data.VDW_D3[numbers.unsqueeze(-1), numbers.unsqueeze(-2)]
-r4r2 = d3.data.R4R2[numbers]
+rcov = mctc.data.COV_D3()[numbers]
+rvdw = mctc.data.VDW_PAIRWISE()[numbers.unsqueeze(-1), numbers.unsqueeze(-2)]
+r4r2 = d3.data.R4R2()[numbers]
 param = {
     "a1": torch.tensor(0.49484001),
     "s8": torch.tensor(0.78981345),
     "a2": torch.tensor(5.73083694),
 }
 
-cn = mctc.ncoord.cn_d3(
-    numbers, positions, counting_function=mctc.ncoord.exp_count, rcov=rcov
-)
-weights = d3.model.weight_references(numbers, cn, ref, d3.model.gaussian_weight)
-c6 = d3.model.atomic_c6(numbers, weights, ref)
-energy = d3.disp.dispersion(
-    numbers,
-    positions,
-    param,
-    c6,
-    rvdw,
-    r4r2,
-    d3.disp.rational_damping,
-)
+energy = d3.dftd3(numbers, positions, param)
 
 torch.set_printoptions(precision=10)
-print(torch.sum(energy, dim=-1))
+print("Expected:", torch.tensor([-0.0014092578, -0.0057840119]))
+print("Actual  :", torch.sum(energy, dim=-1))
 # tensor([-0.0014092578, -0.0057840119])

examples/hessian.py

@@ -0,0 +1,105 @@
+# SPDX-Identifier: CC0-1.0
+from __future__ import annotations
+
+from typing import Callable
+
+import tad_mctc as mctc
+import torch
+
+import tad_dftd3 as d3
+
+sample1 = dict(
+    numbers=mctc.convert.symbol_to_number("Pb H H H H Bi H H H".split()),
+    positions=torch.tensor(
+        [
+            [-0.00000020988889, -4.98043478877778, +0.00000000000000],
+            [+3.06964045311111, -6.06324400177778, +0.00000000000000],
+            [-1.53482054188889, -6.06324400177778, -2.65838526500000],
+            [-1.53482054188889, -6.06324400177778, +2.65838526500000],
+            [-0.00000020988889, -1.72196703577778, +0.00000000000000],
+            [-0.00000020988889, +4.77334244722222, +0.00000000000000],
+            [+1.35700257511111, +6.70626379422222, -2.35039772300000],
+            [-2.71400388988889, +6.70626379422222, +0.00000000000000],
+            [+1.35700257511111, +6.70626379422222, +2.35039772300000],
+        ]
+    ),
+)
+sample2 = dict(
+    numbers=mctc.convert.symbol_to_number(
+        "C C C C C C I H H H H H S H C H H H".split(" ")
+    ),
+    positions=torch.tensor(
+        [
+            [-1.42754169820131, -1.50508961850828, -1.93430551124333],
+            [+1.19860572924150, -1.66299114873979, -2.03189643761298],
+            [+2.65876001301880, +0.37736955363609, -1.23426391650599],
+            [+1.50963368042358, +2.57230374419743, -0.34128058818180],
+            [-1.12092277855371, +2.71045691257517, -0.25246348639234],
+            [-2.60071517756218, +0.67879949508239, -1.04550707592673],
+            [-2.86169588073340, +5.99660765711210, +1.08394899986031],
+            [+2.09930989272956, -3.36144811062374, -2.72237695164263],
+            [+2.64405246349916, +4.15317840474646, +0.27856972788526],
+            [+4.69864865613751, +0.26922271535391, -1.30274048619151],
+            [-4.63786461351839, +0.79856258572808, -0.96906659938432],
+            [-2.57447518692275, -3.08132039046931, -2.54875517521577],
+            [-5.88211879210329, 11.88491819358157, +2.31866455902233],
+            [-8.18022701418703, 10.95619984550779, +1.83940856333092],
+            [-5.08172874482867, 12.66714386256482, -0.92419491629867],
+            [-3.18311711399702, 13.44626574330220, -0.86977613647871],
+            [-5.07177399637298, 10.99164969235585, -2.10739192258756],
+            [-6.35955320518616, 14.08073002965080, -1.68204314084441],
+        ]
+    ),
+)
+numbers = mctc.batch.pack(
+    (
+        sample1["numbers"],
+        sample2["numbers"],
+    )
+)
+positions = mctc.batch.pack(
+    (
+        sample1["positions"],
+        sample2["positions"],
+    )
+)
+
+param = {
+    "a1": torch.tensor(0.49484001),
+    "s8": torch.tensor(0.78981345),
+    "a2": torch.tensor(5.73083694),
+}
+
+
+def _energy(numbers: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+    """
+    Closure over non-tensor argument `param` for `dftd3` function.
+
+    Returns the energy as a scalar, which is required for Hessian computation
+    to obtain the correct shape of ``(..., nat, 3, nat, 3)``.
+    """
+    return d3.dftd3(numbers, positions, param).sum(-1)
+
+
+def hessian(
+    fn: Callable[..., torch.Tensor], argnums: tuple[int] | int = 0
+) -> Callable:
+    """
+    Compute the Hessian using reverse-mode autodiff twice.
+    (Functorch's `hessian` uses forward and backward mode, but forward is
+    not implemented for the custom autograd functions in DFT-D3.)
+    """
+    return torch.func.jacrev(
+        torch.func.jacrev(fn, argnums=argnums), argnums=argnums
+    )
+
+
+hess_fn_single = hessian(_energy, argnums=1)
+hess_fn_batch = torch.func.vmap(hess_fn_single, in_dims=(0, 0))
+
+pos = positions.clone().requires_grad_(True)
+hess = hess_fn_batch(numbers, pos)
+
+print(f"Shape of numbers  : {numbers.shape}")
+print(f"Shape of positions: {positions.shape}")
+print(f"Shape of Hessian  : {hess.shape}")

examples/single.py

@@ -4,7 +4,9 @@
 
 import tad_dftd3 as d3
 
-numbers = mctc.convert.symbol_to_number(symbols="C C C C N C S H H H H H".split())
+numbers = mctc.convert.symbol_to_number(
+    symbols="C C C C N C S H H H H H".split()
+)
 positions = torch.tensor(
     [
         [-2.56745685564671, -0.02509985979910, 0.00000000000000],

setup.cfg

@@ -34,7 +34,7 @@ project_urls =
 packages = find:
 install_requires =
     numpy<2
-    tad-mctc
+    tad-mctc==0.5.3
     torch>=1.11,<3
 python_requires = >=3.8, <3.13
 include_package_data = True

src/tad_dftd3/__version__.py

@@ -15,4 +15,4 @@
 """
 Version module for *tad-dftd3*.
 """
-__version__ = "0.4.1"
+__version__ = "0.5.0"

src/tad_dftd3/damping/atm.py

@@ -33,9 +33,9 @@
 import torch
 from tad_mctc import storch
 from tad_mctc.batch import real_pairs, real_triples
+from tad_mctc.typing import DD, Tensor
 
 from .. import defaults
-from ..typing import DD, Tensor
 
 __all__ = ["dispersion_atm"]
 
@@ -143,7 +143,10 @@ def dispersion_atm(
     )
 
     ang = torch.where(
-        mask_triples * (r2ij <= cutoff2) * (r2jk <= cutoff2) * (r2jk <= cutoff2),
+        mask_triples
+        * (r2ij <= cutoff2)
+        * (r2jk <= cutoff2)
+        * (r2jk <= cutoff2),
         0.375 * s / r5 + 1.0 / r3,
         torch.tensor(0.0, **dd),
     )

src/tad_dftd3/damping/rational.py

@@ -25,12 +25,12 @@
     \dfrac{R^n_{\text{AB}}}{R^n_{\text{AB}} +
     \left( a_1 R_0^{\text{AB}} + a_2 \right)^n}
 """
-from typing import Dict
+from __future__ import annotations
 
 import torch
+from tad_mctc.typing import DD, Tensor
 
 from .. import defaults
-from ..typing import DD, Tensor
 
 __all__ = ["rational_damping"]
 
@@ -39,7 +39,7 @@ def rational_damping(
     order: int,
     distances: Tensor,
     qq: Tensor,
-    param: Dict[str, Tensor],
+    param: dict[str, Tensor],
 ) -> Tensor:
     """
     Rational damped dispersion interaction between pairs.

src/tad_dftd3/data/__init__.py

@@ -16,7 +16,6 @@
 Atomic data
 ===========
 
-Data arrays for atomic constants like covalent radii or van-der-Waals radii.
+Data arrays for atomic constants.
 """
 from .r4r2 import *
-from .radii import *

src/tad_dftd3/data/r4r2.py

@@ -21,40 +21,65 @@
 Also new super heavies Cn, Nh, Fl, Lv, Og and Am-Rg calculated at
 4c-PBE/Dyall-AE4Z (Dirac 2022).
 """
+from __future__ import annotations
+
 import torch
 
 __all__ = ["R4R2"]
 
 
-# fmt: off
-r4_over_r2 = torch.tensor([
-    0.0000,  # None
-    8.0589, 3.4698,  # H,He
-    29.0974,14.8517,11.8799, 7.8715, 5.5588, 4.7566, 3.8025, 3.1036,  # Li-Ne
-    26.1552,17.2304,17.7210,12.7442, 9.5361, 8.1652, 6.7463, 5.6004,  # Na-Ar
-    29.2012,22.3934,  # K,Ca
-    19.0598,16.8590,15.4023,12.5589,13.4788,  # Sc-
-    12.2309,11.2809,10.5569,10.1428, 9.4907,  # -Zn
-    13.4606,10.8544, 8.9386, 8.1350, 7.1251, 6.1971,  # Ga-Kr
-    30.0162,24.4103,  # Rb,Sr
-    20.3537,17.4780,13.5528,11.8451,11.0355,  # Y-
-    10.1997, 9.5414, 9.0061, 8.6417, 8.9975,  # -Cd
-    14.0834,11.8333,10.0179, 9.3844, 8.4110, 7.5152,  # In-Xe
-    32.7622,27.5708,  # Cs,Ba
-    23.1671,21.6003,20.9615,20.4562,20.1010,19.7475,19.4828,  # La-Eu
-    15.6013,19.2362,17.4717,17.8321,17.4237,17.1954,17.1631,  # Gd-Yb
-    14.5716,15.8758,13.8989,12.4834,11.4421,  # Lu-
-    10.2671, 8.3549, 7.8496, 7.3278, 7.4820,  # -Hg
-    13.5124,11.6554,10.0959, 9.7340, 8.8584, 8.0125,  # Tl-Rn
-    29.8135,26.3157,  # Fr,Ra
-    19.1885,15.8542,16.1305,15.6161,15.1226,16.1576,14.6510,  # Ac-Am
-    14.7178,13.9108,13.5623,13.2326,12.9189,12.6133,12.3142,  # Cm-No
-    14.8326,12.3771,10.6378, 9.3638, 8.2297,  # Lr-
-    7.5667, 6.9456, 6.3946, 5.9159, 5.4929,  # -Cn
-    6.7286, 6.5144,10.9169,10.3600, 9.4723, 8.6641,  # Nh-Og
-])  # fmt: on
-"""Actually calculated r⁴ over r² expectation values."""
-# fmt: on
-
-R4R2 = torch.sqrt(0.5 * (r4_over_r2 * torch.sqrt(torch.arange(r4_over_r2.shape[0]))))
-"""r⁴ over r² expectation values."""
+def R4R2(
+    dtype: torch.dtype | None = None, device: torch.device | None = None
+) -> torch.Tensor:
+    """
+    Returns the r⁴ over r² expectation values as a tensor.
+
+    Parameters
+    ----------
+    dtype : torch.dtype | None, optional
+        The desired data type of the returned tensor. Defaults to None.
+    device : torch.device | None, optional
+        The desired device of the returned tensor. Defaults to None.
+
+    Returns
+    -------
+    Tensor
+        A tensor containing the r⁴ over r² expectation values.
+    """
+
+    # Actually calculated r⁴ over r² expectation values
+    # fmt: off
+    _r4_over_r2 = [
+        0.0000,  # None
+        8.0589, 3.4698,  # H,He
+        29.0974,14.8517,11.8799, 7.8715, 5.5588, 4.7566, 3.8025, 3.1036,  # Li-Ne
+        26.1552,17.2304,17.7210,12.7442, 9.5361, 8.1652, 6.7463, 5.6004,  # Na-Ar
+        29.2012,22.3934,  # K,Ca
+        19.0598,16.8590,15.4023,12.5589,13.4788,  # Sc-
+        12.2309,11.2809,10.5569,10.1428, 9.4907,  # -Zn
+        13.4606,10.8544, 8.9386, 8.1350, 7.1251, 6.1971,  # Ga-Kr
+        30.0162,24.4103,  # Rb,Sr
+        20.3537,17.4780,13.5528,11.8451,11.0355,  # Y-
+        10.1997, 9.5414, 9.0061, 8.6417, 8.9975,  # -Cd
+        14.0834,11.8333,10.0179, 9.3844, 8.4110, 7.5152,  # In-Xe
+        32.7622,27.5708,  # Cs,Ba
+        23.1671,21.6003,20.9615,20.4562,20.1010,19.7475,19.4828,  # La-Eu
+        15.6013,19.2362,17.4717,17.8321,17.4237,17.1954,17.1631,  # Gd-Yb
+        14.5716,15.8758,13.8989,12.4834,11.4421,  # Lu-
+        10.2671, 8.3549, 7.8496, 7.3278, 7.4820,  # -Hg
+        13.5124,11.6554,10.0959, 9.7340, 8.8584, 8.0125,  # Tl-Rn
+        29.8135,26.3157,  # Fr,Ra
+        19.1885,15.8542,16.1305,15.6161,15.1226,16.1576,14.6510,  # Ac-Am
+        14.7178,13.9108,13.5623,13.2326,12.9189,12.6133,12.3142,  # Cm-No
+        14.8326,12.3771,10.6378, 9.3638, 8.2297,  # Lr-
+        7.5667, 6.9456, 6.3946, 5.9159, 5.4929,  # -Cn
+        6.7286, 6.5144,10.9169,10.3600, 9.4723, 8.6641,  # Nh-Og
+    ]
+    # fmt: on
+
+    sqrtz = torch.sqrt(
+        torch.arange(len(_r4_over_r2), device=device, dtype=dtype)
+    )
+    return torch.sqrt(
+        0.5 * (torch.tensor(_r4_over_r2, device=device, dtype=dtype) * sqrtz)
+    )