From 711086b3a8354cedaeca5888b964262abfe6b172 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Wed, 30 Oct 2024 18:14:54 +0100 Subject: [PATCH] Update to metatensor-torch v0.6.0 This adds a `pbc` argument to System (for now only all True or all False is supported) and a `strict` option to the neighbor list. We request non-strict NL since we will make a copy & re-filter the neighbor lists anyway. --- docs/requirements.txt | 4 +- pyproject.toml | 2 +- .../rascaline-torch/build-backend/backend.py | 2 +- .../rascaline/torch/calculator_base.py | 3 ++ .../rascaline-torch/rascaline/torch/system.py | 6 +++ python/rascaline-torch/setup.py | 2 +- python/rascaline-torch/tests/autograd.py | 51 ++++++++++--------- python/rascaline-torch/tests/calculator.py | 2 +- python/rascaline-torch/tests/export.py | 3 +- .../rascaline-torch/tests/utils/cg_product.py | 7 ++- .../tests/utils/density_correlations.py | 1 + .../tests/utils/power_spectrum.py | 1 + .../tests/utils/density_correlations.py | 11 ++-- rascaline-c-api/CMakeLists.txt | 4 +- rascaline-torch/CMakeLists.txt | 6 +-- rascaline-torch/src/system.cpp | 21 ++++---- rascaline-torch/tests/calculator.cpp | 4 +- .../tests/cmake-project/src/main.cpp | 3 +- tox.ini | 5 +- 19 files changed, 78 insertions(+), 60 deletions(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index 0cf9c8fe5..3fe557a58 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -10,8 +10,8 @@ myst-parser # markdown => rst translation, used in extensions/rascaline_json # dependencies for the tutorials --extra-index-url https://download.pytorch.org/whl/cpu -metatensor -metatensor-torch >= 0.5.0,<0.6.0 +metatensor-operations >=0.3.0,<0.4.0 +metatensor-torch >= 0.6.0,<0.7.0 torch chemfiles matplotlib diff --git a/pyproject.toml b/pyproject.toml index e5f67df50..ff3326d3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ classifiers = [ dependencies = [ "metatensor-core >=0.1.0,<0.2.0", - "metatensor-operations >=0.2.0,<0.3.0", + "metatensor-operations >=0.3.0,<0.4.0", "wigners", ] diff --git a/python/rascaline-torch/build-backend/backend.py b/python/rascaline-torch/build-backend/backend.py index d90ac3dee..72c8bf849 100644 --- a/python/rascaline-torch/build-backend/backend.py +++ b/python/rascaline-torch/build-backend/backend.py @@ -30,7 +30,7 @@ def get_requires_for_build_wheel(config_settings=None): defaults = build_meta.get_requires_for_build_wheel(config_settings) return defaults + [ "torch >= 1.12", - "metatensor-torch >=0.5.0,<0.6.0", + "metatensor-torch >=0.6.0,<0.7.0", RASCALINE_DEP, ] diff --git a/python/rascaline-torch/rascaline/torch/calculator_base.py b/python/rascaline-torch/rascaline/torch/calculator_base.py index 33ddfd386..47015d13b 100644 --- a/python/rascaline-torch/rascaline/torch/calculator_base.py +++ b/python/rascaline-torch/rascaline/torch/calculator_base.py @@ -90,6 +90,9 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]: NeighborListOptions( cutoff=cutoff, full_list=False, + # we will re-filter the NL when converting to rascaline internal + # type, so we don't need the engine to pre-filter it for us + strict=False, requestor="rascaline", ) ) diff --git a/python/rascaline-torch/rascaline/torch/system.py b/python/rascaline-torch/rascaline/torch/system.py index 49c2f6456..f0bba1599 100644 --- a/python/rascaline-torch/rascaline/torch/system.py +++ b/python/rascaline-torch/rascaline/torch/system.py @@ -1,5 +1,6 @@ from typing import List, Optional, Sequence, overload +import numpy as np import torch from metatensor.torch.atomistic import System from packaging import version @@ -65,6 +66,11 @@ def _system_to_torch(system, positions_requires_grad, cell_requires_grad): types=torch.tensor(system.types()), positions=torch.tensor(system.positions()), cell=torch.tensor(system.cell()), + pbc=( + torch.tensor([False, False, False]) + if np.all(system.cell() == 0.0) + else torch.tensor([True, True, True]) + ), ) if positions_requires_grad is not None: diff --git a/python/rascaline-torch/setup.py b/python/rascaline-torch/setup.py index cbce9ef4a..b6fe7a009 100644 --- a/python/rascaline-torch/setup.py +++ b/python/rascaline-torch/setup.py @@ -243,7 +243,7 @@ def git_extra_version(): install_requires = [ "torch >= 1.12", - "metatensor-torch >=0.5.0,<0.6.0", + "metatensor-torch >=0.6.0,<0.7.0", ] if os.path.exists(RASCALINE_C_API): # we are building from a git checkout diff --git a/python/rascaline-torch/tests/autograd.py b/python/rascaline-torch/tests/autograd.py index 2f4bf8845..05ce80129 100644 --- a/python/rascaline-torch/tests/autograd.py +++ b/python/rascaline-torch/tests/autograd.py @@ -29,12 +29,13 @@ def _create_random_system(n_atoms, cell_size): cell = torch.tensor(cell[:], dtype=torch.float64) positions = torch.rand((n_atoms, 3), dtype=torch.float64) @ cell + pbc = torch.tensor([True, True, True]) - return types, positions, cell + return types, positions, cell, pbc -def _compute_spherical_expansion(types, positions, cell): - system = System(types=types, positions=positions, cell=cell) +def _compute_spherical_expansion(types, positions, cell, pbc): + system = System(types=types, positions=positions, cell=cell, pbc=pbc) calculator = SphericalExpansion(**HYPERS) descriptor = calculator(system) @@ -47,8 +48,8 @@ def _compute_spherical_expansion(types, positions, cell): return descriptor.block(0).values -def _compute_power_spectrum(types, positions, cell): - system = System(types=types, positions=positions, cell=cell) +def _compute_power_spectrum(types, positions, cell, pbc): + system = System(types=types, positions=positions, cell=cell, pbc=pbc) calculator = SoapPowerSpectrum(**HYPERS) descriptor = calculator(system) @@ -59,56 +60,57 @@ def _compute_power_spectrum(types, positions, cell): def test_spherical_expansion_positions_grad(): - types, positions, cell = _create_random_system(n_atoms=75, cell_size=5.0) + types, positions, cell, pbc = _create_random_system(n_atoms=75, cell_size=5.0) positions.requires_grad = True assert torch.autograd.gradcheck( _compute_spherical_expansion, - (types, positions, cell), + (types, positions, cell, pbc), fast_mode=True, ) def test_spherical_expansion_cell_grad(): - types, positions, cell = _create_random_system(n_atoms=75, cell_size=5.0) + types, positions, cell, pbc = _create_random_system(n_atoms=75, cell_size=5.0) cell.requires_grad = True assert torch.autograd.gradcheck( _compute_spherical_expansion, - (types, positions, cell), + (types, positions, cell, pbc), fast_mode=True, ) def test_power_spectrum_positions_grad(): - types, positions, cell = _create_random_system(n_atoms=75, cell_size=5.0) + types, positions, cell, pbc = _create_random_system(n_atoms=75, cell_size=5.0) positions.requires_grad = True assert torch.autograd.gradcheck( _compute_power_spectrum, - (types, positions, cell), + (types, positions, cell, pbc), fast_mode=True, ) def test_power_spectrum_cell_grad(): - types, positions, cell = _create_random_system(n_atoms=75, cell_size=5.0) + types, positions, cell, pbc = _create_random_system(n_atoms=75, cell_size=5.0) cell.requires_grad = True assert torch.autograd.gradcheck( _compute_power_spectrum, - (types, positions, cell), + (types, positions, cell, pbc), fast_mode=True, ) def test_power_spectrum_register_autograd(): # check autograd when registering the graph after pre-computing a representation - types, positions, cell = _create_random_system(n_atoms=75, cell_size=5.0) + types, positions, cell, pbc = _create_random_system(n_atoms=75, cell_size=5.0) calculator = SoapPowerSpectrum(**HYPERS) precomputed = calculator( - System(types, positions, cell), gradients=["positions", "cell"] + System(types, positions, cell, pbc), + gradients=["positions", "cell"], ) # no grad_fn for now @@ -118,7 +120,7 @@ def compute(new_positions, new_cell): same_positions = (new_positions - positions).norm() < 1e-30 same_cell = (new_cell - cell).norm() < 1e-30 - system = System(types=types, positions=positions, cell=cell) + system = System(types, positions, cell, pbc) if same_positions and same_cell: # we can only re-use the calculation when working with the same input @@ -146,14 +148,14 @@ def compute(new_positions, new_cell): def test_power_spectrum_positions_grad_grad(): - types, positions, cell = _create_random_system(n_atoms=75, cell_size=5.0) + types, positions, cell, pbc = _create_random_system(n_atoms=75, cell_size=5.0) positions.requires_grad = True - X = _compute_power_spectrum(types, positions, cell) + X = _compute_power_spectrum(types, positions, cell, pbc) weights = torch.rand((X.shape[-1], 1), requires_grad=True, dtype=torch.float64) def compute(weights): - X = _compute_power_spectrum(types, positions, cell) + X = _compute_power_spectrum(types, positions, cell, pbc) A = X @ weights return torch.autograd.grad( @@ -184,14 +186,14 @@ def compute(weights): def test_power_spectrum_cell_grad_grad(): - types, positions, cell = _create_random_system(n_atoms=75, cell_size=5.0) + types, positions, cell, pbc = _create_random_system(n_atoms=75, cell_size=5.0) cell.requires_grad = True - X = _compute_power_spectrum(types, positions, cell) + X = _compute_power_spectrum(types, positions, cell, pbc) weights = torch.rand((X.shape[-1], 1), requires_grad=True, dtype=torch.float64) def compute(weights): - X = _compute_power_spectrum(types, positions, cell) + X = _compute_power_spectrum(types, positions, cell, pbc) A = X @ weights return torch.autograd.grad( @@ -244,7 +246,7 @@ def test_different_device_dtype(): options.append((torch.device("cuda:0"), torch.float64)) for device, dtype in options: - types, positions, cell = _create_random_system(n_atoms=10, cell_size=3.0) + types, positions, cell, pbc = _create_random_system(n_atoms=10, cell_size=3.0) positions = positions.to(dtype=dtype, device=device, copy=True) positions.requires_grad = True assert positions.grad is None @@ -254,11 +256,12 @@ def test_different_device_dtype(): assert cell.grad is None types = types.to(device=device, copy=True) + pbc = pbc.to(device=device, copy=True) with warnings.catch_warnings(): warnings.filterwarnings("ignore") - X = _compute_power_spectrum(types, positions, cell) + X = _compute_power_spectrum(types, positions, cell, pbc) assert X.dtype == dtype assert X.device == device diff --git a/python/rascaline-torch/tests/calculator.py b/python/rascaline-torch/tests/calculator.py index 8691369e6..392866229 100644 --- a/python/rascaline-torch/tests/calculator.py +++ b/python/rascaline-torch/tests/calculator.py @@ -16,6 +16,7 @@ def system(): types=torch.tensor([1, 1, 8, 8]), positions=torch.tensor([[0.0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), cell=torch.tensor([[10.0, 0, 0], [0, 10, 0], [0, 0, 10]]), + pbc=torch.tensor([True, True, True]), ) @@ -188,7 +189,6 @@ def test_different_device_dtype_errors(system): custom_device = torch.device("cuda:0") if custom_device is not None: - device_system = system.to(device=custom_device) torch.set_warn_always(True) diff --git a/python/rascaline-torch/tests/export.py b/python/rascaline-torch/tests/export.py index 3d61b84d0..b23fa192a 100644 --- a/python/rascaline-torch/tests/export.py +++ b/python/rascaline-torch/tests/export.py @@ -87,9 +87,10 @@ def test_export_as_metatensor_model(tmpdir): model = Model(types=[1, 6, 8]) model.eval() - energy_output = ModelOutput() + energy_output = ModelOutput(quantity="energy", unit="eV") capabilities = ModelCapabilities( supported_devices=["cpu"], + length_unit="A", interaction_range=HYPERS["cutoff"], atomic_types=[1, 6, 8], dtype="float64", diff --git a/python/rascaline-torch/tests/utils/cg_product.py b/python/rascaline-torch/tests/utils/cg_product.py index f81921d51..9263da8a5 100644 --- a/python/rascaline-torch/tests/utils/cg_product.py +++ b/python/rascaline-torch/tests/utils/cg_product.py @@ -32,11 +32,10 @@ def system(): [2.56633400, 2.50000000, 2.50370100], [1.97361700, 1.73067300, 2.47063400], [1.97361700, 3.26932700, 2.47063400], - ], - ), - cell=torch.zeros( - (3, 3), + ] ), + cell=torch.zeros((3, 3)), + pbc=torch.tensor([False, False, False]), ) diff --git a/python/rascaline-torch/tests/utils/density_correlations.py b/python/rascaline-torch/tests/utils/density_correlations.py index 171d1fe41..0063ec5fb 100644 --- a/python/rascaline-torch/tests/utils/density_correlations.py +++ b/python/rascaline-torch/tests/utils/density_correlations.py @@ -35,6 +35,7 @@ def system(): ] ), cell=torch.zeros((3, 3)), + pbc=torch.tensor([False, False, False]), ) diff --git a/python/rascaline-torch/tests/utils/power_spectrum.py b/python/rascaline-torch/tests/utils/power_spectrum.py index ebccc5599..de14cff94 100644 --- a/python/rascaline-torch/tests/utils/power_spectrum.py +++ b/python/rascaline-torch/tests/utils/power_spectrum.py @@ -11,6 +11,7 @@ def system(): types=torch.tensor([1, 1, 8, 8]), positions=torch.tensor([[0.0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), cell=torch.tensor([[10.0, 0, 0], [0, 10, 0], [0, 0, 10]]), + pbc=torch.tensor([True, True, True]), ) diff --git a/python/rascaline/tests/utils/density_correlations.py b/python/rascaline/tests/utils/density_correlations.py index ab0c43068..879b6fda3 100644 --- a/python/rascaline/tests/utils/density_correlations.py +++ b/python/rascaline/tests/utils/density_correlations.py @@ -97,7 +97,6 @@ def h2o_isolated(): def h2o_periodic(): - return [ ase.Atoms( symbols=["O", "H", "H"], @@ -353,7 +352,7 @@ def test_correlate_density_norm(): unique_samples = metatensor.unique_metadata( ps, "samples", names=["system", "atom", "center_type"] ) - grouped_labels = [ + selections = [ Labels(names=ps.sample_names, values=unique_samples.values[i].reshape(1, 3)) for i in range(len(unique_samples)) ] @@ -362,11 +361,11 @@ def test_correlate_density_norm(): norm_nu1 = 0.0 norm_ps = 0.0 norm_ps_sorted = 0.0 - for sample in grouped_labels: + for selection in selections: # Slice the TensorMaps - nu1_sliced = metatensor.slice(density, "samples", labels=sample) - ps_sliced = metatensor.slice(ps, "samples", labels=sample) - ps_sorted_sliced = metatensor.slice(ps_sorted, "samples", labels=sample) + nu1_sliced = metatensor.slice(density, "samples", selection=selection) + ps_sliced = metatensor.slice(ps, "samples", selection=selection) + ps_sorted_sliced = metatensor.slice(ps_sorted, "samples", selection=selection) # Calculate norms norm_nu1 += get_norm(nu1_sliced) ** (n_correlations + 1) diff --git a/rascaline-c-api/CMakeLists.txt b/rascaline-c-api/CMakeLists.txt index 41916f865..74a2e2db7 100644 --- a/rascaline-c-api/CMakeLists.txt +++ b/rascaline-c-api/CMakeLists.txt @@ -222,7 +222,7 @@ endif() # ============================================================================ # # Setup metatensor -set(METATENSOR_FETCH_VERSION "0.1.10") +set(METATENSOR_FETCH_VERSION "0.1.11") set(METATENSOR_REQUIRED_VERSION "0.1") if (RASCALINE_FETCH_METATENSOR) message(STATUS "Fetching metatensor @ ${METATENSOR_FETCH_VERSION} from github") @@ -232,7 +232,7 @@ if (RASCALINE_FETCH_METATENSOR) FetchContent_Declare( metatensor URL ${URL_ROOT}/metatensor-core-v${METATENSOR_FETCH_VERSION}/metatensor-core-cxx-${METATENSOR_FETCH_VERSION}.tar.gz - URL_HASH SHA256=3ec0775da67bb0eb3246b81770426e612f83b6591442a39eb17aad6969b5f9d9 + URL_HASH SHA256=b79435dbeb59a95361b4a5881574dfde4a4ed13f12b05a225df161dcaa7ea61e ) if (CMAKE_VERSION VERSION_GREATER 3.18) diff --git a/rascaline-torch/CMakeLists.txt b/rascaline-torch/CMakeLists.txt index b199ddd97..4a42a60f6 100644 --- a/rascaline-torch/CMakeLists.txt +++ b/rascaline-torch/CMakeLists.txt @@ -58,8 +58,8 @@ find_package(Torch 1.12 REQUIRED) # ============================================================================ # # Setup metatensor_torch -set(METATENSOR_FETCH_VERSION "0.5.5") -set(REQUIRED_METATENSOR_TORCH_VERSION "0.5") +set(METATENSOR_FETCH_VERSION "0.6.0") +set(REQUIRED_METATENSOR_TORCH_VERSION "0.6") if (RASCALINE_TORCH_FETCH_METATENSOR_TORCH) message(STATUS "Fetching metatensor-torch @ ${METATENSOR_FETCH_VERSION} from github") @@ -68,7 +68,7 @@ if (RASCALINE_TORCH_FETCH_METATENSOR_TORCH) FetchContent_Declare( metatensor_torch URL ${URL_ROOT}/metatensor-torch-v${METATENSOR_FETCH_VERSION}/metatensor-torch-cxx-${METATENSOR_FETCH_VERSION}.tar.gz - URL_HASH SHA256=dac306ab59ac8b59167827405f468397dbf0d4a69988fce7b9f4285f2816a57c + URL_HASH SHA256=f050743662ece38948b2087dd025d60110645716840dbfc5370c059e1275d0cf ) if (CMAKE_VERSION VERSION_GREATER 3.18) diff --git a/rascaline-torch/src/system.cpp b/rascaline-torch/src/system.cpp index 55dbe3b04..706ee609c 100644 --- a/rascaline-torch/src/system.cpp +++ b/rascaline-torch/src/system.cpp @@ -14,6 +14,7 @@ SystemAdapter::SystemAdapter(metatensor_torch::System system): system_(std::move for (const auto& options: system_->known_neighbor_lists()) { for (const auto& requestor: options->requestors()) { if (requestor == "rascaline") { + auto cutoff = options->cutoff(); auto neighbors = system_->get_neighbor_list(options); auto samples_values = neighbors->samples()->values().to(torch::kCPU).contiguous(); auto samples = samples_values.accessor(); @@ -24,7 +25,6 @@ SystemAdapter::SystemAdapter(metatensor_torch::System system): system_(std::move auto n_pairs = samples.size(0); auto pairs = std::vector(); - pairs.reserve(static_cast(n_pairs)); for (int64_t i=0; isize()); pair.distance = std::sqrt(x*x + y*y + z*z); - pair.vector[0] = x; - pair.vector[1] = y; - pair.vector[2] = z; - pair.cell_shift_indices[0] = samples[i][2]; - pair.cell_shift_indices[1] = samples[i][3]; - pair.cell_shift_indices[2] = samples[i][4]; + if (pair.distance < cutoff) { + pair.vector[0] = x; + pair.vector[1] = y; + pair.vector[2] = z; - pairs.emplace_back(pair); + pair.cell_shift_indices[0] = samples[i][2]; + pair.cell_shift_indices[1] = samples[i][3]; + pair.cell_shift_indices[2] = samples[i][4]; + + pairs.emplace_back(pair); + } } - this->set_precomputed_pairs(options->cutoff(), std::move(pairs)); + this->set_precomputed_pairs(cutoff, std::move(pairs)); continue; } } diff --git a/rascaline-torch/tests/calculator.cpp b/rascaline-torch/tests/calculator.cpp index 2c3078a49..c86e1818a 100644 --- a/rascaline-torch/tests/calculator.cpp +++ b/rascaline-torch/tests/calculator.cpp @@ -299,5 +299,7 @@ metatensor_torch::System test_system(bool positions_grad, bool cell_grad) { auto cell = 10 * torch::eye(3); cell.requires_grad_(cell_grad); - return torch::make_intrusive(types, positions, cell); + auto pbc = torch::ones(3, torch::TensorOptions().dtype(torch::kBool)); + + return torch::make_intrusive(types, positions, cell, pbc); } diff --git a/rascaline-torch/tests/cmake-project/src/main.cpp b/rascaline-torch/tests/cmake-project/src/main.cpp index b9e778918..02f1e32c8 100644 --- a/rascaline-torch/tests/cmake-project/src/main.cpp +++ b/rascaline-torch/tests/cmake-project/src/main.cpp @@ -9,7 +9,8 @@ int main() { auto system = torch::make_intrusive( torch::zeros({5}, torch::kI32), torch::rand({5, 3}, torch::kF64), - torch::zeros({3, 3}, torch::kF64) + torch::zeros({3, 3}, torch::kF64), + torch::zeros({3}, torch::kBool) ); const auto* HYPERS_JSON = R"({ diff --git a/tox.ini b/tox.ini index 69cbebfed..c4a5c555d 100644 --- a/tox.ini +++ b/tox.ini @@ -24,7 +24,7 @@ metatensor-core-requirement = metatensor-core >=0.1.0,<0.2.0 metatensor-torch-requirement = - metatensor-torch >=0.5.0,<0.6.0 + metatensor-torch >=0.6.0,<0.7.0 build-single-wheel = --no-deps --no-build-isolation --check-build-dependencies warning_options = @@ -77,8 +77,7 @@ deps = torch pyscf;platform_system!="Windows" wigners - # TODO: add mops once it becomes stable enough (and potentially supports windows) - #mops@git+https://github.com/lab-cosmo/mops ; platform_system!="Windows" + commands = pytest {[testenv]test_options} {posargs}