Skip slow calculation of reference C6

src/tad_dftd3/__init__.py

@@ -74,7 +74,7 @@
 """
 import torch
 
-from . import damping, data, defaults, disp, model, ncoord, reference
+from . import constants, damping, data, defaults, disp, model, ncoord, reference
 from ._typing import (
     DD,
     CountingFunction,
@@ -134,6 +134,12 @@ def dftd3(
     """
     dd: DD = {"device": positions.device, "dtype": positions.dtype}
 
+    if torch.max(numbers) >= constants.MAX_ELEMENT:
+        raise ValueError(
+            f"No D3 parameters available for Z > {constants.MAX_ELEMENT-1} "
+            f"({constants.PSE_Z2S[constants.MAX_ELEMENT]})."
+        )
+
     if cutoff is None:
         cutoff = torch.tensor(defaults.D3_DISP_CUTOFF, **dd)
     if ref is None:

src/tad_dftd3/constants.py

@@ -16,13 +16,16 @@
 This module contains fundamental constants and conversion factors.
 """
 
+MAX_ELEMENT = 104
+"""Atomic number (+1 for dummy) of last element supported by DFT-D3."""
+
 BOHR_TO_ANGSTROM = 0.529177210903
 """Bohr radius in Angstroms."""
 
 ANGSTROM_TO_BOHR = 1.0 / BOHR_TO_ANGSTROM
 """Conversion factor from Angstrom to Bohr."""
 
-PSE = {
+PSE_S2Z = {
     "H": 1,
     "He": 2,
     "Li": 3,
@@ -142,3 +145,7 @@
     "Ts": 117,
     "Og": 118,
 }
+"""PSE with mapping from symbol to atomic number."""
+
+PSE_Z2S = {v: k for k, v in PSE_S2Z.items()}
+"""PSE with mapping from atomic number to symbol."""

src/tad_dftd3/disp.py

@@ -54,7 +54,7 @@
 """
 import torch
 
-from . import data, defaults
+from . import constants, data, defaults
 from ._typing import DD, Any, DampingFunction, Dict, Optional, Tensor
 from .damping import dispersion_atm, rational_damping
 from .utils import cdist, real_pairs
