From 913e69843d6b8b95a2275e9b258f3b807ceea858 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Mon, 4 Sep 2023 17:07:56 +0200 Subject: [PATCH] Update metatensor --- docs/requirements.txt | 2 +- pyproject.toml | 2 +- python/rascaline-torch/pyproject.toml | 2 +- .../utils/power_spectrum/calculator.py | 8 +++--- rascaline-c-api/CMakeLists.txt | 4 +-- rascaline-c-api/Cargo.toml | 4 +-- rascaline-torch/CMakeLists.txt | 4 +-- rascaline-torch/src/autograd.cpp | 9 ++++--- rascaline-torch/src/calculator.cpp | 11 ++++---- rascaline-torch/tests/calculator.cpp | 26 ++++++++++--------- rascaline/Cargo.toml | 2 +- tox.ini | 4 +-- 12 files changed, 40 insertions(+), 38 deletions(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index b721d6ad8..d5253562f 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -10,7 +10,7 @@ myst-parser # markdown => rst translation, used in extensions/rascaline_json # dependencies for the tutorials --extra-index-url https://download.pytorch.org/whl/cpu -metatensor[torch] @ https://github.com/lab-cosmo/metatensor/archive/32ad5bb.zip +metatensor[torch] @ https://github.com/lab-cosmo/metatensor/archive/d97ea65.zip torch chemfiles matplotlib diff --git a/pyproject.toml b/pyproject.toml index 01f012eb7..66f43c20a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] dependencies = [ - "metatensor-core @ https://github.com/lab-cosmo/metatensor/archive/32ad5bb.zip#subdirectory=python/metatensor-core", + "metatensor-core @ https://github.com/lab-cosmo/metatensor/archive/d97ea65.zip#subdirectory=python/metatensor-core", ] [project.urls] diff --git a/python/rascaline-torch/pyproject.toml b/python/rascaline-torch/pyproject.toml index f7fd6c9b8..b588cb504 100644 --- a/python/rascaline-torch/pyproject.toml +++ b/python/rascaline-torch/pyproject.toml @@ -38,7 +38,7 @@ requires = [ "wheel >=0.38", "cmake", "torch >= 1.11", - "metatensor-torch @ https://github.com/lab-cosmo/metatensor/archive/32ad5bb.zip#subdirectory=python/metatensor-torch", + "metatensor-torch @ https://github.com/lab-cosmo/metatensor/archive/d97ea65.zip#subdirectory=python/metatensor-torch", ] # use a custom build backend to add a dependency on the right version of rascaline diff --git a/python/rascaline/rascaline/utils/power_spectrum/calculator.py b/python/rascaline/rascaline/utils/power_spectrum/calculator.py index ad522337f..c6b872ecd 100644 --- a/python/rascaline/rascaline/utils/power_spectrum/calculator.py +++ b/python/rascaline/rascaline/utils/power_spectrum/calculator.py @@ -176,10 +176,8 @@ def compute( "species_neighbor", ] - # TODO: re-enable once we update metatensor with - # https://github.com/lab-cosmo/metatensor/pull/322 - # assert spherical_expansion_1.keys.names == expected_key_names - # assert spherical_expansion_1.property_names == ["n"] + assert spherical_expansion_1.keys.names == expected_key_names + assert spherical_expansion_1.properties_names == ["n"] # Fill blocks with `species_neighbor` from ALL blocks. If we don't do this # merging blocks along the ``sample`` direction might be not possible. @@ -200,7 +198,7 @@ def compute( use_native_system=use_native_system, ) assert spherical_expansion_2.keys.names == expected_key_names - assert spherical_expansion_2.property_names == ["n"] + assert spherical_expansion_2.properties_names == ["n"] array = spherical_expansion_2.keys.column("species_neighbor") keys_to_move = Labels( diff --git a/rascaline-c-api/CMakeLists.txt b/rascaline-c-api/CMakeLists.txt index a6fac72a3..db90e4b95 100644 --- a/rascaline-c-api/CMakeLists.txt +++ b/rascaline-c-api/CMakeLists.txt @@ -217,7 +217,7 @@ endif() # ============================================================================ # # Setup metatensor -set(METATENSOR_GIT_VERSION "32ad5bb") +set(METATENSOR_GIT_VERSION "d97ea65") set(METATENSOR_REQUIRED_VERSION "0.1") if (RASCALINE_FETCH_METATENSOR) message(STATUS "Fetching metatensor @ ${METATENSOR_GIT_VERSION} from github") @@ -226,7 +226,7 @@ if (RASCALINE_FETCH_METATENSOR) FetchContent_Declare( metatensor URL https://github.com/lab-cosmo/metatensor/archive/${METATENSOR_GIT_VERSION}.zip - URL_HASH MD5=cbc7bd27e9e2307638405d1613fa7f89 + URL_HASH MD5=6a6899779591ae15861bb3547c7c354b SOURCE_SUBDIR metatensor-core VERBOSE ) diff --git a/rascaline-c-api/Cargo.toml b/rascaline-c-api/Cargo.toml index a515f471e..59e195a0f 100644 --- a/rascaline-c-api/Cargo.toml +++ b/rascaline-c-api/Cargo.toml @@ -18,7 +18,7 @@ chemfiles = ["rascaline/chemfiles"] [dependencies] rascaline = {path = "../rascaline", version = "0.1.0", default-features = false} -metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "32ad5bb"} +metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "d97ea65"} ndarray = "0.15" log = { version = "0.4", features = ["std"] } @@ -29,7 +29,7 @@ libc = "0.2" [build-dependencies] cbindgen = { version = "0.24", default-features = false } fs_extra = "1" -metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "32ad5bb"} +metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "d97ea65"} [dev-dependencies] which = "4" diff --git a/rascaline-torch/CMakeLists.txt b/rascaline-torch/CMakeLists.txt index ed7fe8019..ad23ae3cd 100644 --- a/rascaline-torch/CMakeLists.txt +++ b/rascaline-torch/CMakeLists.txt @@ -58,7 +58,7 @@ find_package(Torch 1.11 REQUIRED) # ============================================================================ # # Setup metatensor_torch -set(METATENSOR_GIT_VERSION "32ad5bb") +set(METATENSOR_GIT_VERSION "d97ea65") set(REQUIRED_METATENSOR_TORCH_VERSION "0.1") if (RASCALINE_TORCH_FETCH_METATENSOR_TORCH) message(STATUS "Fetching metatensor_torch @ ${METATENSOR_GIT_VERSION} from github") @@ -67,7 +67,7 @@ if (RASCALINE_TORCH_FETCH_METATENSOR_TORCH) FetchContent_Declare( metatensor_torch URL https://github.com/lab-cosmo/metatensor/archive/${METATENSOR_GIT_VERSION}.zip - URL_HASH MD5=cbc7bd27e9e2307638405d1613fa7f89 + URL_HASH MD5=6a6899779591ae15861bb3547c7c354b SOURCE_SUBDIR metatensor-torch VERBOSE ) diff --git a/rascaline-torch/src/autograd.cpp b/rascaline-torch/src/autograd.cpp index 6092793bf..91b79151b 100644 --- a/rascaline-torch/src/autograd.cpp +++ b/rascaline-torch/src/autograd.cpp @@ -1,6 +1,7 @@ #include #include +#include "metatensor/torch/tensor.hpp" #include "rascaline/torch/autograd.hpp" using namespace metatensor_torch; @@ -73,8 +74,8 @@ static std::vector extract_gradient_blocks( ) { auto gradients = std::vector(); for (int64_t i=0; ikeys()->count(); i++) { - auto block = tensor->block_by_id(i); - auto gradient = block->gradient(parameter); + auto block = TensorMapHolder::block_by_id(tensor, i); + auto gradient = TensorBlockHolder::gradient(block, parameter); gradients.push_back(torch::make_intrusive( gradient->values(), @@ -103,7 +104,7 @@ std::vector RascalineAutograd::forward( if (all_positions.requires_grad()) { ctx->saved_data.emplace("structures_start", structures_start); - auto gradient = block->gradient("positions"); + auto gradient = TensorBlockHolder::gradient(block, "positions"); ctx->saved_data["positions_gradients"] = torch::make_intrusive( gradient->values(), gradient->samples(), @@ -115,7 +116,7 @@ std::vector RascalineAutograd::forward( if (all_cells.requires_grad()) { ctx->saved_data["samples"] = block->samples(); - auto gradient = block->gradient("cell"); + auto gradient = TensorBlockHolder::gradient(block, "cell"); ctx->saved_data["cell_gradients"] = torch::make_intrusive( gradient->values(), gradient->samples(), diff --git a/rascaline-torch/src/calculator.cpp b/rascaline-torch/src/calculator.cpp index 6771232df..4be2ca7b0 100644 --- a/rascaline-torch/src/calculator.cpp +++ b/rascaline-torch/src/calculator.cpp @@ -1,4 +1,5 @@ #include "rascaline/torch/calculator.hpp" +#include "metatensor/torch/block.hpp" #include "metatensor/torch/tensor.hpp" #include "rascaline/torch/autograd.hpp" #include @@ -86,7 +87,7 @@ static TorchTensorMap remove_other_gradients( ) { auto new_blocks = std::vector(); for (int64_t i=0; ikeys()->count(); i++) { - auto block = tensor->block_by_id(i); + auto block = TensorMapHolder::block_by_id(tensor, i); auto new_block = torch::make_intrusive( block->values(), block->samples(), @@ -95,7 +96,7 @@ static TorchTensorMap remove_other_gradients( ); for (const auto& parameter: gradients_to_keep) { - auto gradient = block->gradient(parameter); + auto gradient = TensorBlockHolder::gradient(block, parameter); new_block->add_gradient(parameter, gradient); } @@ -196,7 +197,7 @@ metatensor_torch::TorchTensorMap CalculatorHolder::compute( } for (int64_t block_i=0; block_ikeys()->count(); block_i++) { - auto block = torch_descriptor->block_by_id(block_i); + auto block = TensorMapHolder::block_by_id(torch_descriptor, block_i); // see `RascalineAutograd::forward` for an explanation of what's happening auto _ = RascalineAutograd::apply( all_positions, @@ -228,7 +229,7 @@ metatensor_torch::TorchTensorMap rascaline_torch::register_autograd( auto all_cells = stack_all_cells(systems); auto structures_start_ivalue = torch::IValue(); - auto precomputed_gradients = precomputed->block_by_id(0)->gradients_list(); + auto precomputed_gradients = TensorMapHolder::block_by_id(precomputed, 0)->gradients_list(); if (all_positions.requires_grad()) { if (!contains(precomputed_gradients, "positions")) { @@ -271,7 +272,7 @@ metatensor_torch::TorchTensorMap rascaline_torch::register_autograd( } for (int64_t block_i=0; block_ikeys()->count(); block_i++) { - auto block = precomputed->block_by_id(block_i); + auto block = TensorMapHolder::block_by_id(precomputed, block_i); auto _ = RascalineAutograd::apply( all_positions, all_cells, diff --git a/rascaline-torch/tests/calculator.cpp b/rascaline-torch/tests/calculator.cpp index 6e4bc470d..bf94903ab 100644 --- a/rascaline-torch/tests/calculator.cpp +++ b/rascaline-torch/tests/calculator.cpp @@ -1,4 +1,5 @@ +#include "metatensor/torch/block.hpp" #include #include @@ -7,6 +8,7 @@ #include using namespace rascaline_torch; +using namespace metatensor_torch; static TorchSystem test_system(bool positions_grad, bool cell_grad); @@ -35,7 +37,7 @@ TEST_CASE("Calculator") { )); // H block - auto block = descriptor->block_by_id(0); + auto block = TensorMapHolder::block_by_id(descriptor, 0); CHECK(*block->samples() == metatensor::Labels( {"structure", "center"}, {{0, 1}, {0, 2}, {0, 3}} @@ -56,7 +58,7 @@ TEST_CASE("Calculator") { CHECK(block->gradients_list().empty()); // C block - block = descriptor->block_by_id(1); + block = TensorMapHolder::block_by_id(descriptor, 1); CHECK(*block->samples() == metatensor::Labels( {"structure", "center"}, {{0, 0}} @@ -85,7 +87,7 @@ TEST_CASE("Calculator") { )); // H block - auto block = descriptor->block_by_id(0); + auto block = TensorMapHolder::block_by_id(descriptor, 0); auto values = block->values(); CHECK(values.requires_grad() == true); @@ -95,7 +97,7 @@ TEST_CASE("Calculator") { CHECK_THAT(grad_fn->name(), Catch::Matchers::Contains("rascaline_torch::RascalineAutograd")); // forward gradients - auto gradient = block->gradient("positions"); + auto gradient = TensorBlockHolder::gradient(block, "positions"); CHECK(*gradient->samples() == metatensor::Labels( {"sample", "structure", "atom"}, { @@ -117,7 +119,7 @@ TEST_CASE("Calculator") { CHECK(torch::all(gradient->values() == expected).item()); // C block - block = descriptor->block_by_id(1); + block = TensorMapHolder::block_by_id(descriptor, 1); values = block->values(); CHECK(values.requires_grad() == true); @@ -127,7 +129,7 @@ TEST_CASE("Calculator") { CHECK_THAT(grad_fn->name(), Catch::Matchers::Contains("rascaline_torch::RascalineAutograd")); // forward gradients - gradient = block->gradient("positions"); + gradient = TensorBlockHolder::gradient(block, "positions"); CHECK(*gradient->samples() == metatensor::Labels( {"sample", "structure", "atom"}, {{0, 0, 0}, {0, 0, 1}} @@ -149,7 +151,7 @@ TEST_CASE("Calculator") { )); // H block - auto block = descriptor->block_by_id(0); + auto block = TensorMapHolder::block_by_id(descriptor, 0); auto values = block->values(); CHECK(values.requires_grad() == true); @@ -162,7 +164,7 @@ TEST_CASE("Calculator") { CHECK(block->gradients_list().empty()); // C block - block = descriptor->block_by_id(1); + block = TensorMapHolder::block_by_id(descriptor, 1); values = block->values(); CHECK(values.requires_grad() == true); @@ -185,25 +187,25 @@ TEST_CASE("Calculator") { )); // H block - auto block = descriptor->block_by_id(0); + auto block = TensorMapHolder::block_by_id(descriptor, 0); auto values = block->values(); CHECK(values.requires_grad() == false); CHECK(values.grad_fn() == nullptr); // forward gradients - auto gradient = block->gradient("positions"); + auto gradient = TensorBlockHolder::gradient(block, "positions"); CHECK(gradient->samples()->count() == 8); // C block - block = descriptor->block_by_id(1); + block = TensorMapHolder::block_by_id(descriptor, 1); values = block->values(); CHECK(values.requires_grad() == false); CHECK(values.grad_fn() == nullptr); // forward gradients - gradient = block->gradient("positions"); + gradient = TensorBlockHolder::gradient(block, "positions"); CHECK(gradient->samples()->count() == 2); } } diff --git a/rascaline/Cargo.toml b/rascaline/Cargo.toml index ff039ae18..fecf212a9 100644 --- a/rascaline/Cargo.toml +++ b/rascaline/Cargo.toml @@ -36,7 +36,7 @@ name = "soap-power-spectrum" harness = false [dependencies] -metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "32ad5bb", features = ["rayon"]} +metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "d97ea65", features = ["rayon"]} ndarray = {version = "0.15", features = ["approx-0_5", "rayon", "serde"]} num-traits = "0.2" diff --git a/tox.ini b/tox.ini index 75012152f..29e8a69ad 100644 --- a/tox.ini +++ b/tox.ini @@ -19,10 +19,10 @@ lint-folders = "{toxinidir}/python" "{toxinidir}/setup.py" # we need to manually install dependencies for rascaline, since tox will install # the fresh wheel with `--no-deps` after building it. metatensor-core-requirement = - metatensor-core @ https://github.com/lab-cosmo/metatensor/archive/32ad5bb.zip\#subdirectory=python/metatensor-core + metatensor-core @ https://github.com/lab-cosmo/metatensor/archive/d97ea65.zip\#subdirectory=python/metatensor-core metatensor-torch-requirement = - metatensor-torch @ https://github.com/lab-cosmo/metatensor/archive/32ad5bb.zip\#subdirectory=python/metatensor-torch + metatensor-torch @ https://github.com/lab-cosmo/metatensor/archive/d97ea65.zip\#subdirectory=python/metatensor-torch build-single-wheel = --no-deps --no-build-isolation --check-build-dependencies