From 4d4baf694c8412dcd53bd73b8b3394c31534c9c3 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 13 Feb 2026 17:16:56 +0100 Subject: [PATCH 1/5] Use vesin-metatomic in `ASECalculator` --- .../metatomic/torch/ase_calculator.py | 34 +++++++------------ 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/python/metatomic_torch/metatomic/torch/ase_calculator.py b/python/metatomic_torch/metatomic/torch/ase_calculator.py index 76964aaa..0c0997ba 100644 --- a/python/metatomic_torch/metatomic/torch/ase_calculator.py +++ b/python/metatomic_torch/metatomic/torch/ase_calculator.py @@ -7,7 +7,7 @@ import metatensor.torch as mts import numpy as np import torch -import vesin +import vesin.metatomic from metatensor.torch import Labels, TensorBlock, TensorMap from torch.profiler import record_function @@ -385,16 +385,12 @@ def run_model( ) system = System(types, positions, cell, pbc) # Compute the neighbors lists requested by the model - for options in self._model.requested_neighbor_lists(): - neighbors = _compute_ase_neighbors( - atoms, options, dtype=self._dtype, device=self._device - ) - register_autograd_neighbors( - system, - neighbors, - check_consistency=self.parameters["check_consistency"], - ) - system.add_neighbor_list(options, neighbors) + vesin.metatomic.compute_requested_neighbors( + systems=systems, + system_length_unit="angstrom", + model=self._model, + check_consistency=self.parameters["check_consistency"], + ) # Get the additional inputs requested by the model for name, option in self._model.requested_inputs().items(): input_tensormap = _get_ase_input( @@ -537,16 +533,12 @@ def calculate( with record_function("MetatomicCalculator::compute_neighbors"): # convert from ase.Atoms to metatomic.torch.System system = System(types, positions, cell, pbc) - for options in self._model.requested_neighbor_lists(): - neighbors = _compute_ase_neighbors( - atoms, options, dtype=self._dtype, device=self._device - ) - register_autograd_neighbors( - system, - neighbors, - check_consistency=self.parameters["check_consistency"], - ) - system.add_neighbor_list(options, neighbors) + vesin.metatomic.compute_requested_neighbors( + systems=[system], + system_length_unit="angstrom", + model=self._model, + check_consistency=self.parameters["check_consistency"], + ) with record_function("MetatomicCalculator::get_model_inputs"): for name, option in self._model.requested_inputs().items(): From a39c1cf3b80849805898efa16d4cd8e0d95fb54d Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 13 Feb 2026 17:22:41 +0100 Subject: [PATCH 2/5] Fix bugs --- .../metatomic/torch/ase_calculator.py | 92 ++----------------- 1 file changed, 7 insertions(+), 85 deletions(-) diff --git a/python/metatomic_torch/metatomic/torch/ase_calculator.py b/python/metatomic_torch/metatomic/torch/ase_calculator.py index 0c0997ba..298a1958 100644 --- a/python/metatomic_torch/metatomic/torch/ase_calculator.py +++ b/python/metatomic_torch/metatomic/torch/ase_calculator.py @@ -20,7 +20,6 @@ load_atomistic_model, pick_device, pick_output, - register_autograd_neighbors, ) @@ -386,7 +385,7 @@ def run_model( system = System(types, positions, cell, pbc) # Compute the neighbors lists requested by the model vesin.metatomic.compute_requested_neighbors( - systems=systems, + systems=[system], system_length_unit="angstrom", model=self._model, check_consistency=self.parameters["check_consistency"], @@ -719,16 +718,12 @@ def compute_energy( strains.append(strain) system = System(types, positions, cell, pbc) # Compute the neighbors lists requested by the model - for options in self._model.requested_neighbor_lists(): - neighbors = _compute_ase_neighbors( - atoms, options, dtype=self._dtype, device=self._device - ) - register_autograd_neighbors( - system, - neighbors, - check_consistency=self.parameters["check_consistency"], - ) - system.add_neighbor_list(options, neighbors) + vesin.metatomic.compute_requested_neighbors( + systems=[system], + system_length_unit="angstrom", + model=self._model, + check_consistency=self.parameters["check_consistency"], + ) systems.append(system) options = ModelEvaluationOptions( @@ -893,79 +888,6 @@ def _ase_properties_to_metatensor_outputs( return metatensor_outputs -def _compute_ase_neighbors(atoms, options, dtype, device): - # options.strict is ignored by this function, since `ase.neighborlist.neighbor_list` - # only computes strict NL, and these are valid even with `strict=False` - - if np.all(atoms.pbc) or np.all(~atoms.pbc): - nl_i, nl_j, nl_S, nl_D = vesin.ase_neighbor_list( - "ijSD", - atoms, - cutoff=options.engine_cutoff(engine_length_unit="angstrom"), - ) - else: - nl_i, nl_j, nl_S, nl_D = ase.neighborlist.neighbor_list( - "ijSD", - atoms, - cutoff=options.engine_cutoff(engine_length_unit="angstrom"), - ) - - if not options.full_list: - # The pair selection code here below avoids a relatively slow loop over - # all pairs to improve performance - reject_condition = ( - # we want a half neighbor list, so drop all duplicated neighbors - (nl_j < nl_i) - | ( - (nl_i == nl_j) - & ( - # only create pairs with the same atom twice if the pair spans more - # than one unit cell - ((nl_S[:, 0] == 0) & (nl_S[:, 1] == 0) & (nl_S[:, 2] == 0)) - # When creating pairs between an atom and one of its periodic - # images, the code generates multiple redundant pairs - # (e.g. with shifts 0 1 1 and 0 -1 -1); and we want to only keep one - # of these. We keep the pair in the positive half plane of shifts. - | ( - (nl_S.sum(axis=1) < 0) - | ( - (nl_S.sum(axis=1) == 0) - & ( - (nl_S[:, 2] < 0) - | ((nl_S[:, 2] == 0) & (nl_S[:, 1] < 0)) - ) - ) - ) - ) - ) - ) - selected = np.logical_not(reject_condition) - nl_i = nl_i[selected] - nl_j = nl_j[selected] - nl_S = nl_S[selected] - nl_D = nl_D[selected] - - samples = np.concatenate([nl_i[:, None], nl_j[:, None], nl_S], axis=1) - distances = torch.from_numpy(nl_D).to(dtype=dtype, device=device) - - return TensorBlock( - values=distances.reshape(-1, 3, 1), - samples=Labels( - names=[ - "first_atom", - "second_atom", - "cell_shift_a", - "cell_shift_b", - "cell_shift_c", - ], - values=torch.from_numpy(samples).to(dtype=torch.int32, device=device), - assume_unique=True, - ), - components=[Labels.range("xyz", 3).to(device)], - properties=Labels.range("distance", 1).to(device), - ) - - def _get_ase_input( atoms: ase.Atoms, name: str, From 552a61e190b1a1c05697ae95fc565fef5ed8f4bc Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 13 Feb 2026 17:49:00 +0100 Subject: [PATCH 3/5] Require torch 2.3 (like vesin-torch) --- docs/src/installation.rst | 4 ++-- metatomic-torch/CMakeLists.txt | 2 +- python/metatomic_torch/build-backend/backend.py | 2 +- python/metatomic_torch/setup.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/src/installation.rst b/docs/src/installation.rst index 8555cf16..38b0bf21 100644 --- a/docs/src/installation.rst +++ b/docs/src/installation.rst @@ -23,7 +23,7 @@ and use it will depend on the programming language you are using. We provide pre-compiled wheels on PyPI that are compatible with all the supported ``torch`` versions at the time of the ``metatomic-torch`` release. - Currently PyTorch version 2.1 and above is supported. + Currently PyTorch version 2.3 and above is supported. If you want to use the code with an unsupported PyTorch version, or a new release of PyTorch which did not exist yet when we released @@ -105,7 +105,7 @@ and use it will depend on the programming language you are using. metatensor-torch. - the C++ part of PyTorch, which you can install `on its own `_. We are compatible with - ``libtorch`` version 2.1 or above. You can also use the same library as + ``libtorch`` version 2.3 or above. You can also use the same library as the Python version of torch by adding the output of the command below to ``CMAKE_PREFIX_PATH``: diff --git a/metatomic-torch/CMakeLists.txt b/metatomic-torch/CMakeLists.txt index 57fd9d46..8f253344 100644 --- a/metatomic-torch/CMakeLists.txt +++ b/metatomic-torch/CMakeLists.txt @@ -91,7 +91,7 @@ endif() # fixed version in `cmake/FindCUDNN.cmake` set(CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake;${CMAKE_MODULE_PATH}") -find_package(Torch 2.1 REQUIRED) +find_package(Torch 2.3 REQUIRED) set(METATOMIC_TORCH_HEADERS "include/metatomic/torch/system.hpp" diff --git a/python/metatomic_torch/build-backend/backend.py b/python/metatomic_torch/build-backend/backend.py index b350ac20..c762d91e 100644 --- a/python/metatomic_torch/build-backend/backend.py +++ b/python/metatomic_torch/build-backend/backend.py @@ -11,7 +11,7 @@ if FORCED_TORCH_VERSION is not None: TORCH_DEP = f"torch =={FORCED_TORCH_VERSION}" else: - TORCH_DEP = "torch >=2.1" + TORCH_DEP = "torch >=2.3" # ==================================================================================== # # Build backend functions definition # diff --git a/python/metatomic_torch/setup.py b/python/metatomic_torch/setup.py index 28a2cc55..defd269f 100644 --- a/python/metatomic_torch/setup.py +++ b/python/metatomic_torch/setup.py @@ -312,7 +312,7 @@ def create_version_number(version): torch_version = f"== {torch_v_major}.{torch_v_minor}.*" except ImportError: # otherwise we are building a sdist - torch_version = ">= 2.1" + torch_version = ">= 2.3" install_requires = [ f"torch {torch_version}", From 49ba2305795101dc9f8d5fcafaf8479622a76c88 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 13 Feb 2026 17:50:53 +0100 Subject: [PATCH 4/5] Call NLs outside the loop for `run_model()` and `compute_energy()` --- .../metatomic/torch/ase_calculator.py | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/python/metatomic_torch/metatomic/torch/ase_calculator.py b/python/metatomic_torch/metatomic/torch/ase_calculator.py index 298a1958..2c8422b0 100644 --- a/python/metatomic_torch/metatomic/torch/ase_calculator.py +++ b/python/metatomic_torch/metatomic/torch/ase_calculator.py @@ -383,13 +383,6 @@ def run_model( atoms=atoms, dtype=self._dtype, device=self._device ) system = System(types, positions, cell, pbc) - # Compute the neighbors lists requested by the model - vesin.metatomic.compute_requested_neighbors( - systems=[system], - system_length_unit="angstrom", - model=self._model, - check_consistency=self.parameters["check_consistency"], - ) # Get the additional inputs requested by the model for name, option in self._model.requested_inputs().items(): input_tensormap = _get_ase_input( @@ -398,6 +391,14 @@ def run_model( system.add_data(name, input_tensormap) systems.append(system) + # Compute the neighbors lists requested by the model + vesin.metatomic.compute_requested_neighbors( + systems=systems, + system_length_unit="angstrom", + model=self._model, + check_consistency=self.parameters["check_consistency"], + ) + available_outputs = self._model.capabilities().outputs for key in outputs: if key not in available_outputs: @@ -717,15 +718,16 @@ def compute_energy( cell = cell @ strain strains.append(strain) system = System(types, positions, cell, pbc) - # Compute the neighbors lists requested by the model - vesin.metatomic.compute_requested_neighbors( - systems=[system], - system_length_unit="angstrom", - model=self._model, - check_consistency=self.parameters["check_consistency"], - ) systems.append(system) + # Compute the neighbors lists requested by the model + vesin.metatomic.compute_requested_neighbors( + systems=systems, + system_length_unit="angstrom", + model=self._model, + check_consistency=self.parameters["check_consistency"], + ) + options = ModelEvaluationOptions( length_unit="angstrom", outputs=outputs, From 903d616feaf02dee6277715b0ec79d1deeda22cb Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 13 Feb 2026 17:57:09 +0100 Subject: [PATCH 5/5] Set minimum version for torch to 2.3 in the CI too --- .github/workflows/build-wheels.yml | 4 +--- .github/workflows/torch-tests.yml | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml index af230581..467493e9 100644 --- a/.github/workflows/build-wheels.yml +++ b/.github/workflows/build-wheels.yml @@ -31,7 +31,7 @@ jobs: matrix: os: ['ubuntu-24.04', 'ubuntu-24.04-arm', 'macos-15-intel', 'macos-15', 'windows-2022'] arch: ['arm64', 'x86_64'] - torch-version: ['2.1', '2.2', '2.3', '2.4', '2.5', '2.6', '2.7', '2.8', '2.9', '2.10'] + torch-version: ['2.3', '2.4', '2.5', '2.6', '2.7', '2.8', '2.9', '2.10'] exclude: # remove mismatched arch-os combinations - {os: macos-15-intel, arch: arm64} @@ -76,8 +76,6 @@ jobs: rust-target: x86_64-pc-windows-msvc cibw-arch: AMD64 # add the right python version image for each torch version - - {torch-version: '2.1', cibw-python: 'cp311-*'} - - {torch-version: '2.2', cibw-python: 'cp312-*'} - {torch-version: '2.3', cibw-python: 'cp312-*'} - {torch-version: '2.4', cibw-python: 'cp312-*'} - {torch-version: '2.5', cibw-python: 'cp312-*'} diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index 994b461b..4aa30a9a 100644 --- a/.github/workflows/torch-tests.yml +++ b/.github/workflows/torch-tests.yml @@ -19,7 +19,7 @@ jobs: include: - os: ubuntu-24.04 python-version: "3.10" - torch-version: "2.1" + torch-version: "2.3" numpy-version-pin: "<2.0" - os: ubuntu-24.04 python-version: "3.10"