@@ -103,6 +103,7 @@ def dispersion(
         cutoff = torch.tensor(defaults.D3_DISP_CUTOFF, **dd)
     if r4r2 is None:
         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.",
@@ -111,6 +112,11 @@ def dispersion(
         raise ValueError(
             "Shape of expectation values is not consistent with atomic numbers.",
         )
+    if torch.max(numbers) >= constants.MAX_ELEMENT:
+        raise ValueError(
+            f"No D3 parameters available for Z > {constants.MAX_ELEMENT-1} "
+            f"({constants.PSE_Z2S[constants.MAX_ELEMENT]})."
+        )
 
     # two-body dispersion
     energy = dispersion2(

src/tad_dftd3/reference.py

@@ -30,6 +30,21 @@
 def _load_cn(
     dtype: torch.dtype = torch.double, device: Optional[torch.device] = None
 ) -> Tensor:
+    """
+    Load reference coordination numbers.
+
+    Parameters
+    ----------
+    dtype : torch.dtype, optional
+        Floating point precision for tensor. Defaults to `torch.double`.
+    device : Optional[torch.device], optional
+        Device of tensor. Defaults to None.
+
+    Returns
+    -------
+    Tensor
+        Reference coordination numbers.
+    """
     return torch.tensor(
         [
             [-1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000],  # None
@@ -146,31 +161,22 @@ def _load_c6(
     dtype: torch.dtype = torch.double, device: Optional[torch.device] = None
 ) -> Tensor:
     """
-    Load reference C6 coefficients from file and fill them into a tensor
+    Load reference C6 coefficients from file.
+
+    Parameters
+    ----------
+    dtype : torch.dtype, optional
+        Floating point precision for tensor. Defaults to `torch.double`.
+    device : Optional[torch.device], optional
+        Device of tensor. Defaults to None.
+
+    Returns
+    -------
+    Tensor
+        Reference C6 coefficients.
     """
-
-    # pylint: disable=import-outside-toplevel
-    import math
-
-    import numpy as np
-
-    path = op.join(op.dirname(__file__), "reference-c6.npy")
-    ref = torch.from_numpy(np.load(path)).type(dtype).to(device)
-
-    n_element = (math.isqrt(8 * ref.shape[0] + 1) - 1) // 2 + 1
-    n_reference = ref.shape[-1]
-    c6 = torch.zeros(
-        (n_element, n_element, n_reference, n_reference),
-        dtype=ref.dtype,
-        device=ref.device,
-    )
-
-    for i in range(1, n_element):
-        for j in range(1, n_element):
-            ij = i * (i - 1) // 2 + j - 1 if j < i else j * (j - 1) // 2 + i - 1
-            c6[i, j, :, :] = ref[ij, :, :].T if j < i else ref[ij, :, :]
-
-    return c6
+    path = op.join(op.dirname(__file__), "reference-c6.pt")
+    return torch.load(path).type(dtype).to(device)
 
 
 class Reference:

src/tad_dftd3/utils/misc.py

@@ -22,7 +22,7 @@
 import torch
 
 from .._typing import List, Optional, Size, Tensor, TensorOrTensors, Union
-from ..constants import PSE
+from ..constants import PSE_S2Z
 
 __all__ = [
     "real_atoms",
@@ -174,7 +174,7 @@ def to_number(symbols: List[str]) -> Tensor:
     Obtain atomic numbers from element symbols.
     """
     return torch.flatten(
-        torch.tensor([PSE.get(symbol.capitalize(), 0) for symbol in symbols])
+        torch.tensor([PSE_S2Z.get(symbol.capitalize(), 0) for symbol in symbols])
     )
 
 

tests/test_dftd3.py

@@ -25,6 +25,28 @@
 from .samples import samples
 
 
+def test_fail() -> None:
+    positions = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+
+    # TPSS0-D3BJ-ATM parameters
+    param = {
+        "s6": torch.tensor(1.0000),
+        "s8": torch.tensor(1.2576),
+        "s9": torch.tensor(1.0000),
+        "alp": torch.tensor(14.00),
+        "a1": torch.tensor(0.3768),
+        "a2": torch.tensor(4.5865),
+    }
+
+    # unsupported element
+    with pytest.raises(ValueError):
+        dftd3(torch.tensor([1, 105]), positions, param)
+
+    # wrong numbers
+    with pytest.raises(ValueError):
+        dftd3(torch.tensor([1]), positions, param)
+
+
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
 @pytest.mark.parametrize("name", ["LiH", "SiH4", "PbH4-BiH3"])
 def test_single(dtype: torch.dtype, name: str) -> None:

tests/test_disp.py

@@ -54,8 +54,11 @@ def test_fail() -> None:
 
     # wrong numbers
     with pytest.raises(ValueError):
-        numbers = torch.tensor([1])
-        disp.dispersion(numbers, positions, param, c6)
+        disp.dispersion(torch.tensor([1]), positions, param, c6)
+
+    # unsupported element
+    with pytest.raises(ValueError):
+        disp.dispersion(torch.tensor([1, 105]), positions, param, c6)
 
 
 @pytest.mark.parametrize("dtype", [torch.float, torch.double])

tests/test_model/test_load.py

@@ -0,0 +1,27 @@
+# This file is part of tad-dftd3.
+# SPDX-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Test loading C6 coefficients.
+"""
+import torch
+
+from tad_dftd3 import constants, reference
+
+
+def test_ref() -> None:
+    c6 = reference._load_c6(dtype=torch.double)
+    assert c6.shape == torch.Size(
+        (constants.MAX_ELEMENT, constants.MAX_ELEMENT, 7, 7),
+    )