Skip slow calculation of reference C6 (#39)

The loops in the `_load_c6` function are actually really slow and create
a bottleneck for the whole calculation. Since the output of this
function is not dependent on the system, we only use it to store the
tensor (".pt" file) and now skip the calculation and directly load the
tensor from the "reference-c6.pt" file.

This gives a speed-up of more than one order of magnitude for small
systems.

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),
+    )