Refactoring and better GPU tests

.vscode/settings.json

@@ -1,18 +1,17 @@
 {
   "[python]": {
+    "editor.defaultFormatter": "ms-python.black-formatter",
     "editor.detectIndentation": false,
     "editor.formatOnSave": true,
     "editor.formatOnPaste": false, // not supported by black
     "editor.insertSpaces": true,
     "editor.tabSize": 4
   },
+  "pylint.args": ["--indent-string='    '"],
   "python.analysis.diagnosticSeverityOverrides": {
     "reportPrivateImportUsage": "information"
   },
-  "python.formatting.provider": "black",
-  "python.linting.enabled": true,
-  "python.linting.pylintEnabled": true,
-  "python.linting.pylintArgs": ["--indent-string='    '"],
+  "python.defaultInterpreterPath": "${env:CONDA_PREFIX}/envs/dxtb/bin/python",
   "python.testing.pytestArgs": [],
   "python.testing.unittestEnabled": false,
   "python.testing.pytestEnabled": true

README.rst

@@ -229,7 +229,7 @@ The next example shows the calculation of dispersion energies for a batch of str
        "a2": torch.tensor(5.73083694),
    }
 
-   cn = d3.ncoord.coordination_number(numbers, positions, rcov, d3.ncoord.exp_count)
+   cn = d3.ncoord.coordination_number(numbers, positions, d3.ncoord.exp_count, 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(

examples/batch.py

@@ -4,7 +4,7 @@
 import tad_dftd3 as d3
 
 sample1 = dict(
-    numbers=d3.util.to_number("Pb H H H H Bi H H H".split()),
+    numbers=d3.utils.to_number("Pb H H H H Bi H H H".split()),
     positions=torch.tensor(
         [
             [-0.00000020988889, -4.98043478877778, +0.00000000000000],
@@ -20,7 +20,7 @@
     ),
 )
 sample2 = dict(
-    numbers=d3.util.to_number("C C C C C C I H H H H H S H C H H H".split(" ")),
+    numbers=d3.utils.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],
@@ -44,13 +44,13 @@
         ]
     ),
 )
