diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 9028cd2..813bb58 100644 --- a/.github/workflows/release.yaml +++ b/.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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 81965b1..667bdbb 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/docs/source/modules/data/index.rst b/docs/source/modules/data/index.rst index 0593c57..d25a5b6 100644 --- a/docs/source/modules/data/index.rst +++ b/docs/source/modules/data/index.rst @@ -5,4 +5,3 @@ .. toctree:: r4r2 - radii diff --git a/docs/source/modules/data/radii.rst b/docs/source/modules/data/radii.rst deleted file mode 100644 index e0e9ded..0000000 --- a/docs/source/modules/data/radii.rst +++ /dev/null @@ -1,2 +0,0 @@ -.. automodule:: tad_dftd3.data.radii - :members: diff --git a/docs/source/modules/typing/builtin.rst b/docs/source/modules/typing/builtin.rst deleted file mode 100644 index 8f79631..0000000 --- a/docs/source/modules/typing/builtin.rst +++ /dev/null @@ -1,2 +0,0 @@ -.. automodule:: tad_dftd3.typing.builtin - :members: diff --git a/docs/source/modules/typing/index.rst b/docs/source/modules/typing/index.rst index 6c66f2e..b9a2dff 100644 --- a/docs/source/modules/typing/index.rst +++ b/docs/source/modules/typing/index.rst @@ -4,6 +4,5 @@ .. toctree:: - builtin d3 pytorch diff --git a/examples/batch.py b/examples/batch.py index 935ba27..e87d8a2 100644 --- a/examples/batch.py +++ b/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]) diff --git a/examples/hessian.py b/examples/hessian.py new file mode 100644 index 0000000..f77b68f --- /dev/null +++ b/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}") diff --git a/examples/single.py b/examples/single.py index 24c2c3f..f0f14c3 100644 --- a/examples/single.py +++ b/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], diff --git a/setup.cfg b/setup.cfg index 952b8da..1cdd4a3 100644 --- a/setup.cfg +++ b/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 diff --git a/src/tad_dftd3/__version__.py b/src/tad_dftd3/__version__.py index a1c88f7..b6f2a57 100644 --- a/src/tad_dftd3/__version__.py +++ b/src/tad_dftd3/__version__.py @@ -15,4 +15,4 @@ """ Version module for *tad-dftd3*. """ -__version__ = "0.4.1" +__version__ = "0.5.0" diff --git a/src/tad_dftd3/damping/atm.py b/src/tad_dftd3/damping/atm.py index 52647a6..cd4d38a 100644 --- a/src/tad_dftd3/damping/atm.py +++ b/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), ) diff --git a/src/tad_dftd3/damping/rational.py b/src/tad_dftd3/damping/rational.py index 5d08877..739e7ef 100644 --- a/src/tad_dftd3/damping/rational.py +++ b/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. diff --git a/src/tad_dftd3/data/__init__.py b/src/tad_dftd3/data/__init__.py index f5645e9..f9ce090 100644 --- a/src/tad_dftd3/data/__init__.py +++ b/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 * diff --git a/src/tad_dftd3/data/r4r2.py b/src/tad_dftd3/data/r4r2.py index 559567c..ec6f86a 100644 --- a/src/tad_dftd3/data/r4r2.py +++ b/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) + ) diff --git a/src/tad_dftd3/data/radii.py b/src/tad_dftd3/data/radii.py deleted file mode 100644 index ac729c3..0000000 --- a/src/tad_dftd3/data/radii.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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. -""" -Data: Radii -=========== - -Data arrays for atomic constants like covalent radii or van-der-Waals radii. - -The `vdw_rad_d3` were previously stored explicitly in one list and then -reshaped to the required `(MAX_ELEMENT, MAX_ELEMENT)` tensor. For the old -version, see older commits (e.g. https://github.com/dftd3/tad-dftd3/blob/ecc50f19adb8aa8baa38a188d04228c4f26975d6/src/tad_dftd3/data/radii.py) -""" -import os.path as op -from typing import Optional - -import torch -from tad_mctc._version import __tversion__ -from tad_mctc.data.radii import COV_D3 - -from ..typing import Any, Tensor - -__all__ = ["COV_D3", "VDW_D3"] - - -def _load_vdw_rad_d3( - dtype: torch.dtype = torch.double, device: Optional[torch.device] = None -) -> Tensor: - """ - Load reference VDW radii 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 - VDW radii. - """ - kwargs: dict[str, Any] = {"map_location": device} - if __tversion__ > (1, 12, 1): # pragma: no cover - kwargs["weights_only"] = True - - path = op.join(op.dirname(__file__), "vdw-d3.pt") - return torch.load(path, **kwargs).type(dtype=dtype) - - -VDW_D3 = _load_vdw_rad_d3() diff --git a/src/tad_dftd3/data/vdw-d3.pt b/src/tad_dftd3/data/vdw-d3.pt deleted file mode 100644 index 32626cd..0000000 Binary files a/src/tad_dftd3/data/vdw-d3.pt and /dev/null differ diff --git a/src/tad_dftd3/disp.py b/src/tad_dftd3/disp.py index e686267..1beca60 100644 --- a/src/tad_dftd3/disp.py +++ b/src/tad_dftd3/disp.py @@ -54,17 +54,19 @@ """ from __future__ import annotations +from typing import Any + import torch from tad_mctc import storch +from tad_mctc.autograd import is_functorch_tensor from tad_mctc.batch import real_pairs -from tad_mctc.data import pse +from tad_mctc.data import pse, radii from . import data, defaults, model, ncoord from .damping import dispersion_atm, rational_damping from .reference import Reference from .typing import ( DD, - Any, CountingFunction, DampingFunction, Tensor, @@ -125,22 +127,25 @@ def dftd3( """ dd: DD = {"device": positions.device, "dtype": positions.dtype} - if torch.max(numbers) >= defaults.MAX_ELEMENT: - raise ValueError( - f"No D3 parameters available for Z > {defaults.MAX_ELEMENT-1} " - f"({pse.Z2S[defaults.MAX_ELEMENT]})." - ) + if not is_functorch_tensor(numbers): + if torch.max(numbers) >= defaults.MAX_ELEMENT: + raise ValueError( + f"No D3 parameters available for Z > {defaults.MAX_ELEMENT-1} " + f"({pse.Z2S[defaults.MAX_ELEMENT]})." + ) if cutoff is None: cutoff = torch.tensor(defaults.D3_DISP_CUTOFF, **dd) if ref is None: ref = Reference(**dd) if rcov is None: - rcov = data.COV_D3.to(**dd)[numbers] + rcov = radii.COV_D3(**dd)[numbers] if rvdw is None: - rvdw = data.VDW_D3.to(**dd)[numbers.unsqueeze(-1), numbers.unsqueeze(-2)] + rvdw = radii.VDW_PAIRWISE(**dd)[ + numbers.unsqueeze(-1), numbers.unsqueeze(-2) + ] if r4r2 is None: - r4r2 = data.R4R2.to(**dd)[numbers] + r4r2 = data.R4R2(**dd)[numbers] cn = ncoord.cn_d3( numbers, positions, counting_function=counting_function, rcov=rcov @@ -202,7 +207,7 @@ def dispersion( if cutoff is None: cutoff = torch.tensor(defaults.D3_DISP_CUTOFF, **dd) if r4r2 is None: - r4r2 = data.R4R2.to(**dd)[numbers] + r4r2 = data.R4R2(**dd)[numbers] if numbers.shape != positions.shape[:-1]: raise ValueError( @@ -212,11 +217,13 @@ def dispersion( raise ValueError( "Shape of expectation values is not consistent with atomic numbers.", ) - if torch.max(numbers) >= defaults.MAX_ELEMENT: - raise ValueError( - f"No D3 parameters available for Z > {defaults.MAX_ELEMENT - 1} " - f"({pse.Z2S[defaults.MAX_ELEMENT]})." - ) + + if not is_functorch_tensor(numbers): + if torch.max(numbers) >= defaults.MAX_ELEMENT: + raise ValueError( + f"No D3 parameters available for Z > {defaults.MAX_ELEMENT-1} " + f"({pse.Z2S[defaults.MAX_ELEMENT]})." + ) # two-body dispersion energy = dispersion2( @@ -226,7 +233,9 @@ def dispersion( # three-body dispersion if "s9" in param and param["s9"] != 0.0: if rvdw is None: - rvdw = data.VDW_D3.to(**dd)[numbers.unsqueeze(-1), numbers.unsqueeze(-2)] + rvdw = radii.VDW_PAIRWISE(**dd)[ + numbers.unsqueeze(-1), numbers.unsqueeze(-2) + ] energy += dispersion3(numbers, positions, param, c6, rvdw, cutoff) diff --git a/src/tad_dftd3/model/c6.py b/src/tad_dftd3/model/c6.py index a146202..78805a6 100644 --- a/src/tad_dftd3/model/c6.py +++ b/src/tad_dftd3/model/c6.py @@ -23,13 +23,15 @@ """ from __future__ import annotations +from typing import NamedTuple, Protocol + import torch from tad_mctc._version import __tversion__ from tad_mctc.math import einsum from tad_mctc.tools import memory +from tad_mctc.typing import Callable, Tensor from ..reference import Reference -from ..typing import Callable, Protocol, Tensor __all__ = ["atomic_c6"] @@ -260,13 +262,18 @@ def _atomic_c6_chunked( # (..., n1, n2, r1, r2) * (..., n1, r1) * (..., n2, r2) -> (..., n1, n2) contribution = _einsum(rc6_chunk, weights_chunk, weights) - # Add contributions to the correct slice of the output tensor - c6_output[..., start:end, :] += contribution + # Add contributions to the correct slice of the output tensor with + # out-of-place aggregation along the "i" axis. In-place aggregation: + # c6_output[..., start:end, :] += contribution + idx = torch.arange(start, end, device=numbers.device) + c6_output = torch.index_add( + c6_output, dim=-2, index=idx, source=contribution + ) return c6_output -# custom autograd functions +# typing class CTX(Protocol): @@ -276,6 +283,14 @@ class CTX(Protocol): reference: Reference +class VmapInfo(NamedTuple): + batch_size: int + randomness: str + + +# custom autograd functions + + class AtomicC6Base(torch.autograd.Function): """ Base class for the version-specific autograd function for atomic C6. @@ -283,7 +298,9 @@ class AtomicC6Base(torch.autograd.Function): """ @staticmethod - def backward(ctx: CTX, grad_out: Tensor) -> tuple[None, Tensor, None, None]: + def backward( + ctx: CTX, grad_outputs: Tensor + ) -> tuple[None, Tensor, None, None]: numbers, weights = ctx.saved_tensors chunk_size = ctx.chunk_size ref = ctx.reference @@ -304,14 +321,14 @@ def backward(ctx: CTX, grad_out: Tensor) -> tuple[None, Tensor, None, None]: g_jb = einsum("...ijab,...ia->...ijb", rc6, weights) # vjp: (..., n1, n2) * (..., n1, n2, r2) -> (..., n2, r2) - _gj = einsum("...ij,...ijb->...jb", grad_out, g_jb) + _gj = einsum("...ij,...ijb->...jb", grad_outputs, g_jb) # ∂c_ij/∂w_ia = ∑b w_jb * c_ijab # (..., n1, n2, r1, r2) * (..., n2, r2) -> (..., n1, n2, r1) g_ia = einsum("...ijab,...jb->...ija", rc6, weights) # vjp: (..., n1, n2) * (..., n1, n2, r1) -> (..., n1, r1) - _gi = einsum("...ij,...ija->...ia", grad_out, g_ia) + _gi = einsum("...ij,...ija->...ia", grad_outputs, g_ia) weights_bar = _gi + _gj @@ -322,13 +339,14 @@ def backward(ctx: CTX, grad_out: Tensor) -> tuple[None, Tensor, None, None]: ####################### nat = weights.shape[-2] - weights_bar = torch.zeros_like(weights) + gi_accum = torch.zeros_like(weights) + gj_accum = torch.zeros_like(weights) for start in range(0, nat, chunk_size): end = min(start + chunk_size, nat) # Numbers and derivatives for this chunk - grad_chunk = grad_out[..., start:end, :] # (..., chunk_size, nat) + grad_chunk = grad_outputs[..., start:end, :] # (..., c_size, nat) num_chunk = numbers[..., start:end] # (..., chunk_size) # Chunked indexing into reference.c6: (..., chunk_size, nat, 7, 7) @@ -346,11 +364,14 @@ def backward(ctx: CTX, grad_out: Tensor) -> tuple[None, Tensor, None, None]: g_jb = einsum("...ijab,...ia->...ijb", rc6_chunk, weights_chunk) _gj = einsum("...ij,...ijb->...jb", grad_chunk, g_jb) - # Accumulate gradients for the current chunk - weights_bar[..., start:end, :] += _gi - weights_bar += _gj + # Accumulate gradients for current chunk with using out-of-place ops + # to allow vmap. Old version: weights_bar[..., start:end, :] += _gi + idx = torch.arange(start, end, device=weights.device) + gi_accum = torch.index_add(gi_accum, dim=-2, index=idx, source=_gi) - return None, weights_bar, None, None + gj_accum = gj_accum + _gj + + return None, gi_accum + gj_accum, None, None class AtomicC6_V1(AtomicC6Base): @@ -383,16 +404,17 @@ class AtomicC6_V2(AtomicC6Base): This is supposed to reduce memory usage. """ - generate_vmap_rule = True - # https://pytorch.org/docs/master/notes/extending.func.html#automatically-generate-a-vmap-rule - # should work since we only use PyTorch operations + generate_vmap_rule = False + # Auto-generation should work since we only use PyTorch operations, + # however, it does not: PyTorch throws an internal error when indexing + # `reference.c6` with the `numbers` tensor. @staticmethod def forward( numbers: Tensor, weights: Tensor, reference: Reference, - chunk_size: None | int = None, + chunk_size: int | None = None, ) -> Tensor: if chunk_size is None: return _atomic_c6_full(numbers, weights, reference) @@ -410,3 +432,49 @@ def setup_context( ctx.save_for_backward(numbers, weights) ctx.chunk_size = chunk_size ctx.reference = reference + + @staticmethod + def vmap( + info: VmapInfo, + in_dims: tuple[int | None, ...], + numbers: Tensor, + weights: Tensor, + reference: Reference, + chunk_size: int | None, + ) -> tuple[Tensor, int]: + bd_n, bd_w, bd_ref, bd_cs = in_dims + + # Only numbers and weights are batched + if bd_ref is not None or bd_cs is not None: # pragma: no cover + raise ValueError( + "`Reference` and `chunk_size` must be static under vmap." + ) + + # Move batch dimensions to the front if necessary + if info.batch_size != numbers.shape[0]: + raise ValueError( + f"Batch size mismatch: expected {info.batch_size}, got " + f"{numbers.shape[0]} in `numbers`. The first dimension " + "should be the batch dimension." + ) + + if info.batch_size != weights.shape[0]: + raise ValueError( + f"Batch size mismatch: expected {info.batch_size}, got " + f"{weights.shape[0]} in `weights`. The first dimension " + "should be the batch dimension." + ) + + if bd_n not in (0, None): # pragma: no cover + raise ValueError( + f"Batch dimension for `numbers` must be 0 (first dimension) " + f"or `None`, got {bd_n}." + ) + if bd_w not in (0, None): # pragma: no cover + raise ValueError( + f"Batch dimension for `weights` must be 0 (first dimension) " + f"or `None`, got {bd_w}." + ) + + out = AtomicC6_V2.forward(numbers, weights, reference, chunk_size) + return out, 0 diff --git a/src/tad_dftd3/model/weights.py b/src/tad_dftd3/model/weights.py index c6b7b83..d493aa2 100644 --- a/src/tad_dftd3/model/weights.py +++ b/src/tad_dftd3/model/weights.py @@ -42,11 +42,15 @@ """ from __future__ import annotations +from typing import Any + import torch from tad_mctc import storch +from tad_mctc.autograd import is_functorch_tensor +from tad_mctc.typing import Tensor from ..reference import Reference -from ..typing import Any, Tensor, WeightingFunction +from ..typing import WeightingFunction __all__ = ["gaussian_weight", "weight_references"] @@ -141,7 +145,10 @@ def weight_references( # back to real dtype gw_temp = storch.divide(weights, norm, eps=small).type(cn.dtype) - assert torch.isnan(gw_temp).sum() == 0 + + # If the tensor is not a grad tracking tensor, we can check for NaN's + if not is_functorch_tensor(gw_temp): + assert torch.isnan(gw_temp).sum() == 0 # The following section handles cases with large CNs that lead to zeros in # after the exponential in the weighting function. If this happens all diff --git a/src/tad_dftd3/reference.py b/src/tad_dftd3/reference.py index ba277bf..80b556e 100644 --- a/src/tad_dftd3/reference.py +++ b/src/tad_dftd3/reference.py @@ -20,12 +20,11 @@ C6 dispersion coefficients. """ import os.path as op -from typing import Optional +from typing import Any, NoReturn, Optional import torch from tad_mctc._version import __tversion__ - -from .typing import Any, NoReturn, Tensor, get_default_device, get_default_dtype +from tad_mctc.typing import Tensor, get_default_device, get_default_dtype __all__ = ["Reference"] diff --git a/src/tad_dftd3/typing/__init__.py b/src/tad_dftd3/typing/__init__.py index 45c1d5b..51bbf70 100644 --- a/src/tad_dftd3/typing/__init__.py +++ b/src/tad_dftd3/typing/__init__.py @@ -18,6 +18,5 @@ All type annotations for this project. """ -from .builtin import * from .d3 import * from .pytorch import * diff --git a/src/tad_dftd3/typing/builtin.py b/src/tad_dftd3/typing/builtin.py deleted file mode 100644 index aab741b..0000000 --- a/src/tad_dftd3/typing/builtin.py +++ /dev/null @@ -1,24 +0,0 @@ -# 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. -""" -Type annotations: Built-ins -=========================== - -Built-in type annotations are imported from the *tad-mctc* library, which -handles some version checking. -""" -from tad_mctc.typing import Any, Callable, NoReturn, Protocol, TypedDict - -__all__ = ["Any", "Callable", "NoReturn", "Protocol", "TypedDict"] diff --git a/test/conftest.py b/test/conftest.py index 6f2c6fa..4fd143f 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -96,7 +96,8 @@ def pytest_addoption(parser: pytest.Parser) -> None: default=6, type=int, help=( - "Number of digits of precision for floating point output " "(default = 4)." + "Number of digits of precision for floating point output " + "(default = 4)." ), ) @@ -176,4 +177,6 @@ def pytest_runtest_setup(item: pytest.Function) -> None: for _ in item.iter_markers(name="cuda"): if not torch.cuda.is_available(): - pytest.skip("Torch not compiled with CUDA or no CUDA device available.") + pytest.skip( + "Torch not compiled with CUDA or no CUDA device available." + ) diff --git a/test/test_disp/test_dftd3.py b/test/test_disp/test_dftd3.py index 778f9c8..3d0d69d 100644 --- a/test/test_disp/test_dftd3.py +++ b/test/test_disp/test_dftd3.py @@ -18,6 +18,7 @@ import pytest import torch from tad_mctc.batch import pack +from tad_mctc.data import radii from tad_dftd3 import damping, data, dftd3, model, reference from tad_dftd3.ncoord import exp_count @@ -59,9 +60,11 @@ def test_single(dtype: torch.dtype, name: str) -> None: positions = sample["positions"].to(**dd) ref = (sample["disp2"] + sample["disp3"]).to(**dd) - rcov = data.COV_D3.to(**dd)[numbers] - rvdw = data.VDW_D3.to(**dd)[numbers.unsqueeze(-1), numbers.unsqueeze(-2)] - r4r2 = data.R4R2.to(**dd)[numbers] + rcov = radii.COV_D3(**dd)[numbers] + rvdw = radii.VDW_PAIRWISE(**dd)[ + numbers.unsqueeze(-1), numbers.unsqueeze(-2) + ] + r4r2 = data.R4R2(**dd)[numbers] cutoff = torch.tensor(50, **dd) param = { diff --git a/test/test_disp/test_disp.py b/test/test_disp/test_disp.py index 5820b2c..4e67dc7 100644 --- a/test/test_disp/test_disp.py +++ b/test/test_disp/test_disp.py @@ -20,6 +20,7 @@ import pytest import torch from tad_mctc.batch import pack +from tad_mctc.data import radii from tad_dftd3 import damping, data, disp from tad_dftd3.typing import DD @@ -40,7 +41,9 @@ } # TPSS0-D3BJ parameters -param_noatm = {k: torch.tensor(0.0) if k == "s9" else v for k, v in param.items()} +param_noatm = { + k: torch.tensor(0.0) if k == "s9" else v for k, v in param.items() +} def test_fail() -> None: @@ -73,8 +76,10 @@ def test_disp2_single(dtype: torch.dtype, name: str) -> None: positions = sample["positions"].to(**dd) ref = sample["disp2"].to(**dd) c6 = sample["c6"].to(**dd) - rvdw = data.VDW_D3.to(**dd)[numbers.unsqueeze(-1), numbers.unsqueeze(-2)] - r4r2 = data.R4R2.to(**dd)[numbers] + rvdw = radii.VDW_PAIRWISE(**dd)[ + numbers.unsqueeze(-1), numbers.unsqueeze(-2) + ] + r4r2 = data.R4R2(**dd)[numbers] cutoff = torch.tensor(50.0, **dd) par = {k: v.to(**dd) for k, v in param_noatm.items()} @@ -147,7 +152,9 @@ def test_atm_single(dtype: torch.dtype, name: str) -> None: c6 = sample["c6"].to(**dd) ref = sample["disp3"].to(**dd) - rvdw = data.VDW_D3.to(**dd)[numbers.unsqueeze(-1), numbers.unsqueeze(-2)] + rvdw = radii.VDW_PAIRWISE(**dd)[ + numbers.unsqueeze(-1), numbers.unsqueeze(-2) + ] par = {k: v.to(**dd) for k, v in param.items()} @@ -200,7 +207,9 @@ def test_atm_batch(dtype: torch.dtype, name1: str, name2: str) -> None: par = {k: v.to(**dd) for k, v in param.items()} - rvdw = data.VDW_D3.to(**dd)[numbers.unsqueeze(-1), numbers.unsqueeze(-2)] + rvdw = radii.VDW_PAIRWISE(**dd)[ + numbers.unsqueeze(-1), numbers.unsqueeze(-2) + ] energy = damping.dispersion_atm( numbers, diff --git a/test/test_disp/test_special.py b/test/test_disp/test_special.py index 7fc697e..31d4f69 100644 --- a/test/test_disp/test_special.py +++ b/test/test_disp/test_special.py @@ -19,6 +19,7 @@ import pytest import torch from tad_mctc.batch import pack +from tad_mctc.data import radii from tad_dftd3 import damping, data, dftd3, model, reference from tad_dftd3.ncoord import exp_count @@ -38,9 +39,11 @@ def test_single(dtype: torch.dtype, name: str) -> None: positions = sample["positions"].to(**dd) ref = sample["disp2"].to(**dd) - rcov = data.COV_D3.to(**dd)[numbers] - rvdw = data.VDW_D3.to(**dd)[numbers.unsqueeze(-1), numbers.unsqueeze(-2)] - r4r2 = data.R4R2.to(**dd)[numbers] + rcov = radii.COV_D3(**dd)[numbers] + rvdw = radii.VDW_PAIRWISE(**dd)[ + numbers.unsqueeze(-1), numbers.unsqueeze(-2) + ] + r4r2 = data.R4R2(**dd)[numbers] cutoff = torch.tensor(50, **dd) # GFN1-xTB parameters diff --git a/test/test_grad/test_hessian.py b/test/test_grad/test_hessian.py index abc0c90..829b8dc 100644 --- a/test/test_grad/test_hessian.py +++ b/test/test_grad/test_hessian.py @@ -19,7 +19,8 @@ import pytest import torch -from tad_mctc.autograd import hessian +from tad_mctc._version import __tversion__ +from tad_mctc.autograd import hess_fn_rev, hessian from tad_mctc.batch import pack from tad_mctc.convert import reshape_fortran @@ -88,11 +89,59 @@ def test_single(dtype: torch.dtype, name: str) -> None: positions.detach_() -# TODO: Figure out batched Hessian computation +@pytest.mark.skipif(__tversion__ < (2, 0, 0), reason="Requires PyTorch>=2.0.0") +@pytest.mark.parametrize("dtype", [torch.double]) +@pytest.mark.parametrize("name", sample_list) +def test_single_v2(dtype: torch.dtype, name: str) -> None: + dd: DD = {"device": DEVICE, "dtype": dtype} + + sample = samples[name] + numbers = sample["numbers"].to(DEVICE) + positions = sample["positions"].to(**dd) + + # GFN1-xTB parameters + param = { + "s6": torch.tensor(1.00000000, **dd), + "s8": torch.tensor(2.40000000, **dd), + "s9": torch.tensor(0.00000000, **dd), + "a1": torch.tensor(0.63000000, **dd), + "a2": torch.tensor(5.00000000, **dd), + } + + ref = reshape_fortran( + sample["hessian"].to(**dd), + torch.Size(2 * (numbers.shape[-1], 3)), + ) + + # variable to be differentiated + positions.requires_grad_(True) + + def _energy(numbers: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """ + Closure over non-tensor argument `param` for `dftd3` function. + + Returns energy as scalar, which is required for Hessian computation + to obtain the correct shape of ``(..., nat, 3, nat, 3)``. + """ + return dftd3(numbers, positions, param).sum(-1) + + pos = positions.clone().requires_grad_(True) + hess = hess_fn_rev(_energy, argnums=1)(numbers, pos) + assert isinstance(hess, Tensor) + + assert pytest.approx(ref.cpu(), abs=tol, rel=tol) == hess.detach().cpu() + + positions.detach_() + + +@pytest.mark.skipif(__tversion__ < (2, 0, 0), reason="Requires PyTorch>=2.0.0") @pytest.mark.parametrize("dtype", [torch.double]) @pytest.mark.parametrize("name1", ["LiH"]) @pytest.mark.parametrize("name2", sample_list) -def skip_test_batch(dtype: torch.dtype, name1: str, name2: str) -> None: +@pytest.mark.parametrize("chunk_size", [None, 2]) +def test_batch( + dtype: torch.dtype, name1: str, name2: str, chunk_size: int | None +) -> None: dd: DD = {"device": DEVICE, "dtype": dtype} sample1, sample2 = samples[name1], samples[name2] @@ -134,7 +183,20 @@ def skip_test_batch(dtype: torch.dtype, name1: str, name2: str) -> None: # variable to be differentiated positions.requires_grad_(True) - hess = hessian(dftd3, (numbers, positions, param), argnums=1) + def _energy(numbers: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """ + Closure over non-tensor argument `param` for `dftd3` function. + + Returns energy as scalar, which is required for Hessian computation + to obtain the correct shape of ``(..., nat, 3, nat, 3)``. + """ + return dftd3(numbers, positions, param, chunk_size=chunk_size).sum(-1) + + hess_fn = hess_fn_rev(_energy, argnums=1) + hess_fn_batch = torch.func.vmap(hess_fn, in_dims=(0, 0)) + + hess = hess_fn_batch(numbers, positions) + assert isinstance(hess, Tensor) assert pytest.approx(ref.cpu(), abs=tol, rel=tol) == hess.detach().cpu() positions.detach_() diff --git a/test/test_grad/test_param.py b/test/test_grad/test_param.py index dfc166d..a9d6715 100644 --- a/test/test_grad/test_param.py +++ b/test/test_grad/test_param.py @@ -22,9 +22,9 @@ from tad_mctc.autograd import dgradcheck, dgradgradcheck from tad_mctc.batch import pack from tad_mctc.data.molecules import mols as samples +from tad_mctc.typing import DD, Callable, Tensor from tad_dftd3 import dftd3 -from tad_dftd3.typing import DD, Callable, Tensor from ..conftest import DEVICE, FAST_MODE @@ -136,7 +136,9 @@ def test_gradcheck_batch(dtype: torch.dtype, name1: str, name2: str) -> None: @pytest.mark.parametrize("dtype", [torch.double]) @pytest.mark.parametrize("name1", ["LiH"]) @pytest.mark.parametrize("name2", sample_list) -def test_gradgradcheck_batch(dtype: torch.dtype, name1: str, name2: str) -> None: +def test_gradgradcheck_batch( + dtype: torch.dtype, name1: str, name2: str +) -> None: """ Check a single analytical gradient of parameters against numerical gradient from `torch.autograd.gradgradcheck`. diff --git a/test/test_grad/test_pos.py b/test/test_grad/test_pos.py index e1693b4..1be8ccc 100644 --- a/test/test_grad/test_pos.py +++ b/test/test_grad/test_pos.py @@ -21,9 +21,9 @@ import torch from tad_mctc.autograd import dgradcheck, dgradgradcheck, jacrev from tad_mctc.batch import pack +from tad_mctc.typing import DD, Callable, Tensor from tad_dftd3 import dftd3 -from tad_dftd3.typing import DD, Callable, Tensor from ..conftest import DEVICE, FAST_MODE from .samples import samples @@ -137,7 +137,9 @@ def test_gradcheck_batch(dtype: torch.dtype, name1: str, name2: str) -> None: @pytest.mark.parametrize("dtype", [torch.double]) @pytest.mark.parametrize("name1", ["LiH"]) @pytest.mark.parametrize("name2", sample_list) -def test_gradgradcheck_batch(dtype: torch.dtype, name1: str, name2: str) -> None: +def test_gradgradcheck_batch( + dtype: torch.dtype, name1: str, name2: str +) -> None: """ Check a single analytical gradient of parameters against numerical gradient from `torch.autograd.gradgradcheck`. diff --git a/test/test_model/samples.py b/test/test_model/samples.py index c4bdfff..c8dd9cf 100644 --- a/test/test_model/samples.py +++ b/test/test_model/samples.py @@ -19,8 +19,7 @@ import torch from tad_mctc.data.molecules import merge_nested_dicts, mols - -from tad_dftd3.typing import Molecule, Tensor, TypedDict +from tad_mctc.typing import Molecule, Tensor, TypedDict class Refs(TypedDict): diff --git a/test/test_model/test_c6.py b/test/test_model/test_c6.py index f27092d..a66db6b 100644 --- a/test/test_model/test_c6.py +++ b/test/test_model/test_c6.py @@ -17,13 +17,16 @@ """ from __future__ import annotations +from typing import Protocol + import pytest import torch +from tad_mctc._version import __tversion__ from tad_mctc.autograd import dgradcheck, dgradgradcheck from tad_mctc.batch import pack +from tad_mctc.typing import DD, Callable, Tensor from tad_dftd3 import model, ncoord, reference -from tad_dftd3.typing import DD, Callable, Protocol, Tensor from ..conftest import DEVICE, FAST_MODE from .samples import samples @@ -109,6 +112,53 @@ def test_chunked(dtype: torch.dtype, size: int, chunk_size: int) -> None: assert pytest.approx(c6.cpu(), abs=tol, rel=tol) == c6_chunked.cpu() +@pytest.mark.skipif(__tversion__ < (2, 1, 0), reason="Requires PyTorch>=2.1.0") +def test_fail() -> None: + dd: DD = {"device": DEVICE, "dtype": torch.float64} + size = 10 + nbatch = 2 + + ref = reference.Reference(**dd) + numbers = torch.randint(1, 86, (nbatch, size), device=DEVICE) + positions = torch.rand((nbatch, size, 3), **dd) * 10 + + cn = ncoord.cn_d3(numbers, positions) + weights = model.weight_references(numbers, cn, ref) + + def _c6(nums: Tensor, ws: Tensor) -> Tensor: + return model.atomic_c6(nums, ws, ref) + + jac = torch.func.jacrev(_c6, argnums=1) + + # Correct batch dimensions + with pytest.raises(ValueError) as excinfo: + vjac = torch.func.vmap(jac, in_dims=(0, None)) + _ = vjac(numbers, weights.moveaxis(1, 0)) + + assert "Batch size mismatch" in str(excinfo.value) + assert "weights" in str(excinfo.value) + + # Correct batch dimensions + with pytest.raises(ValueError) as excinfo: + vjac = torch.func.vmap(jac, in_dims=(None, 0)) + _ = vjac(numbers.moveaxis(1, 0), weights) + + assert "Batch size mismatch" in str(excinfo.value) + assert "numbers" in str(excinfo.value) + + # Internal vmap errors + + # Batch dimensions is always 0 for numbers and weights + with pytest.raises(ValueError) as excinfo: + vjac = torch.func.vmap(jac, in_dims=(1, 0)) + _ = vjac(numbers, weights) + + # Batch dimensions is always 0 for numbers and weights + with pytest.raises(ValueError): + vjac = torch.func.vmap(jac, in_dims=(0, 1)) + _ = vjac(numbers, weights) + + ############################################################################### diff --git a/test/test_model/test_load.py b/test/test_model/test_load.py index 14637a2..c7a7dc4 100644 --- a/test/test_model/test_load.py +++ b/test/test_model/test_load.py @@ -17,7 +17,7 @@ """ import torch -from tad_dftd3 import data, defaults, reference +from tad_dftd3 import defaults, reference def test_ref() -> None: @@ -25,8 +25,3 @@ def test_ref() -> None: assert c6.shape == torch.Size( (defaults.MAX_ELEMENT, defaults.MAX_ELEMENT, 7, 7), ) - - -def test_vdw() -> None: - vdw = data.radii._load_vdw_rad_d3(dtype=torch.double) - assert vdw.shape == torch.Size((defaults.MAX_ELEMENT, defaults.MAX_ELEMENT)) diff --git a/test/test_model/test_reference.py b/test/test_model/test_reference.py index b92c7c6..f498461 100644 --- a/test/test_model/test_reference.py +++ b/test/test_model/test_reference.py @@ -15,16 +15,15 @@ """ Test the reference. """ -from typing import Optional, Union +from typing import Any, Optional, TypedDict, Union from unittest.mock import patch import pytest import torch from tad_mctc.convert import str_to_device -from tad_mctc.typing import MockTensor +from tad_mctc.typing import DD, MockTensor, Tensor from tad_dftd3 import reference -from tad_dftd3.typing import DD, Any, Tensor, TypedDict from ..conftest import DEVICE diff --git a/test/test_model/test_weights.py b/test/test_model/test_weights.py index 38fe4ba..72bfa47 100644 --- a/test/test_model/test_weights.py +++ b/test/test_model/test_weights.py @@ -18,9 +18,9 @@ import pytest import torch from tad_mctc.batch import pack +from tad_mctc.typing import DD from tad_dftd3 import model, reference -from tad_dftd3.typing import DD from ..conftest import DEVICE from .samples import samples