-numbers = d3.util.pack(
+numbers = d3.utils.pack(
     (
         sample1["numbers"],
         sample2["numbers"],
     )
 )
-positions = d3.util.pack(
+positions = d3.utils.pack(
     (
         sample1["positions"],
         sample2["positions"],
@@ -66,7 +66,7 @@
     "a2": torch.tensor(5.73083694),
 }
 
-cn = d3.ncoord.coordination_number(numbers, positions, rcov, d3.ncoord.exp_count)
+cn = d3.ncoord.coordination_number(numbers, positions, d3.ncoord.exp_count, 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(

examples/single.py

@@ -3,7 +3,7 @@
 
 import tad_dftd3 as d3
 
-numbers = d3.util.to_number(symbols="C C C C N C S H H H H H".split())
+numbers = d3.utils.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],

pyproject.toml

@@ -40,4 +40,4 @@ plugins = ["covdefaults"]
 source = ["./src"]
 
 [tool.coverage.report]
-fail_under = 80
+fail_under = 90

setup.cfg

@@ -14,6 +14,7 @@ classifiers =
     Programming Language :: Python :: 3.9
     Programming Language :: Python :: 3.10
     Programming Language :: Python :: 3.11
+    Programming Language :: Python :: Implementation :: CPython
 project_urls =
     Documentation = https://tad-dftd3.readthedocs.io/en/latest/
     Source Code = https://github.com/dftd3/tad-dftd3
@@ -36,13 +37,20 @@ where = src
 dev =
     black
     covdefaults
-    coverage
     mypy
     pre-commit
     pylint
     pytest
+    pytest-cov
     pytest-random-order
+    pytest-xdist
     tox
+tox =
+    covdefaults
+    pytest
+    pytest-cov
+    pytest-random-order
+    pytest-xdist
 
 [options.package_data]
 tad_dftd3 =

src/tad_dftd3/__init__.py

@@ -74,7 +74,7 @@
 """
 import torch
 
-from . import damping, data, disp, model, ncoord, reference
+from . import damping, data, defaults, disp, model, ncoord, reference
 from ._typing import (
     DD,
     CountingFunction,
@@ -84,7 +84,7 @@
     Tensor,
     WeightingFunction,
 )
-from .util import misc
+from .utils import misc
 
 
 def dftd3(
@@ -135,20 +135,17 @@ def dftd3(
     dd: DD = {"device": positions.device, "dtype": positions.dtype}
 
     if cutoff is None:
-        cutoff = torch.tensor(50.0, **dd)
+        cutoff = torch.tensor(defaults.D3_DISP_CUTOFF, **dd)
     if ref is None:
         ref = reference.Reference(**dd)
     if rcov is None:
-        rcov = data.covalent_rad_d3[numbers].to(**dd)
+        rcov = data.covalent_rad_d3.to(**dd)[numbers]
     if rvdw is None:
-        rvdw = data.vdw_rad_d3[
-            numbers.unsqueeze(-1),
-            numbers.unsqueeze(-2),
-        ].to(**dd)
+        rvdw = data.vdw_rad_d3.to(**dd)[numbers.unsqueeze(-1), numbers.unsqueeze(-2)]
     if r4r2 is None:
-        r4r2 = data.sqrt_z_r4_over_r2[numbers].to(**dd)
+        r4r2 = data.sqrt_z_r4_over_r2.to(**dd)[numbers]
 
-    cn = ncoord.coordination_number(numbers, positions, rcov, counting_function)
+    cn = ncoord.coordination_number(numbers, positions, counting_function, rcov)
     weights = model.weight_references(numbers, cn, ref, weighting_function)
     c6 = model.atomic_c6(numbers, weights, ref)
     energy = disp.dispersion(

src/tad_dftd3/constants.py

@@ -21,3 +21,124 @@
 
 ANGSTROM_TO_BOHR = 1.0 / BOHR_TO_ANGSTROM
 """Conversion factor from Angstrom to Bohr."""
+
+PSE = {
+    "H": 1,
+    "He": 2,
+    "Li": 3,
+    "Be": 4,
+    "B": 5,
+    "C": 6,
+    "N": 7,
+    "O": 8,
+    "F": 9,
+    "Ne": 10,
+    "Na": 11,
+    "Mg": 12,
+    "Al": 13,
+    "Si": 14,
+    "P": 15,
+    "S": 16,
+    "Cl": 17,
+    "Ar": 18,
+    "K": 19,
+    "Ca": 20,
+    "Sc": 21,
+    "Ti": 22,
+    "V": 23,
+    "Cr": 24,
+    "Mn": 25,
+    "Fe": 26,
+    "Co": 27,
+    "Ni": 28,
+    "Cu": 29,
+    "Zn": 30,
+    "Ga": 31,
+    "Ge": 32,
+    "As": 33,
+    "Se": 34,
+    "Br": 35,
+    "Kr": 36,
+    "Rb": 37,
+    "Sr": 38,
+    "Y": 39,
+    "Zr": 40,
+    "Nb": 41,
+    "Mo": 42,
+    "Tc": 43,
+    "Ru": 44,
+    "Rh": 45,
+    "Pd": 46,
+    "Ag": 47,
+    "Cd": 48,
+    "In": 49,
+    "Sn": 50,
+    "Sb": 51,
+    "Te": 52,
+    "I": 53,
+    "Xe": 54,
+    "Cs": 55,
+    "Ba": 56,
+    "La": 57,
+    "Ce": 58,
+    "Pr": 59,
+    "Nd": 60,
+    "Pm": 61,
+    "Sm": 62,
+    "Eu": 63,
+    "Gd": 64,
+    "Tb": 65,
+    "Dy": 66,
+    "Ho": 67,
+    "Er": 68,
+    "Tm": 69,
+    "Yb": 70,
+    "Lu": 71,
+    "Hf": 72,
+    "Ta": 73,
+    "W": 74,
+    "Re": 75,
+    "Os": 76,
+    "Ir": 77,
+    "Pt": 78,
+    "Au": 79,
+    "Hg": 80,
+    "Tl": 81,
+    "Pb": 82,
+    "Bi": 83,
+    "Po": 84,
+    "At": 85,
+    "Rn": 86,
+    "Fr": 87,
+    "Ra": 88,
+    "Ac": 89,
+    "Th": 90,
+    "Pa": 91,
+    "U": 92,
+    "Np": 93,
+    "Pu": 94,
+    "Am": 95,
+    "Cm": 96,
+    "Bk": 97,
+    "Cf": 98,
+    "Es": 99,
+    "Fm": 100,
+    "Md": 101,
+    "No": 102,
+    "Lr": 103,
+    "Rf": 104,
+    "Db": 105,
+    "Sg": 106,
+    "Bh": 107,
+    "Hs": 108,
+    "Mt": 109,
+    "Ds": 110,
+    "Rg": 111,
+    "Cn": 112,
+    "Nh": 113,
+    "Fl": 114,
+    "Mc": 115,
+    "Lv": 116,
+    "Ts": 117,
+    "Og": 118,
+}

src/tad_dftd3/damping/atm.py

@@ -34,7 +34,7 @@
 
 from .. import defaults
 from .._typing import DD, Tensor
-from ..util import cdist, real_pairs, real_triples
+from ..utils import cdist, real_pairs, real_triples
 
 
 def dispersion_atm(

src/tad_dftd3/defaults.py

@@ -19,6 +19,19 @@
 This module defines the default values for all parameters within DFT-D3.
 """
 
+# DFT-D3
+
+D3_CN_CUTOFF = 25.0
+"""Coordination number cutoff (25.0)."""
+
+D3_DISP_CUTOFF = 50.0
+"""Two/three-body interaction cutoff (50.0)."""
+
+D3_KCN = 16.0
+"""Steepness of counting function (16.0)."""
+
+# DFT-D3 damping parameters
+
 A1 = 0.4
 """Scaling for the C8 / C6 ratio in the critical radius (0.4)."""
 

src/tad_dftd3/disp.py

@@ -57,7 +57,7 @@
 from . import data, defaults
 from ._typing import DD, Any, DampingFunction, Dict, Optional, Tensor
 from .damping import dispersion_atm, rational_damping
-from .util import cdist, real_pairs
+from .utils import cdist, real_pairs
 
 
 def dispersion(
@@ -100,9 +100,9 @@ def dispersion(
     dd: DD = {"device": positions.device, "dtype": positions.dtype}
 
     if cutoff is None:
-        cutoff = torch.tensor(50.0, **dd)
+        cutoff = torch.tensor(defaults.D3_DISP_CUTOFF, **dd)
     if r4r2 is None:
-        r4r2 = data.sqrt_z_r4_over_r2[numbers].to(**dd)
+        r4r2 = data.sqrt_z_r4_over_r2.to(**dd)[numbers]
     if numbers.shape != positions.shape[:-1]:
         raise ValueError(
             "Shape of positions is not consistent with atomic numbers.",
@@ -120,10 +120,9 @@ def dispersion(
     # three-body dispersion
     if "s9" in param and param["s9"] != 0.0:
         if rvdw is None:
-            rvdw = data.vdw_rad_d3[
-                numbers.unsqueeze(-1),
-                numbers.unsqueeze(-2),
-            ].to(**dd)
+            rvdw = data.vdw_rad_d3.to(**dd)[
+                numbers.unsqueeze(-1), numbers.unsqueeze(-2)
+            ]
 
         energy += dispersion3(numbers, positions, param, c6, rvdw, cutoff)
 

src/tad_dftd3/model.py

@@ -43,7 +43,7 @@
 
 from ._typing import Any, Tensor, WeightingFunction
 from .reference import Reference
-from .util import real_atoms
+from .utils import real_atoms
 
 
 def atomic_c6(numbers: Tensor, weights: Tensor, reference: Reference) -> Tensor: