From f2299e50be42d8a58c1d8a902738e8df57f8257a Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 6 Oct 2023 10:46:08 +0200 Subject: [PATCH 01/30] Fix new clippy warnings for rustc 1.73 --- rascaline-c-api/src/calculator.rs | 3 +-- rascaline-c-api/src/status.rs | 2 +- rascaline/src/calculators/neighbor_list.rs | 2 +- rascaline/src/calculators/soap/spherical_expansion.rs | 4 ++-- 4 files changed, 5 insertions(+), 6 deletions(-) diff --git a/rascaline-c-api/src/calculator.rs b/rascaline-c-api/src/calculator.rs index 8a512f05a..234a2aba0 100644 --- a/rascaline-c-api/src/calculator.rs +++ b/rascaline-c-api/src/calculator.rs @@ -54,14 +54,13 @@ pub unsafe extern fn rascal_calculator(name: *const c_char, parameters: *const c let mut raw = std::ptr::null_mut(); let unwind_wrapper = std::panic::AssertUnwindSafe(&mut raw); let status = catch_unwind(move || { - let unwind_wrapper = unwind_wrapper; - check_pointers!(name, parameters); let name = CStr::from_ptr(name).to_str()?; let parameters = CStr::from_ptr(parameters).to_str()?; let calculator = Calculator::new(name, parameters.to_owned())?; let boxed = Box::new(rascal_calculator_t(calculator)); + let _ = &unwind_wrapper; *unwind_wrapper.0 = Box::into_raw(boxed); Ok(()) }); diff --git a/rascaline-c-api/src/status.rs b/rascaline-c-api/src/status.rs index a861323bc..2dd255a67 100644 --- a/rascaline-c-api/src/status.rs +++ b/rascaline-c-api/src/status.rs @@ -84,7 +84,7 @@ impl From for rascal_status_t { /// the error into `rascal_status_t`. pub fn catch_unwind(function: F) -> rascal_status_t where F: FnOnce() -> Result<(), Error> + UnwindSafe { match std::panic::catch_unwind(function) { - Ok(Ok(_)) => rascal_status_t(RASCAL_SUCCESS), + Ok(Ok(())) => rascal_status_t(RASCAL_SUCCESS), Ok(Err(error)) => error.into(), Err(error) => Error::from(error).into() } diff --git a/rascaline/src/calculators/neighbor_list.rs b/rascaline/src/calculators/neighbor_list.rs index c092b4753..fc19c22cb 100644 --- a/rascaline/src/calculators/neighbor_list.rs +++ b/rascaline/src/calculators/neighbor_list.rs @@ -500,7 +500,7 @@ impl FullNeighborList { system.compute_neighbors(self.cutoff)?; let species = system.species()?; - for pair in system.pairs()?.iter() { + for pair in system.pairs()? { let first_block_i = descriptor.keys().position(&[ species[pair.first].into(), species[pair.second].into() ]); diff --git a/rascaline/src/calculators/soap/spherical_expansion.rs b/rascaline/src/calculators/soap/spherical_expansion.rs index 1601e2f35..b62cfa37d 100644 --- a/rascaline/src/calculators/soap/spherical_expansion.rs +++ b/rascaline/src/calculators/soap/spherical_expansion.rs @@ -189,7 +189,7 @@ impl SphericalExpansion { let neighbor_i = pair.second; result.pair_to_pair_ids.entry((pair.first, pair.second)) - .or_insert_with(Vec::new) + .or_default() .push(pair_id); let species_neighbor_i = result.species_mapping[&species[neighbor_i]]; @@ -253,7 +253,7 @@ impl SphericalExpansion { let neighbor_i = pair.first; result.pair_to_pair_ids.entry((pair.second, pair.first)) - .or_insert_with(Vec::new) + .or_default() .push(pair_id); From e723afd2a7c6e5a3c1f9bac34d7574eb029fadc3 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 6 Oct 2023 11:17:09 +0200 Subject: [PATCH 02/30] Update metatensor --- docs/requirements.txt | 2 +- pyproject.toml | 2 +- python/rascaline-torch/pyproject.toml | 2 +- rascaline-c-api/CMakeLists.txt | 4 ++-- rascaline-c-api/Cargo.toml | 4 ++-- rascaline-torch/CMakeLists.txt | 4 ++-- rascaline/Cargo.toml | 2 +- tox.ini | 4 ++-- 8 files changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index 6047c6435..7d63ac552 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/2248a3c.zip +metatensor[torch] @ git+https://github.com/lab-cosmo/metatensor@ae10944 torch chemfiles matplotlib diff --git a/pyproject.toml b/pyproject.toml index f436d2495..e088915fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] dependencies = [ - "metatensor-core @ https://github.com/lab-cosmo/metatensor/archive/2248a3c.zip#subdirectory=python/metatensor-core", + "metatensor-core @ git+https://github.com/lab-cosmo/metatensor@ae10944#subdirectory=python/metatensor-core", ] [project.urls] diff --git a/python/rascaline-torch/pyproject.toml b/python/rascaline-torch/pyproject.toml index 2769faa93..3dc56acd0 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/2248a3c.zip#subdirectory=python/metatensor-torch", + "metatensor-torch @ https://github.com/lab-cosmo/metatensor/archive/ae10944.zip#subdirectory=python/metatensor-torch", ] # use a custom build backend to add a dependency on the right version of rascaline diff --git a/rascaline-c-api/CMakeLists.txt b/rascaline-c-api/CMakeLists.txt index e4e340076..b7073c400 100644 --- a/rascaline-c-api/CMakeLists.txt +++ b/rascaline-c-api/CMakeLists.txt @@ -217,7 +217,7 @@ endif() # ============================================================================ # # Setup metatensor -set(METATENSOR_GIT_VERSION "2248a3c") +set(METATENSOR_GIT_VERSION "ae10944") 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=50e70a2e74246a985e1520260a74205c + URL_HASH MD5=1de1208bd96fa9acda0b835d1c9e2c5d SOURCE_SUBDIR metatensor-core VERBOSE ) diff --git a/rascaline-c-api/Cargo.toml b/rascaline-c-api/Cargo.toml index 0521fec19..e82a58bac 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 = "2248a3c"} +metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "ae10944"} 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 = "2248a3c"} +metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "ae10944"} [dev-dependencies] which = "4" diff --git a/rascaline-torch/CMakeLists.txt b/rascaline-torch/CMakeLists.txt index b14fb7210..156c727f2 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 "2248a3c") +set(METATENSOR_GIT_VERSION "ae10944") 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=50e70a2e74246a985e1520260a74205c + URL_HASH MD5=1de1208bd96fa9acda0b835d1c9e2c5d SOURCE_SUBDIR metatensor-torch VERBOSE ) diff --git a/rascaline/Cargo.toml b/rascaline/Cargo.toml index 885007351..26b975bf6 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 = "2248a3c", features = ["rayon"]} +metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "ae10944", 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 9331a771e..351f8da60 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/2248a3c.zip\#subdirectory=python/metatensor-core + metatensor-core @ git+https://github.com/lab-cosmo/metatensor@ae10944\#subdirectory=python/metatensor-core metatensor-torch-requirement = - metatensor-torch @ https://github.com/lab-cosmo/metatensor/archive/2248a3c.zip\#subdirectory=python/metatensor-torch + metatensor-torch @ git+https://github.com/lab-cosmo/metatensor@ae10944\#subdirectory=python/metatensor-torch build-single-wheel = --no-deps --no-build-isolation --check-build-dependencies From 29d8a7b4837c7a859728b72e0a489dc3b454814d Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 6 Oct 2023 10:41:51 +0200 Subject: [PATCH 03/30] Update to C++17 for torch>=2.1 --- .github/workflows/torch-tests.yml | 6 +++--- python/rascaline-torch/pyproject.toml | 2 +- rascaline-c-api/tests/catch/CMakeLists.txt | 2 +- rascaline-torch/CMakeLists.txt | 2 +- rascaline-torch/tests/utils/mod.rs | 2 +- 5 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index e362acb83..79ec93e2d 100644 --- a/.github/workflows/torch-tests.yml +++ b/.github/workflows/torch-tests.yml @@ -23,18 +23,18 @@ jobs: cargo-test-flags: --release - os: ubuntu-20.04 - torch-version: 2.0.* + torch-version: 2.1.* python-version: "3.11" cargo-test-flags: --release do-valgrind: true - os: macos-11 - torch-version: 2.0.* + torch-version: 2.1.* python-version: "3.11" cargo-test-flags: --release - os: windows-2019 - torch-version: 2.0.* + torch-version: 2.1.* python-version: "3.11" cargo-test-flags: --release diff --git a/python/rascaline-torch/pyproject.toml b/python/rascaline-torch/pyproject.toml index 3dc56acd0..5545f32bd 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/ae10944.zip#subdirectory=python/metatensor-torch", + "metatensor-torch @git+https://github.com/lab-cosmo/metatensor@ae10944#subdirectory=python/metatensor-torch", ] # use a custom build backend to add a dependency on the right version of rascaline diff --git a/rascaline-c-api/tests/catch/CMakeLists.txt b/rascaline-c-api/tests/catch/CMakeLists.txt index 4c6184820..c783b584e 100644 --- a/rascaline-c-api/tests/catch/CMakeLists.txt +++ b/rascaline-c-api/tests/catch/CMakeLists.txt @@ -1,3 +1,3 @@ add_library(catch STATIC catch.cpp) target_include_directories(catch PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}) -target_compile_features(catch PUBLIC cxx_std_11) +target_compile_features(catch PUBLIC cxx_std_17) diff --git a/rascaline-torch/CMakeLists.txt b/rascaline-torch/CMakeLists.txt index 156c727f2..f1f343809 100644 --- a/rascaline-torch/CMakeLists.txt +++ b/rascaline-torch/CMakeLists.txt @@ -106,7 +106,7 @@ add_library(rascaline_torch SHARED ) target_link_libraries(rascaline_torch PUBLIC torch metatensor_torch rascaline) -target_compile_features(rascaline_torch PUBLIC cxx_std_14) +target_compile_features(rascaline_torch PUBLIC cxx_std_17) target_include_directories(rascaline_torch PUBLIC $ $ diff --git a/rascaline-torch/tests/utils/mod.rs b/rascaline-torch/tests/utils/mod.rs index fb2213c34..008e06b72 100644 --- a/rascaline-torch/tests/utils/mod.rs +++ b/rascaline-torch/tests/utils/mod.rs @@ -84,7 +84,7 @@ pub fn setup_pytorch(build_dir: PathBuf) -> PathBuf { .expect("failed to run python"); assert!(status.success(), "failed to run `python -m pip install --upgrade pip`"); - let torch_version = std::env::var("RASCALINE_TORCH_TEST_VERSION").unwrap_or("2.0.*".into()); + let torch_version = std::env::var("RASCALINE_TORCH_TEST_VERSION").unwrap_or("2.1.*".into()); let status = Command::new(&python) .arg("-m") .arg("pip") From 365e79fa10c168de091e19c7268f8229e72f6628 Mon Sep 17 00:00:00 2001 From: Andrea Grisafi Date: Fri, 6 Oct 2023 16:16:30 +0200 Subject: [PATCH 04/30] Exploited subtraction theorem in LODE structure factor (#242) --- .../calculators/lode/spherical_expansion.rs | 29 +++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/rascaline/src/calculators/lode/spherical_expansion.rs b/rascaline/src/calculators/lode/spherical_expansion.rs index afaed528e..a10666e51 100644 --- a/rascaline/src/calculators/lode/spherical_expansion.rs +++ b/rascaline/src/calculators/lode/spherical_expansion.rs @@ -126,14 +126,31 @@ fn compute_structure_factors(positions: &[Vector3D], species: &[i32], k_vectors: .map(|s| (s, Array2::from_elem((n_atoms, n_k_vectors), 0.0))) .collect::>(); + // Precompute species-dependent sum over neighbours j of sines and cosines + let mut sumjcos = all_species.iter().copied() + .map(|s| (s, Array1::from_elem(n_k_vectors, 0.0))) + .collect::>(); + let mut sumjsin = all_species.iter().copied() + .map(|s| (s, Array1::from_elem(n_k_vectors, 0.0))) + .collect::>(); + for j in 0..n_atoms { + let sumjcos = sumjcos.get_mut(&species[j]).unwrap(); + let sumjsin = sumjsin.get_mut(&species[j]).unwrap(); + for k in 0..n_k_vectors { + sumjcos[k] += cosines[[j, k]]; + sumjsin[k] += sines[[j, k]]; + } + } + + // Compute Sum_j cos(k*r_ij) and Sum_j sin(k*r_ij) using the subtraction theorem for i in 0..n_atoms { - for j in 0..n_atoms { + for (species, real_per_center) in &mut real_per_center { + let sumjcos = sumjcos.get_mut(&species).unwrap(); + let sumjsin = sumjsin.get_mut(&species).unwrap(); + let imag_per_center = imag_per_center.get_mut(&species).unwrap(); for k in 0..n_k_vectors { - let real = cosines[[i, k]] * cosines[[j, k]] + sines[[i, k]] * sines[[j, k]]; - let imag = sines[[i, k]] * cosines[[j, k]] - cosines[[i, k]] * sines[[j, k]]; - - let real_per_center = real_per_center.get_mut(&species[j]).unwrap(); - let imag_per_center = imag_per_center.get_mut(&species[j]).unwrap(); + let real = cosines[[i, k]] * sumjcos[k] + sines[[i, k]] * sumjsin[k]; + let imag = sines[[i, k]] * sumjcos[k] - cosines[[i, k]] * sumjsin[k]; real_per_center[[i, k]] += 2.0 * real; imag_per_center[[i, k]] += 2.0 * imag; } From 7107376fe828316860bc9c6420e2f7bfd3dc313a Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 6 Oct 2023 16:21:14 +0200 Subject: [PATCH 05/30] Update Python installation instruction Suggest updating pip & also point people to the nightly wheels --- docs/src/get-started/installation.rst | 29 ++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/docs/src/get-started/installation.rst b/docs/src/get-started/installation.rst index 034165dc1..3c01322b9 100644 --- a/docs/src/get-started/installation.rst +++ b/docs/src/get-started/installation.rst @@ -10,9 +10,33 @@ install using `rustup `_ or your OS package manager. Installing the Python module ---------------------------- +For building and using the Python package, clone the repository using `git +`_ and install rascaline using `pip +`_. + +From source: + +.. code-block:: bash + + # Make sure you are using the latest version of pip + pip install --upgrade pip + + git clone https://github.com/Luthaf/rascaline + cd rascaline + pip install . + + # alternatively, the same thing in a single command + pip install git+https://github.com/Luthaf/rascaline + + +Rascaline is also provided as prebuilt wheel which avoids the intermediate step +of building the package with a Rust compiler from the source code. + .. code-block:: bash - pip install git+https://github.com/Luthaf/rascaline.git + pip install --upgrade pip + pip install --extra-index-url https://luthaf.fr/temporary-wheels/ rascaline + .. _install-c-lib: @@ -96,6 +120,9 @@ Building from source: cd rascaline/python/rascaline-torch pip install . + # Make sure you are using the latest version of pip + pip install --upgrade pip + # alternatively, the same thing in a single command pip install git+https://github.com/luthaf/rascaline#subdirectory=python/rascaline-torch From b3e095b5f7219d269020b220470406337b4ccfc7 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 6 Oct 2023 16:25:53 +0200 Subject: [PATCH 06/30] Allow to disable valgrind when for C/C++ API tests --- CONTRIBUTING.rst | 33 ++++++++++++++++++++++------ rascaline-c-api/tests/CMakeLists.txt | 14 +++++++----- rascaline-torch/tests/CMakeLists.txt | 2 +- 3 files changed, 35 insertions(+), 14 deletions(-) diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index fee04112d..74a2fa45c 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -123,9 +123,29 @@ Github CI workflows. You can also run only a subset of tests with one of these commands: - ``cargo test`` runs everything -- ``cargo test --test=`` to run only the tests in ``tests/.rs``; - - ``cargo test --test=python-api`` (or ``tox`` directly) to run Python tests only; - - ``cargo test --test=c-api`` to run the C API tests only; +- ``cargo test --package=rascaline`` to run the calculators tests; +- ``cargo test --package=rascaline-c-api`` to run the C/C++ tests only; + + - ``cargo test --test=run-cxx-tests`` will run the unit tests for the C/C++ + API. If `valgrind`_ is installed, it will be used to check for memory + errors. You can disable this by setting the `RASCALINE_DISABLE_VALGRIND` + environment variable to 1 (`export RASCALINE_DISABLE_VALGRIND=1` for most + Linux/macOS shells); + - ``cargo test --test=check-cxx-install`` will build the C/C++ interfaces, + install them and the associated CMake files and then try to build a basic + project depending on this interface with CMake; + +- ``cargo test --package=rascaline-torch`` to run the C++ TorchScript extension + tests only; + + - ``cargo test --test=run-torch-tests`` will run the unit tests for the + TorchScript C++ extension; + - ``cargo test --test=check-cxx-install`` will build the C++ TorchScript + extension, install it and then try to build a basic project depending on + this extension with CMake; + +- ``cargo test --package=rascaline-python`` (or ``tox`` directly, see below) to + run Python tests only; - ``cargo test --lib`` to run unit tests; - ``cargo test --doc`` to run documentation tests; - ``cargo bench --test`` compiles and run the benchmarks once, to quickly ensure @@ -135,10 +155,8 @@ You can add some flags to any of above commands to further refine which tests should run: - ``--release`` to run tests in release mode (default is to run tests in debug mode) -- ``-- `` to only run tests whose name contains filter, for example ``cargo test -- spherical_harmonics`` -- ``--package rascaline`` to run tests defined in the rascaline crate (the core implementation) -- ``--package rascaline-c-api`` to run tests defined in the rascaline-c-api - crate (the C API implementation) +- ``-- `` to only run tests whose name contains filter, for example + ``cargo test -- spherical_harmonics`` Also, you can run individual python tests using `tox`_ if you wish to test only specific functionalities, for example: @@ -156,6 +174,7 @@ The latter command ``tox -e format`` will use tox to do actual formatting instea of just testing it. .. _`cargo` : https://doc.rust-lang.org/cargo/ +.. _valgrind: https://valgrind.org/ Writing your own calculator --------------------------- diff --git a/rascaline-c-api/tests/CMakeLists.txt b/rascaline-c-api/tests/CMakeLists.txt index ef81e4740..bbe08bec4 100644 --- a/rascaline-c-api/tests/CMakeLists.txt +++ b/rascaline-c-api/tests/CMakeLists.txt @@ -70,12 +70,14 @@ target_link_libraries(tests_helpers rascaline::shared) find_program(VALGRIND valgrind) if (VALGRIND) - message(STATUS "Running tests using valgrind") - set(TEST_COMMAND - "${VALGRIND}" "--tool=memcheck" "--dsymutil=yes" "--error-exitcode=125" - "--leak-check=full" "--show-leak-kinds=definite,indirect,possible" "--track-origins=yes" - "--suppressions=${CMAKE_CURRENT_SOURCE_DIR}/valgrind.supp" "--gen-suppressions=all" - ) + if (NOT "$ENV{RASCALINE_DISABLE_VALGRIND}" EQUAL "1") + message(STATUS "Running tests using valgrind") + set(TEST_COMMAND + "${VALGRIND}" "--tool=memcheck" "--dsymutil=yes" "--error-exitcode=125" + "--leak-check=full" "--show-leak-kinds=definite,indirect,possible" "--track-origins=yes" + "--suppressions=${CMAKE_CURRENT_SOURCE_DIR}/valgrind.supp" "--gen-suppressions=all" + ) + endif() else() set(TEST_COMMAND "") endif() diff --git a/rascaline-torch/tests/CMakeLists.txt b/rascaline-torch/tests/CMakeLists.txt index 07ad38238..874738c63 100644 --- a/rascaline-torch/tests/CMakeLists.txt +++ b/rascaline-torch/tests/CMakeLists.txt @@ -8,7 +8,7 @@ target_link_libraries(catch torch) find_program(VALGRIND valgrind) if (VALGRIND) - if (NOT "$ENV{METATENSOR_DISABLE_VALGRIND}" EQUAL "1") + if (NOT "$ENV{RASCALINE_DISABLE_VALGRIND}" EQUAL "1") message(STATUS "Running tests using valgrind") set(TEST_COMMAND "${VALGRIND}" "--tool=memcheck" "--dsymutil=yes" "--error-exitcode=125" From e8265ec987efece1fa3188e3d6ce268a426d34ce Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Tue, 10 Oct 2023 11:59:39 +0200 Subject: [PATCH 07/30] Bump rustc to 1.65 Co-authored-by: Guillaume Fraux --- .github/workflows/rust-tests.yml | 2 +- .readthedocs.yml | 2 +- CONTRIBUTING.rst | 2 +- docs/rascaline-json-schema/Cargo.toml | 2 +- python/scripts/rustc-manylinux2014_x86_64/Dockerfile | 2 +- rascaline-c-api/Cargo.toml | 2 +- rascaline/Cargo.toml | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/rust-tests.yml b/.github/workflows/rust-tests.yml index 73bd6113e..20d7f83a0 100644 --- a/.github/workflows/rust-tests.yml +++ b/.github/workflows/rust-tests.yml @@ -40,7 +40,7 @@ jobs: # check the build on a stock Ubuntu 20.04, including cmake 3.16 - os: ubuntu-20.04 - rust-version: "1.64" + rust-version: "1.65" container: ubuntu:20.04 rust-target: x86_64-unknown-linux-gnu build-type: debug diff --git a/.readthedocs.yml b/.readthedocs.yml index a8c7d2575..9586d0bdc 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -12,7 +12,7 @@ build: - cmake tools: python: "3.11" - rust: "1.64" + rust: "1.70" jobs: post_install: # install rascaline-torch with the CPU version of PyTorch. We can not use diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 74a2fa45c..7b61cb772 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -57,7 +57,7 @@ on rascaline: - **the rust compiler**: you will need both ``rustc`` (the compiler) and ``cargo`` (associated build tool). You can install both using `rustup`_, or use a version provided by your operating system. We need at least Rust version - 1.64 to build rascaline. + 1.65 to build rascaline. - **Python**: you can install ``Python`` and ``pip`` from your operating system. We require a Python version of at least 3.6. - **tox**: a Python test runner, cf https://tox.readthedocs.io/en/latest/. You diff --git a/docs/rascaline-json-schema/Cargo.toml b/docs/rascaline-json-schema/Cargo.toml index c6fea5aa1..f4bb1f017 100644 --- a/docs/rascaline-json-schema/Cargo.toml +++ b/docs/rascaline-json-schema/Cargo.toml @@ -4,7 +4,7 @@ version = "0.1.0" authors = ["Luthaf "] edition = "2021" publish = false -rust-version = "1.64" +rust-version = "1.65" [[bin]] name = "rascaline-json-schema" diff --git a/python/scripts/rustc-manylinux2014_x86_64/Dockerfile b/python/scripts/rustc-manylinux2014_x86_64/Dockerfile index f4fdc5e3b..026efd13a 100644 --- a/python/scripts/rustc-manylinux2014_x86_64/Dockerfile +++ b/python/scripts/rustc-manylinux2014_x86_64/Dockerfile @@ -5,7 +5,7 @@ RUN yum install git -y RUN git config --global --add safe.directory /code # Download rustup-init asn install -RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --profile minimal --default-toolchain 1.64 +RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --profile minimal --default-toolchain 1.65 ENV PATH="/root/.cargo/bin:${PATH}" ENV RUST_BUILD_TARGET="x86_64-unknown-linux-gnu" diff --git a/rascaline-c-api/Cargo.toml b/rascaline-c-api/Cargo.toml index e82a58bac..36b44812e 100644 --- a/rascaline-c-api/Cargo.toml +++ b/rascaline-c-api/Cargo.toml @@ -3,7 +3,7 @@ name = "rascaline-c-api" version = "0.1.0" authors = ["Luthaf "] edition = "2021" -rust-version = "1.64" +rust-version = "1.65" [lib] # when https://github.com/rust-lang/cargo/pull/8789 lands, use it here! diff --git a/rascaline/Cargo.toml b/rascaline/Cargo.toml index 26b975bf6..89651b43d 100644 --- a/rascaline/Cargo.toml +++ b/rascaline/Cargo.toml @@ -3,7 +3,7 @@ name = "rascaline" version = "0.1.0" authors = ["Luthaf "] edition = "2021" -rust-version = "1.64" +rust-version = "1.65" [lib] bench = false From 449ac3710e49b9f0c60fed17b1ddb9e848f4b78c Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Tue, 10 Oct 2023 18:05:32 +0200 Subject: [PATCH 08/30] Update metatensor --- docs/requirements.txt | 2 +- pyproject.toml | 2 +- python/rascaline-torch/pyproject.toml | 2 +- .../rascaline/examples/property-selection.py | 2 +- python/rascaline/examples/sample-selection.py | 2 +- .../utils/power_spectrum/calculator.py | 12 ++--- rascaline-c-api/CMakeLists.txt | 4 +- rascaline-c-api/Cargo.toml | 4 +- rascaline-c-api/tests/calculator.cpp | 52 +++++++++---------- rascaline-torch/CMakeLists.txt | 4 +- rascaline-torch/src/autograd.cpp | 32 ++++++------ rascaline/Cargo.toml | 2 +- rascaline/src/calculator.rs | 4 +- .../src/calculators/atomic_composition.rs | 8 +-- rascaline/src/calculators/dummy_calculator.rs | 8 +-- .../calculators/lode/spherical_expansion.rs | 14 ++--- rascaline/src/calculators/mod.rs | 4 +- rascaline/src/calculators/neighbor_list.rs | 6 +-- .../src/calculators/soap/power_spectrum.rs | 8 +-- .../src/calculators/soap/radial_spectrum.rs | 8 +-- .../calculators/soap/spherical_expansion.rs | 8 +-- .../soap/spherical_expansion_pair.rs | 6 +-- rascaline/src/calculators/sorted_distances.rs | 8 +-- rascaline/src/labels/samples/atom_centered.rs | 6 +-- rascaline/src/labels/samples/long_range.rs | 4 +- rascaline/src/labels/samples/mod.rs | 2 +- rascaline/src/tutorials/moments/moments.rs | 8 +-- .../src/tutorials/moments/s1_scaffold.rs | 4 +- .../src/tutorials/moments/s2_metadata.rs | 8 +-- .../src/tutorials/moments/s3_compute_1.rs | 4 +- .../src/tutorials/moments/s3_compute_2.rs | 4 +- .../src/tutorials/moments/s3_compute_3.rs | 4 +- .../src/tutorials/moments/s3_compute_4.rs | 4 +- .../src/tutorials/moments/s3_compute_5.rs | 4 +- tox.ini | 4 +- 35 files changed, 129 insertions(+), 129 deletions(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index 7d63ac552..cb9112693 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] @ git+https://github.com/lab-cosmo/metatensor@ae10944 +metatensor[torch] @ git+https://github.com/lab-cosmo/metatensor@ad3eda5 torch chemfiles matplotlib diff --git a/pyproject.toml b/pyproject.toml index e088915fd..47dd87747 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] dependencies = [ - "metatensor-core @ git+https://github.com/lab-cosmo/metatensor@ae10944#subdirectory=python/metatensor-core", + "metatensor-core @ git+https://github.com/lab-cosmo/metatensor@ad3eda5#subdirectory=python/metatensor-core", ] [project.urls] diff --git a/python/rascaline-torch/pyproject.toml b/python/rascaline-torch/pyproject.toml index 5545f32bd..00fe173a1 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 @git+https://github.com/lab-cosmo/metatensor@ae10944#subdirectory=python/metatensor-torch", + "metatensor-torch @git+https://github.com/lab-cosmo/metatensor@ad3eda5#subdirectory=python/metatensor-torch", ] # use a custom build backend to add a dependency on the right version of rascaline diff --git a/python/rascaline/examples/property-selection.py b/python/rascaline/examples/property-selection.py index 4353af90c..3b1093017 100644 --- a/python/rascaline/examples/property-selection.py +++ b/python/rascaline/examples/property-selection.py @@ -47,7 +47,7 @@ # of the labels must be a subset of the names of the properties produced by the # calculator. You can see the default set of names with: -print("properties names:", descriptor.properties_names) +print("property names:", descriptor.property_names) # %% # diff --git a/python/rascaline/examples/sample-selection.py b/python/rascaline/examples/sample-selection.py index 6622eef47..d9f9576d9 100644 --- a/python/rascaline/examples/sample-selection.py +++ b/python/rascaline/examples/sample-selection.py @@ -47,7 +47,7 @@ # of the labels must be a subset of the names of the samples produced by the # calculator. You can see the default set of names with: -print("samples names:", descriptor.samples_names) +print("sample names:", descriptor.sample_names) # %% # diff --git a/python/rascaline/rascaline/utils/power_spectrum/calculator.py b/python/rascaline/rascaline/utils/power_spectrum/calculator.py index c6b872ecd..6815f4b16 100644 --- a/python/rascaline/rascaline/utils/power_spectrum/calculator.py +++ b/python/rascaline/rascaline/utils/power_spectrum/calculator.py @@ -177,7 +177,7 @@ def compute( ] assert spherical_expansion_1.keys.names == expected_key_names - assert spherical_expansion_1.properties_names == ["n"] + assert spherical_expansion_1.property_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. @@ -198,7 +198,7 @@ def compute( use_native_system=use_native_system, ) assert spherical_expansion_2.keys.names == expected_key_names - assert spherical_expansion_2.properties_names == ["n"] + assert spherical_expansion_2.property_names == ["n"] array = spherical_expansion_2.keys.column("species_neighbor") keys_to_move = Labels( @@ -239,7 +239,7 @@ def compute( properties_1.values.shape[1] + properties_2.values.shape[1] ) - new_properties_values = _dispatch.empty_like( + new_property_values = _dispatch.empty_like( array=properties_1.values, shape=[ properties_1.values.shape[0], @@ -250,12 +250,12 @@ def compute( for i, values_1 in enumerate(properties_1.values): for j, values_2 in enumerate(properties_2.values): - new_properties_values[i, j, : len(values_1)] = values_1 - new_properties_values[i, j, len(values_1) :] = values_2 + new_property_values[i, j, : len(values_1)] = values_1 + new_property_values[i, j, len(values_1) :] = values_2 properties = Labels( names=["species_neighbor_1", "n1", "species_neighbor_2", "n2"], - values=new_properties_values.reshape(-1, n_keys_dimensions), + values=new_property_values.reshape(-1, n_keys_dimensions), ) # Compute the invariants by summation and store the results this is diff --git a/rascaline-c-api/CMakeLists.txt b/rascaline-c-api/CMakeLists.txt index b7073c400..2ed0cd486 100644 --- a/rascaline-c-api/CMakeLists.txt +++ b/rascaline-c-api/CMakeLists.txt @@ -217,7 +217,7 @@ endif() # ============================================================================ # # Setup metatensor -set(METATENSOR_GIT_VERSION "ae10944") +set(METATENSOR_GIT_VERSION "ad3eda5") 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=1de1208bd96fa9acda0b835d1c9e2c5d + URL_HASH MD5=0ca550a82d587c4426eddf82ca8deae2 SOURCE_SUBDIR metatensor-core VERBOSE ) diff --git a/rascaline-c-api/Cargo.toml b/rascaline-c-api/Cargo.toml index 36b44812e..19170f99a 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 = "ae10944"} +metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "ad3eda5"} 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 = "ae10944"} +metatensor = {git = "https://github.com/lab-cosmo/metatensor", rev = "ad3eda5"} [dev-dependencies] which = "4" diff --git a/rascaline-c-api/tests/calculator.cpp b/rascaline-c-api/tests/calculator.cpp index a1709896a..965a4146b 100644 --- a/rascaline-c-api/tests/calculator.cpp +++ b/rascaline-c-api/tests/calculator.cpp @@ -212,16 +212,16 @@ TEST_CASE("Compute descriptor") { } SECTION("Partial compute -- samples") { - auto selected_samples_values = std::vector{ + auto selected_sample_values = std::vector{ 0, 1, /**/ 0, 3, }; - auto selected_samples_names = std::vector{ + auto selected_sample_names = std::vector{ "structure", "center" }; mts_labels_t selected_samples = {0}; - selected_samples.names = selected_samples_names.data(); - selected_samples.values = selected_samples_values.data(); + selected_samples.names = selected_sample_names.data(); + selected_samples.values = selected_sample_values.data(); selected_samples.count = 2; selected_samples.size = 2; @@ -290,16 +290,16 @@ TEST_CASE("Compute descriptor") { } SECTION("Partial compute -- features") { - auto selected_properties_values = std::vector{ + auto selected_property_values = std::vector{ 0, 1, }; - auto selected_properties_names = std::vector{ + auto selected_property_names = std::vector{ "index_delta", "x_y_z" }; mts_labels_t selected_properties = {0}; - selected_properties.names = selected_properties_names.data(); - selected_properties.values = selected_properties_values.data(); + selected_properties.names = selected_property_names.data(); + selected_properties.values = selected_property_values.data(); selected_properties.count = 1; selected_properties.size = 2; @@ -380,17 +380,17 @@ TEST_CASE("Compute descriptor") { } SECTION("Partial compute -- preselected") { - auto samples_names = std::vector{ + auto sample_names = std::vector{ "structure", "center" }; - auto properties_names = std::vector{ + auto property_names = std::vector{ "index_delta", "x_y_z" }; - auto h_samples_values = std::vector{ + auto h_sample_values = std::vector{ 0, 3, }; - auto h_properties_values = std::vector{ + auto h_property_values = std::vector{ 0, 1, }; @@ -398,37 +398,37 @@ TEST_CASE("Compute descriptor") { mts_labels_t h_samples = {0}; h_samples.size = 2; - h_samples.names = samples_names.data(); + h_samples.names = sample_names.data(); h_samples.count = 1; - h_samples.values = h_samples_values.data(); + h_samples.values = h_sample_values.data(); mts_labels_t h_properties = {0}; h_properties.size = 2; - h_properties.names = properties_names.data(); + h_properties.names = property_names.data(); h_properties.count = 1; - h_properties.values = h_properties_values.data(); + h_properties.values = h_property_values.data(); blocks[0] = mts_block(empty_array({1, 1}), h_samples, nullptr, 0, h_properties); REQUIRE(blocks[0] != nullptr); - auto c_samples_values = std::vector{ + auto c_sample_values = std::vector{ 0, 0, }; - auto c_properties_values = std::vector{ + auto c_property_values = std::vector{ 1, 0, }; mts_labels_t c_samples = {0}; c_samples.size = 2; - c_samples.names = samples_names.data(); + c_samples.names = sample_names.data(); c_samples.count = 1; - c_samples.values = c_samples_values.data(); + c_samples.values = c_sample_values.data(); mts_labels_t c_properties = {0}; c_properties.size = 2; - c_properties.names = properties_names.data(); + c_properties.names = property_names.data(); c_properties.count = 1; - c_properties.values = c_properties_values.data(); + c_properties.values = c_property_values.data(); blocks[1] = mts_block(empty_array({1, 1}), c_samples, nullptr, 0, c_properties); REQUIRE(blocks[1] != nullptr); @@ -521,12 +521,12 @@ TEST_CASE("Compute descriptor") { // in a different order than what would be the default (6, 12). mts_labels_t selected_keys = {0}; - const char* samples_names[] = {"species_center"}; - selected_keys.names = samples_names; + const char* sample_names[] = {"species_center"}; + selected_keys.names = sample_names; selected_keys.size = 1; - int32_t samples_values[] = {12, 6}; - selected_keys.values = samples_values; + int32_t sample_values[] = {12, 6}; + selected_keys.values = sample_values; selected_keys.count = 2; auto system = simple_system(); diff --git a/rascaline-torch/CMakeLists.txt b/rascaline-torch/CMakeLists.txt index f1f343809..b8b2c7dc2 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 "ae10944") +set(METATENSOR_GIT_VERSION "ad3eda5") 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=1de1208bd96fa9acda0b835d1c9e2c5d + URL_HASH MD5=0ca550a82d587c4426eddf82ca8deae2 SOURCE_SUBDIR metatensor-torch VERBOSE ) diff --git a/rascaline-torch/src/autograd.cpp b/rascaline-torch/src/autograd.cpp index 9f6af7b9a..4d03929ad 100644 --- a/rascaline-torch/src/autograd.cpp +++ b/rascaline-torch/src/autograd.cpp @@ -214,8 +214,8 @@ std::vector PositionsGrad::forward( auto structures_start = structures_start_ivalue.toIntList(); auto samples = dX_dr->samples(); - auto samples_values = samples->values(); - auto* samples_values_ptr = samples_values.data_ptr(); + auto sample_values = samples->values(); + auto* sample_values_ptr = sample_values.data_ptr(); always_assert(samples->names().size() == 3); always_assert(samples->names()[0] == "sample"); always_assert(samples->names()[1] == "structure"); @@ -242,9 +242,9 @@ std::vector PositionsGrad::forward( // =========================== compute dA_dr ============================ // for (int64_t grad_sample_i=0; grad_sample_icount(); grad_sample_i++) { - auto sample_i = samples_values_ptr[grad_sample_i * 3 + 0]; - auto structure_i = samples_values_ptr[grad_sample_i * 3 + 1]; - auto atom_i = samples_values_ptr[grad_sample_i* 3 + 2]; + auto sample_i = sample_values_ptr[grad_sample_i * 3 + 0]; + auto structure_i = sample_values_ptr[grad_sample_i * 3 + 1]; + auto atom_i = sample_values_ptr[grad_sample_i* 3 + 2]; auto global_atom_i = structures_start[structure_i] + atom_i; @@ -283,8 +283,8 @@ std::vector PositionsGrad::backward( auto dB_d_dA_dr = grad_outputs[0]; // gradient of B w.r.t. dA/dr (output of forward) auto samples = dX_dr->samples(); - auto samples_values = samples->values(); - auto* samples_values_ptr = samples_values.data_ptr(); + auto sample_values = samples->values(); + auto* sample_values_ptr = sample_values.data_ptr(); always_assert(samples->names().size() == 3); always_assert(samples->names()[0] == "sample"); always_assert(samples->names()[1] == "structure"); @@ -329,9 +329,9 @@ std::vector PositionsGrad::backward( // dB_d_dA_dr.shape == [n_atoms, 3] // dB_d_dA_dX.shape == [samples, features...] for (int64_t grad_sample_i=0; grad_sample_icount(); grad_sample_i++) { - auto sample_i = samples_values_ptr[grad_sample_i * 3 + 0]; - auto structure_i = samples_values_ptr[grad_sample_i * 3 + 1]; - auto atom_i = samples_values_ptr[grad_sample_i* 3 + 2]; + auto sample_i = sample_values_ptr[grad_sample_i * 3 + 0]; + auto structure_i = sample_values_ptr[grad_sample_i * 3 + 1]; + auto atom_i = sample_values_ptr[grad_sample_i* 3 + 2]; auto global_atom_i = structures_start[structure_i] + atom_i; @@ -375,8 +375,8 @@ std::vector CellGrad::forward( auto* cell_grad_ptr = cell_grad.data_ptr(); auto samples = dX_dH->samples(); - auto samples_values = samples->values(); - auto* samples_values_ptr = samples_values.data_ptr(); + auto sample_values = samples->values(); + auto* sample_values_ptr = sample_values.data_ptr(); always_assert(samples->names().size() == 1); always_assert(samples->names()[0] == "sample"); @@ -397,7 +397,7 @@ std::vector CellGrad::forward( // =========================== compute dA_dH ============================ // for (int64_t grad_sample_i=0; grad_sample_icount(); grad_sample_i++) { - auto sample_i = samples_values_ptr[grad_sample_i]; + auto sample_i = sample_values_ptr[grad_sample_i]; // we get the structure from the samples of the values auto structure_i = static_cast(structures[sample_i].item()); @@ -439,8 +439,8 @@ std::vector CellGrad::backward( auto dB_d_dA_dH = grad_outputs[0]; // gradient of B w.r.t. dA/dH (output of forward) auto samples = dX_dH->samples(); - auto samples_values = samples->values(); - auto* samples_values_ptr = samples_values.data_ptr(); + auto sample_values = samples->values(); + auto* sample_values_ptr = sample_values.data_ptr(); always_assert(samples->names().size() == 1); always_assert(samples->names()[0] == "sample"); @@ -483,7 +483,7 @@ std::vector CellGrad::backward( // dB_d_dA_dH.shape == [structures, 3, 3] // dB_d_dA_dX.shape == [samples, features...] for (int64_t grad_sample_i=0; grad_sample_icount(); grad_sample_i++) { - auto sample_i = samples_values_ptr[grad_sample_i]; + auto sample_i = sample_values_ptr[grad_sample_i]; auto structure_i = static_cast(structures[sample_i].item()); for (int64_t i=0; i Vec<&str> { + fn sample_names(&self) -> Vec<&str> { if self.per_structure { return vec!["structure"]; } @@ -51,7 +51,7 @@ impl CalculatorBase for AtomicComposition { assert_eq!(keys.names(), ["species_center"]); let mut samples = Vec::new(); for [species_center_key] in keys.iter_fixed_size() { - let mut builder = LabelsBuilder::new(self.samples_names()); + let mut builder = LabelsBuilder::new(self.sample_names()); for (system_i, system) in systems.iter_mut().enumerate() { if self.per_structure { @@ -96,12 +96,12 @@ impl CalculatorBase for AtomicComposition { return vec![Vec::new(); keys.count()]; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { return vec!["count"]; } fn properties(&self, keys: &Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); properties.add(&[0]); let properties = properties.finish(); diff --git a/rascaline/src/calculators/dummy_calculator.rs b/rascaline/src/calculators/dummy_calculator.rs index 0674db16b..de2d19b0a 100644 --- a/rascaline/src/calculators/dummy_calculator.rs +++ b/rascaline/src/calculators/dummy_calculator.rs @@ -47,8 +47,8 @@ impl CalculatorBase for DummyCalculator { return CenterSpeciesKeys.keys(systems); } - fn samples_names(&self) -> Vec<&str> { - AtomCenteredSamples::samples_names() + fn sample_names(&self) -> Vec<&str> { + AtomCenteredSamples::sample_names() } fn samples(&self, keys: &Labels, systems: &mut [Box]) -> Result, Error> { @@ -96,12 +96,12 @@ impl CalculatorBase for DummyCalculator { return vec![Vec::new(); keys.count()]; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["index_delta", "x_y_z"] } fn properties(&self, keys: &Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); properties.add(&[1, 0]); properties.add(&[0, 1]); let properties = properties.finish(); diff --git a/rascaline/src/calculators/lode/spherical_expansion.rs b/rascaline/src/calculators/lode/spherical_expansion.rs index a10666e51..40e2e4629 100644 --- a/rascaline/src/calculators/lode/spherical_expansion.rs +++ b/rascaline/src/calculators/lode/spherical_expansion.rs @@ -126,7 +126,7 @@ fn compute_structure_factors(positions: &[Vector3D], species: &[i32], k_vectors: .map(|s| (s, Array2::from_elem((n_atoms, n_k_vectors), 0.0))) .collect::>(); - // Precompute species-dependent sum over neighbours j of sines and cosines + // Precompute species-dependent sum over neighbours j of sines and cosines let mut sumjcos = all_species.iter().copied() .map(|s| (s, Array1::from_elem(n_k_vectors, 0.0))) .collect::>(); @@ -141,7 +141,7 @@ fn compute_structure_factors(positions: &[Vector3D], species: &[i32], k_vectors: sumjsin[k] += sines[[j, k]]; } } - + // Compute Sum_j cos(k*r_ij) and Sum_j sin(k*r_ij) using the subtraction theorem for i in 0..n_atoms { for (species, real_per_center) in &mut real_per_center { @@ -149,7 +149,7 @@ fn compute_structure_factors(positions: &[Vector3D], species: &[i32], k_vectors: let sumjsin = sumjsin.get_mut(&species).unwrap(); let imag_per_center = imag_per_center.get_mut(&species).unwrap(); for k in 0..n_k_vectors { - let real = cosines[[i, k]] * sumjcos[k] + sines[[i, k]] * sumjsin[k]; + let real = cosines[[i, k]] * sumjcos[k] + sines[[i, k]] * sumjsin[k]; let imag = sines[[i, k]] * sumjcos[k] - cosines[[i, k]] * sumjsin[k]; real_per_center[[i, k]] += 2.0 * real; imag_per_center[[i, k]] += 2.0 * imag; @@ -467,8 +467,8 @@ impl CalculatorBase for LodeSphericalExpansion { return Ok(builder.finish()); } - fn samples_names(&self) -> Vec<&str> { - LongRangeSamplesPerAtom::samples_names() + fn sample_names(&self) -> Vec<&str> { + LongRangeSamplesPerAtom::sample_names() } fn samples(&self, keys: &Labels, systems: &mut [Box]) -> Result, Error> { @@ -556,12 +556,12 @@ impl CalculatorBase for LodeSphericalExpansion { return result; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["n"] } fn properties(&self, keys: &Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); for n in 0..self.parameters.max_radial { properties.add(&[n]); } diff --git a/rascaline/src/calculators/mod.rs b/rascaline/src/calculators/mod.rs index 5603b1d94..e10dd828d 100644 --- a/rascaline/src/calculators/mod.rs +++ b/rascaline/src/calculators/mod.rs @@ -25,7 +25,7 @@ pub trait CalculatorBase: std::panic::RefUnwindSafe { fn keys(&self, systems: &mut [Box]) -> Result; /// Get the names used for sample labels by this calculator - fn samples_names(&self) -> Vec<&str>; + fn sample_names(&self) -> Vec<&str>; /// Get the full list of samples this calculator would create for the given /// systems. This function should return one set of samples for each key. @@ -48,7 +48,7 @@ pub trait CalculatorBase: std::panic::RefUnwindSafe { fn components(&self, keys: &Labels) -> Vec>; /// Get the names used for property labels by this calculator - fn properties_names(&self) -> Vec<&str>; + fn property_names(&self) -> Vec<&str>; /// Get the properties this calculator computes for each key. fn properties(&self, keys: &Labels) -> Vec; diff --git a/rascaline/src/calculators/neighbor_list.rs b/rascaline/src/calculators/neighbor_list.rs index fc19c22cb..0834713c0 100644 --- a/rascaline/src/calculators/neighbor_list.rs +++ b/rascaline/src/calculators/neighbor_list.rs @@ -81,7 +81,7 @@ impl CalculatorBase for NeighborList { } } - fn samples_names(&self) -> Vec<&str> { + fn sample_names(&self) -> Vec<&str> { return vec!["structure", "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"]; } @@ -134,12 +134,12 @@ impl CalculatorBase for NeighborList { return vec![components; keys.count()]; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["distance"] } fn properties(&self, keys: &Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); properties.add(&[LabelValue::new(1)]); let properties = properties.finish(); diff --git a/rascaline/src/calculators/soap/power_spectrum.rs b/rascaline/src/calculators/soap/power_spectrum.rs index a879c8bc6..229f23834 100644 --- a/rascaline/src/calculators/soap/power_spectrum.rs +++ b/rascaline/src/calculators/soap/power_spectrum.rs @@ -435,8 +435,8 @@ impl CalculatorBase for SoapPowerSpectrum { return builder.keys(systems); } - fn samples_names(&self) -> Vec<&str> { - AtomCenteredSamples::samples_names() + fn sample_names(&self) -> Vec<&str> { + AtomCenteredSamples::sample_names() } fn samples(&self, keys: &metatensor::Labels, systems: &mut [Box]) -> Result, Error> { @@ -498,12 +498,12 @@ impl CalculatorBase for SoapPowerSpectrum { return vec![vec![]; keys.count()]; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["l", "n1", "n2"] } fn properties(&self, keys: &metatensor::Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); for l in 0..=self.parameters.max_angular { for n1 in 0..self.parameters.max_radial { for n2 in 0..self.parameters.max_radial { diff --git a/rascaline/src/calculators/soap/radial_spectrum.rs b/rascaline/src/calculators/soap/radial_spectrum.rs index a6a353a92..621ec82f0 100644 --- a/rascaline/src/calculators/soap/radial_spectrum.rs +++ b/rascaline/src/calculators/soap/radial_spectrum.rs @@ -137,8 +137,8 @@ impl CalculatorBase for SoapRadialSpectrum { return builder.keys(systems); } - fn samples_names(&self) -> Vec<&str> { - AtomCenteredSamples::samples_names() + fn sample_names(&self) -> Vec<&str> { + AtomCenteredSamples::sample_names() } fn samples( @@ -192,12 +192,12 @@ impl CalculatorBase for SoapRadialSpectrum { return vec![vec![]; keys.count()]; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["n"] } fn properties(&self, keys: &metatensor::Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); for n in 0..self.parameters.max_radial { properties.add(&[n]); } diff --git a/rascaline/src/calculators/soap/spherical_expansion.rs b/rascaline/src/calculators/soap/spherical_expansion.rs index b62cfa37d..63a8ddf19 100644 --- a/rascaline/src/calculators/soap/spherical_expansion.rs +++ b/rascaline/src/calculators/soap/spherical_expansion.rs @@ -609,8 +609,8 @@ impl CalculatorBase for SphericalExpansion { return Ok(builder.finish()); } - fn samples_names(&self) -> Vec<&str> { - AtomCenteredSamples::samples_names() + fn sample_names(&self) -> Vec<&str> { + AtomCenteredSamples::sample_names() } fn samples(&self, keys: &Labels, systems: &mut [Box]) -> Result, Error> { @@ -703,12 +703,12 @@ impl CalculatorBase for SphericalExpansion { return result; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["n"] } fn properties(&self, keys: &Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); for n in 0..self.by_pair.parameters().max_radial { properties.add(&[n]); } diff --git a/rascaline/src/calculators/soap/spherical_expansion_pair.rs b/rascaline/src/calculators/soap/spherical_expansion_pair.rs index 897ecd7cf..737b68797 100644 --- a/rascaline/src/calculators/soap/spherical_expansion_pair.rs +++ b/rascaline/src/calculators/soap/spherical_expansion_pair.rs @@ -551,7 +551,7 @@ impl CalculatorBase for SphericalExpansionByPair { return Ok(keys.finish()); } - fn samples_names(&self) -> Vec<&str> { + fn sample_names(&self) -> Vec<&str> { return vec!["structure", "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"]; } @@ -665,12 +665,12 @@ impl CalculatorBase for SphericalExpansionByPair { return result; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["n"] } fn properties(&self, keys: &Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); for n in 0..self.parameters.max_radial { properties.add(&[n]); } diff --git a/rascaline/src/calculators/sorted_distances.rs b/rascaline/src/calculators/sorted_distances.rs index cfd27a86d..21621cef7 100644 --- a/rascaline/src/calculators/sorted_distances.rs +++ b/rascaline/src/calculators/sorted_distances.rs @@ -53,8 +53,8 @@ impl CalculatorBase for SortedDistances { return CenterSpeciesKeys.keys(systems); } - fn samples_names(&self) -> Vec<&str> { - AtomCenteredSamples::samples_names() + fn sample_names(&self) -> Vec<&str> { + AtomCenteredSamples::sample_names() } fn samples(&self, keys: &Labels, systems: &mut [Box]) -> Result, Error> { @@ -100,12 +100,12 @@ impl CalculatorBase for SortedDistances { return vec![Vec::new(); keys.count()]; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["neighbor"] } fn properties(&self, keys: &Labels) -> Vec { - let mut properties = LabelsBuilder::new(self.properties_names()); + let mut properties = LabelsBuilder::new(self.property_names()); for i in 0..self.max_neighbors { properties.add(&[i]); } diff --git a/rascaline/src/labels/samples/atom_centered.rs b/rascaline/src/labels/samples/atom_centered.rs index 58445f374..544a2ec7d 100644 --- a/rascaline/src/labels/samples/atom_centered.rs +++ b/rascaline/src/labels/samples/atom_centered.rs @@ -8,7 +8,7 @@ use super::{SamplesBuilder, SpeciesFilter}; /// `SampleBuilder` for atom-centered representation. This will create one /// sample for each atom, optionally filtering on the central atom species. The -/// samples names are ("structure", "center"). +/// sample names are ("structure", "center"). /// /// Positions gradient samples include all atoms within a spherical cutoff, /// optionally filtering on the neighbor atom species. @@ -24,13 +24,13 @@ pub struct AtomCenteredSamples { } impl SamplesBuilder for AtomCenteredSamples { - fn samples_names() -> Vec<&'static str> { + fn sample_names() -> Vec<&'static str> { vec!["structure", "center"] } fn samples(&self, systems: &mut [Box]) -> Result { assert!(self.cutoff > 0.0 && self.cutoff.is_finite(), "cutoff must be positive for AtomCenteredSamples"); - let mut builder = LabelsBuilder::new(Self::samples_names()); + let mut builder = LabelsBuilder::new(Self::sample_names()); for (system_i, system) in systems.iter_mut().enumerate() { system.compute_neighbors(self.cutoff)?; let species = system.species()?; diff --git a/rascaline/src/labels/samples/long_range.rs b/rascaline/src/labels/samples/long_range.rs index ad61d2cfb..ccd8c27b8 100644 --- a/rascaline/src/labels/samples/long_range.rs +++ b/rascaline/src/labels/samples/long_range.rs @@ -18,14 +18,14 @@ pub struct LongRangeSamplesPerAtom { } impl SamplesBuilder for LongRangeSamplesPerAtom { - fn samples_names() -> Vec<&'static str> { + fn sample_names() -> Vec<&'static str> { vec!["structure", "center"] } fn samples(&self, systems: &mut [Box]) -> Result { assert!(self.self_pairs, "self.self_pairs = false is not implemented"); - let mut builder = LabelsBuilder::new(Self::samples_names()); + let mut builder = LabelsBuilder::new(Self::sample_names()); for (system_i, system) in systems.iter_mut().enumerate() { let species = system.species()?; diff --git a/rascaline/src/labels/samples/mod.rs b/rascaline/src/labels/samples/mod.rs index fe7ada49f..5b7d77ff6 100644 --- a/rascaline/src/labels/samples/mod.rs +++ b/rascaline/src/labels/samples/mod.rs @@ -39,7 +39,7 @@ impl SpeciesFilter { /// and can optionally implement samples for gradients calculation. pub trait SamplesBuilder { /// Get the names used by the samples - fn samples_names() -> Vec<&'static str>; + fn sample_names() -> Vec<&'static str>; /// Create `Labels` containing all the samples corresponding to the given /// list of systems. diff --git a/rascaline/src/tutorials/moments/moments.rs b/rascaline/src/tutorials/moments/moments.rs index 5e4b2d964..283c10500 100644 --- a/rascaline/src/tutorials/moments/moments.rs +++ b/rascaline/src/tutorials/moments/moments.rs @@ -33,8 +33,8 @@ impl CalculatorBase for GeometricMoments { return builder.keys(systems); } - fn samples_names(&self) -> Vec<&str> { - AtomCenteredSamples::samples_names() + fn sample_names(&self) -> Vec<&str> { + AtomCenteredSamples::sample_names() } fn samples(&self, keys: &Labels, systems: &mut [Box]) -> Result, Error> { @@ -85,12 +85,12 @@ impl CalculatorBase for GeometricMoments { return vec![vec![]; keys.count()]; } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["k"] } fn properties(&self, keys: &Labels) -> Vec { - let mut builder = LabelsBuilder::new(self.properties_names()); + let mut builder = LabelsBuilder::new(self.property_names()); for k in 0..=self.max_moment { builder.add(&[k]); } diff --git a/rascaline/src/tutorials/moments/s1_scaffold.rs b/rascaline/src/tutorials/moments/s1_scaffold.rs index 20cae5f71..1bcdf94b9 100644 --- a/rascaline/src/tutorials/moments/s1_scaffold.rs +++ b/rascaline/src/tutorials/moments/s1_scaffold.rs @@ -33,7 +33,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn samples_names(&self) -> Vec<&str> { + fn sample_names(&self) -> Vec<&str> { todo!() } @@ -53,7 +53,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { todo!() } diff --git a/rascaline/src/tutorials/moments/s2_metadata.rs b/rascaline/src/tutorials/moments/s2_metadata.rs index 336cc6901..4e152b550 100644 --- a/rascaline/src/tutorials/moments/s2_metadata.rs +++ b/rascaline/src/tutorials/moments/s2_metadata.rs @@ -46,8 +46,8 @@ impl CalculatorBase for GeometricMoments { // [CalculatorBase::keys] // [CalculatorBase::samples] - fn samples_names(&self) -> Vec<&str> { - AtomCenteredSamples::samples_names() + fn sample_names(&self) -> Vec<&str> { + AtomCenteredSamples::sample_names() } fn samples(&self, keys: &Labels, systems: &mut [Box]) -> Result, Error> { @@ -112,12 +112,12 @@ impl CalculatorBase for GeometricMoments { // [CalculatorBase::components] // [CalculatorBase::properties] - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { vec!["k"] } fn properties(&self, keys: &Labels) -> Vec { - let mut builder = LabelsBuilder::new(self.properties_names()); + let mut builder = LabelsBuilder::new(self.property_names()); for k in 0..=self.max_moment { builder.add(&[k]); } diff --git a/rascaline/src/tutorials/moments/s3_compute_1.rs b/rascaline/src/tutorials/moments/s3_compute_1.rs index 3eca33d44..89aacce22 100644 --- a/rascaline/src/tutorials/moments/s3_compute_1.rs +++ b/rascaline/src/tutorials/moments/s3_compute_1.rs @@ -29,7 +29,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn samples_names(&self) -> Vec<&str> { + fn sample_names(&self) -> Vec<&str> { todo!() } @@ -49,7 +49,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { todo!() } diff --git a/rascaline/src/tutorials/moments/s3_compute_2.rs b/rascaline/src/tutorials/moments/s3_compute_2.rs index e9f43d0d5..92e7cdd20 100644 --- a/rascaline/src/tutorials/moments/s3_compute_2.rs +++ b/rascaline/src/tutorials/moments/s3_compute_2.rs @@ -29,7 +29,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn samples_names(&self) -> Vec<&str> { + fn sample_names(&self) -> Vec<&str> { todo!() } @@ -49,7 +49,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { todo!() } diff --git a/rascaline/src/tutorials/moments/s3_compute_3.rs b/rascaline/src/tutorials/moments/s3_compute_3.rs index 6f8a929b3..27324c2af 100644 --- a/rascaline/src/tutorials/moments/s3_compute_3.rs +++ b/rascaline/src/tutorials/moments/s3_compute_3.rs @@ -29,7 +29,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn samples_names(&self) -> Vec<&str> { + fn sample_names(&self) -> Vec<&str> { todo!() } @@ -49,7 +49,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { todo!() } diff --git a/rascaline/src/tutorials/moments/s3_compute_4.rs b/rascaline/src/tutorials/moments/s3_compute_4.rs index b7f98c18e..5c3288892 100644 --- a/rascaline/src/tutorials/moments/s3_compute_4.rs +++ b/rascaline/src/tutorials/moments/s3_compute_4.rs @@ -36,7 +36,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn samples_names(&self) -> Vec<&str> { + fn sample_names(&self) -> Vec<&str> { todo!() } @@ -56,7 +56,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { todo!() } diff --git a/rascaline/src/tutorials/moments/s3_compute_5.rs b/rascaline/src/tutorials/moments/s3_compute_5.rs index 0e853a3cd..b782f7a8a 100644 --- a/rascaline/src/tutorials/moments/s3_compute_5.rs +++ b/rascaline/src/tutorials/moments/s3_compute_5.rs @@ -37,7 +37,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn samples_names(&self) -> Vec<&str> { + fn sample_names(&self) -> Vec<&str> { todo!() } @@ -57,7 +57,7 @@ impl CalculatorBase for GeometricMoments { todo!() } - fn properties_names(&self) -> Vec<&str> { + fn property_names(&self) -> Vec<&str> { todo!() } diff --git a/tox.ini b/tox.ini index 351f8da60..af80a8ed0 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 @ git+https://github.com/lab-cosmo/metatensor@ae10944\#subdirectory=python/metatensor-core + metatensor-core @ git+https://github.com/lab-cosmo/metatensor@ad3eda5\#subdirectory=python/metatensor-core metatensor-torch-requirement = - metatensor-torch @ git+https://github.com/lab-cosmo/metatensor@ae10944\#subdirectory=python/metatensor-torch + metatensor-torch @ git+https://github.com/lab-cosmo/metatensor@ad3eda5\#subdirectory=python/metatensor-torch build-single-wheel = --no-deps --no-build-isolation --check-build-dependencies From 434f92389f6b68bacb34e951093921d3d564d8c4 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Wed, 11 Oct 2023 12:06:42 +0200 Subject: [PATCH 09/30] Apply linter suggestions for C++ and Rust code I used clang-tidy for C++ code, and clippy for Rust code --- rascaline-c-api/src/calculator.rs | 2 +- rascaline-c-api/tests/calculator.cpp | 78 +++++++++++-------- .../tests/cmake-project/src/main.c | 1 - .../tests/cmake-project/src/main.cpp | 4 +- rascaline-c-api/tests/cxx/systems.cpp | 2 +- rascaline-c-api/tests/helpers.cpp | 12 +-- rascaline-c-api/tests/logging.cpp | 4 +- rascaline-c-api/tests/systems.cpp | 4 +- .../include/rascaline/torch/system.hpp | 8 +- rascaline-torch/src/calculator.cpp | 12 +-- rascaline-torch/src/register.cpp | 6 +- rascaline-torch/src/system.cpp | 2 +- rascaline-torch/tests/calculator.cpp | 6 +- rascaline-torch/tests/system.cpp | 6 +- .../calculators/lode/spherical_expansion.rs | 6 +- rascaline/src/lib.rs | 2 +- scripts/clean-python.sh | 6 +- scripts/package-torch.sh | 4 +- 18 files changed, 88 insertions(+), 77 deletions(-) diff --git a/rascaline-c-api/src/calculator.rs b/rascaline-c-api/src/calculator.rs index 234a2aba0..34607d50b 100644 --- a/rascaline-c-api/src/calculator.rs +++ b/rascaline-c-api/src/calculator.rs @@ -257,7 +257,7 @@ fn convert_labels_selection<'a>( } (true, false) => { let tensor = unsafe { - TensorMap::from_raw(selection.predefined as *mut mts_tensormap_t) + TensorMap::from_raw(selection.predefined.cast_mut()) }; match tensor.try_clone() { diff --git a/rascaline-c-api/tests/calculator.cpp b/rascaline-c-api/tests/calculator.cpp index 965a4146b..d19a810da 100644 --- a/rascaline-c-api/tests/calculator.cpp +++ b/rascaline-c-api/tests/calculator.cpp @@ -9,11 +9,11 @@ static void check_block( mts_tensormap_t* descriptor, size_t block_id, - std::vector samples, - std::vector properties, - std::vector values, - std::vector gradient_samples, - std::vector gradients + const std::vector& samples, + const std::vector& properties, + const std::vector& values, + const std::vector& gradient_samples, + const std::vector& gradients ); TEST_CASE("calculator name") { @@ -26,7 +26,7 @@ TEST_CASE("calculator name") { auto* calculator = rascal_calculator("dummy_calculator", HYPERS_JSON); REQUIRE(calculator != nullptr); - char buffer[256] = {0}; + char buffer[256] = {}; CHECK_SUCCESS(rascal_calculator_name(calculator, buffer, sizeof(buffer))); CHECK(buffer == std::string("dummy test calculator with cutoff: 3.5 - delta: 25 - name: bar")); @@ -68,7 +68,7 @@ TEST_CASE("calculator parameters") { auto* calculator = rascal_calculator("dummy_calculator", HYPERS_JSON.c_str()); REQUIRE(calculator != nullptr); - char buffer[256] = {0}; + char buffer[256] = {}; CHECK_SUCCESS(rascal_calculator_parameters(calculator, buffer, sizeof(buffer))); CHECK(buffer == HYPERS_JSON); @@ -138,7 +138,9 @@ TEST_CASE("Compute descriptor") { SECTION("Full compute") { auto system = simple_system(); - rascal_calculation_options_t options = {0}; + rascal_calculation_options_t options; + std::memset(&options, 0, sizeof(rascal_calculation_options_t)); + const char* gradients_list[] = {"positions"}; options.gradients = gradients_list; options.gradients_count = 1; @@ -151,7 +153,9 @@ TEST_CASE("Compute descriptor") { ); CHECK_SUCCESS(status); - mts_labels_t keys = {0}; + mts_labels_t keys; + std::memset(&keys, 0, sizeof(mts_labels_t)); + status = mts_tensormap_keys(descriptor, &keys); CHECK_SUCCESS(status); @@ -219,7 +223,9 @@ TEST_CASE("Compute descriptor") { "structure", "center" }; - mts_labels_t selected_samples = {0}; + mts_labels_t selected_samples; + std::memset(&selected_samples, 0, sizeof(mts_labels_t)); + selected_samples.names = selected_sample_names.data(); selected_samples.values = selected_sample_values.data(); selected_samples.count = 2; @@ -227,7 +233,9 @@ TEST_CASE("Compute descriptor") { auto system = simple_system(); - rascal_calculation_options_t options = {0}; + rascal_calculation_options_t options; + std::memset(&options, 0, sizeof(rascal_calculation_options_t)); + const char* gradients_list[] = {"positions"}; options.gradients = gradients_list; options.gradients_count = 1; @@ -242,7 +250,9 @@ TEST_CASE("Compute descriptor") { CHECK_SUCCESS(status); - mts_labels_t keys = {0}; + mts_labels_t keys; + std::memset(&keys, 0, sizeof(mts_labels_t)); + status = mts_tensormap_keys(descriptor, &keys); CHECK_SUCCESS(status); @@ -297,7 +307,9 @@ TEST_CASE("Compute descriptor") { "index_delta", "x_y_z" }; - mts_labels_t selected_properties = {0}; + mts_labels_t selected_properties; + std::memset(&selected_properties, 0, sizeof(mts_labels_t)); + selected_properties.names = selected_property_names.data(); selected_properties.values = selected_property_values.data(); selected_properties.count = 1; @@ -305,7 +317,9 @@ TEST_CASE("Compute descriptor") { auto system = simple_system(); - rascal_calculation_options_t options = {0}; + rascal_calculation_options_t options; + std::memset(&options, 0, sizeof(rascal_calculation_options_t)); + const char* gradients_list[] = {"positions"}; options.gradients = gradients_list; options.gradients_count = 1; @@ -319,7 +333,7 @@ TEST_CASE("Compute descriptor") { ); CHECK_SUCCESS(status); - mts_labels_t keys = {0}; + mts_labels_t keys = {}; status = mts_tensormap_keys(descriptor, &keys); CHECK_SUCCESS(status); @@ -396,13 +410,13 @@ TEST_CASE("Compute descriptor") { mts_block_t* blocks[2] = {nullptr, nullptr}; - mts_labels_t h_samples = {0}; + mts_labels_t h_samples = {}; h_samples.size = 2; h_samples.names = sample_names.data(); h_samples.count = 1; h_samples.values = h_sample_values.data(); - mts_labels_t h_properties = {0}; + mts_labels_t h_properties = {}; h_properties.size = 2; h_properties.names = property_names.data(); h_properties.count = 1; @@ -418,13 +432,13 @@ TEST_CASE("Compute descriptor") { 1, 0, }; - mts_labels_t c_samples = {0}; + mts_labels_t c_samples = {}; c_samples.size = 2; c_samples.names = sample_names.data(); c_samples.count = 1; c_samples.values = c_sample_values.data(); - mts_labels_t c_properties = {0}; + mts_labels_t c_properties = {}; c_properties.size = 2; c_properties.names = property_names.data(); c_properties.count = 1; @@ -435,17 +449,17 @@ TEST_CASE("Compute descriptor") { auto keys_names = std::vector{"species_center"}; auto keys_values = std::vector{1, 6}; - mts_labels_t keys = {0}; + mts_labels_t keys = {}; keys.size = 1; keys.names = keys_names.data(); keys.count = 2; keys.values = keys_values.data(); - auto predefined = mts_tensormap(keys, blocks, 2); + auto* predefined = mts_tensormap(keys, blocks, 2); REQUIRE(predefined != nullptr); auto system = simple_system(); - rascal_calculation_options_t options = {0}; + rascal_calculation_options_t options = {}; const char* gradients_list[] = {"positions"}; options.gradients = gradients_list; options.gradients_count = 1; @@ -520,7 +534,7 @@ TEST_CASE("Compute descriptor") { // existing one (1) from the default set of keys. We also put the keys // in a different order than what would be the default (6, 12). - mts_labels_t selected_keys = {0}; + mts_labels_t selected_keys = {}; const char* sample_names[] = {"species_center"}; selected_keys.names = sample_names; selected_keys.size = 1; @@ -531,7 +545,7 @@ TEST_CASE("Compute descriptor") { auto system = simple_system(); - rascal_calculation_options_t options = {0}; + rascal_calculation_options_t options = {}; const char* gradients_list[] = {"positions"}; options.gradients = gradients_list; options.gradients_count = 1; @@ -546,7 +560,7 @@ TEST_CASE("Compute descriptor") { ); CHECK_SUCCESS(status); - mts_labels_t keys = {0}; + mts_labels_t keys = {}; status = mts_tensormap_keys(descriptor, &keys); CHECK_SUCCESS(status); @@ -596,11 +610,11 @@ TEST_CASE("Compute descriptor") { void check_block( mts_tensormap_t* descriptor, size_t block_id, - std::vector samples, - std::vector properties, - std::vector values, - std::vector gradient_samples, - std::vector gradients + const std::vector& samples, + const std::vector& properties, + const std::vector& values, + const std::vector& gradient_samples, + const std::vector& gradients ) { mts_block_t* block = nullptr; @@ -608,7 +622,7 @@ void check_block( CHECK_SUCCESS(status); /**************************************************************************/ - mts_labels_t labels = {0}; + mts_labels_t labels = {}; status = mts_block_labels(block, 0, &labels); CHECK_SUCCESS(status); @@ -639,7 +653,7 @@ void check_block( mts_labels_free(&labels); /**************************************************************************/ - mts_array_t array = {0}; + mts_array_t array = {}; status = mts_block_data(block, &array); CHECK_SUCCESS(status); diff --git a/rascaline-c-api/tests/cmake-project/src/main.c b/rascaline-c-api/tests/cmake-project/src/main.c index 309de1342..84c6093b5 100644 --- a/rascaline-c-api/tests/cmake-project/src/main.c +++ b/rascaline-c-api/tests/cmake-project/src/main.c @@ -1,5 +1,4 @@ #include - #include diff --git a/rascaline-c-api/tests/cmake-project/src/main.cpp b/rascaline-c-api/tests/cmake-project/src/main.cpp index 3de19f8ea..986fcb69f 100644 --- a/rascaline-c-api/tests/cmake-project/src/main.cpp +++ b/rascaline-c-api/tests/cmake-project/src/main.cpp @@ -1,5 +1,3 @@ -#include - #include @@ -7,7 +5,7 @@ int main() { try { auto calculator = rascaline::Calculator( "dummy_calculator", - "{\"cutoff\": 3.4, \"delta\": -3, \"name\": \"testing\", \"gradients\": true}" + R"({"cutoff": 3.4, "delta": -3, "name": "testing", "gradients": true})" ); return 0; } catch (const rascaline::RascalineError& e) { diff --git a/rascaline-c-api/tests/cxx/systems.cpp b/rascaline-c-api/tests/cxx/systems.cpp index 76ff8638f..03f878fe5 100644 --- a/rascaline-c-api/tests/cxx/systems.cpp +++ b/rascaline-c-api/tests/cxx/systems.cpp @@ -9,7 +9,7 @@ TEST_CASE("basic systems") { CHECK(systems.count() == 30); - auto system = systems.systems(); + auto* system = systems.systems(); uintptr_t size = 0; system->size(system->user_data, &size); CHECK(size == 54); diff --git a/rascaline-c-api/tests/helpers.cpp b/rascaline-c-api/tests/helpers.cpp index 69fa6710f..b240227de 100644 --- a/rascaline-c-api/tests/helpers.cpp +++ b/rascaline-c-api/tests/helpers.cpp @@ -8,7 +8,7 @@ #define SQRT_3 1.73205080756887729352 rascal_system_t simple_system() { - rascal_system_t system = {0}; + rascal_system_t system = {}; system.size = [](const void* _, uintptr_t* size) { *size = 4; @@ -104,25 +104,25 @@ rascal_system_t simple_system() { } mts_array_t empty_array(std::vector array_shape) { - mts_array_t array = {0}; + mts_array_t array = {}; - array.ptr = new std::vector(array_shape); + array.ptr = new std::vector(std::move(array_shape)); array.origin = [](const void *array, mts_data_origin_t *origin){ mts_register_data_origin("c-tests-empty-array", origin); return MTS_SUCCESS; }; array.shape = [](const void *array, const uintptr_t** shape, uintptr_t* shape_count){ - auto array_shape = static_cast*>(array); + const auto* array_shape = static_cast*>(array); *shape = array_shape->data(); *shape_count = array_shape->size(); return MTS_SUCCESS; }; array.destroy = [](void *array){ - auto array_shape = static_cast*>(array); + auto* array_shape = static_cast*>(array); delete array_shape; }; array.copy = [](const void *array, mts_array_t* new_array){ - auto array_shape = static_cast*>(array); + const auto* array_shape = static_cast*>(array); *new_array = empty_array(*array_shape); return MTS_SUCCESS; }; diff --git a/rascaline-c-api/tests/logging.cpp b/rascaline-c-api/tests/logging.cpp index b3391fb27..61e00860f 100644 --- a/rascaline-c-api/tests/logging.cpp +++ b/rascaline-c-api/tests/logging.cpp @@ -13,7 +13,7 @@ static void run_calculation(const char* hypers) { auto* calculator = rascal_calculator("dummy_calculator", hypers); REQUIRE(calculator != nullptr); auto system = simple_system(); - rascal_calculation_options_t options = {0}; + rascal_calculation_options_t options = {}; mts_tensormap_t* descriptor = nullptr; CHECK_SUCCESS(rascal_calculator_compute( @@ -26,7 +26,7 @@ static void run_calculation(const char* hypers) { TEST_CASE("Logging") { auto record_log_events = [](int level, const char* message) { - RECORDED_LOG_EVENTS.push_back(std::make_tuple(level, std::string(message))); + RECORDED_LOG_EVENTS.emplace_back(level, std::string(message)); }; CHECK_SUCCESS(rascal_set_logging_callback(record_log_events)); diff --git a/rascaline-c-api/tests/systems.cpp b/rascaline-c-api/tests/systems.cpp index e185efc56..344182c7d 100644 --- a/rascaline-c-api/tests/systems.cpp +++ b/rascaline-c-api/tests/systems.cpp @@ -57,8 +57,8 @@ TEST_CASE("systems errors") { auto* calculator = rascal_calculator("dummy_calculator", HYPERS_JSON); REQUIRE(calculator != nullptr); - rascal_system_t system = {0}; - rascal_calculation_options_t options = {0}; + rascal_system_t system = {}; + rascal_calculation_options_t options = {}; mts_tensormap_t* descriptor = nullptr; auto status = rascal_calculator_compute( diff --git a/rascaline-torch/include/rascaline/torch/system.hpp b/rascaline-torch/include/rascaline/torch/system.hpp index fefc583ab..22f810758 100644 --- a/rascaline-torch/include/rascaline/torch/system.hpp +++ b/rascaline-torch/include/rascaline/torch/system.hpp @@ -43,7 +43,7 @@ class RASCALINE_TORCH_EXPORT SystemHolder final: public rascaline::System, publi /// SystemHolder can be move assigned SystemHolder& operator=(SystemHolder&&) = default; - virtual ~SystemHolder() override = default; + ~SystemHolder() override = default; /*========================================================================*/ /* Functions to implement rascaline::System */ @@ -66,7 +66,7 @@ class RASCALINE_TORCH_EXPORT SystemHolder final: public rascaline::System, publi /// @private CellMatrix cell() const override { - auto data = cell_.data_ptr(); + auto* data = cell_.data_ptr(); return CellMatrix{{ {{data[0], data[1], data[2]}}, {{data[3], data[4], data[5]}}, @@ -117,12 +117,12 @@ class RASCALINE_TORCH_EXPORT SystemHolder final: public rascaline::System, publi } /// @private implementation of __len__ for TorchScript - int64_t __len__() const { + int64_t len() const { return species_.size(0); } /// @private implementation of __str__ for TorchScript - std::string __str__() const; + std::string str() const; // TODO: convert from a Dict[str, TorchTensorMap] for the interface with LAMMPS // static TorchSystem from_metatensor_dict(); diff --git a/rascaline-torch/src/calculator.cpp b/rascaline-torch/src/calculator.cpp index 71859ce9a..18ea2f86d 100644 --- a/rascaline-torch/src/calculator.cpp +++ b/rascaline-torch/src/calculator.cpp @@ -22,8 +22,8 @@ static metatensor::TensorBlock block_to_torch( auto torch_values = torch::from_blob( values.data(), - std::move(sizes), - [tensor](void*){ + sizes, + [tensor](void*) mutable { // this function holds a copy of `tensor`, which will make sure that // (a) the TensorMap is kept alive for as long as the values // returned by `torch::from_blob` is; and (b) the TensorMap will be @@ -41,9 +41,9 @@ static metatensor::TensorBlock block_to_torch( block.properties() ); - for (auto parameter: block.gradients_list()) { + for (const auto& parameter: block.gradients_list()) { auto gradient = block_to_torch(tensor, block.gradient(parameter)); - new_block.add_gradient(std::move(parameter), std::move(gradient)); + new_block.add_gradient(parameter, std::move(gradient)); } return new_block; @@ -146,7 +146,7 @@ metatensor_torch::TorchTensorMap CalculatorHolder::compute( int64_t current_start = 0; for (auto& system: systems) { structures_start.push_back(current_start); - current_start += system->size(); + current_start += static_cast(system->size()); } structures_start_ivalue = torch::IValue(std::move(structures_start)); } @@ -252,7 +252,7 @@ metatensor_torch::TorchTensorMap rascaline_torch::register_autograd( int64_t current_start = 0; for (auto& system: systems) { structures_start.push_back(current_start); - current_start += system->size(); + current_start += static_cast(system->size()); } structures_start_ivalue = torch::IValue(std::move(structures_start)); } diff --git a/rascaline-torch/src/register.cpp b/rascaline-torch/src/register.cpp index 703a5757b..965b10865 100644 --- a/rascaline-torch/src/register.cpp +++ b/rascaline-torch/src/register.cpp @@ -13,9 +13,9 @@ TORCH_LIBRARY(rascaline, module) { DOCSTRING, {torch::arg("species"), torch::arg("positions"), torch::arg("cell")} ) - .def("__str__", &SystemHolder::__str__) - .def("__repr__", &SystemHolder::__str__) - .def("__len__", &SystemHolder::__len__) + .def("__str__", &SystemHolder::str) + .def("__repr__", &SystemHolder::str) + .def("__len__", &SystemHolder::len) .def_property("species", &SystemHolder::get_species) .def_property("positions", &SystemHolder::get_positions) .def_property("cell", &SystemHolder::get_cell) diff --git a/rascaline-torch/src/system.cpp b/rascaline-torch/src/system.cpp index 119916428..b1bd999a2 100644 --- a/rascaline-torch/src/system.cpp +++ b/rascaline-torch/src/system.cpp @@ -125,7 +125,7 @@ const std::vector& SystemHolder::pairs_containing(uintptr_t cente } -std::string SystemHolder::__str__() const { +std::string SystemHolder::str() const { auto result = std::ostringstream(); result << "System with " << this->size() << " atoms, "; if (torch::all(cell_ == torch::zeros({3, 3})).item()) { diff --git a/rascaline-torch/tests/calculator.cpp b/rascaline-torch/tests/calculator.cpp index fe510d934..d841d603a 100644 --- a/rascaline-torch/tests/calculator.cpp +++ b/rascaline-torch/tests/calculator.cpp @@ -16,7 +16,7 @@ static TorchSystem test_system(bool positions_grad, bool cell_grad); TEST_CASE("Calculator") { - auto HYPERS_JSON = R"({ + const auto* HYPERS_JSON = R"({ "cutoff": 3.0, "delta": 4, "name": "bar" @@ -152,7 +152,7 @@ TEST_CASE("Calculator") { auto system = test_system(true, false); auto options = torch::make_intrusive(); - options->gradients.push_back("positions"); + options->gradients.emplace_back("positions"); auto descriptor = calculator.compute({system}, options); CHECK(*descriptor->keys() == metatensor::Labels( @@ -255,7 +255,7 @@ TEST_CASE("Calculator") { auto system = test_system(false, false); auto options = torch::make_intrusive(); - options->gradients.push_back("positions"); + options->gradients.emplace_back("positions"); auto descriptor = calculator.compute({system}, options); CHECK(*descriptor->keys() == metatensor::Labels( diff --git a/rascaline-torch/tests/system.cpp b/rascaline-torch/tests/system.cpp index 07f7ce58a..4226bffce 100644 --- a/rascaline-torch/tests/system.cpp +++ b/rascaline-torch/tests/system.cpp @@ -13,7 +13,7 @@ TEST_CASE("Systems") { auto system = SystemHolder(species, positions, cell); CHECK(system.size() == 5); - CHECK(system.__len__() == 5); + CHECK(system.len() == 5); CHECK(system.use_native_system() == true); @@ -47,10 +47,10 @@ TEST_CASE("Systems") { auto cell = torch::zeros({3, 3}); auto system = SystemHolder(species, positions, cell); - CHECK(system.__str__() == "System with 5 atoms, non periodic"); + CHECK(system.str() == "System with 5 atoms, non periodic"); cell = torch::eye(3); system = SystemHolder(species, positions, cell); - CHECK(system.__str__() == "System with 5 atoms, periodic cell: [1, 0, 0, 0, 1, 0, 0, 0, 1]"); + CHECK(system.str() == "System with 5 atoms, periodic cell: [1, 0, 0, 0, 1, 0, 0, 0, 1]"); } } diff --git a/rascaline/src/calculators/lode/spherical_expansion.rs b/rascaline/src/calculators/lode/spherical_expansion.rs index 40e2e4629..341067c35 100644 --- a/rascaline/src/calculators/lode/spherical_expansion.rs +++ b/rascaline/src/calculators/lode/spherical_expansion.rs @@ -145,9 +145,9 @@ fn compute_structure_factors(positions: &[Vector3D], species: &[i32], k_vectors: // Compute Sum_j cos(k*r_ij) and Sum_j sin(k*r_ij) using the subtraction theorem for i in 0..n_atoms { for (species, real_per_center) in &mut real_per_center { - let sumjcos = sumjcos.get_mut(&species).unwrap(); - let sumjsin = sumjsin.get_mut(&species).unwrap(); - let imag_per_center = imag_per_center.get_mut(&species).unwrap(); + let sumjcos = sumjcos.get_mut(species).unwrap(); + let sumjsin = sumjsin.get_mut(species).unwrap(); + let imag_per_center = imag_per_center.get_mut(species).unwrap(); for k in 0..n_k_vectors { let real = cosines[[i, k]] * sumjcos[k] + sines[[i, k]] * sumjsin[k]; let imag = sines[[i, k]] * sumjcos[k] - cosines[[i, k]] * sumjsin[k]; diff --git a/rascaline/src/lib.rs b/rascaline/src/lib.rs index fc7f33108..41acd16dc 100644 --- a/rascaline/src/lib.rs +++ b/rascaline/src/lib.rs @@ -6,7 +6,7 @@ #![allow(clippy::unreadable_literal, clippy::option_if_let_else, clippy::range_plus_one)] #![allow(clippy::missing_errors_doc, clippy::missing_panics_doc, clippy::module_name_repetitions)] #![allow(clippy::manual_assert, clippy::return_self_not_must_use, clippy::match_like_matches_macro)] -#![allow(clippy::needless_range_loop, clippy::uninlined_format_args)] +#![allow(clippy::needless_range_loop, clippy::uninlined_format_args, clippy::manual_let_else)] #![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] #![allow(clippy::cast_possible_wrap, clippy::cast_lossless, clippy::cast_sign_loss)] diff --git a/scripts/clean-python.sh b/scripts/clean-python.sh index 564e88196..650c09747 100755 --- a/scripts/clean-python.sh +++ b/scripts/clean-python.sh @@ -1,12 +1,12 @@ -#/usr/bin/env bash +#!/usr/bin/env bash # This script removes all temporary files created by Python during # installation and tests running. set -eux -ROOT_DIR=$(cd $(dirname "${BASH_SOURCE[0]}")/.. && pwd) -cd $ROOT_DIR +ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) +cd "$ROOT_DIR" rm -rf dist rm -rf build diff --git a/scripts/package-torch.sh b/scripts/package-torch.sh index db62ca4d9..39a97eadb 100755 --- a/scripts/package-torch.sh +++ b/scripts/package-torch.sh @@ -4,10 +4,10 @@ # rascaline-torch, and copy it to be included in the rascaline-torch python # package sdist. -ROOT_DIR=$(cd $(dirname "${BASH_SOURCE[0]}")/.. && pwd) +ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) set -eux -cd $ROOT_DIR +cd "$ROOT_DIR" tar cf rascaline-torch.tar rascaline-torch gzip -9 rascaline-torch.tar From 14d2a72ef60ce52be97450aad0f773da49637eea Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Thu, 12 Oct 2023 14:19:32 +0200 Subject: [PATCH 10/30] Use the released version of metatensor --- docs/requirements.txt | 2 +- pyproject.toml | 2 +- python/rascaline-torch/pyproject.toml | 2 +- rascaline-c-api/CMakeLists.txt | 13 ++++++------- rascaline-c-api/Cargo.toml | 4 ++-- rascaline-torch/CMakeLists.txt | 13 ++++++------- rascaline/Cargo.toml | 2 +- tox.ini | 4 ++-- 8 files changed, 20 insertions(+), 22 deletions(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index cb9112693..943d98e85 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] @ git+https://github.com/lab-cosmo/metatensor@ad3eda5 +metatensor[torch] >=0.1.0,<0.2.0 torch chemfiles matplotlib diff --git a/pyproject.toml b/pyproject.toml index 47dd87747..924103358 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] dependencies = [ - "metatensor-core @ git+https://github.com/lab-cosmo/metatensor@ad3eda5#subdirectory=python/metatensor-core", + "metatensor-core >=0.1.0,<0.2.0", ] [project.urls] diff --git a/python/rascaline-torch/pyproject.toml b/python/rascaline-torch/pyproject.toml index 00fe173a1..4c4daacf4 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 @git+https://github.com/lab-cosmo/metatensor@ad3eda5#subdirectory=python/metatensor-torch", + "metatensor-torch >=0.1.0,<0.2.0", ] # use a custom build backend to add a dependency on the right version of rascaline diff --git a/rascaline-c-api/CMakeLists.txt b/rascaline-c-api/CMakeLists.txt index 2ed0cd486..27de85d69 100644 --- a/rascaline-c-api/CMakeLists.txt +++ b/rascaline-c-api/CMakeLists.txt @@ -217,18 +217,17 @@ endif() # ============================================================================ # # Setup metatensor -set(METATENSOR_GIT_VERSION "ad3eda5") +set(METATENSOR_FETCH_VERSION "0.1.0") set(METATENSOR_REQUIRED_VERSION "0.1") if (RASCALINE_FETCH_METATENSOR) - message(STATUS "Fetching metatensor @ ${METATENSOR_GIT_VERSION} from github") + message(STATUS "Fetching metatensor @ ${METATENSOR_FETCH_VERSION} from github") include(FetchContent) + set(URL_ROOT "https://github.com/lab-cosmo/metatensor/releases/download") FetchContent_Declare( metatensor - URL https://github.com/lab-cosmo/metatensor/archive/${METATENSOR_GIT_VERSION}.zip - URL_HASH MD5=0ca550a82d587c4426eddf82ca8deae2 - SOURCE_SUBDIR metatensor-core - VERBOSE + URL ${URL_ROOT}/metatensor-core-v${METATENSOR_FETCH_VERSION}/metatensor-core-cxx-${METATENSOR_FETCH_VERSION}.tar.gz + URL_HASH MD5=3b48fdd8b6ca72e92a7785326dd74f6a ) if (CMAKE_VERSION VERSION_GREATER 3.18) @@ -238,7 +237,7 @@ if (RASCALINE_FETCH_METATENSOR) FetchContent_Populate(metatensor) endif() - add_subdirectory(${metatensor_SOURCE_DIR}/metatensor-core ${metatensor_BINARY_DIR}) + add_subdirectory(${metatensor_SOURCE_DIR} ${metatensor_BINARY_DIR}) endif() # metatensor will be installed in the same place as rascaline, so set diff --git a/rascaline-c-api/Cargo.toml b/rascaline-c-api/Cargo.toml index 19170f99a..7c22cf3ef 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 = "ad3eda5"} +metatensor = "0.1" 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 = "ad3eda5"} +metatensor = "0.1" [dev-dependencies] which = "4" diff --git a/rascaline-torch/CMakeLists.txt b/rascaline-torch/CMakeLists.txt index b8b2c7dc2..79bbb3402 100644 --- a/rascaline-torch/CMakeLists.txt +++ b/rascaline-torch/CMakeLists.txt @@ -58,18 +58,17 @@ find_package(Torch 1.11 REQUIRED) # ============================================================================ # # Setup metatensor_torch -set(METATENSOR_GIT_VERSION "ad3eda5") +set(METATENSOR_FETCH_VERSION "0.1.0") set(REQUIRED_METATENSOR_TORCH_VERSION "0.1") if (RASCALINE_TORCH_FETCH_METATENSOR_TORCH) - message(STATUS "Fetching metatensor_torch @ ${METATENSOR_GIT_VERSION} from github") + message(STATUS "Fetching metatensor_torch @ ${METATENSOR_FETCH_VERSION} from github") + set(URL_ROOT "https://github.com/lab-cosmo/metatensor/releases/download") include(FetchContent) FetchContent_Declare( metatensor_torch - URL https://github.com/lab-cosmo/metatensor/archive/${METATENSOR_GIT_VERSION}.zip - URL_HASH MD5=0ca550a82d587c4426eddf82ca8deae2 - SOURCE_SUBDIR metatensor-torch - VERBOSE + URL ${URL_ROOT}/metatensor-torch-v${METATENSOR_FETCH_VERSION}/metatensor-torch-cxx-${METATENSOR_FETCH_VERSION}.tar.gz + URL_HASH MD5=f4355be47f849fedcaf9bb2abd66d10c ) if (CMAKE_VERSION VERSION_GREATER 3.18) @@ -79,7 +78,7 @@ if (RASCALINE_TORCH_FETCH_METATENSOR_TORCH) FetchContent_Populate(metatensor_torch) endif() - add_subdirectory(${metatensor_torch_SOURCE_DIR}/metatensor-torch ${metatensor_torch_BINARY_DIR}) + add_subdirectory(${metatensor_torch_SOURCE_DIR} ${metatensor_torch_BINARY_DIR}) endif() else() find_package(metatensor_torch ${REQUIRED_METATENSOR_TORCH_VERSION} REQUIRED CONFIG) diff --git a/rascaline/Cargo.toml b/rascaline/Cargo.toml index dd6ccbafa..b45ac15a1 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 = "ad3eda5", features = ["rayon"]} +metatensor = {version = "0.1", 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 af80a8ed0..c4df209aa 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 @ git+https://github.com/lab-cosmo/metatensor@ad3eda5\#subdirectory=python/metatensor-core + metatensor-core >=0.1.0,<0.2.0 metatensor-torch-requirement = - metatensor-torch @ git+https://github.com/lab-cosmo/metatensor@ad3eda5\#subdirectory=python/metatensor-torch + metatensor-torch >=0.1.0,<0.2.0 build-single-wheel = --no-deps --no-build-isolation --check-build-dependencies From 581d0ca4dece424b4594c8c3cac40a7381e13ae5 Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Tue, 17 Oct 2023 10:54:09 +0200 Subject: [PATCH 11/30] Added real and Fourier space spliner base classes (#235) Co-authored-by: Guillaume Fraux --- docs/src/conf.py | 1 + docs/src/devdoc/explanations/index.rst | 1 + .../devdoc/explanations/radial-integral.rst | 275 ++++++ .../api/python/utils/atomic-density.rst | 1 + .../src/references/api/python/utils/index.rst | 2 + .../api/python/utils/radial-basis.rst | 1 + .../references/api/python/utils/splines.rst | 2 +- python/rascaline-torch/tests/calculator.py | 1 - python/rascaline/rascaline/utils/__init__.py | 15 +- python/rascaline/rascaline/utils/splines.py | 403 --------- .../rascaline/utils/splines/__init__.py | 18 + .../rascaline/utils/splines/atomic_density.py | 166 ++++ .../rascaline/utils/splines/radial_basis.py | 331 +++++++ .../rascaline/utils/splines/splines.py | 849 ++++++++++++++++++ python/rascaline/tests/utils/radial_basis.py | 91 ++ python/rascaline/tests/utils/splines.py | 228 ++++- tox.ini | 1 + 17 files changed, 1979 insertions(+), 407 deletions(-) create mode 100644 docs/src/devdoc/explanations/radial-integral.rst create mode 100644 docs/src/references/api/python/utils/atomic-density.rst create mode 100644 docs/src/references/api/python/utils/radial-basis.rst delete mode 100644 python/rascaline/rascaline/utils/splines.py create mode 100644 python/rascaline/rascaline/utils/splines/__init__.py create mode 100644 python/rascaline/rascaline/utils/splines/atomic_density.py create mode 100644 python/rascaline/rascaline/utils/splines/radial_basis.py create mode 100644 python/rascaline/rascaline/utils/splines/splines.py create mode 100644 python/rascaline/tests/utils/radial_basis.py diff --git a/docs/src/conf.py b/docs/src/conf.py index 3b45ea3f6..a8d8a21da 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -140,6 +140,7 @@ def setup(app): "metatensor": ("https://lab-cosmo.github.io/metatensor/latest/", None), "matplotlib": ("https://matplotlib.org/stable/", None), "numpy": ("https://numpy.org/doc/stable/", None), + "scipy": ("https://docs.scipy.org/doc/scipy/", None), "skmatter": ("https://scikit-matter.readthedocs.io/en/latest/", None), "torch": ("https://pytorch.org/docs/stable/", None), "python": ("https://docs.python.org/3", None), diff --git a/docs/src/devdoc/explanations/index.rst b/docs/src/devdoc/explanations/index.rst index 8d054cc20..9ccd23922 100644 --- a/docs/src/devdoc/explanations/index.rst +++ b/docs/src/devdoc/explanations/index.rst @@ -12,3 +12,4 @@ to give you more clarity and understanding of what rascaline is all about. architecture interfaces + radial-integral diff --git a/docs/src/devdoc/explanations/radial-integral.rst b/docs/src/devdoc/explanations/radial-integral.rst new file mode 100644 index 000000000..c17f5a901 --- /dev/null +++ b/docs/src/devdoc/explanations/radial-integral.rst @@ -0,0 +1,275 @@ +.. _radial-integral: + +SOAP and LODE radial integrals +=================================== + +On this page, we describe the exact mathematical expression that are implemented in the +radial integral and the splined radial integral classes i.e. +:ref:`python-splined-radial-integral`. Note that this page assumes knowledge of +spherical expansion & friends and currently serves as a reference page for +the developers to support the implementation. + +Preliminaries +------------- + +In this subsection, we briefly provide all the preliminary knowledge that is needed to +understand what the radial integral class is doing. The actual explanation for what is +computed in the radial integral class can be found in the next subsection (1.2). The +spherical expansion coefficients :math:`\langle anlm | \rho_i \rangle` are completely +determined by specifying two ingredients: + +- the atomic density function :math:`g(r)` as implemented in + :ref:`python-atomic-density`, often chosen to be a Gaussian or Delta function, that + defined the type of density under consideration. For a given center atom :math:`i` in + the structure, the total density function :math:`\rho_i(\boldsymbol{r})` around is + then defined as :math:`\rho_i(\boldsymbol{r}) = \sum_{j} g(\boldsymbol{r} - + \boldsymbol{r}_{ij})`. + +- the radial basis functions :math:`R_{nl}(r)` as implementated + :ref:`python-radial-basis`, on which the density :math:`\rho_i` is projected. To be + more precise, the actual basis functions are of the form + + .. math:: + + B_{nlm}(\boldsymbol{r}) = R_{nl}(r)Y_{lm}(\hat{r}), + + where :math:`Y_{lm}(\hat{r})` are the real spherical harmonics evaluated at the point + :math:`\hat{r}`, i.e. at the spherical angles :math:`(\theta, \phi)` that determine + the orientation of the unit vector :math:`\hat{r} = \boldsymbol{r}/r`. + +The spherical expansion coefficient :math:`\langle nlm | \rho_i \rangle` (we ommit the +chemical species index :math:`a` for simplicity) is then defined as + +.. math:: + + \begin{aligned} + \langle nlm | \rho_i \rangle & = \int \mathrm{d}^3\boldsymbol{r} + B_{nlm}(\boldsymbol{r}) \rho_i(\boldsymbol{r}) \\ \label{expansion_coeff_def} & = + \int \mathrm{d}^3\boldsymbol{r} R_{nl}(r)Y_{lm}(\hat{r})\rho_i(\boldsymbol{r}). + \end{aligned} + +In practice, the atom centered density :math:`\rho_i` is a superposition of the neighbor +contributions, namely :math:`\rho_i(\boldsymbol{r}) = \sum_{j} g(\boldsymbol{r} - +\boldsymbol{r}_{ij})`. Due to linearity of integration, evaluating the integral can then +be simplified to + +.. math:: + + \begin{aligned} + \langle nlm | \rho_i \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})\rho_i(\boldsymbol{r}) \\ & = \int + \mathrm{d}^3\boldsymbol{r} R_{nl}(r)Y_{lm}(\hat{r})\left( \sum_{j} + g(\boldsymbol{r} - \boldsymbol{r}_{ij})\right) \\ & = \sum_{j} \int + \mathrm{d}^3\boldsymbol{r} R_{nl}(r)Y_{lm}(\hat{r}) g(\boldsymbol{r} - + \boldsymbol{r}_{ij}) \\ & = \sum_j \langle nlm | g;\boldsymbol{r}_{ij} \rangle. + \end{aligned} + +Thus, instead of having to compute integrals for arbitrary densities :math:`\rho_i`, we +have reduced our problem to the evaluation of integrals of the form + +.. math:: + + \begin{aligned} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})g(\boldsymbol{r} - \boldsymbol{r}_{ij}), + \end{aligned} + +which are completely specified by + +- the density function :math:`g(\boldsymbol{r})` + +- the radial basis :math:`R_{nl}(r)` + +- the position of the neighbor atom :math:`\boldsymbol{r}_{ij}` relative to the center + atom + +The radial integral class +------------------------- + +In the previous subsection, we have explained how the computation of the spherical +expansion coefficients can be reduced to integrals of the form + +.. math:: + + \begin{aligned} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})g(\boldsymbol{r} - \boldsymbol{r}_{ij}). + \end{aligned} + +If the atomic density is spherically symmetric, i.e. if :math:`g(\boldsymbol{r}) = g(r)` +this integral can always be written in the following form: + +.. math:: + + \begin{aligned} \label{expansion_coeff_spherical_symmetric} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = + Y_{lm}(\hat{r}_{ij})I_{nl}(r_{ij}). + \end{aligned} + +The key point is that the dependence on the vectorial position +:math:`\boldsymbol{r}_{ij}` is split into a factor that contains information about the +orientation of this vector, namely :math:`Y_{lm}(\hat{r}_{ij})`, which is just the +spherical harmonic evaluated at :math:`\hat{r}_{ij}`, and a remaining part that captures +the dependence on the distance of atom :math:`j` from the center atom :math:`i`, namely +:math:`I_{nl}(r_{ij})`, which we shall call the radial integral. The radial integral +class computes and outputs this radial part :math:`I_{nl}(r_{ij})`. Since the angular +part is just the usual spherical harmonic, this is the part that also depends on the +choice of atomic density :math:`g(r)`, as well as the radial basis :math:`R_{nl}(r)`. In +the following, for users only interested in a specific type of density, we provide the +explicit expressions of :math:`I_{nl}(r)` for the Delta and Gaussian densities, followed +by the general expression. + +Delta Densities +~~~~~~~~~~~~~~~ + +Here, we consider the especially simple special case where the atomic density function +:math:`g(\boldsymbol{r}) = \delta(\boldsymbol{r})`. Then: + +.. math:: + + \begin{aligned} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})g(\boldsymbol{r} - \boldsymbol{r}_{ij}) \\ & = \int + \mathrm{d}^3\boldsymbol{r} R_{nl}(r)Y_{lm}(\hat{r})\delta(\boldsymbol{r} - + \boldsymbol{r}_{ij}) \\ & = R_{nl}(r) Y_{lm}(\hat{r}_{ij}) = + B_{nlm}(\boldsymbol{r}_{ij}). + \end{aligned} + +Thus, in this particularly simple case, the radial integral is simply the radial basis +function evaluated at the pair distance :math:`r_{ij}`, and we see that the integrals +have indeed the form presented above. + +Gaussian Densities +~~~~~~~~~~~~~~~~~~ + +Here, we consider another popular use case, where the atomic density function is a +Gaussian. In rascaline, we use the convention + +.. math:: + + g(r) = \frac{1}{(\pi \sigma^2)^{3/4}}e^{-\frac{r^2}{2\sigma^2}}. + +The prefactor was chosen such that the “L2-norm” of the Gaussian + +.. math:: + + \begin{aligned} + \|g\|^2 = \int \mathrm{d}^3\boldsymbol{r} |g(r)|^2 = 1, + \end{aligned} + +but does not affect the following calculations in any way. With these conventions, it +can be shown that the integral has the desired form + +.. math:: + + \begin{aligned} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})g(\boldsymbol{r} - \boldsymbol{r}_{ij}) \\ & = + Y_{lm}(\hat{r}_{ij}) \cdot I_{nl}(r_{ij}) + \end{aligned} + +with + +.. math:: + + I_{nl}(r_{ij}) = \frac{1}{(\pi \sigma^2)^{3/4}}4\pi e^{-\frac{r_{ij}^2}{2\sigma^2}} + \int_0^\infty \mathrm{d}r r^2 R_{nl}(r) e^{-\frac{r^2}{2\sigma^2}} + i_l\left(\frac{rr_{ij}}{\sigma^2}\right), + +where :math:`i_l` is a modified spherical Bessel function. The first factor, of course, +is just the normalization factor of the Gaussian density. See the next two subsections +for a derivation of this formula. + +Derivation of the General Case +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +We now derive an explicit formula for radial integral that works for any density. Let +:math:`g(r)` be a generic spherically symmetric density function. Our goal will be to +show that + +.. math:: + + \langle nlm | g;\boldsymbol{r}_{ij} \rangle = Y_{lm}(\hat{r}_{ij}) \left[2\pi + \int_0^\infty \mathrm{d}r r^2 R_{nl}(r) \int_{-1}^1 \mathrm{d}(\cos\theta) + P_l(\cos\theta) g(\sqrt{r^2+r_{ij}^2-2rr_{ij}\cos\theta}) \right] + +and thus we have the desired form :math:`\langle nlm | g;\boldsymbol{r}_{ij} \rangle = +Y_{lm}(\hat{r}_{ij}) I_{nl}(r_{ij})` with + +.. math:: + + \begin{aligned} + I_{nl}(r_{ij}) = 2\pi \int_0^\infty \mathrm{d}r r^2 R_{nl}(r) \int_{-1}^1 + \mathrm{d}u P_l(u) g(\sqrt{r^2+r_{ij}^2-2rr_{ij}u}), + \end{aligned} + +where :math:`P_l(x)` is the :math:`l`-th Legendre polynomial. + +Derivation of the explicit radial integral for Gaussian densities +----------------------------------------------------------------- + +Denoting by :math:`\theta(\boldsymbol{r},\boldsymbol{r}_{ij})` the angle between a +generic position vector :math:`\boldsymbol{r}` and the vector +:math:`\boldsymbol{r}_{ij}`, we can write + +.. math:: + + \begin{aligned} + g(\boldsymbol{r}- \boldsymbol{r}_{ij}) & = \frac{1}{(\pi + \sigma^2)^{3/4}}e^{-\frac{(\boldsymbol{r}- \boldsymbol{r}_{ij})^2}{2\sigma^2}} \\ + & = \frac{1}{(\pi + \sigma^2)^{3/4}}e^{-\frac{(r_{ij})^2}{2\sigma^2}}e^{-\frac{(\boldsymbol{r}^2- + 2\boldsymbol{r}\boldsymbol{r}_{ij})}{2\sigma^2}}, + \end{aligned} + +where the first factor no longer depends on the integration variable :math:`r`. + +Analytical Expressions for the GTO Basis +---------------------------------------- + +While the above integrals are hard to compute in general, the GTO basis is one of the +few sets of basis functions for which many of the integrals can be evaluated +analytically. This is also useful to test the correctness of more numerical +implementations. + +The primitive basis functions are defined as + +.. math:: + + \begin{aligned} + R_{nl}(r) = R_n(r) = r^n e^{-\frac{r^2}{2\sigma_n^2}} + \end{aligned} + +In this form, the basis functions are not yet orthonormal, which requires an extra +linear transformation. Since this transformation can also be applied after computing the +integrals, we simply evaluate the radial integral with respect to these primitive basis +functions. + +Real Space Integral for Gaussian Densities +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +We now evaluate + +.. math:: + + \begin{aligned} + I_{nl}(r_{ij}) & = \frac{1}{(\pi \sigma^2)^{3/4}}4\pi + e^{-\frac{r_{ij}^2}{2\sigma^2}} \int_0^\infty \mathrm{d}r r^2 R_{nl}(r) + e^{-\frac{r^2}{2\sigma^2}} i_l\left(\frac{rr_{ij}}{\sigma^2}\right) \\ & = + \frac{1}{(\pi \sigma^2)^{3/4}}4\pi e^{-\frac{r_{ij}^2}{2\sigma^2}} \int_0^\infty + \mathrm{d}r r^2 r^n e^{-\frac{r^2}{2\sigma_n^2}} e^{-\frac{r^2}{2\sigma^2}} + i_l\left(\frac{rr_{ij}}{\sigma^2}\right), + \end{aligned} + +the result of which can be conveniently expressed using :math:`a=\frac{1}{2\sigma^2}`, +:math:`b_n = \frac{1}{2\sigma_n^2}`, :math:`n_\mathrm{eff}=\frac{n+l+3}{2}` and +:math:`l_\mathrm{eff}=l+\frac{3}{2}` as + +.. math:: + + \begin{aligned} + I_{nl}(r_{ij}) = \frac{1}{(\pi \sigma^2)^{3/4}} \cdot + \pi^{\frac{3}{2}}\frac{\Gamma\left(n_\mathrm{eff}\right)}{\Gamma\left(l_\mathrm{eff}\right)}\frac{(ar_{ij})^l}{(a+b)^{n_\mathrm{eff}}}M\left(n_\mathrm{eff},l_\mathrm{eff},\frac{a^2r_{ij}^2}{a^2+b^2}\right), + \end{aligned} + +where :math:`M(a,b,z)` is the confluent hypergeometric function (hyp1f1). diff --git a/docs/src/references/api/python/utils/atomic-density.rst b/docs/src/references/api/python/utils/atomic-density.rst new file mode 100644 index 000000000..4c0fc5591 --- /dev/null +++ b/docs/src/references/api/python/utils/atomic-density.rst @@ -0,0 +1 @@ +.. automodule:: rascaline.utils.splines.atomic_density diff --git a/docs/src/references/api/python/utils/index.rst b/docs/src/references/api/python/utils/index.rst index abe7e11a0..3d88c4da9 100644 --- a/docs/src/references/api/python/utils/index.rst +++ b/docs/src/references/api/python/utils/index.rst @@ -7,5 +7,7 @@ Utility functions and classes that extend the core usage of rascaline. .. toctree:: :maxdepth: 1 + atomic-density + radial-basis power-spectrum splines diff --git a/docs/src/references/api/python/utils/radial-basis.rst b/docs/src/references/api/python/utils/radial-basis.rst new file mode 100644 index 000000000..fa8bacf62 --- /dev/null +++ b/docs/src/references/api/python/utils/radial-basis.rst @@ -0,0 +1 @@ +.. automodule:: rascaline.utils.splines.radial_basis diff --git a/docs/src/references/api/python/utils/splines.rst b/docs/src/references/api/python/utils/splines.rst index 7e93ba9db..4a3b0a912 100644 --- a/docs/src/references/api/python/utils/splines.rst +++ b/docs/src/references/api/python/utils/splines.rst @@ -1 +1 @@ -.. automodule:: rascaline.utils.splines +.. automodule:: rascaline.utils.splines.splines diff --git a/python/rascaline-torch/tests/calculator.py b/python/rascaline-torch/tests/calculator.py index 4f6d4017a..28c2c6999 100644 --- a/python/rascaline-torch/tests/calculator.py +++ b/python/rascaline-torch/tests/calculator.py @@ -62,7 +62,6 @@ def test_compute(system): assert torch.all(gradient.values[i, 2, :] == torch.tensor([0, 1])) assert len(gradient.samples) == 8 - print(gradient.samples.values) assert gradient.samples.names == ["sample", "structure", "atom"] assert tuple(gradient.samples[0]) == (0, 0, 0) assert tuple(gradient.samples[1]) == (0, 0, 1) diff --git a/python/rascaline/rascaline/utils/__init__.py b/python/rascaline/rascaline/utils/__init__.py index 88d7e9e90..b5e6221e0 100644 --- a/python/rascaline/rascaline/utils/__init__.py +++ b/python/rascaline/rascaline/utils/__init__.py @@ -3,7 +3,20 @@ import metatensor from .power_spectrum import PowerSpectrum # noqa -from .splines import RadialIntegralFromFunction, RadialIntegralSplinerBase # noqa +from .splines import ( # noqa + AtomicDensityBase, + DeltaDensity, + GaussianDensity, + GtoBasis, + LodeDensity, + LodeSpliner, + MonomialBasis, + RadialBasisBase, + RadialIntegralFromFunction, + RadialIntegralSplinerBase, + SoapSpliner, + SphericalBesselBasis, +) # path that can be used with cmake to access the rascaline library and headers diff --git a/python/rascaline/rascaline/utils/splines.py b/python/rascaline/rascaline/utils/splines.py deleted file mode 100644 index af8534b5e..000000000 --- a/python/rascaline/rascaline/utils/splines.py +++ /dev/null @@ -1,403 +0,0 @@ -""" -Splined radial integrals -======================== - -Classes for generating splines which can be used as tabulated radial integrals in the -various SOAP and LODE calculators. For an complete example of how to use these classes -see :ref:`example-splines`. - -.. autoclass:: rascaline.utils.RadialIntegralSplinerBase - :members: - :show-inheritance: - -.. autoclass:: rascaline.utils.RadialIntegralFromFunction - :members: - :show-inheritance: -""" - -from abc import ABC, abstractmethod -from typing import Callable, Dict, Optional, Union - -import numpy as np - - -class RadialIntegralSplinerBase(ABC): - """Base class for splining arbitrary radial integrals. - - If ``_radial_integral_derivative`` is not implemented in a child class it will - computed based on finite differences. - - :parameter max_angular: number of radial components - :parameter max_radial: number of angular components - :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also - the maximal value that can be interpolated. - :parameter accuracy: accuracy of the numerical integration and the splining. - Accuracy is reached when either the mean absolute error or the mean relative - error gets below the ``accuracy`` threshold. - """ - - def __init__( - self, - max_radial: int, - max_angular: int, - spline_cutoff: float, - accuracy: float, - ): - self.max_radial = max_radial - self.max_angular = max_angular - self.spline_cutoff = spline_cutoff - self.accuracy = accuracy - - @abstractmethod - def _radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: - """Method calculating the radial integral.""" - ... - - @property - def _center_contribution(self) -> Union[None, np.ndarray]: - r"""Contribution of the central atom required for LODE calculations.""" - - return None - - def _radial_integral_derivative( - self, n: int, ell: int, positions: np.ndarray - ) -> np.ndarray: - """Method calculating the derivatice of the radial integral.""" - displacement = 1e-6 - mean_abs_positions = np.abs(positions).mean() - - if mean_abs_positions <= 1.0: - raise ValueError( - "Numerically derivative of the radial integral can not be performed " - "since positions are too small. Mean of the absolute positions is " - f"{mean_abs_positions:.1e} but should be at least 1." - ) - - radial_integral_pos = self._radial_integral( - n, ell, positions + displacement / 2 - ) - radial_integral_neg = self._radial_integral( - n, ell, positions - displacement / 2 - ) - - return (radial_integral_pos - radial_integral_neg) / displacement - - def _value_evaluator_3D( - self, - positions: np.ndarray, - derivative: bool, - ): - values = np.zeros([len(positions), self.max_angular + 1, self.max_radial]) - for ell in range(self.max_angular + 1): - for n in range(self.max_radial): - if derivative: - values[:, ell, n] = self._radial_integral_derivative( - n, ell, positions - ) - else: - values[:, ell, n] = self._radial_integral(n, ell, positions) - - return values - - def compute( - self, - n_spline_points: Optional[int] = None, - ) -> Dict: - """Compute the spline for rascaline's tabulated radial integrals. - - :parameter n_spline_points: Use fixed number of spline points instead of find - the number based on the provided ``accuracy``. - :returns dict: dictionary for the input as the ``radial_basis`` parameter of a - rascaline calculator. - """ - - def value_evaluator_3D(positions): - return self._value_evaluator_3D(positions, derivative=False) - - def derivative_evaluator_3D(positions): - return self._value_evaluator_3D(positions, derivative=True) - - if n_spline_points is not None: - positions = np.linspace(0, self.spline_cutoff, n_spline_points) - values = value_evaluator_3D(positions) - derivatives = derivative_evaluator_3D(positions) - else: - dynamic_spliner = DynamicSpliner( - 0, - self.spline_cutoff, - value_evaluator_3D, - derivative_evaluator_3D, - self.accuracy, - ) - positions, values, derivatives = dynamic_spliner.spline() - - # Convert positions, values, derivatives into the appropriate json formats: - spline_points = [] - for position, value, derivative in zip(positions, values, derivatives): - spline_points.append( - { - "position": position, - "values": { - "v": 1, - "dim": value.shape, - "data": value.flatten().tolist(), - }, - "derivatives": { - "v": 1, - "dim": derivative.shape, - "data": derivative.flatten().tolist(), - }, - } - ) - - parameters = {"points": spline_points} - - center_contribution = self._center_contribution - if center_contribution is not None: - parameters["center_contribution"] = center_contribution - - return {"TabulatedRadialIntegral": parameters} - - -class DynamicSpliner: - def __init__( - self, - start: float, - stop: float, - values_fn: Callable[[np.ndarray], np.ndarray], - derivatives_fn: Callable[[np.ndarray], np.ndarray], - accuracy: float = 1e-8, - ) -> None: - """Dynamic spline generator. - - This class can be used to spline any set of functions defined within the - start-stop interval. Cubic Hermite splines - (https://en.wikipedia.org/wiki/Cubic_Hermite_spline) are used. The same spline - points will be used for all functions, and more will be added until either the - relative error or the absolute error fall below the requested accuracy on - average across all functions. The functions are specified via values_fn and - derivatives_fn. These must be able to take a numpy 1D array of positions as - their input, and they must output a numpy array where the first dimension - corresponds to the input positions, while other dimensions are arbitrary and can - correspond to any way in which the target functions can be classified. The - splines can be obtained via the spline method. - """ - - self.start = start - self.stop = stop - self.values_fn = values_fn - self.derivatives_fn = derivatives_fn - self.requested_accuracy = accuracy - - # initialize spline with 11 points - positions = np.linspace(start, stop, 11) - self.spline_positions = positions - self.spline_values = values_fn(positions) - self.spline_derivatives = derivatives_fn(positions) - - self.number_of_custom_axes = len(self.spline_values.shape) - 1 - - def spline(self): - """Calculates and outputs the splines. - - The outputs of this function are, respectively: - A numpy 1D array containing - the spline positions. These are equally spaced in the start-stop interval. - - A numpy ndarray containing the values of the splined functions at the spline - positions. The first dimension corresponds to the spline positions, while all - subsequent dimensions are consistent with the values_fn and - `get_function_derivative` provided during initialization of the class. - - A numpy ndarray containing the derivatives of the splined functions at the - spline positions, with the same structure as that of the ndarray of values. - """ - - while True: - n_intermediate_positions = len(self.spline_positions) - 1 - - if n_intermediate_positions >= 50000: - raise ValueError( - "Maximum number of spline points reached. \ - There might be a problem with the functions to be splined" - ) - - half_step = (self.spline_positions[1] - self.spline_positions[0]) / 2 - intermediate_positions = np.linspace( - self.start + half_step, self.stop - half_step, n_intermediate_positions - ) - - estimated_values = self._compute_from_spline(intermediate_positions) - new_values = self.values_fn(intermediate_positions) - - mean_absolute_error = np.mean(np.abs(estimated_values - new_values)) - with np.errstate(divide="ignore"): # Ignore divide-by-zero warnings - mean_relative_error = np.mean( - np.abs((estimated_values - new_values) / new_values) - ) - - if ( - mean_absolute_error < self.requested_accuracy - or mean_relative_error < self.requested_accuracy - ): - break - - new_derivatives = self.derivatives_fn(intermediate_positions) - - concatenated_positions = np.concatenate( - [self.spline_positions, intermediate_positions], axis=0 - ) - concatenated_values = np.concatenate( - [self.spline_values, new_values], axis=0 - ) - concatenated_derivatives = np.concatenate( - [self.spline_derivatives, new_derivatives], axis=0 - ) - - sort_indices = np.argsort(concatenated_positions, axis=0) - - self.spline_positions = concatenated_positions[sort_indices] - self.spline_values = concatenated_values[sort_indices] - self.spline_derivatives = concatenated_derivatives[sort_indices] - - return self.spline_positions, self.spline_values, self.spline_derivatives - - def _compute_from_spline(self, positions): - x = positions - delta_x = self.spline_positions[1] - self.spline_positions[0] - n = (np.floor(x / delta_x)).astype(np.int32) - - t = (x - n * delta_x) / delta_x - t_2 = t**2 - t_3 = t**3 - - h00 = 2.0 * t_3 - 3.0 * t_2 + 1.0 - h10 = t_3 - 2.0 * t_2 + t - h01 = -2.0 * t_3 + 3.0 * t_2 - h11 = t_3 - t_2 - - p_k = self.spline_values[n] - p_k_1 = self.spline_values[n + 1] - - m_k = self.spline_derivatives[n] - m_k_1 = self.spline_derivatives[n + 1] - - new_shape = (-1,) + (1,) * self.number_of_custom_axes - h00 = h00.reshape(new_shape) - h10 = h10.reshape(new_shape) - h01 = h01.reshape(new_shape) - h11 = h11.reshape(new_shape) - - interpolated_values = ( - h00 * p_k + h10 * delta_x * m_k + h01 * p_k_1 + h11 * delta_x * m_k_1 - ) - - return interpolated_values - - -class RadialIntegralFromFunction(RadialIntegralSplinerBase): - r"""Compute the radial integral spline points based on a provided function. - - :parameter radial_integral: Function to compute the radial integral. Function must - take ``n``, ``l``, and ``positions`` as inputs, where ``n`` and ``l`` are - integers and ``positions`` is a numpy 1-D array that contains the spline points - at which the radial integral will be evaluated. The function must return a numpy - 1-D array containing the values of the radial integral. - :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also - the maximal value that can be interpolated. - :parameter max_radial: number of angular componentss - :parameter max_angular: number of radial components - :parameter radial_integral_derivative: The derivative of the radial integral taking - the same paramaters as ``radial_integral``. If it is :py:obj:`None` (default), - finite differences are used to calculate the derivative of the radial integral. - It is recommended to provide this parameter if possible. Derivatives from finite - differences can cause problems when evaluating at the edges of the domain (i.e., - at ``0`` and ``spline_cutoff``) because the function might not be defined - outside of the domain. - :parameter accuracy: accuracy of the numerical integration and the splining. - Accuracy is reached when either the mean absolute error or the mean relative - error gets below the ``accuracy`` threshold. - :parameter center_contribution: Contribution of the central atom required for LODE - calculations. The ``center_contribution`` is defined as - - .. math:: - c_n = \sqrt{4π}\int_0^\infty dr r^2 R_n(r) g(r) - - where :math:`g(r)` is the radially symmetric density function, `R_n(r)` the - radial basis function and :math:`n` the current radial channel. This should be - pre-computed and provided as a separate parameter. - - Example - ------- - First define a ``radial_integral`` function - - >>> def radial_integral(n, ell, r): - ... return np.sin(r) - ... - - and provide this as input to the spline generator - - >>> spliner = RadialIntegralFromFunction( - ... radial_integral=radial_integral, - ... max_radial=12, - ... max_angular=9, - ... spline_cutoff=8.0, - ... ) - - Finally, we can use the ``spliner`` directly in the ``radial_integral`` section of a - calculator - - >>> from rascaline import SoapPowerSpectrum - >>> calculator = SoapPowerSpectrum( - ... cutoff=8.0, - ... max_radial=12, - ... max_angular=9, - ... center_atom_weight=1.0, - ... radial_basis=spliner.compute(), - ... atomic_gaussian_width=1.0, # ignored - ... cutoff_function={"Step": {}}, - ... ) - - The ``atomic_gaussian_width`` paramater is required by the calculator but will be - will be ignored during the feature computation. - - A more in depth example using a "rectangular" Laplacian eigenstate basis - is provided in the :ref:`example section`. - """ - - def __init__( - self, - radial_integral: Callable[[int, int, np.ndarray], np.ndarray], - spline_cutoff: float, - max_radial: int, - max_angular: int, - radial_integral_derivative: Optional[ - Callable[[int, int, np.ndarray], np.ndarray] - ] = None, - center_contribution: Optional[np.ndarray] = None, - accuracy: float = 1e-8, - ): - self.radial_integral_function = radial_integral - self.radial_integral_derivative_funcion = radial_integral_derivative - self.center_contribution = center_contribution - - super().__init__( - max_radial=max_radial, - max_angular=max_angular, - spline_cutoff=spline_cutoff, - accuracy=accuracy, - ) - - def _radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: - return self.radial_integral_function(n, ell, positions) - - @property - def _center_contribution(self) -> Union[None, np.ndarray]: - # Test that ``len(self.center_contribution) == max_radial`` is performed by the - # calculator. - return self.center_contribution - - def _radial_integral_derivative( - self, n: int, ell: int, positions: np.ndarray - ) -> np.ndarray: - if self.radial_integral_derivative_funcion is None: - return super()._radial_integral_derivative(n, ell, positions) - else: - return self.radial_integral_derivative_funcion(n, ell, positions) diff --git a/python/rascaline/rascaline/utils/splines/__init__.py b/python/rascaline/rascaline/utils/splines/__init__.py new file mode 100644 index 000000000..247aaa83e --- /dev/null +++ b/python/rascaline/rascaline/utils/splines/__init__.py @@ -0,0 +1,18 @@ +from .atomic_density import ( # noqa + AtomicDensityBase, + DeltaDensity, + GaussianDensity, + LodeDensity, +) +from .radial_basis import ( # noqa + GtoBasis, + MonomialBasis, + RadialBasisBase, + SphericalBesselBasis, +) +from .splines import ( # noqa + LodeSpliner, + RadialIntegralFromFunction, + RadialIntegralSplinerBase, + SoapSpliner, +) diff --git a/python/rascaline/rascaline/utils/splines/atomic_density.py b/python/rascaline/rascaline/utils/splines/atomic_density.py new file mode 100644 index 000000000..b1e074391 --- /dev/null +++ b/python/rascaline/rascaline/utils/splines/atomic_density.py @@ -0,0 +1,166 @@ +r""" +.. _python-atomic-density: + +Atomic Density +============== + +the atomic density function :math:`g(r)`, often chosen to be a Gaussian or Delta +function, that defined the type of density under consideration. For a given center atom +:math:`i` in the structure, the total density function :math:`\rho_i(\boldsymbol{r})` +around is then defined as :math:`\rho_i(\boldsymbol{r}) = \sum_{j} g(\boldsymbol{r} - +\boldsymbol{r}_{ij})`. + +Atomic densities are represented as different child class of +:py:class:`rascaline.utils.AtomicDensityBase`: :py:class:`rascaline.utils.DeltaDensity`, +:py:class:`rascaline.utils.GaussianDensity`, and :py:class:`rascaline.utils.LodeDensity` +are provided, and you can implement your own by defining a new class. + +.. autoclass:: rascaline.utils.AtomicDensityBase + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.DeltaDensity + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.GaussianDensity + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.LodeDensity + :members: + :show-inheritance: + +""" +from abc import ABC, abstractmethod +from typing import Union + +import numpy as np + + +try: + from scipy.special import gamma, gammainc + + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + + +class AtomicDensityBase(ABC): + """Base class representing atomic densities.""" + + @abstractmethod + def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]: + """Compute the atomic density arising from atoms at ``positions`` + + :param positions: positions to evaluate the atomic densities + :returns: evaluated atomic density + """ + + +class DeltaDensity(AtomicDensityBase): + r"""Delta atomic densities of the form :math:`g(r)=\delta(r)`.""" + + def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]: + raise ValueError( + "Compute function of the delta density should never called directly." + ) + + +class GaussianDensity(AtomicDensityBase): + r"""Gaussian atomic density function. + + In rascaline, we use the convention + + .. math:: + + g(r) = \frac{1}{(\pi \sigma^2)^{3/4}}e^{-\frac{r^2}{2\sigma^2}} \,. + + The prefactor was chosen such that the "L2-norm" of the Gaussian + + .. math:: + + \|g\|^2 = \int \mathrm{d}^3\boldsymbol{r} |g(r)|^2 = 1\,, + + :param atomic_gaussian_width: Width of the atom-centered gaussian used to create the + atomic density + """ + + def __init__(self, atomic_gaussian_width: float): + self.atomic_gaussian_width = atomic_gaussian_width + + def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]: + atomic_gaussian_width_sq = self.atomic_gaussian_width**2 + return np.exp(-0.5 * positions**2 / atomic_gaussian_width_sq) / ( + np.pi * atomic_gaussian_width_sq + ) ** (3 / 4) + + +class LodeDensity(AtomicDensityBase): + r"""Smeared power law density, as used in LODE. + + It is defined as + + .. math:: + + g(r) = \frac{1}{\Gamma\left(\frac{p}{2}\right)} + \frac{\gamma\left( \frac{p}{2}, \frac{r^2}{2\sigma^2} \right)} + {r^p}, + + where :math:`\Gamma(z)` is the Gamma function and :math:`\gamma(a, x)` is the + incomplete lower Gamma function. However its evaluation at :math:`r=0` is + problematic because :math:`g(r)` is of the form :math:`0/0`. For practical + implementations, it is thus more convenient to rewrite the density as + + .. math:: + + g(r) = \frac{1}{\Gamma(a)}\frac{1}{\left(2 \sigma^2\right)^a} + \begin{cases} + \frac{1}{a} - \frac{x}{a+1} + \frac{x^2}{2(a+2)} + \mathcal{O}(x^3) + & x < 10^{-5} \\ + \frac{\gamma(a,x)}{x^a} + & x \geq 10^{-5} + \end{cases} + + It is convenient to use the expression for sufficiently small :math:`x` since the + relative weight of the first neglected term is on the order of :math:`1/6x^3`. + Therefore, the threshold :math:`x = 10^{-5}` leads to relative errors on the order + of the machine epsilon. + + :param atomic_gaussian_width: Width of the atom-centered gaussian used to create the + atomic density + :param potential_exponent: Potential exponent of the decorated atom density. + Currently only implemented for potential_exponent < 10. Some exponents can be + connected to SOAP or physics-based quantities: p=0 uses Gaussian densities as in + SOAP, p=1 uses 1/r Coulomb like densities, p=6 uses 1/r^6 dispersion like + densities. + """ + + def __init__(self, atomic_gaussian_width: float, potential_exponent: int): + if not HAS_SCIPY: + raise ValueError("LodeDensity requires scipy to be installed") + + self.potential_exponent = potential_exponent + self.atomic_gaussian_width = atomic_gaussian_width + + def _short_range(self, a, x): + return 1 / a - x / (a + 1) + x**2 / (2 * (a + 2)) + + def _long_range(self, a, x): + return gamma(a) * gammainc(a, x) / x**a + + def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]: + if self.potential_exponent == 0: + return GaussianDensity.compute(self, positions=positions) + else: + atomic_gaussian_width_sq = self.atomic_gaussian_width**2 + a = self.potential_exponent / 2 + x = positions**2 / (2 * atomic_gaussian_width_sq) + + prefac = 1 / gamma(a) / (2 * atomic_gaussian_width_sq) ** a + + return prefac * np.where( + x < 1e-5, + self._short_range(a, x), + self._long_range(a, x), + ) diff --git a/python/rascaline/rascaline/utils/splines/radial_basis.py b/python/rascaline/rascaline/utils/splines/radial_basis.py new file mode 100644 index 000000000..1b68396b3 --- /dev/null +++ b/python/rascaline/rascaline/utils/splines/radial_basis.py @@ -0,0 +1,331 @@ +r""" +.. _python-radial-basis: + +Radial Basis +============ + +Radial basis functions :math:`R_{nl}(\boldsymbol{r})` are besides :ref:`atomic densities +` :math:`\rho_i` the central ingredients to compute spherical +expansion coefficients :math:`\langle anlm\vert\rho_i\rangle`. Radial basis functions, +define how which the atomic density is projected. To be more precise, the actual basis +functions are of + +.. math:: + + B_{nlm}(\boldsymbol{r}) = R_{nl}(r)Y_{lm}(\hat{r}) \,, + +where :math:`Y_{lm}(\hat{r})` are the real spherical harmonics evaluated at the point +:math:`\hat{r}`, i.e. at the spherical angles :math:`(\theta, \phi)` that determine the +orientation of the unit vector :math:`\hat{r} = \boldsymbol{r}/r`. + +Radial basis are represented as different child class of +:py:class:`rascaline.utils.RadialBasisBase`: :py:class:`rascaline.utils.GtoBasis`, +:py:class:`rascaline.utils.MonomialBasis`, and +:py:class:`rascaline.utils.SphericalBesselBasis` are provided, and you can implement +your own by defining a new class. + +.. autoclass:: rascaline.utils.RadialBasisBase + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.GtoBasis + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.MonomialBasis + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.SphericalBesselBasis + :members: + :show-inheritance: +""" + +from abc import ABC, abstractmethod +from typing import Union + +import numpy as np + + +try: + import scipy.integrate + import scipy.optimize + import scipy.special + + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + + +class RadialBasisBase(ABC): + r""" + Base class to define radial basis and their evaluation. + + The class provides methods to evaluate the radial basis :math:`R_{nl}(r)` as well as + its (numerical) derivative with respect to positions :math:`r`. + + :parameter integration_radius: Value up to which the radial integral should be + performed. The usual value is :math:`\infty`. + """ + + def __init__(self, integration_radius: float): + self.integration_radius = integration_radius + + @abstractmethod + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + """Compute the ``n``/``l`` radial basis at all given ``integrand_positions`` + + :param n: radial channel + :param ell: angular channel + :param integrand_positions: positions to evaluate the radial basis + :returns: evaluated radial basis + """ + + def compute_derivative( + self, n: float, ell: float, integrand_positions: np.ndarray + ) -> np.ndarray: + """Compute the derivative of the ``n``/``l`` radial basis at all given + ``integrand_positions`` + + This is used for radial integrals with delta-like atomic densities. If not + defined in a child class, a numerical derivative based on finite differences of + ``integrand_positions`` will be used instead. + + :param n: radial channel + :param ell: angular channel + :param integrand_positions: positions to evaluate the radial basis + :returns: evaluated derivative of the radial basis + """ + displacement = 1e-6 + mean_abs_positions = np.abs(integrand_positions).mean() + + if mean_abs_positions < 1.0: + raise ValueError( + "Numerically derivative of the radial integral can not be performed " + "since positions are too small. Mean of the absolute positions is " + f"{mean_abs_positions:.1e} but should be at least 1." + ) + + radial_basis_pos = self.compute(n, ell, integrand_positions + displacement / 2) + radial_basis_neg = self.compute(n, ell, integrand_positions - displacement / 2) + + return (radial_basis_pos - radial_basis_neg) / displacement + + def compute_gram_matrix( + self, + max_radial: int, + max_angular: int, + ) -> np.ndarray: + """Gram matrix of the current basis. + + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :returns: orthonormalization matrix of shape + ``(max_angular + 1, max_radial, max_radial)`` + """ + + if not HAS_SCIPY: + raise ValueError("Orthonormalization requires scipy!") + + # Gram matrix (also called overlap matrix or inner product matrix) + gram_matrix = np.zeros((max_angular + 1, max_radial, max_radial)) + + def integrand( + integrand_positions: np.ndarray, + n1: int, + n2: int, + ell: int, + ) -> np.ndarray: + return ( + integrand_positions**2 + * self.compute(n1, ell, integrand_positions) + * self.compute(n2, ell, integrand_positions) + ) + + for ell in range(max_angular + 1): + for n1 in range(max_radial): + for n2 in range(max_radial): + gram_matrix[ell, n1, n2] = scipy.integrate.quad( + func=integrand, + a=0, + b=self.integration_radius, + args=(n1, n2, ell), + )[0] + + return gram_matrix + + def compute_orthonormalization_matrix( + self, + max_radial: int, + max_angular: int, + ) -> np.ndarray: + """Compute orthonormalization matrix + + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :returns: orthonormalization matrix of shape (max_angular + 1, max_radial, + max_radial) + """ + + gram_matrix = self.compute_gram_matrix(max_radial, max_angular) + + # Get the normalization constants from the diagonal entries + normalizations = np.zeros((max_angular + 1, max_radial)) + + for ell in range(max_angular + 1): + for n in range(max_radial): + normalizations[ell, n] = 1 / np.sqrt(gram_matrix[ell, n, n]) + + # Rescale orthonormalization matrix to be defined + # in terms of the normalized (but not yet orthonormalized) + # basis functions + gram_matrix[ell, n, :] *= normalizations[ell, n] + gram_matrix[ell, :, n] *= normalizations[ell, n] + + orthonormalization_matrix = np.zeros_like(gram_matrix) + for ell in range(max_angular + 1): + eigvals, eigvecs = np.linalg.eigh(gram_matrix[ell]) + orthonormalization_matrix[ell] = ( + eigvecs @ np.diag(np.sqrt(1.0 / eigvals)) @ eigvecs.T + ) + + # Rescale the orthonormalization matrix so that it + # works with respect to the primitive (not yet normalized) + # radial basis functions + for ell in range(max_angular + 1): + for n in range(max_radial): + orthonormalization_matrix[ell, :, n] *= normalizations[ell, n] + + return orthonormalization_matrix + + +class GtoBasis(RadialBasisBase): + r"""Primitive (not normalized nor orthonormalized) GTO radial basis. + + It is defined as + + .. math:: + + R_{nl}(r) = R_n(r) = r^n e^{-\frac{r^2}{2\sigma_n^2}}, + + where :math:`\sigma_n = \sqrt{n} r_\mathrm{cut}/n_\mathrm{max}` with + :math:`r_\mathrm{cut}` being the ``cutoff`` and :math:`n_\mathrm{max}` the maximal + number of radial components. + + :parameter cutoff: spherical cutoff for the radial basis + :parameter max_radial: number of radial components + """ + + def __init__(self, cutoff, max_radial): + # choosing infinity leads to problems when calculating the radial integral with + # `quad`! + super().__init__(integration_radius=5 * cutoff) + self.max_radial = max_radial + self.cutoff = cutoff + self.sigmas = np.ones(self.max_radial, dtype=float) + + for n in range(1, self.max_radial): + self.sigmas[n] = np.sqrt(n) + self.sigmas *= self.cutoff / self.max_radial + + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return integrand_positions**n * np.exp( + -0.5 * (integrand_positions / self.sigmas[n]) ** 2 + ) + + def compute_derivative( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return n / integrand_positions * self.compute( + n, ell, integrand_positions + ) - integrand_positions / self.sigmas[n] ** 2 * self.compute( + n, ell, integrand_positions + ) + + +class MonomialBasis(RadialBasisBase): + r"""Monomial basis. + + Basis is consisting of functions + + .. math:: + R_{nl}(r) = r^{l+2n}, + + where :math:`n` runs from :math:`0,1,...,n_\mathrm{max}-1`. These capture precisely + the radial dependence if we compute the Taylor expansion of a generic function + defined in 3D space. + + :parameter cutoff: spherical cutoff for the radial basis + """ + + def __init__(self, cutoff): + super().__init__(integration_radius=cutoff) + + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return integrand_positions ** (ell + 2 * n) + + def compute_derivative( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return (ell + 2 * n) * integrand_positions ** (ell + 2 * n - 1) + + +def _spherical_jn_swapped(z, n): + """spherical_jn with swapped arguments for usage in `fsolve`.""" + return scipy.special.spherical_jn(n=n, z=z) + + +class SphericalBesselBasis(RadialBasisBase): + """Spherical Bessel functions used in the Laplacian eigenstate (LE) basis. + + :parameter cutoff: spherical cutoff for the radial basis + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + """ + + def __init__(self, cutoff, max_radial, max_angular): + if not HAS_SCIPY: + raise ValueError("SphericalBesselBasis requires scipy!") + + super().__init__(integration_radius=cutoff) + + self.max_radial = max_radial + self.max_angular = max_angular + self.roots = np.zeros([max_angular + 1, self.max_radial]) + + # Define target function and the estimated location of roots obtained from the + # asymptotic expansion of the spherical Bessel functions for large arguments x + for ell in range(self.max_angular + 1): + roots_guesses = np.pi * (np.arange(1, self.max_radial + 1) + ell / 2) + # Compute roots from initial guess using Newton method + for n, root_guess in enumerate(roots_guesses): + self.roots[ell, n] = scipy.optimize.fsolve( + func=_spherical_jn_swapped, x0=root_guess, args=(ell,) + )[0] + + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return scipy.special.spherical_jn( + ell, + integrand_positions * self.roots[ell, n] / self.integration_radius, + ) + + def compute_derivative( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return ( + self.roots[ell, n] + / self.integration_radius + * scipy.special.spherical_jn( + ell, + integrand_positions * self.roots[ell, n] / self.integration_radius, + derivative=True, + ) + ) diff --git a/python/rascaline/rascaline/utils/splines/splines.py b/python/rascaline/rascaline/utils/splines/splines.py new file mode 100644 index 000000000..42ad3c968 --- /dev/null +++ b/python/rascaline/rascaline/utils/splines/splines.py @@ -0,0 +1,849 @@ +""" +.. _python-splined-radial-integral: + +Splined radial integrals +======================== + +Classes for generating splines which can be used as tabulated radial integrals in the +various SOAP and LODE calculators. + +All classes are based on :py:class:`rascaline.utils.RadialIntegralSplinerBase`. We +provides several ways to compute a radial integral: you may chose and initialize a pre +defined atomic density and radial basis and provide them to +:py:class:`rascaline.utils.SoapSpliner` or :py:class:`rascaline.utils.LodeSpliner`. Both +classes require `scipy`_ to be installed in order to perform the numerical integrals. + +Alternatively, you can also explicitly provide functions for the radial integral and its +derivative and passing them to :py:class:`rascaline.utils.RadialIntegralFromFunction`. + +.. autoclass:: rascaline.utils.RadialIntegralSplinerBase + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.SoapSpliner + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.LodeSpliner + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.RadialIntegralFromFunction + :members: + :show-inheritance: + + +.. _`scipy`: https://scipy.org +""" + +from abc import ABC, abstractmethod +from typing import Callable, Dict, Optional, Union + +import numpy as np + + +try: + from scipy.integrate import quad, quad_vec + from scipy.special import spherical_in, spherical_jn + + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + +from .atomic_density import AtomicDensityBase, DeltaDensity, GaussianDensity +from .radial_basis import RadialBasisBase + + +class RadialIntegralSplinerBase(ABC): + """Base class for splining arbitrary radial integrals. + + If :py:meth:`RadialIntegralSplinerBase.radial_integral_derivative` is not + implemented in a child class it will computed based on finite differences. + + :parameter max_angular: number of radial components + :parameter max_radial: number of angular components + :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also + the maximal value that can be interpolated. + :parameter basis: Provide a :class:`RadialBasisBase` instance to orthonormalize the + radial integral. + :parameter accuracy: accuracy of the numerical integration and the splining. + Accuracy is reached when either the mean absolute error or the mean relative + error gets below the ``accuracy`` threshold. + """ + + def __init__( + self, + max_radial: int, + max_angular: int, + spline_cutoff: float, + basis: Optional[RadialBasisBase], + accuracy: float, + ): + self.max_radial = max_radial + self.max_angular = max_angular + self.spline_cutoff = spline_cutoff + self.basis = basis + self.accuracy = accuracy + + def compute( + self, + n_spline_points: Optional[int] = None, + ) -> Dict: + """Compute the spline for rascaline's tabulated radial integrals. + + :parameter n_spline_points: Use fixed number of spline points instead of find + the number based on the provided ``accuracy``. + :returns dict: dictionary for the input as the ``radial_basis`` parameter of a + rascaline calculator. + """ + + if self.basis is not None: + orthonormalization_matrix = self.basis.compute_orthonormalization_matrix( + self.max_radial, self.max_angular + ) + else: + orthonormalization_matrix = None + + def value_evaluator_3D(positions): + return self._value_evaluator_3D( + positions, orthonormalization_matrix, derivative=False + ) + + def derivative_evaluator_3D(positions): + return self._value_evaluator_3D( + positions, orthonormalization_matrix, derivative=True + ) + + if n_spline_points is not None: + positions = np.linspace(0, self.spline_cutoff, n_spline_points) + values = value_evaluator_3D(positions) + derivatives = derivative_evaluator_3D(positions) + else: + dynamic_spliner = DynamicSpliner( + 0, + self.spline_cutoff, + value_evaluator_3D, + derivative_evaluator_3D, + self.accuracy, + ) + positions, values, derivatives = dynamic_spliner.spline() + + # Convert positions, values, derivatives into the appropriate json formats: + spline_points = [] + for position, value, derivative in zip(positions, values, derivatives): + spline_points.append( + { + "position": position, + "values": { + "v": 1, + "dim": value.shape, + "data": value.flatten().tolist(), + }, + "derivatives": { + "v": 1, + "dim": derivative.shape, + "data": derivative.flatten().tolist(), + }, + } + ) + + parameters = {"points": spline_points} + + center_contribution = self.center_contribution + if center_contribution is not None: + if self.basis is not None: + # consider only `l=0` component of the `orthonormalization_matrix` + parameters["center_contribution"] = list( + orthonormalization_matrix[0] @ center_contribution + ) + else: + parameters["center_contribution"] = center_contribution + + return {"TabulatedRadialIntegral": parameters} + + @abstractmethod + def radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + """evaluate the radial integral""" + ... + + @property + def center_contribution(self) -> Union[None, np.ndarray]: + r"""Pre-computed value for the contribution of the central atom. + + Required for LODE calculations. The central atom contribution will be + orthonormalized in the same way as the radial integral. + """ + + return None + + def radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + """evaluate the derivative of the radial integral""" + displacement = 1e-6 + mean_abs_positions = np.mean(np.abs(positions)) + + if mean_abs_positions < 1.0: + raise ValueError( + "Numerically derivative of the radial integral can not be performed " + "since positions are too small. Mean of the absolute positions is " + f"{mean_abs_positions:.1e} but should be at least 1." + ) + + radial_integral_pos = self.radial_integral(n, ell, positions + displacement / 2) + radial_integral_neg = self.radial_integral(n, ell, positions - displacement / 2) + + return (radial_integral_pos - radial_integral_neg) / displacement + + def _value_evaluator_3D( + self, + positions: np.ndarray, + orthonormalization_matrix: Optional[np.ndarray], + derivative: bool, + ): + values = np.zeros([len(positions), self.max_angular + 1, self.max_radial]) + for ell in range(self.max_angular + 1): + for n in range(self.max_radial): + if derivative: + values[:, ell, n] = self.radial_integral_derivative( + n, ell, positions + ) + else: + values[:, ell, n] = self.radial_integral(n, ell, positions) + + if orthonormalization_matrix is not None: + # For each l channel we do a dot product of the orthonormalization_matrix of + # shape (n, n) with the values which should have the shape (n, n_positions). + # To achieve the correct broadcasting we have to transpose twice. + for ell in range(self.max_angular + 1): + values[:, ell, :] = ( + orthonormalization_matrix[ell] @ values[:, ell, :].T + ).T + + return values + + +class DynamicSpliner: + def __init__( + self, + start: float, + stop: float, + values_fn: Callable[[np.ndarray], np.ndarray], + derivatives_fn: Callable[[np.ndarray], np.ndarray], + accuracy: float = 1e-8, + ) -> None: + """Dynamic spline generator. + + This class can be used to spline any set of functions defined within the + start-stop interval. Cubic Hermite splines + (https://en.wikipedia.org/wiki/Cubic_Hermite_spline) are used. The same spline + points will be used for all functions, and more will be added until either the + relative error or the absolute error fall below the requested accuracy on + average across all functions. The functions are specified via values_fn and + derivatives_fn. These must be able to take a numpy 1D array of positions as + their input, and they must output a numpy array where the first dimension + corresponds to the input positions, while other dimensions are arbitrary and can + correspond to any way in which the target functions can be classified. The + splines can be obtained via the spline method. + """ + + self.start = start + self.stop = stop + self.values_fn = values_fn + self.derivatives_fn = derivatives_fn + self.requested_accuracy = accuracy + + # initialize spline with 11 points + positions = np.linspace(start, stop, 11) + self.spline_positions = positions + self.spline_values = values_fn(positions) + self.spline_derivatives = derivatives_fn(positions) + + self.number_of_custom_axes = len(self.spline_values.shape) - 1 + + def spline(self): + """Calculates and outputs the splines. + + The outputs of this function are, respectively: + + - A numpy 1D array containing the spline positions. These are equally spaced in + the start-stop interval. + - A numpy ndarray containing the values of the splined functions at the spline + positions. The first dimension corresponds to the spline positions, while all + subsequent dimensions are consistent with the values_fn and + `get_function_derivative` provided during initialization of the class. + - A numpy ndarray containing the derivatives of the splined functions at the + spline positions, with the same structure as that of the ndarray of values. + """ + + while True: + n_intermediate_positions = len(self.spline_positions) - 1 + + if n_intermediate_positions >= 50000: + raise ValueError( + "Maximum number of spline points reached. \ + There might be a problem with the functions to be splined" + ) + + half_step = (self.spline_positions[1] - self.spline_positions[0]) / 2 + intermediate_positions = np.linspace( + self.start + half_step, self.stop - half_step, n_intermediate_positions + ) + + estimated_values = self._compute_from_spline(intermediate_positions) + new_values = self.values_fn(intermediate_positions) + + mean_absolute_error = np.mean(np.abs(estimated_values - new_values)) + with np.errstate(divide="ignore"): # Ignore divide-by-zero warnings + mean_relative_error = np.mean( + np.abs((estimated_values - new_values) / new_values) + ) + + if ( + mean_absolute_error < self.requested_accuracy + or mean_relative_error < self.requested_accuracy + ): + break + + new_derivatives = self.derivatives_fn(intermediate_positions) + + concatenated_positions = np.concatenate( + [self.spline_positions, intermediate_positions], axis=0 + ) + concatenated_values = np.concatenate( + [self.spline_values, new_values], axis=0 + ) + concatenated_derivatives = np.concatenate( + [self.spline_derivatives, new_derivatives], axis=0 + ) + + sort_indices = np.argsort(concatenated_positions, axis=0) + + self.spline_positions = concatenated_positions[sort_indices] + self.spline_values = concatenated_values[sort_indices] + self.spline_derivatives = concatenated_derivatives[sort_indices] + + return self.spline_positions, self.spline_values, self.spline_derivatives + + def _compute_from_spline(self, positions): + x = positions + delta_x = self.spline_positions[1] - self.spline_positions[0] + n = (np.floor(x / delta_x)).astype(np.int32) + + t = (x - n * delta_x) / delta_x + t_2 = t**2 + t_3 = t**3 + + h00 = 2.0 * t_3 - 3.0 * t_2 + 1.0 + h10 = t_3 - 2.0 * t_2 + t + h01 = -2.0 * t_3 + 3.0 * t_2 + h11 = t_3 - t_2 + + p_k = self.spline_values[n] + p_k_1 = self.spline_values[n + 1] + + m_k = self.spline_derivatives[n] + m_k_1 = self.spline_derivatives[n + 1] + + new_shape = (-1,) + (1,) * self.number_of_custom_axes + h00 = h00.reshape(new_shape) + h10 = h10.reshape(new_shape) + h01 = h01.reshape(new_shape) + h11 = h11.reshape(new_shape) + + interpolated_values = ( + h00 * p_k + h10 * delta_x * m_k + h01 * p_k_1 + h11 * delta_x * m_k_1 + ) + + return interpolated_values + + +class RadialIntegralFromFunction(RadialIntegralSplinerBase): + r"""Compute radial integral spline points based on a provided function. + + :parameter radial_integral: Function to compute the radial integral. Function must + take ``n``, ``l``, and ``positions`` as inputs, where ``n`` and ``l`` are + integers and ``positions`` is a numpy 1-D array that contains the spline points + at which the radial integral will be evaluated. The function must return a numpy + 1-D array containing the values of the radial integral. + :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also + the maximal value that can be interpolated. + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :parameter radial_integral_derivative: The derivative of the radial integral taking + the same parameters as ``radial_integral``. If it is :py:obj:`None` (default), + finite differences are used to calculate the derivative of the radial integral. + It is recommended to provide this parameter if possible. Derivatives from finite + differences can cause problems when evaluating at the edges of the domain (i.e., + at ``0`` and ``spline_cutoff``) because the function might not be defined + outside of the domain. + :parameter accuracy: accuracy of the numerical integration and the splining. + Accuracy is reached when either the mean absolute error or the mean relative + error gets below the ``accuracy`` threshold. + :parameter center_contribution: Contribution of the central atom required for LODE + calculations. The ``center_contribution`` is defined as + + .. math:: + c_n = \sqrt{4π}\int_0^\infty dr r^2 R_n(r) g(r) + + where :math:`g(r)` is the radially symmetric density function, `R_n(r)` the + radial basis function and :math:`n` the current radial channel. This should be + pre-computed and provided as a separate parameter. + + Example + ------- + First define a ``radial_integral`` function + + >>> def radial_integral(n, ell, r): + ... return np.sin(r) + ... + + and provide this as input to the spline generator + + >>> spliner = RadialIntegralFromFunction( + ... radial_integral=radial_integral, + ... max_radial=12, + ... max_angular=9, + ... spline_cutoff=8.0, + ... ) + + Finally, we can use the ``spliner`` directly in the ``radial_integral`` section of a + calculator + + >>> from rascaline import SoapPowerSpectrum + >>> calculator = SoapPowerSpectrum( + ... cutoff=8.0, + ... max_radial=12, + ... max_angular=9, + ... center_atom_weight=1.0, + ... radial_basis=spliner.compute(), + ... atomic_gaussian_width=1.0, # ignored + ... cutoff_function={"Step": {}}, + ... ) + + The ``atomic_gaussian_width`` parameter is required by the calculator but will be + will be ignored during the feature computation. + + A more in depth example using a "rectangular" Laplacian eigenstate basis is provided + in the :ref:`example section`. + """ + + def __init__( + self, + radial_integral: Callable[[int, int, np.ndarray], np.ndarray], + spline_cutoff: float, + max_radial: int, + max_angular: int, + radial_integral_derivative: Optional[ + Callable[[int, int, np.ndarray], np.ndarray] + ] = None, + center_contribution: Optional[np.ndarray] = None, + accuracy: float = 1e-8, + ): + self.radial_integral_function = radial_integral + self.radial_integral_derivative_function = radial_integral_derivative + + if center_contribution is not None and len(center_contribution) != max_radial: + raise ValueError( + f"center contribution has {len(center_contribution)} entries but " + f"should be the same as max_radial ({max_radial})" + ) + + self._center_contribution = center_contribution + + super().__init__( + max_radial=max_radial, + max_angular=max_angular, + spline_cutoff=spline_cutoff, + basis=None, # do no orthonormalize the radial integral + accuracy=accuracy, + ) + + def radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + return self.radial_integral_function(n, ell, positions) + + @property + def center_contribution(self) -> Union[None, np.ndarray]: + return self._center_contribution + + def radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + if self.radial_integral_derivative_function is None: + return super().radial_integral_derivative(n, ell, positions) + else: + return self.radial_integral_derivative_function(n, ell, positions) + + +class SoapSpliner(RadialIntegralSplinerBase): + """Compute radial integral spline points for real space calculators. + + Use only in combination with a real space calculators like + :class:`rascaline.SphericalExpansion` or :class:`rascaline.SoapPowerSpectrum`. For + k-space spherical expansions use :class:`LodeSpliner`. + + If ``density`` is either :class:`rascaline.utils.DeltaDensity` or + :class:`rascaline.utils.GaussianDensity` the radial integral will be partly solved + analytical. These simpler expressions result in a faster and more stable evaluation. + + :parameter cutoff: spherical cutoff for the radial basis + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :parameter basis: definition of the radial basis + :parameter density: definition of the atomic density + :parameter accuracy: accuracy of the numerical integration and the splining. + Accuracy is reached when either the mean absolute error or the mean relative + error gets below the ``accuracy`` threshold. + :raise ValueError: if `scipy`_ is not available + + Example + ------- + + First import the necessary classed and define hyper parameters for the spherical + expansions. + + >>> from rascaline import SphericalExpansion + >>> from rascaline.utils import GaussianDensity, GtoBasis + + >>> cutoff = 2 + >>> max_radial = 6 + >>> max_angular = 4 + >>> atomic_gaussian_width = 1.0 + + Next we initialize our radial basis and the density + + >>> basis = GtoBasis(cutoff=cutoff, max_radial=max_radial) + >>> density = GaussianDensity(atomic_gaussian_width=atomic_gaussian_width) + + And finally the actual spliner instance + + >>> spliner = SoapSpliner( + ... cutoff=cutoff, + ... max_radial=max_radial, + ... max_angular=max_angular, + ... basis=basis, + ... density=density, + ... accuracy=1e-3, + ... ) + + Above we reduced ``accuracy`` from the default value of ``1e-8`` to ``1e-3`` to + speed up calculations. + + As for all spliner classes you can use the output + :meth:`RadialIntegralSplinerBase.compute` method directly as the ``radial_basis`` + parameter. + + >>> calculator = SphericalExpansion( + ... cutoff=cutoff, + ... max_radial=max_radial, + ... max_angular=max_angular, + ... center_atom_weight=1.0, + ... atomic_gaussian_width=atomic_gaussian_width, + ... radial_basis=spliner.compute(), + ... cutoff_function={"Step": {}}, + ... ) + + You can now use ``calculator`` to obtain the spherical expansion coefficients of + your systems. Note that the the spliner based used here will produce the same + coefficients as if ``radial_basis={"Gto": {}}`` would be used. + + .. seealso:: + :class:`LodeSpliner` for a spliner class that works with + :class:`rascaline.LodeSphericalExpansion` + """ + + def __init__( + self, + cutoff: float, + max_radial: int, + max_angular: int, + basis: RadialBasisBase, + density: AtomicDensityBase, + accuracy: float = 1e-8, + ): + if not HAS_SCIPY: + raise ValueError("Spliner class requires scipy!") + + self.density = density + + super().__init__( + max_radial=max_radial, + max_angular=max_angular, + spline_cutoff=cutoff, + basis=basis, + accuracy=accuracy, + ) + + def radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + if type(self.density) is DeltaDensity: + return self._radial_integral_delta(n, ell, positions) + elif type(self.density) is GaussianDensity: + return self._radial_integral_gaussian(n, ell, positions) + else: + return self._radial_integral_custom(n, ell, positions) + + def radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + if type(self.density) is DeltaDensity: + return self._radial_integral_delta_derivative(n, ell, positions) + elif type(self.density) is GaussianDensity: + return self._radial_integral_gaussian_derivative(n, ell, positions) + else: + return self._radial_integral_custom_derivative(n, ell, positions) + + def _radial_integral_delta( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return self.basis.compute(n, ell, positions) + + def _radial_integral_delta_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return self.basis.compute_derivative(n, ell, positions) + + def _radial_integral_gaussian( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + atomic_gaussian_width_sq = self.density.atomic_gaussian_width**2 + + prefactor = ( + (4 * np.pi) + / (np.pi * atomic_gaussian_width_sq) ** (3 / 4) + * np.exp(-0.5 * positions**2 / atomic_gaussian_width_sq) + ) + + def integrand( + integrand_position: float, n: int, ell: int, positions: np.array + ) -> np.ndarray: + return ( + integrand_position**2 + * self.basis.compute(n, ell, integrand_position) + * np.exp(-0.5 * integrand_position**2 / atomic_gaussian_width_sq) + * spherical_in( + ell, + integrand_position * positions / atomic_gaussian_width_sq, + ) + ) + + return ( + prefactor + * quad_vec( + f=lambda x: integrand(x, n, ell, positions), + a=0, + b=self.basis.integration_radius, + )[0] + ) + + def _radial_integral_gaussian_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + atomic_gaussian_width_sq = self.density.atomic_gaussian_width**2 + + prefactor = ( + (4 * np.pi) + / (np.pi * atomic_gaussian_width_sq) ** (3 / 4) + * np.exp(-0.5 * positions**2 / atomic_gaussian_width_sq) + ) + + def integrand( + integrand_position: float, n: int, ell: int, positions: np.array + ) -> np.ndarray: + return ( + integrand_position**3 + * self.basis.compute(n, ell, integrand_position) + * np.exp(-(integrand_position**2) / (2 * atomic_gaussian_width_sq)) + * spherical_in( + ell, + integrand_position * positions / atomic_gaussian_width_sq, + derivative=True, + ) + ) + + return atomic_gaussian_width_sq**-1 * ( + prefactor + * quad_vec( + f=lambda x: integrand(x, n, ell, positions), + a=0, + b=self.basis.integration_radius, + )[0] + - positions * self._radial_integral_gaussian(n, ell, positions) + ) + + def _radial_integral_custom( + self, n: int, ell: int, positions: np.ndarray, derivative: bool + ) -> np.ndarray: + raise NotImplementedError( + "Radial integral with custom atomic densities is not implemented yet!" + ) + + def _radial_integral_custom_derivative( + self, n: int, ell: int, positions: np.ndarray, derivative: bool + ) -> np.ndarray: + raise NotImplementedError( + "Radial integral with custom atomic densities is not implemented yet!" + ) + + +class LodeSpliner(RadialIntegralSplinerBase): + r"""Compute radial integral spline points for k-space calculators. + + Use only in combination with a k-space/Fourier-space calculators like + :class:`rascaline.LodeSphericalExpansion`. For real space spherical expansions use + :class:`SoapSpliner`. + + :parameter k_cutoff: spherical reciprocal cutoff + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :parameter basis: definition of the radial basis + :parameter density: definition of the atomic density + :parameter accuracy: accuracy of the numerical integration and the splining. + Accuracy is reached when either the mean absolute error or the mean relative + error gets below the ``accuracy`` threshold. + :raise ValueError: if `scipy`_ is not available + + Example + ------- + + First import the necessary classed and define hyper parameters for the spherical + expansions. + + >>> from rascaline import LodeSphericalExpansion + >>> from rascaline.utils import GaussianDensity, GtoBasis + + Note that ``cutoff`` defined below denotes the maximal distance for the projection + of the density. In contrast to SOAP, LODE also takes atoms outside of this + ``cutoff`` into account for the density. + + >>> cutoff = 2 + >>> max_radial = 6 + >>> max_angular = 4 + >>> atomic_gaussian_width = 1.0 + + :math:`1.2 \, \pi \, \sigma` where :math:`\sigma` is the ``atomic_gaussian_width`` + which is a reasonable value for most systems. + + >>> k_cutoff = 1.2 * np.pi / atomic_gaussian_width + + Next we initialize our radial basis and the density + + >>> basis = GtoBasis(cutoff=cutoff, max_radial=max_radial) + >>> density = GaussianDensity(atomic_gaussian_width=atomic_gaussian_width) + + And finally the actual spliner instance + + >>> spliner = LodeSpliner( + ... k_cutoff=k_cutoff, + ... max_radial=max_radial, + ... max_angular=max_angular, + ... basis=basis, + ... density=density, + ... ) + + As for all spliner classes you can use the output + :meth:`RadialIntegralSplinerBase.compute` method directly as the ``radial_basis`` + parameter. + + >>> calculator = LodeSphericalExpansion( + ... cutoff=cutoff, + ... max_radial=max_radial, + ... max_angular=max_angular, + ... center_atom_weight=1.0, + ... atomic_gaussian_width=atomic_gaussian_width, + ... potential_exponent=1, + ... radial_basis=spliner.compute(), + ... ) + + You can now use ``calculator`` to obtain the spherical expansion coefficients of + your systems. Note that the the spliner based used here will produce the same + coefficients as if ``radial_basis={"Gto": {}}`` would be used. + + .. seealso:: + :class:`SoapSpliner` for a spliner class that works with + :class:`rascaline.SphericalExpansion` + """ + + def __init__( + self, + k_cutoff: float, + max_radial: int, + max_angular: int, + basis: RadialBasisBase, + density: AtomicDensityBase, + accuracy: float = 1e-8, + ): + if not HAS_SCIPY: + raise ValueError("Spliner class requires scipy!") + + self.density = density + + super().__init__( + max_radial=max_radial, + max_angular=max_angular, + basis=basis, + spline_cutoff=k_cutoff, # use k_cutoff here because we spline in k-space + accuracy=accuracy, + ) + + def radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + def integrand( + integrand_position: float, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return ( + integrand_position**2 + * self.basis.compute(n, ell, integrand_position) + * spherical_jn(ell, integrand_position * positions) + ) + + return quad_vec( + f=lambda x: integrand(x, n, ell, positions), + a=0, + b=self.basis.integration_radius, + )[0] + + def radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + def integrand( + integrand_position: float, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return ( + integrand_position**3 + * self.basis.compute(n, ell, integrand_position) + * spherical_jn(ell, integrand_position * positions, derivative=True) + ) + + return quad_vec( + f=lambda x: integrand(x, n, ell, positions), + a=0, + b=self.basis.integration_radius, + )[0] + + @property + def center_contribution(self) -> np.ndarray: + if type(self.density) is DeltaDensity: + center_contrib = self._center_contribution_delta + else: + center_contrib = self._center_contribution_custom + + return [np.sqrt(4 * np.pi) * center_contrib(n) for n in range(self.max_radial)] + + def _center_contribution_delta(self, n: int): + raise NotImplementedError( + "center contribution for delta distributions is not implemented yet." + ) + + def _center_contribution_custom(self, n: int): + def integrand(integrand_position: float, n: int) -> np.ndarray: + return ( + integrand_position**2 + * self.basis.compute(n, 0, integrand_position) + * self.density.compute(integrand_position) + ) + + return quad( + func=integrand, + a=0, + b=self.basis.integration_radius, + args=(n), + )[0] diff --git a/python/rascaline/tests/utils/radial_basis.py b/python/rascaline/tests/utils/radial_basis.py new file mode 100644 index 000000000..986717fd0 --- /dev/null +++ b/python/rascaline/tests/utils/radial_basis.py @@ -0,0 +1,91 @@ +from typing import Union + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +from rascaline.utils import ( + GtoBasis, + MonomialBasis, + RadialBasisBase, + SphericalBesselBasis, +) + + +pytest.importorskip("scipy") + + +class RtoNRadialBasis(RadialBasisBase): + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return integrand_positions**n + + +def test_radial_basis_gram(): + """Test that quad integration of the gram matrix is the same as an analytical.""" + + integration_radius = 1 + max_radial = 4 + max_angular = 2 + + test_basis = RtoNRadialBasis(integration_radius=integration_radius) + + numerical_gram = test_basis.compute_gram_matrix(max_radial, max_angular) + analytical_gram = np.zeros_like(numerical_gram) + + for ell in range(max_angular + 1): + for n1 in range(max_radial): + for n2 in range(max_radial): + exp = 3 + n1 + n2 + analytical_gram[ell, n1, n2] = integration_radius**exp / exp + + assert_allclose(numerical_gram, analytical_gram) + + +def test_radial_basis_orthornormalization(): + integration_radius = 1 + max_radial = 4 + max_angular = 2 + + test_basis = RtoNRadialBasis(integration_radius=integration_radius) + + gram = test_basis.compute_gram_matrix(max_radial, max_angular) + ortho = test_basis.compute_orthonormalization_matrix(max_radial, max_angular) + + for ell in range(max_angular): + eye = ortho[ell] @ gram[ell] @ ortho[ell].T + assert_allclose(eye, np.eye(max_radial, max_radial), atol=1e-11) + + +@pytest.mark.parametrize( + "analytical_basis", + [ + GtoBasis(cutoff=4, max_radial=6), + MonomialBasis(cutoff=4), + SphericalBesselBasis(cutoff=4, max_radial=6, max_angular=4), + ], +) +def test_derivative(analytical_basis: RadialBasisBase): + """Finite difference test for testing the derivative of a radial basis""" + + class NumericalRadialBasis(RadialBasisBase): + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return analytical_basis.compute(n, ell, integrand_positions) + + numerical_basis = NumericalRadialBasis(integration_radius=np.inf) + + cutoff = 4 + max_radial = 6 + max_angular = 4 + positions = np.linspace(2, cutoff) + + for n in range(max_radial): + for ell in range(max_angular): + assert_allclose( + numerical_basis.compute_derivative(n, ell, positions), + analytical_basis.compute_derivative(n, ell, positions), + atol=1e-9, + ) diff --git a/python/rascaline/tests/utils/splines.py b/python/rascaline/tests/utils/splines.py index a08cc54e3..9ff918e99 100644 --- a/python/rascaline/tests/utils/splines.py +++ b/python/rascaline/tests/utils/splines.py @@ -2,7 +2,22 @@ import pytest from numpy.testing import assert_allclose, assert_equal -from rascaline.utils import RadialIntegralFromFunction +from rascaline import LodeSphericalExpansion, SphericalExpansion +from rascaline.utils import ( + DeltaDensity, + GaussianDensity, + GtoBasis, + LodeDensity, + LodeSpliner, + RadialIntegralFromFunction, + SoapSpliner, +) + +from ..test_systems import SystemForTests + + +pytest.importorskip("scipy") +from scipy.special import gamma, hyp1f1 # noqa def sine(n: int, ell: int, positions: np.ndarray) -> np.ndarray: @@ -102,3 +117,214 @@ def test_splines_numerical_derivative_error(): match = "Numerically derivative of the radial integral can not be performed" with pytest.raises(ValueError, match=match): RadialIntegralFromFunction(**kwargs).compute() + + +def test_kspace_radial_integral(): + """Test against analytical integral with Gaussian densities and GTOs""" + + cutoff = 2 + max_radial = 6 + max_angular = 3 + atomic_gaussian_width = 1.0 + k_cutoff = 1.2 * np.pi / atomic_gaussian_width + + basis = GtoBasis(cutoff=cutoff, max_radial=max_radial) + + spliner = LodeSpliner( + max_radial=max_radial, + max_angular=max_angular, + k_cutoff=k_cutoff, + basis=basis, + density=DeltaDensity(), # density does not enter in a Kspace radial integral + accuracy=1e-8, + ) + + Neval = 100 + kk = np.linspace(0, k_cutoff, Neval) + + sigma = np.ones(max_radial, dtype=float) + for i in range(1, max_radial): + sigma[i] = np.sqrt(i) + sigma *= cutoff / max_radial + + factors = np.sqrt(np.pi) * np.ones((max_radial, max_angular + 1)) + + coeffs_num = np.zeros([max_radial, max_angular + 1, Neval]) + coeffs_exact = np.zeros_like(coeffs_num) + + for ell in range(max_angular + 1): + for n in range(max_radial): + i1 = 0.5 * (3 + n + ell) + i2 = 1.5 + ell + factors[n, ell] *= ( + 2 ** (0.5 * (n - ell - 1)) + * gamma(i1) + / gamma(i2) + * sigma[n] ** (2 * i1) + ) + coeffs_exact[n, ell] = ( + factors[n, ell] + * kk**ell + * hyp1f1(i1, i2, -0.5 * (kk * sigma[n]) ** 2) + ) + + coeffs_num[n, ell] = spliner.radial_integral(n, ell, kk) + + assert_allclose(coeffs_num, coeffs_exact) + + +def test_rspace_delta(): + cutoff = 2 + max_radial = 6 + max_angular = 3 + + basis = GtoBasis(cutoff=cutoff, max_radial=max_radial) + density = DeltaDensity() + + spliner = SoapSpliner( + max_radial=max_radial, + max_angular=max_angular, + cutoff=cutoff, + basis=basis, + density=density, + accuracy=1e-8, + ) + + positions = np.linspace(0, cutoff) + + for ell in range(max_angular + 1): + for n in range(max_radial): + assert_equal( + spliner.radial_integral(n, ell, positions), + basis.compute(n, ell, positions), + ) + assert_equal( + spliner.radial_integral_derivative(n, ell, positions), + basis.compute_derivative(n, ell, positions), + ) + + +def test_real_space_spliner(): + """Compare splined spherical expansion with GTOs and a Gaussian density to + analytical implementation.""" + cutoff = 8.0 + max_radial = 12 + max_angular = 9 + atomic_gaussian_width = 1.2 + + # We choose an accuracy that is larger then the default one (1e-8) to limit the time + # consumption of the test. + accuracy = 1e-4 + + spliner = SoapSpliner( + cutoff=cutoff, + max_radial=max_radial, + max_angular=max_angular, + basis=GtoBasis(cutoff=cutoff, max_radial=max_radial), + density=GaussianDensity(atomic_gaussian_width=atomic_gaussian_width), + accuracy=accuracy, + ) + + hypers_spherical_expansion = { + "cutoff": cutoff, + "max_radial": max_radial, + "max_angular": max_angular, + "center_atom_weight": 1.0, + "atomic_gaussian_width": atomic_gaussian_width, + "cutoff_function": {"Step": {}}, + } + + analytic = SphericalExpansion( + radial_basis={"Gto": {}}, **hypers_spherical_expansion + ).compute(SystemForTests()) + splined = SphericalExpansion( + radial_basis=spliner.compute(), **hypers_spherical_expansion + ).compute(SystemForTests()) + + for key, block_analytic in analytic.items(): + block_splined = splined.block(key) + assert_allclose( + block_splined.values, block_analytic.values, rtol=5e-4, atol=2e-5 + ) + + +@pytest.mark.parametrize("center_atom_weight", [1.0, 0.0]) +@pytest.mark.parametrize("potential_exponent", [0, 1]) +def test_fourier_space_spliner(center_atom_weight, potential_exponent): + """Compare splined LODE spherical expansion with GTOs and a Gaussian density to + analytical implementation.""" + + cutoff = 2 + max_radial = 6 + max_angular = 4 + atomic_gaussian_width = 0.8 + k_cutoff = 1.2 * np.pi / atomic_gaussian_width + + spliner = LodeSpliner( + k_cutoff=k_cutoff, + max_radial=max_radial, + max_angular=max_angular, + basis=GtoBasis(cutoff=cutoff, max_radial=max_radial), + density=LodeDensity( + atomic_gaussian_width=atomic_gaussian_width, + potential_exponent=potential_exponent, + ), + ) + + hypers_spherical_expansion = { + "cutoff": cutoff, + "max_radial": max_radial, + "max_angular": max_angular, + "center_atom_weight": center_atom_weight, + "atomic_gaussian_width": atomic_gaussian_width, + "potential_exponent": potential_exponent, + } + + analytic = LodeSphericalExpansion( + radial_basis={"Gto": {}}, **hypers_spherical_expansion + ).compute(SystemForTests()) + splined = LodeSphericalExpansion( + radial_basis=spliner.compute(), **hypers_spherical_expansion + ).compute(SystemForTests()) + + for key, block_analytic in analytic.items(): + block_splined = splined.block(key) + assert_allclose(block_splined.values, block_analytic.values, atol=1e-14) + + +def test_center_contribution_gto_gaussian(): + cutoff = 2.0 + max_radial = 6 + max_angular = 4 + atomic_gaussian_width = 0.8 + k_cutoff = 1.2 * np.pi / atomic_gaussian_width + + # Numerical evaluation of center contributions + spliner = LodeSpliner( + k_cutoff=k_cutoff, + max_radial=max_radial, + max_angular=max_angular, + basis=GtoBasis(cutoff=cutoff, max_radial=max_radial), + density=GaussianDensity(atomic_gaussian_width=atomic_gaussian_width), + ) + + # Analytical evaluation of center contributions + center_contr_analytical = np.zeros((max_radial)) + + normalization = 1.0 / (np.pi * atomic_gaussian_width**2) ** (3 / 4) + sigma_radial = np.ones(max_radial, dtype=float) + + for n in range(1, max_radial): + sigma_radial[n] = np.sqrt(n) + sigma_radial *= cutoff / max_radial + + for n in range(max_radial): + sigmatemp_sq = 1.0 / ( + 1.0 / atomic_gaussian_width**2 + 1.0 / sigma_radial[n] ** 2 + ) + neff = 0.5 * (3 + n) + center_contr_analytical[n] = (2 * sigmatemp_sq) ** neff * gamma(neff) + + center_contr_analytical *= normalization * 2 * np.pi / np.sqrt(4 * np.pi) + + assert_allclose(spliner.center_contribution, center_contr_analytical, rtol=1e-14) diff --git a/tox.ini b/tox.ini index c4df209aa..52e95df30 100644 --- a/tox.ini +++ b/tox.ini @@ -53,6 +53,7 @@ deps = chemfiles pytest pytest-cov + scipy commands = pytest --cov={env_site_packages_dir}/rascaline --cov-report xml:.tox/coverage.xml --import-mode=append {posargs} From 7384f7a244a31df81f5fa3158fe24a6887a43382 Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Wed, 18 Oct 2023 13:23:18 +0200 Subject: [PATCH 12/30] Add coverage for Python torch bindings (#234) --- .codecov.yml | 12 ++++-------- .github/workflows/coverage.yml | 20 ++++++++++++++++---- .gitignore | 4 +++- CONTRIBUTING.rst | 18 ++++++++++++++++++ tox.ini | 27 +++++++++++++++++++++++---- 5 files changed, 64 insertions(+), 17 deletions(-) diff --git a/.codecov.yml b/.codecov.yml index 3b8bba5e6..94bab7aef 100644 --- a/.codecov.yml +++ b/.codecov.yml @@ -1,15 +1,11 @@ coverage: - ignore: - - tests/.* status: project: default: target: 80% patch: off + ignore: + - "tests/.*" + - "examples/.*" -comment: off - -fixes: - # map coverage collected inside tox virtual environments - # to the source dir in git - - ".tox/all-deps/lib/*/site-packages/::python/" +comment: false diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 6949a6c0d..006530972 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -38,7 +38,7 @@ jobs: - name: install python dependencies run: | python -m pip install --upgrade pip - python -m pip install tox + python -m pip install tox coverage - name: install lcov run: sudo apt install -y lcov @@ -47,7 +47,7 @@ jobs: uses: actions/cache@v3 with: path: .tox - key: tox-${{ matrix.os }}-${{ hashFiles('pyproject.toml', 'setup.cfg', 'tox.ini') }} + key: tox-${{ hashFiles('pyproject.toml', 'setup.cfg', 'tox.ini') }} - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 @@ -69,10 +69,22 @@ jobs: lcov --remove coverage.info '/usr/*' "$(pwd)/rascaline-c-api/tests/*" "$(pwd)/rascaline-c-api/examples/*" --output-file coverage.info - name: collect Python coverage - run: tox -e all-deps + run: | + tox -e all-deps + tox -e torch-tests + env: + # Use the CPU only version of torch when building/running the code + PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu + + - name: combine Python coverage files + run: | + coverage combine --append \ + ./.coverage \ + ./python/rascaline-torch/.coverage + coverage xml - name: upload to codecov.io uses: codecov/codecov-action@v3 with: fail_ci_if_error: true - files: target/tarpaulin/cobertura.xml,.tox/coverage.xml,coverage.info + files: target/tarpaulin/cobertura.xml,coverage.xml,coverage.info diff --git a/.gitignore b/.gitignore index 7bcc9343a..4f077544b 100644 --- a/.gitignore +++ b/.gitignore @@ -3,10 +3,12 @@ target/ Cargo.lock .tox/ -.coverage build/ dist/ *.egg-info __pycache__/ .vscode *.DS_Store +htmlcov/ +.coverage +coverage.xml diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 7b61cb772..ead7509a3 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -176,6 +176,24 @@ of just testing it. .. _`cargo` : https://doc.rust-lang.org/cargo/ .. _valgrind: https://valgrind.org/ +Inspecting Python code coverage +------------------------------- + +The code coverage is reported at `codecov`_. You can also inspect the coverage locally. +To get the full coverage first combine all reports and open produced html file in a +browser + +.. code-block:: bash + + tox + coverage combine --append \ + ./.coverage \ + ./python/rascaline-torch/.coverage + coverage html + firefox htmlcov/index.html + +.. _codecov: https://codecov.io/gh/lab-cosmo/metatensor + Writing your own calculator --------------------------- diff --git a/tox.ini b/tox.ini index 52e95df30..0a0a60bbb 100644 --- a/tox.ini +++ b/tox.ini @@ -4,8 +4,8 @@ min_version = 4.0 # execute `tox` in the command-line without anything else envlist = lint - all-deps min-deps + all-deps docs-tests torch-tests @@ -25,6 +25,7 @@ metatensor-torch-requirement = metatensor-torch >=0.1.0,<0.2.0 build-single-wheel = --no-deps --no-build-isolation --check-build-dependencies +test_options = --cov={env_site_packages_dir}/rascaline --cov-append --import-mode=append commands = # error if the user gives a wrong testenv name in `tox -e` @@ -56,7 +57,7 @@ deps = scipy commands = - pytest --cov={env_site_packages_dir}/rascaline --cov-report xml:.tox/coverage.xml --import-mode=append {posargs} + pytest {[testenv]test_options} {posargs} [testenv:min-deps] @@ -64,9 +65,10 @@ commands = deps = {[testenv]metatensor-core-requirement} pytest + pytest-cov commands = - pytest --import-mode=append {posargs} + pytest {[testenv]test_options} {posargs} [testenv:torch-tests] @@ -74,6 +76,7 @@ deps = {[testenv]metatensor-torch-requirement} pytest + pytest-cov numpy torch ase @@ -86,7 +89,7 @@ commands = # install rascaline-torch pip install . {[testenv]build-single-wheel} --force-reinstall # run the unit tests - pytest --import-mode=append --assert=plain {posargs} + pytest {[testenv]test_options} --assert=plain {posargs} [testenv:docs] deps = @@ -184,3 +187,19 @@ commands = [flake8] max_line_length = 88 extend-ignore = E203 + +[coverage:report] +skip_covered = True +show_missing = True +omit = + tests/.* + examples/.* + +[coverage:paths] +rascaline = + python/rascaline/rascaline + .tox/*/lib/python*/site-packages/rascaline + +rascaline_torch = + python/rascaline-torch/rascaline/torch + .tox/*/lib/python*/site-packages/rascaline/torch From c62302bd1524c3be9673a2955a1495e0cb98e4f5 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 8 Sep 2023 18:33:39 +0200 Subject: [PATCH 13/30] rascaline-torch depends on metatensor-torch --- python/rascaline-torch/build-backend/backend.py | 2 +- python/rascaline-torch/setup.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/python/rascaline-torch/build-backend/backend.py b/python/rascaline-torch/build-backend/backend.py index 057d0157c..84f667978 100644 --- a/python/rascaline-torch/build-backend/backend.py +++ b/python/rascaline-torch/build-backend/backend.py @@ -18,7 +18,7 @@ RASCALINE_DEP = f"rascaline @ file://{RASCALINE}?{uuid}" else: # we are building from a sdist - RASCALINE_DEP = "rascaline == 0.1.0" + RASCALINE_DEP = "rascaline >=0.1.0.dev0,<0.2.0" prepare_metadata_for_build_wheel = build_meta.prepare_metadata_for_build_wheel diff --git a/python/rascaline-torch/setup.py b/python/rascaline-torch/setup.py index dea776aee..dcce11623 100644 --- a/python/rascaline-torch/setup.py +++ b/python/rascaline-torch/setup.py @@ -238,10 +238,13 @@ def git_extra_version(): with open(os.path.join(ROOT, authors[0])) as fd: authors = fd.read().splitlines() - install_requires = ["torch >= 1.11"] + install_requires = [ + "torch >= 1.11", + "metatensor-torch", + ] if os.path.exists(RASCALINE_C_API): # we are building from a git checkout - rascaline_path = os.path.join(ROOT, "..", "..") + rascaline_path = os.path.realpath(os.path.join(ROOT, "..", "..")) # add a random uuid to the file url to prevent pip from using a cached # wheel for rascaline, and force it to re-build from scratch From ad8a967807a9c34f8b102319343045b04244c820 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 8 Sep 2023 18:38:51 +0200 Subject: [PATCH 14/30] Add rascaline-torch as optional dependency of rascaline --- pyproject.toml | 2 +- setup.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 924103358..c319086d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "rascaline" -dynamic = ["version", "authors"] +dynamic = ["version", "authors", "optional-dependencies"] requires-python = ">=3.7" readme = "README.rst" diff --git a/setup.py b/setup.py index d26c8a651..4d7fca5f9 100644 --- a/setup.py +++ b/setup.py @@ -3,6 +3,7 @@ import shutil import subprocess import sys +import uuid from setuptools import Extension, setup from setuptools.command.bdist_egg import bdist_egg @@ -12,6 +13,7 @@ ROOT = os.path.realpath(os.path.dirname(__file__)) +RASCALINE_TORCH = os.path.join(ROOT, "python", "rascaline-torch") RASCALINE_BUILD_TYPE = os.environ.get("RASCALINE_BUILD_TYPE", "release") if RASCALINE_BUILD_TYPE not in ["debug", "release"]: @@ -250,9 +252,22 @@ def git_extra_version(): with open(os.path.join(ROOT, "AUTHORS")) as fd: authors = fd.read().splitlines() + extras_require = {} + if os.path.exists(RASCALINE_TORCH): + # we are building from a git checkout + + # add a random uuid to the file url to prevent pip from using a cached + # wheel for rascaline-torch, and force it to re-build from scratch + uuid = uuid.uuid4() + extras_require["torch"] = f"rascaline-torch @ file://{RASCALINE_TORCH}?{uuid}" + else: + # we are building from a sdist/installing from a wheel + extras_require["torch"] = "rascaline-torch >=0.1.0.dev0,<0.2.0" + setup( version=version, author=", ".join(authors), + extras_require=extras_require, ext_modules=[ # only declare the extension, it is built & copied as required by cmake # in the build_ext command From e961e1573edf073f2cc1aea6e971f43a22b2ef53 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Wed, 18 Oct 2023 15:48:13 +0200 Subject: [PATCH 15/30] Add a (-1)^l factor to the power spectrum invariants This make the power spectrum produce the same output as doing a full Clebsch-Gordan product and keeping only the lambda=0 term --- .../utils/power_spectrum/calculator.py | 5 +++- .../src/calculators/soap/power_spectrum.rs | 26 +++++++++++------ rascaline/tests/soap-power-spectrum.rs | 28 +++++++++++++++---- scripts/clean-python.sh | 1 + 4 files changed, 44 insertions(+), 16 deletions(-) diff --git a/python/rascaline/rascaline/utils/power_spectrum/calculator.py b/python/rascaline/rascaline/utils/power_spectrum/calculator.py index 6815f4b16..139a3e091 100644 --- a/python/rascaline/rascaline/utils/power_spectrum/calculator.py +++ b/python/rascaline/rascaline/utils/power_spectrum/calculator.py @@ -217,7 +217,10 @@ def compute( ell = key[0] species_center = key[1] - factor = 1 / sqrt(2 * ell + 1) + # For consistency with a full Clebsch-Gordan product we need to add + # a `-1^l / sqrt(2 l + 1)` factor to the power spectrum invariants + factor = (-1) ** ell / sqrt(2 * ell + 1) + # Find that block indices that have the same spherical_harmonics_l and # species_center selection = Labels( diff --git a/rascaline/src/calculators/soap/power_spectrum.rs b/rascaline/src/calculators/soap/power_spectrum.rs index 229f23834..911eaaed1 100644 --- a/rascaline/src/calculators/soap/power_spectrum.rs +++ b/rascaline/src/calculators/soap/power_spectrum.rs @@ -358,8 +358,19 @@ impl SoapPowerSpectrum { let property_1 = block_1.properties.position(&[n1]).expect("missing n1"); let property_2 = block_2.properties.position(&[n2]).expect("missing n2"); + let spherical_harmonics_l = l.usize(); + + // For consistency with a full Clebsch-Gordan product we need to add + // a `-1^l / sqrt(2 l + 1)` factor to the power spectrum invariants + let normalization = if spherical_harmonics_l % 2 == 0 { + f64::sqrt((2 * spherical_harmonics_l + 1) as f64) + } else { + -f64::sqrt((2 * spherical_harmonics_l + 1) as f64) + }; + SpxPropertiesToCombine { - spherical_harmonics_l: l.usize(), + spherical_harmonics_l, + normalization, property_1, property_2, spx_1: block_1.clone(), @@ -375,6 +386,8 @@ impl SoapPowerSpectrum { struct SpxPropertiesToCombine<'a> { /// value of l spherical_harmonics_l: usize, + /// normalization factor $-1^l * \sqrt{2 l + 1}$ + normalization: f64, /// position of n1 in the first spherical expansion properties property_1: usize, /// position of n2 in the second spherical expansion properties @@ -599,7 +612,7 @@ impl CalculatorBase for SoapPowerSpectrum { } unsafe { - *values.uget_mut(property_i) = sum / f64::sqrt((2 * spx.spherical_harmonics_l + 1) as f64); + *values.uget_mut(property_i) = sum / spx.normalization; } } }); @@ -655,10 +668,9 @@ impl CalculatorBase for SoapPowerSpectrum { } } - let normalization = f64::sqrt((2 * spx.spherical_harmonics_l + 1) as f64); for d in 0..3 { unsafe { - *values.uget_mut([d, property_i]) = sum[d] / normalization; + *values.uget_mut([d, property_i]) = sum[d] / spx.normalization; } } } @@ -694,7 +706,6 @@ impl CalculatorBase for SoapPowerSpectrum { let value_2 = spx_2.values.uget([spx_sample_2, m, spx.property_2]); for d1 in 0..3 { for d2 in 0..3 { - // TODO: ensure that gradient samples are 0..nsamples sum[d1][d2] += value_2 * spx_1_gradient.uget([spx_sample_1, d1, d2, m, spx.property_1]); } } @@ -707,7 +718,6 @@ impl CalculatorBase for SoapPowerSpectrum { let value_1 = spx_1.values.uget([spx_sample_1, m, spx.property_1]); for d1 in 0..3 { for d2 in 0..3 { - // TODO: ensure that gradient samples are 0..nsamples sum[d1][d2] += value_1 * spx_2_gradient.uget([spx_sample_2, d1, d2, m, spx.property_2]); } } @@ -723,12 +733,10 @@ impl CalculatorBase for SoapPowerSpectrum { } } - let normalization = f64::sqrt((2 * spx.spherical_harmonics_l + 1) as f64); - for d1 in 0..3 { for d2 in 0..3 { unsafe { - *values.uget_mut([d1, d2, property_i]) = sum[d1][d2] / normalization; + *values.uget_mut([d1, d2, property_i]) = sum[d1][d2] / spx.normalization; } } } diff --git a/rascaline/tests/soap-power-spectrum.rs b/rascaline/tests/soap-power-spectrum.rs index 42bfdb1da..632b02e6b 100644 --- a/rascaline/tests/soap-power-spectrum.rs +++ b/rascaline/tests/soap-power-spectrum.rs @@ -24,8 +24,10 @@ fn values() { let block = &descriptor.block_by_id(0); let array = block.values().to_array(); - let expected = &data::load_expected_values("soap-power-spectrum-values.npy.gz"); - assert_relative_eq!(array, expected, max_relative=1e-5); + let mut expected = data::load_expected_values("soap-power-spectrum-values.npy.gz"); + correct_factor(&mut expected, block.properties()); + + assert_relative_eq!(array, &expected, max_relative=1e-5); } #[test] @@ -51,13 +53,15 @@ fn gradients() { let gradients = block.gradient("positions").unwrap(); let array = sum_gradients(n_atoms, gradients); - let expected = &data::load_expected_values("soap-power-spectrum-positions-gradient.npy.gz"); - assert_relative_eq!(array, expected, max_relative=1e-6); + let mut expected = data::load_expected_values("soap-power-spectrum-positions-gradient.npy.gz"); + correct_factor(&mut expected, block.properties()); + assert_relative_eq!(array, &expected, max_relative=1e-6); let gradient = block.gradient("cell").unwrap(); let array = gradient.values().to_array(); - let expected = &data::load_expected_values("soap-power-spectrum-cell-gradient.npy.gz"); - assert_relative_eq!(array, expected, max_relative=1e-6); + let mut expected = data::load_expected_values("soap-power-spectrum-cell-gradient.npy.gz"); + correct_factor(&mut expected, block.properties()); + assert_relative_eq!(array, &expected, max_relative=1e-6); } fn sum_gradients(n_atoms: usize, gradients: TensorBlockRef<'_>) -> ArrayD { @@ -72,3 +76,15 @@ fn sum_gradients(n_atoms: usize, gradients: TensorBlockRef<'_>) -> ArrayD { sum } + + +// Add back the missing (-1)^l factor to the reference data +fn correct_factor(reference: &mut ndarray::ArrayD, properties: Labels) { + assert!(properties.names() == [ "species_neighbor_1", "species_neighbor_2", "l", "n1", "n2"]); + + let last_axis = reference.shape().len() - 1; + + for (&[_, _, l, _, _], mut column) in properties.iter_fixed_size().zip(reference.axis_iter_mut(Axis(last_axis))) { + column *= (-1.0_f64).powi(l.i32()); + } +} diff --git a/scripts/clean-python.sh b/scripts/clean-python.sh index 650c09747..992aa6c06 100755 --- a/scripts/clean-python.sh +++ b/scripts/clean-python.sh @@ -10,6 +10,7 @@ cd "$ROOT_DIR" rm -rf dist rm -rf build +rm -rf .coverage rm -rf python/rascaline-torch/dist rm -rf python/rascaline-torch/build From b8bb5cc3a4379272ef1f23f18923e7fa302717e7 Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Wed, 25 Oct 2023 22:40:01 +0200 Subject: [PATCH 16/30] rename cutoff -> k_cutoff --- .../calculators/lode/radial_integral/mod.rs | 4 +-- .../lode/radial_integral/spline.rs | 28 +++++++++---------- 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/rascaline/src/calculators/lode/radial_integral/mod.rs b/rascaline/src/calculators/lode/radial_integral/mod.rs index 5b98cadb3..c76b626df 100644 --- a/rascaline/src/calculators/lode/radial_integral/mod.rs +++ b/rascaline/src/calculators/lode/radial_integral/mod.rs @@ -81,7 +81,7 @@ impl LodeRadialIntegralCache { // the largest value the spline should interpolate is // the k-space cutoff, not the real-space cutoff // associated with the GTO basis - cutoff: parameters.k_cutoff, + k_cutoff: parameters.k_cutoff, }; Box::new(LodeRadialIntegralSpline::with_accuracy( @@ -95,7 +95,7 @@ impl LodeRadialIntegralCache { let parameters = LodeRadialIntegralSplineParameters { max_radial: parameters.max_radial, max_angular: parameters.max_angular, - cutoff: parameters.cutoff, + k_cutoff: parameters.k_cutoff, }; let center_contribution = center_contribution.ok_or(Error::InvalidParameter( diff --git a/rascaline/src/calculators/lode/radial_integral/spline.rs b/rascaline/src/calculators/lode/radial_integral/spline.rs index eabeb00e5..7597de15f 100644 --- a/rascaline/src/calculators/lode/radial_integral/spline.rs +++ b/rascaline/src/calculators/lode/radial_integral/spline.rs @@ -24,8 +24,8 @@ pub struct LodeRadialIntegralSplineParameters { pub max_radial: usize, /// Number of angular components pub max_angular: usize, - /// cutoff radius, this is also the maximal value that can be interpolated - pub cutoff: f64, + /// k-space cutoff radius, this is also the maximal value that can be interpolated + pub k_cutoff: f64, } impl LodeRadialIntegralSpline { @@ -44,7 +44,7 @@ impl LodeRadialIntegralSpline { let parameters = SplineParameters { start: 0.0, - stop: parameters.cutoff, + stop: parameters.k_cutoff, shape: vec![parameters.max_angular + 1, parameters.max_radial], }; @@ -74,7 +74,7 @@ impl LodeRadialIntegralSpline { let spline_parameters = SplineParameters { start: 0.0, - stop: parameters.cutoff, + stop: parameters.k_cutoff, shape: vec![parameters.max_angular + 1, parameters.max_radial], }; @@ -123,18 +123,17 @@ mod tests { #[test] fn high_accuracy() { - // Check that even with high accuracy and large domain MAX_SPLINE_SIZE - // is enough + // Check that even with high accuracy and large domain MAX_SPLINE_SIZE is enough let parameters = LodeRadialIntegralSplineParameters { max_radial: 15, max_angular: 10, - cutoff: 12.0, + k_cutoff: 10.0, }; let gto = LodeRadialIntegralGto::new(LodeRadialIntegralGtoParameters { max_radial: parameters.max_radial, max_angular: parameters.max_angular, - cutoff: parameters.cutoff, + cutoff: 5.0, atomic_gaussian_width: 0.5, potential_exponent: 1, }).unwrap(); @@ -150,20 +149,19 @@ mod tests { let parameters = LodeRadialIntegralSplineParameters { max_radial: max_radial, max_angular: max_angular, - cutoff: 5.0, + k_cutoff: 10.0, }; let gto = LodeRadialIntegralGto::new(LodeRadialIntegralGtoParameters { max_radial: parameters.max_radial, max_angular: parameters.max_angular, - cutoff: parameters.cutoff, + cutoff: 5.0, atomic_gaussian_width: 0.5, potential_exponent: 1, }).unwrap(); - // even with very bad accuracy, we want the gradients of the spline to - // match the values produces by the spline, and not necessarily the - // actual GTO gradients. + // even with very bad accuracy, we want the gradients of the spline to match the + // values produces by the spline, and not necessarily the actual GTO gradients. let spline = LodeRadialIntegralSpline::with_accuracy(parameters, 1e-2, gto).unwrap(); let rij = 3.4; @@ -204,13 +202,13 @@ mod tests { let parameters = LodeRadialIntegralSplineParameters { max_radial: max_radial, max_angular: max_angular, - cutoff: 5.0, + k_cutoff: 10.0, }; let gto = LodeRadialIntegralGto::new(LodeRadialIntegralGtoParameters { max_radial: parameters.max_radial, max_angular: parameters.max_angular, - cutoff: parameters.cutoff, + cutoff: 5.0, atomic_gaussian_width: 0.5, potential_exponent: 1, }).unwrap(); From e5dc30c3f5c5003ba3e3fa8268e542c35cd6f9dc Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 27 Oct 2023 16:20:42 +0200 Subject: [PATCH 17/30] Fix the version of sccache we pull from github --- .github/workflows/coverage.yml | 2 ++ .github/workflows/python-tests.yml | 2 ++ .github/workflows/rust-tests.yml | 4 ++++ .github/workflows/torch-tests.yml | 2 ++ 4 files changed, 10 insertions(+) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 006530972..8ef085c1a 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -51,6 +51,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 7627b9508..55773e416 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -47,6 +47,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | diff --git a/.github/workflows/rust-tests.yml b/.github/workflows/rust-tests.yml index 20d7f83a0..999542e2e 100644 --- a/.github/workflows/rust-tests.yml +++ b/.github/workflows/rust-tests.yml @@ -100,6 +100,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | @@ -140,6 +142,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index 79ec93e2d..4fd2d77cb 100644 --- a/.github/workflows/torch-tests.yml +++ b/.github/workflows/torch-tests.yml @@ -60,6 +60,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | From 30592db335784e37dcb7f9f23be71ff0ce99c216 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Mon, 30 Oct 2023 15:26:41 +0100 Subject: [PATCH 18/30] Disable coverage reporting in terminal (#252) --- tox.ini | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tox.ini b/tox.ini index 0a0a60bbb..2af417108 100644 --- a/tox.ini +++ b/tox.ini @@ -25,7 +25,11 @@ metatensor-torch-requirement = metatensor-torch >=0.1.0,<0.2.0 build-single-wheel = --no-deps --no-build-isolation --check-build-dependencies -test_options = --cov={env_site_packages_dir}/rascaline --cov-append --import-mode=append +test_options = + --cov={env_site_packages_dir}/rascaline \ + --cov-append \ + --cov-report= \ + --import-mode=append commands = # error if the user gives a wrong testenv name in `tox -e` From fca103ec4f975a7fcba521d2b0c27d582d2cbd7f Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Wed, 1 Nov 2023 15:28:34 +0100 Subject: [PATCH 19/30] Fix serialization of calculators in torch We need to store the registration name, not the self-reported calculator name --- python/rascaline-torch/tests/calculator.py | 1 + rascaline-torch/include/rascaline/torch/calculator.hpp | 9 ++++++++- rascaline-torch/src/register.cpp | 8 ++++---- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/python/rascaline-torch/tests/calculator.py b/python/rascaline-torch/tests/calculator.py index 28c2c6999..1c26be1bb 100644 --- a/python/rascaline-torch/tests/calculator.py +++ b/python/rascaline-torch/tests/calculator.py @@ -190,3 +190,4 @@ def forward( with tmpdir.as_cwd(): torch.jit.save(module, "test-save.torch") + module = torch.jit.load("test-save.torch") diff --git a/rascaline-torch/include/rascaline/torch/calculator.hpp b/rascaline-torch/include/rascaline/torch/calculator.hpp index bdc09bb1f..ee54929de 100644 --- a/rascaline-torch/include/rascaline/torch/calculator.hpp +++ b/rascaline-torch/include/rascaline/torch/calculator.hpp @@ -75,7 +75,8 @@ class RASCALINE_TORCH_EXPORT CalculatorHolder: public torch::CustomClassHolder { public: /// Create a new calculator with the given `name` and JSON `parameters` CalculatorHolder(std::string name, std::string parameters): - calculator_(std::move(name), std::move(parameters)) + c_name_(std::move(name)), + calculator_(c_name_, std::move(parameters)) {} /// Get the name of this calculator @@ -83,6 +84,11 @@ class RASCALINE_TORCH_EXPORT CalculatorHolder: public torch::CustomClassHolder { return calculator_.name(); } + /// Get the name used to register this calculator + std::string c_name() const { + return c_name_; + } + /// Get the parameters of this calculator std::string parameters() const { return calculator_.parameters(); @@ -100,6 +106,7 @@ class RASCALINE_TORCH_EXPORT CalculatorHolder: public torch::CustomClassHolder { ); private: + std::string c_name_; rascaline::Calculator calculator_; }; diff --git a/rascaline-torch/src/register.cpp b/rascaline-torch/src/register.cpp index 965b10865..75cf845d4 100644 --- a/rascaline-torch/src/register.cpp +++ b/rascaline-torch/src/register.cpp @@ -52,13 +52,13 @@ TORCH_LIBRARY(rascaline, module) { }) .def_pickle( // __getstate__ - [](const TorchCalculator& self) -> std::vector { - return {self->name(), self->parameters()}; + [](const TorchCalculator& self) -> std::tuple { + return {self->c_name(), self->parameters()}; }, // __setstate__ - [](std::vector state) -> TorchCalculator { + [](std::tuple state) -> TorchCalculator { return c10::make_intrusive( - state[0], state[1] + std::get<0>(state), std::get<1>(state) ); }) ; From 2851d47ab3471057f3794ac2b48b0417add2f748 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 3 Nov 2023 12:54:28 +0100 Subject: [PATCH 20/30] Fix access to rascaline-torch cmake_prefix_path from Python --- .../rascaline-torch/rascaline/torch/utils.py | 21 ------------------- .../rascaline/torch/utils/__init__.py | 9 ++++++++ python/rascaline-torch/setup.py | 2 ++ .../tests/utils/cmake_prefix.py | 7 +++++++ python/rascaline/rascaline/utils/__init__.py | 8 ++----- 5 files changed, 20 insertions(+), 27 deletions(-) delete mode 100644 python/rascaline-torch/rascaline/torch/utils.py create mode 100644 python/rascaline-torch/tests/utils/cmake_prefix.py diff --git a/python/rascaline-torch/rascaline/torch/utils.py b/python/rascaline-torch/rascaline/torch/utils.py deleted file mode 100644 index 96bce84d6..000000000 --- a/python/rascaline-torch/rascaline/torch/utils.py +++ /dev/null @@ -1,21 +0,0 @@ -import os - -import metatensor.torch - -import rascaline - - -_HERE = os.path.realpath(os.path.dirname(__file__)) - -_rascaline_torch_cmake_prefix = os.path.join(os.path.dirname(__file__), "lib", "cmake") - -cmake_prefix_path = ";".join( - [ - _rascaline_torch_cmake_prefix, - rascaline.utils.cmake_prefix_path, - metatensor.torch.utils.cmake_prefix_path, - ] -) -""" -Path containing the CMake configuration files for the underlying C library -""" diff --git a/python/rascaline-torch/rascaline/torch/utils/__init__.py b/python/rascaline-torch/rascaline/torch/utils/__init__.py index bc31628b9..3a73c2d20 100644 --- a/python/rascaline-torch/rascaline/torch/utils/__init__.py +++ b/python/rascaline-torch/rascaline/torch/utils/__init__.py @@ -1,4 +1,13 @@ +import os + from .power_spectrum import PowerSpectrum +_HERE = os.path.dirname(__file__) + +cmake_prefix_path = os.path.realpath(os.path.join(_HERE, "..", "lib", "cmake")) +""" +Path containing the CMake configuration files for the underlying C library +""" + __all__ = ["PowerSpectrum"] diff --git a/python/rascaline-torch/setup.py b/python/rascaline-torch/setup.py index dcce11623..0d357482e 100644 --- a/python/rascaline-torch/setup.py +++ b/python/rascaline-torch/setup.py @@ -38,6 +38,7 @@ class cmake_ext(build_ext): """Build the native library using cmake""" def run(self): + import metatensor import metatensor.torch import torch @@ -52,6 +53,7 @@ def run(self): # Tell CMake where to find rascaline & torch cmake_prefix_path = [ rascaline.utils.cmake_prefix_path, + metatensor.utils.cmake_prefix_path, metatensor.torch.utils.cmake_prefix_path, torch.utils.cmake_prefix_path, ] diff --git a/python/rascaline-torch/tests/utils/cmake_prefix.py b/python/rascaline-torch/tests/utils/cmake_prefix.py new file mode 100644 index 000000000..d6172a3bb --- /dev/null +++ b/python/rascaline-torch/tests/utils/cmake_prefix.py @@ -0,0 +1,7 @@ +import os + +import rascaline.torch + + +def test_cmake_prefix(): + assert os.path.exists(rascaline.torch.utils.cmake_prefix_path) diff --git a/python/rascaline/rascaline/utils/__init__.py b/python/rascaline/rascaline/utils/__init__.py index b5e6221e0..8f0f588a8 100644 --- a/python/rascaline/rascaline/utils/__init__.py +++ b/python/rascaline/rascaline/utils/__init__.py @@ -1,7 +1,5 @@ import os -import metatensor - from .power_spectrum import PowerSpectrum # noqa from .splines import ( # noqa AtomicDensityBase, @@ -19,11 +17,9 @@ ) -# path that can be used with cmake to access the rascaline library and headers -_HERE = os.path.realpath(os.path.dirname(__file__)) +_HERE = os.path.dirname(__file__) -_rascaline_cmake_prefix = os.path.realpath(os.path.join(_HERE, "..", "lib", "cmake")) -cmake_prefix_path = f"{_rascaline_cmake_prefix};{metatensor.utils.cmake_prefix_path}" +cmake_prefix_path = os.path.realpath(os.path.join(_HERE, "..", "lib", "cmake")) """ Path containing the CMake configuration files for the underlying C library """ From 1c0694096bc6832dc059f4f382627c246b4895f9 Mon Sep 17 00:00:00 2001 From: "Joseph W. Abbott" Date: Thu, 30 Nov 2023 16:04:19 +0100 Subject: [PATCH 21/30] Redesigned metatensor-based Clebsch Gordan iterations (#237) Co-authored-by: Alexander Goscinski Co-authored-by: DivyaSuman14 Co-authored-by: Guillaume Fraux Co-authored-by: Jigyasa Nigam Co-authored-by: Edoardo Cignoni Co-authored-by: frostedoyster q --- docs/src/explanations/index.rst | 1 + docs/src/explanations/rotation_adapted.rst | 88 ++ pyproject.toml | 1 + python/rascaline/rascaline/utils/__init__.py | 1 + .../utils/clebsch_gordan/__init__.py | 7 + .../utils/clebsch_gordan/_cg_cache.py | 551 +++++++++++++ .../utils/clebsch_gordan/_dispatch.py | 133 +++ .../utils/clebsch_gordan/clebsch_gordan.py | 772 ++++++++++++++++++ .../rascaline/tests/utils/clebsch_gordan.py | 540 ++++++++++++ .../tests/utils/data/h2_isolated.xyz | 4 + .../tests/utils/data/h2o_isolated.xyz | 5 + .../tests/utils/data/h2o_periodic.xyz | 5 + python/rascaline/tests/utils/rotations.py | 330 ++++++++ tox.ini | 5 + 14 files changed, 2443 insertions(+) create mode 100644 docs/src/explanations/rotation_adapted.rst create mode 100644 python/rascaline/rascaline/utils/clebsch_gordan/__init__.py create mode 100644 python/rascaline/rascaline/utils/clebsch_gordan/_cg_cache.py create mode 100644 python/rascaline/rascaline/utils/clebsch_gordan/_dispatch.py create mode 100644 python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py create mode 100644 python/rascaline/tests/utils/clebsch_gordan.py create mode 100644 python/rascaline/tests/utils/data/h2_isolated.xyz create mode 100644 python/rascaline/tests/utils/data/h2o_isolated.xyz create mode 100644 python/rascaline/tests/utils/data/h2o_periodic.xyz create mode 100644 python/rascaline/tests/utils/rotations.py diff --git a/docs/src/explanations/index.rst b/docs/src/explanations/index.rst index 8643a7235..77eb9ffa0 100644 --- a/docs/src/explanations/index.rst +++ b/docs/src/explanations/index.rst @@ -13,3 +13,4 @@ all about. concepts soap + rotation_adapted \ No newline at end of file diff --git a/docs/src/explanations/rotation_adapted.rst b/docs/src/explanations/rotation_adapted.rst new file mode 100644 index 000000000..e98bb332d --- /dev/null +++ b/docs/src/explanations/rotation_adapted.rst @@ -0,0 +1,88 @@ +Rotation-Adapted Features +========================= + +Equivariance +------------ + +Descriptors like SOAP are translation, rotation, and permutation invariant. +Indeed, such invariances are extremely useful if one wants to learn an invariant target (e.g., the energy). +Being already encoded in the descriptor, the learning algorithm does not have to learn such a physical requirement. + +The situation is different if the target is not invariant. For example, one may want to learn a dipole. The dipole rotates with a rotation of the molecule, and as such, invariant descriptors do not have the required symmetries for this task. + +Instead, one would need a rotation equivariant descriptor. +Rotation equivariance means that, if I first rotate the structure and compute the descriptor, I obtain the same result as first computing the descriptor and then applying the rotation, i.e., the descriptor behaves correctly upon rotation operations. +Denoting a structure as :math:`A`, the function computing the descriptor as :math:`f(\cdot)`, and the rotation operator as :math:`\hat{R}`, rotation equivariance can be expressed as: + +.. math:: + :name: eq:equivariance + + f(\hat{R} A) = \hat{R} f(A) + +Of course, invariance is a special case of equivariance. + + +Rotation Equivariance of the Spherical Expansion +------------------------------------------------ + +The spherical expansion is a rotation equivariant descriptor. +Let's consider the expansion coefficients of :math:`\rho_i(\mathbf{r})`. +We have: + +.. math:: + + \hat{R} \rho_i(\mathbf{r}) &= \sum_{nlm} c_{nlm}^{i} R_n(r) \hat{R} Y_l^m(\hat{\mathbf{r}}) \nonumber \\ + &= \sum_{nlmm'} c_{nlm}^{i} R_n(r) D_{m,m'}^{l}(\hat{R}) Y_l^{m'}(\hat{\mathbf{r}}) \nonumber \\ + &= \sum_{nlm} \left( \sum_{m'} D_{m',m}^l(\hat{R}) c_{nlm'}^{i}\right) B_{nlm}(\mathbf{r}) \nonumber + +and noting that :math:`Y_l^m(\hat{R} \hat{\mathbf{r}}) = \hat{R} Y_l^m(\hat{\mathbf{r}})` and :math:`\hat{R}r = r`, equation :ref:`(1) ` is satisfied and we conclude that the expansion coefficients :math:`c_{nlm}^{i}` are rotation equivariant. +Indeed, each :math:`c_{nlm}^{i}` transforms under rotation as the spherical harmonics :math:`Y_l^m(\hat{\mathbf{r}})`. + +Using the Dirac notation, the coefficient :math:`c_{nlm}^{i}` can be expressed as :math:`\braket{nlm\vert\rho_i}`. +Equivalently, and to stress the fact that this coefficient describes something that transforms under rotation as a spherical harmonics :math:`Y_l^m(\hat{\mathbf{r}})`, it is sometimes written as :math:`\braket{n\vert\rho_i;lm}`, i.e., the atomic density is "tagged" with a label that tells how it transforms under rotations. + + +Completeness Relations of Spherical Harmonics +--------------------------------------------- + +Spherical harmonics can be combined together using rules coming from standard theory of angular momentum: + +.. math:: + :name: eq:cg_coupling + + \ket{lm} \propto \ket{l_1 l_2 l m} = \sum_{m_1 m_2} C_{m_1 m_2 m}^{l_1 l_2 l} \ket{l_1 m_1} \ket{l_2 m_2} + +where :math:`C_{m_1 m_2 m}^{l_1 l_2 l}` is a Clebsch-Gordan (CG) coefficient. + +Thanks to the one-to-one correspondence (under rotation) between :math:`c_{nlm}^{i}` and :math:`Y_l^m`, +:ref:`(2) ` means that one can take products of two spherical expansion coefficients (which amounts to considering density correlations), and combine them with CG coefficients to get new coefficients that transform as a single spherical harmonics. +This process is known as coupling, from the uncoupled basis of angular momentum (formed by the product of rotation eigenstates) to a coupled basis (a single rotation eigenstate). + +One can also write the inverse of :ref:`(2) `: + +.. math:: + :name: eq:cg_decoupling + + \ket{l_1 m_1} \ket{l_2 m_2} = \sum_{l m} C_{m_1 m_2 m}^{l_1 l_2 l m} \ket{l_1 l_2 l m} + +that express the product of two rotation eigenstates in terms of one. This process is known as decoupling. + +Example: :math:`\lambda`-SOAP +----------------------------- + +A straightforward application of :ref:`(2) ` is the construction of :math:`\lambda`-SOAP features. +Indeed, :math:`\lambda`-SOAP was created in order to have a rotation and inversion equivariant version of the 3-body density correlations. +The :math:`\lambda` represents the degree of a spherical harmonics, :math:`Y_{\lambda}^{\mu}(\hat{\mathbf{r}})`, +and it indicates that this descriptor can transform under rotations as a spherical harmonics, i.e., it is rotation equivariant. + +It is then obtained by considering two expansion coefficients of the atomic density, and combining them with a CG iteration to a coupled basis, +as in :ref:`(2) `. +The :math:`\lambda`-SOAP descriptor is then: + +.. math:: + + \braket{n_1 l_1 n_2 l_2\vert\overline{\rho_i^{\otimes 2}, \sigma, \lambda \mu}} = + \frac{\delta_{\sigma, (-1)^{l_1 + l_2 + \lambda}}}{\sqrt{2 \lambda + 1}} + \sum_{m} C_{m (\mu-m) \mu}^{l_1 l_2 \lambda} c_{n_1 l_1 m}^{i} c_{n_2 l_2 (\mu - m)}^{i} + +where we have assumed real spherical harmonics coefficients. \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index c319086d0..36feada37 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ classifiers = [ dependencies = [ "metatensor-core >=0.1.0,<0.2.0", + "wigners", ] [project.urls] diff --git a/python/rascaline/rascaline/utils/__init__.py b/python/rascaline/rascaline/utils/__init__.py index 8f0f588a8..8309ac44c 100644 --- a/python/rascaline/rascaline/utils/__init__.py +++ b/python/rascaline/rascaline/utils/__init__.py @@ -1,5 +1,6 @@ import os +from .clebsch_gordan import * # noqa from .power_spectrum import PowerSpectrum # noqa from .splines import ( # noqa AtomicDensityBase, diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py b/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py new file mode 100644 index 000000000..38777d18e --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py @@ -0,0 +1,7 @@ +from .clebsch_gordan import correlate_density, correlate_density_metadata # noqa + + +__all__ = [ + "correlate_density", + "correlate_density_metadata", +] diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/_cg_cache.py b/python/rascaline/rascaline/utils/clebsch_gordan/_cg_cache.py new file mode 100644 index 000000000..8ad082fb8 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/_cg_cache.py @@ -0,0 +1,551 @@ +""" +Module that stores the ClebschGordanReal class for computing and caching Clebsch +Gordan coefficients for use in CG combinations. +""" +from typing import Union + +import numpy as np +import wigners + +from . import _dispatch + + +try: + from mops import sparse_accumulation_of_products as sap # noqa F401 + + HAS_MOPS = True +except ImportError: + HAS_MOPS = False + +try: + from torch import Tensor as TorchTensor +except ImportError: + + class TorchTensor: + pass + + +UNKNOWN_ARRAY_TYPE = ( + "unknown array type, only numpy arrays and torch tensors are supported" +) + + +# ================================= +# ===== ClebschGordanReal class +# ================================= + + +class ClebschGordanReal: + """ + Class for computing Clebsch-Gordan coefficients for real spherical + harmonics. + + Stores the coefficients in a dictionary in the `self.coeffs` attribute, + which is built at initialization. There are 3 current use cases for the + format of these coefficients. By default, sparse accumulation of products is + performed, whether or not Mops is installed. + + Case 1: standard sparse format. + + Each dictionary entry is a dictionary with entries for each (m1, m2, mu) + combination. + + { + (l1, l2, lambda): { + (m1, m2, mu) : cg_{m1, m2, mu}^{l1, l2, lambda} + for m1 in range(-l1, l1 + 1), + for m2 in range(-l2, l2 + 1), + }, + ... + for l1 in range(0, l1_list) + for l2 in range(0, l2_list) + for lambda in range(0, range(|l1 - l2|, ..., |l1 + l2|)) + } + + Case 2: standard dense format. + + Each dictionary entry is a dense array with shape (2 * l1 + 1, 2 * l2 + 1, 2 + * lambda + 1). + + { + (l1, l2, lambda): + array( + cg_{m1, m2, mu}^{l1, l2, lambda} + ... + for m1 in range(-l1, l1 + 1), + for m2 in range(-l2, l2 + 1), + for mu in range(-lambda, lambda + 1), + + shape=(2 * l1 + 1, 2 * l2 + 1, 2 * lambda + 1), + ) + ... + for l1 in range(0, l1_list) + for l2 in range(0, l2_list) + for lambda in range(0, range(|l1 - l2|, ..., |l1 + l2|)) + } + + Case 3: MOPS sparse format. + + Each dictionary entry contains a tuple with four 1D arrays, corresponding to + the CG coeffs and m1, m2, mu indices respectively. All of these arrays are + sorted according to the mu index. This format is used for Sparse + Accumulation of Products (SAP) as implemented in MOPS. See + https://github.com/lab-cosmo/mops . + + { + (l1, l2, lambda): + ( + [ + cg_{m1, m2, mu}^{l1, l2, lambda} + ... + for m1 in range(-l1, l1 + 1), + for m2 in range(-l2, l2 + 1), + for mu in range(-lambda, lambda + 1) + ], + [ + m1 for m1 in range(-l1, l1 + 1), + ], + [ + m2 for m2 in range(-l2, l2 + 1), + ], + [ + mu for mu in range(-lambda, lambda + 1), + ], + ) + + + } + + where `cg_{m1, m2, mu}^{l1, l2, lambda}` is the Clebsch-Gordan coefficient + that describes the combination of the `m1` irreducible component of the `l1` + angular channel and the `m2` irreducible component of the `l2` angular + channel into the irreducible tensor of order `lambda`. In all cases, these + correspond to the non-zero CG coefficients, i.e. those in the range |-l, + ..., +l| for each angular order l in {l1, l2, lambda}. + + :param lambda_max: maximum lambda value to compute CG coefficients for. + :param sparse: whether to store the CG coefficients in sparse format. + :param use_mops: whether to store the CG coefficients in MOPS sparse format. + This is recommended as the default for sparse accumulation, but can only + be used if Mops is installed. + """ + + def __init__(self, lambda_max: int, sparse: bool = True, use_mops: bool = HAS_MOPS): + self._lambda_max = lambda_max + self._sparse = sparse + + if sparse: + if not HAS_MOPS: + # TODO: provide a warning once Mops is fully ready + # import warnings + # warnings.warn( + # "It is recommended to use MOPS for sparse accumulation. " + # " This can be installed with ``pip install" + # " git+https://github.com/lab-cosmo/mops`." + # " Falling back to numpy for now." + # ) + self._use_mops = False + else: + self._use_mops = True + + else: + # TODO: provide a warning once Mops is fully ready + # if HAS_MOPS: + # import warnings + # warnings.warn( + # "Mops is installed, but not being used" + # " as dense operations chosen." + # ) + self._use_mops = False + + self._coeffs = ClebschGordanReal.build_coeff_dict( + self._lambda_max, + self._sparse, + self._use_mops, + ) + + @property + def lambda_max(self): + return self._lambda_max + + @property + def sparse(self): + return self._sparse + + @property + def use_mops(self): + return self._use_mops + + @property + def coeffs(self): + return self._coeffs + + @staticmethod + def build_coeff_dict(lambda_max: int, sparse: bool, use_mops: bool): + """ + Builds a dictionary of Clebsch-Gordan coefficients for all possible + combination of l1 and l2, up to lambda_max. + """ + # real-to-complex and complex-to-real transformations as matrices + r2c = {} + c2r = {} + coeff_dict = {} + for lambda_ in range(0, lambda_max + 1): + c2r[lambda_] = _complex2real(lambda_) + r2c[lambda_] = _real2complex(lambda_) + + for l1 in range(lambda_max + 1): + for l2 in range(lambda_max + 1): + for lambda_ in range( + max(l1, l2) - min(l1, l2), min(lambda_max, (l1 + l2)) + 1 + ): + complex_cg = _complex_clebsch_gordan_matrix(l1, l2, lambda_) + + real_cg = (r2c[l1].T @ complex_cg.reshape(2 * l1 + 1, -1)).reshape( + complex_cg.shape + ) + + real_cg = real_cg.swapaxes(0, 1) + real_cg = (r2c[l2].T @ real_cg.reshape(2 * l2 + 1, -1)).reshape( + real_cg.shape + ) + real_cg = real_cg.swapaxes(0, 1) + + real_cg = real_cg @ c2r[lambda_].T + + if (l1 + l2 + lambda_) % 2 == 0: + cg_l1l2lam = np.real(real_cg) + else: + cg_l1l2lam = np.imag(real_cg) + + if sparse: + # Find the m1, m2, mu idxs of the nonzero CG coeffs + nonzeros_cg_coeffs_idx = np.where(np.abs(cg_l1l2lam) > 1e-15) + if use_mops: + # Store CG coeffs in a specific format for use in + # MOPS. Here we need the m1, m2, mu, and CG coeffs + # to be stored as separate 1D arrays. + m1_arr, m2_arr, mu_arr, C_arr = [], [], [], [] + for m1, m2, mu in zip(*nonzeros_cg_coeffs_idx): + m1_arr.append(m1) + m2_arr.append(m2) + mu_arr.append(mu) + C_arr.append(cg_l1l2lam[m1, m2, mu]) + + # Reorder the arrays based on sorted mu values + mu_idxs = np.argsort(mu_arr) + m1_arr = np.array(m1_arr)[mu_idxs] + m2_arr = np.array(m2_arr)[mu_idxs] + mu_arr = np.array(mu_arr)[mu_idxs] + C_arr = np.array(C_arr)[mu_idxs] + cg_l1l2lam = (C_arr, m1_arr, m2_arr, mu_arr) + else: + # Otherwise fall back to torch/numpy and store as + # sparse dicts. + cg_l1l2lam = { + (m1, m2, mu): cg_l1l2lam[m1, m2, mu] + for m1, m2, mu in zip(*nonzeros_cg_coeffs_idx) + } + + # Store + coeff_dict[(l1, l2, lambda_)] = cg_l1l2lam + + return coeff_dict + + +# ============================ +# ===== Helper functions +# ============================ + + +def _real2complex(lambda_: int) -> np.ndarray: + """ + Computes a matrix that can be used to convert from real to complex-valued + spherical harmonics(coefficients) of order ``lambda_``. + + This is meant to be applied to the left: ``real2complex @ [-lambda_, ..., + +lambda_]``. + + See https://en.wikipedia.org/wiki/Spherical_harmonics#Real_form for details + on the convention for how these tranformations are defined. + """ + result = np.zeros((2 * lambda_ + 1, 2 * lambda_ + 1), dtype=np.complex128) + inv_sqrt_2 = 1.0 / np.sqrt(2) + i_sqrt_2 = 1j / np.sqrt(2) + for m in range(-lambda_, lambda_ + 1): + if m < 0: + # Positve part + result[lambda_ + m, lambda_ + m] = +i_sqrt_2 + # Negative part + result[lambda_ - m, lambda_ + m] = -i_sqrt_2 * ((-1) ** m) + + if m == 0: + result[lambda_, lambda_] = +1.0 + + if m > 0: + # Negative part + result[lambda_ - m, lambda_ + m] = +inv_sqrt_2 + # Positive part + result[lambda_ + m, lambda_ + m] = +inv_sqrt_2 * ((-1) ** m) + + return result + + +def _complex2real(lambda_: int) -> np.ndarray: + """ + Converts from complex to real spherical harmonics. This is just given by the + conjugate tranpose of the real->complex transformation matrices. + """ + return np.conjugate(_real2complex(lambda_)).T + + +def _complex_clebsch_gordan_matrix(l1, l2, lambda_): + r"""clebsch-gordan matrix + Computes the Clebsch-Gordan (CG) matrix for + transforming complex-valued spherical harmonics. + The CG matrix is computed as a 3D array of elements + < l1 m1 l2 m2 | lambda_ mu > + where the first axis loops over m1, the second loops over m2, + and the third one loops over mu. The matrix is real. + For example, using the relation: + | l1 l2 lambda_ mu > = + \sum_{m1, m2} + | l1 m1 > | l2 m2 > + (https://en.wikipedia.org/wiki/Clebsch–Gordan_coefficients, section + "Formal definition of Clebsch-Gordan coefficients", eq 2) + one can obtain the spherical harmonics lambda_ from two sets of + spherical harmonics with l1 and l2 (up to a normalization factor). + E.g.: + Args: + l1: l number for the first set of spherical harmonics + l2: l number for the second set of spherical harmonics + lambda_: l number For the third set of spherical harmonics + Returns: + cg: CG matrix for transforming complex-valued spherical harmonics + >>> from scipy.special import sph_harm + >>> import numpy as np + >>> import wigners + >>> C_112 = _complex_clebsch_gordan_matrix(1, 1, 2) + >>> comp_sph_1 = np.array([sph_harm(m, 1, 0.2, 0.2) for m in range(-1, 1 + 1)]) + >>> comp_sph_2 = np.array([sph_harm(m, 1, 0.2, 0.2) for m in range(-1, 1 + 1)]) + >>> # obtain the (unnormalized) spherical harmonics + >>> # with l = 2 by contraction over m1 and m2 + >>> comp_sph_2_u = np.einsum("ijk,i,j->k", C_112, comp_sph_1, comp_sph_2) + >>> # we can check that they differ from the spherical harmonics + >>> # by a constant factor + >>> comp_sph_2 = np.array([sph_harm(m, 2, 0.2, 0.2) for m in range(-2, 2 + 1)]) + >>> ratio = comp_sph_2 / comp_sph_2_u + >>> np.allclose(ratio[0], ratio) + True + """ + if np.abs(l1 - l2) > lambda_ or np.abs(l1 + l2) < lambda_: + return np.zeros((2 * l1 + 1, 2 * l2 + 1, 2 * lambda_ + 1), dtype=np.double) + else: + return wigners.clebsch_gordan_array(l1, l2, lambda_) + + +# ================================================= +# ===== Functions for performing CG combinations +# ================================================= + + +def combine_arrays( + arr_1: Union[np.ndarray, TorchTensor], + arr_2: Union[np.ndarray, TorchTensor], + lambda_: int, + cg_cache, + return_empty_array: bool = False, +) -> Union[np.ndarray, TorchTensor]: + """ + Couples arrays `arr_1` and `arr_2` corresponding to the irreducible + spherical components of 2 angular channels l1 and l2 using the appropriate + Clebsch-Gordan coefficients. As l1 and l2 can be combined to form multiple + lambda channels, this function returns the coupling to a single specified + channel `lambda`. The angular channels l1 and l2 are inferred from the size + of the components axis (axis 1) of the input arrays. + + `arr_1` has shape (n_i, 2 * l1 + 1, n_p) and `arr_2` has shape (n_i, 2 * l2 + + 1, n_q). n_i is the number of samples, n_p and n_q are the number of + properties in each array. The number of samples in each array must be the + same. + + The ouput array has shape (n_i, 2 * lambda + 1, n_p * n_q), where lambda is + the input parameter `lambda_`. + + The Clebsch-Gordan coefficients are cached in `cg_cache`. Currently, these + must be produced by the ClebschGordanReal class in this module. These + coefficients can be stored in either sparse dictionaries or dense arrays. + + The combination operation is dispatched such that numpy arrays or torch + tensors are automatically handled. + + `return_empty_array` can be used to return an empty array of the correct + shape, without performing the CG combination step. This can be useful for + probing the outputs of CG iterations in terms of metadata without the + computational cost of performing the CG combinations - i.e. using the + function :py:func:`combine_single_center_to_body_order_metadata_only`. + + :param arr_1: array with the m values for l1 with shape [n_samples, 2 * l1 + + 1, n_q_properties] + :param arr_2: array with the m values for l2 with shape [n_samples, 2 * l2 + + 1, n_p_properties] + :param lambda_: int value of the resulting coupled channel + :param cg_cache: either a sparse dictionary with keys (m1, m2, mu) and array + values being sparse blocks of shape , or a dense array + of shape [(2 * l1 +1) * (2 * l2 +1), (2 * lambda_ + 1)]. + + :returns: array of shape [n_samples, (2*lambda_+1), q_properties * p_properties] + """ + # If just precomputing metadata, return an empty array + if return_empty_array: + return sparse_combine(arr_1, arr_2, lambda_, cg_cache, return_empty_array=True) + + # Otherwise, perform the CG combination + # Spare CG cache + if cg_cache.sparse: + return sparse_combine(arr_1, arr_2, lambda_, cg_cache, return_empty_array=False) + + # Dense CG cache + return dense_combine(arr_1, arr_2, lambda_, cg_cache) + + +def sparse_combine( + arr_1: Union[np.ndarray, TorchTensor], + arr_2: Union[np.ndarray, TorchTensor], + lambda_: int, + cg_cache, + return_empty_array: bool = False, +) -> Union[np.ndarray, TorchTensor]: + """ + Performs a Clebsch-Gordan combination step on 2 arrays using sparse + operations. The angular channel of each block is inferred from the size of + its component axis, and the blocks are combined to the desired output + angular channel `lambda_` using the appropriate Clebsch-Gordan coefficients. + + :param arr_1: array with the m values for l1 with shape [n_samples, 2 * l1 + + 1, n_q_properties] + :param arr_2: array with the m values for l2 with shape [n_samples, 2 * l2 + + 1, n_p_properties] + :param lambda_: int value of the resulting coupled channel + :param cg_cache: sparse dictionary with keys (m1, m2, mu) and array values + being sparse blocks of shape + + :returns: array of shape [n_samples, (2*lambda_+1), q_properties * p_properties] + """ + # Samples dimensions must be the same + assert arr_1.shape[0] == arr_2.shape[0] + + # Infer l1 and l2 from the len of the length of axis 1 of each tensor + l1 = (arr_1.shape[1] - 1) // 2 + l2 = (arr_2.shape[1] - 1) // 2 + + # Define other useful dimensions + n_i = arr_1.shape[0] # number of samples + n_p = arr_1.shape[2] # number of properties in arr_1 + n_q = arr_2.shape[2] # number of properties in arr_2 + + if return_empty_array: # used when only computing metadata + return _dispatch.zeros_like((n_i, 2 * lambda_ + 1, n_p * n_q), like=arr_1) + + if isinstance(arr_1, np.ndarray) and HAS_MOPS: + # Reshape + arr_1 = np.repeat(arr_1[:, :, :, None], n_q, axis=3).reshape( + n_i, 2 * l1 + 1, n_p * n_q + ) + arr_2 = np.repeat(arr_2[:, :, None, :], n_p, axis=2).reshape( + n_i, 2 * l2 + 1, n_p * n_q + ) + + arr_1 = _dispatch.swapaxes(arr_1, 1, 2).reshape(n_i * n_p * n_q, 2 * l1 + 1) + arr_2 = _dispatch.swapaxes(arr_2, 1, 2).reshape(n_i * n_p * n_q, 2 * l2 + 1) + + # Do SAP + arr_out = sap( + arr_1, + arr_2, + *cg_cache._coeffs[(l1, l2, lambda_)], + output_size=2 * lambda_ + 1, + ) + assert arr_out.shape == (n_i * n_p * n_q, 2 * lambda_ + 1) + + # Reshape back + arr_out = arr_out.reshape(n_i, n_p * n_q, 2 * lambda_ + 1) + arr_out = _dispatch.swapaxes(arr_out, 1, 2) + + return arr_out + + if isinstance(arr_1, np.ndarray) or isinstance(arr_1, TorchTensor): + # Initialise output array + arr_out = _dispatch.zeros_like((n_i, 2 * lambda_ + 1, n_p * n_q), like=arr_1) + + # Get the corresponding Clebsch-Gordan coefficients + cg_coeffs = cg_cache.coeffs[(l1, l2, lambda_)] + + # Fill in each mu component of the output array in turn + for m1, m2, mu in cg_coeffs.keys(): + # Broadcast arrays, multiply together and with CG coeff + arr_out[:, mu, :] += ( + arr_1[:, m1, :, None] * arr_2[:, m2, None, :] * cg_coeffs[(m1, m2, mu)] + ).reshape(n_i, n_p * n_q) + + return arr_out + + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def dense_combine( + arr_1: Union[np.ndarray, TorchTensor], + arr_2: Union[np.ndarray, TorchTensor], + lambda_: int, + cg_cache, +) -> Union[np.ndarray, TorchTensor]: + """ + Performs a Clebsch-Gordan combination step on 2 arrays using a dense + operation. The angular channel of each block is inferred from the size of + its component axis, and the blocks are combined to the desired output + angular channel `lambda_` using the appropriate Clebsch-Gordan coefficients. + + :param arr_1: array with the m values for l1 with shape [n_samples, 2 * l1 + + 1, n_q_properties] + :param arr_2: array with the m values for l2 with shape [n_samples, 2 * l2 + + 1, n_p_properties] + :param lambda_: int value of the resulting coupled channel + :param cg_cache: dense array of shape [(2 * l1 +1) * (2 * l2 +1), (2 * lambda_ + + 1)] + + :returns: array of shape [n_samples, (2*lambda_+1), q_properties * p_properties] + """ + if isinstance(arr_1, np.ndarray) or isinstance(arr_1, TorchTensor): + # Infer l1 and l2 from the len of the length of axis 1 of each tensor + l1 = (arr_1.shape[1] - 1) // 2 + l2 = (arr_2.shape[1] - 1) // 2 + cg_coeffs = cg_cache.coeffs[(l1, l2, lambda_)] + + # (samples None None l1_mu q) * (samples l2_mu p None None) + # -> (samples l2_mu p l1_mu q) we broadcast it in this way + # so we only need to do one swapaxes in the next step + arr_out = arr_1[:, None, None, :, :] * arr_2[:, :, :, None, None] + + # (samples l2_mu p l1_mu q) -> (samples q p l1_mu l2_mu) + arr_out = _dispatch.swapaxes(arr_out, 1, 4) + + # samples (q p l1_mu l2_mu) -> (samples (q p) (l1_mu l2_mu)) + arr_out = arr_out.reshape( + -1, + arr_1.shape[2] * arr_2.shape[2], + arr_1.shape[1] * arr_2.shape[1], + ) + + # (l1_mu l2_mu lam_mu) -> ((l1_mu l2_mu) lam_mu) + cg_coeffs = cg_coeffs.reshape(-1, 2 * lambda_ + 1) + + # (samples (q p) (l1_mu l2_mu)) @ ((l1_mu l2_mu) lam_mu) + # -> samples (q p) lam_mu + arr_out = arr_out @ cg_coeffs + + # (samples (q p) lam_mu) -> (samples lam_mu (q p)) + return _dispatch.swapaxes(arr_out, 1, 2) + + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/_dispatch.py b/python/rascaline/rascaline/utils/clebsch_gordan/_dispatch.py new file mode 100644 index 000000000..f44aa3713 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/_dispatch.py @@ -0,0 +1,133 @@ +""" +Module containing dispatch functions for numpy/torch CG combination operations. +""" +from typing import List, Optional + +import numpy as np + + +try: + import torch + from torch import Tensor as TorchTensor +except ImportError: + + class TorchTensor: + pass + + +UNKNOWN_ARRAY_TYPE = ( + "unknown array type, only numpy arrays and torch tensors are supported" +) + + +def unique(array, axis: Optional[int] = None): + """Find the unique elements of an array.""" + if isinstance(array, TorchTensor): + return torch.unique(array, dim=axis) + elif isinstance(array, np.ndarray): + return np.unique(array, axis=axis) + + +def int_range_like(min_val, max_val, like): + """Returns an array of integers from min to max, non-inclusive, based on the + type of `like`""" + if isinstance(like, TorchTensor): + return torch.arange(min_val, max_val, dtype=torch.int64, device=like.device) + elif isinstance(like, np.ndarray): + return np.arange(min_val, max_val).astype(np.int64) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def int_array_like(int_list: List[int], like): + """ + Converts the input list of int to a numpy array or torch tensor + based on the type of `like`. + """ + if isinstance(like, TorchTensor): + return torch.tensor(int_list, dtype=torch.int64, device=like.device) + elif isinstance(like, np.ndarray): + return np.array(int_list).astype(np.int64) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def concatenate(arrays, axis: Optional[int] = 0): + """Concatenate arrays along an axis.""" + if isinstance(arrays[0], TorchTensor): + return torch.cat(arrays, dim=axis) + elif isinstance(arrays[0], np.ndarray): + return np.concatenate(arrays, axis=axis) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def all(array, axis: Optional[int] = None): + """Test whether all array elements along a given axis evaluate to True. + + This function has the same behavior as + ``np.all(array,axis=axis)``. + """ + if isinstance(array, bool): + array = np.array(array) + if isinstance(array, list): + array = np.array(array) + + if isinstance(array, TorchTensor): + # torch.all has two implementation, and picks one depending if more than one + # parameter is given. The second one does not supports setting dim to `None` + if axis is None: + return torch.all(input=array) + else: + return torch.all(input=array, dim=axis) + elif isinstance(array, np.ndarray): + return np.all(a=array, axis=axis) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def any(array): + """Test whether any array elements along a given axis evaluate to True. + + This function has the same behavior as + ``np.any(array)``. + """ + if isinstance(array, bool): + array = np.array(array) + if isinstance(array, list): + array = np.array(array) + if isinstance(array, TorchTensor): + return torch.any(array) + elif isinstance(array, np.ndarray): + return np.any(array) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def zeros_like(shape, like): + """Return an array of zeros with the same shape and type as a given array. + + This function has the same behavior as + ``np.zeros_like(array)``. + """ + if isinstance(like, TorchTensor): + return torch.zeros( + shape, + requires_grad=like.requires_grad, + dtype=like.dtype, + device=like.device, + ) + elif isinstance(like, np.ndarray): + return np.zeros(shape, dtype=like.dtype) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def swapaxes(array, axis0: int, axis1: int): + """Swaps axes of an array.""" + if isinstance(array, TorchTensor): + return torch.swapaxes(array, axis0, axis1) + elif isinstance(array, np.ndarray): + return np.swapaxes(array, axis0, axis1) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py b/python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py new file mode 100644 index 000000000..5572c7598 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py @@ -0,0 +1,772 @@ +""" +Module for computing Clebsch-gordan iterations with metatensor TensorMaps. +""" +import itertools +from typing import List, Optional, Tuple, Union + +from metatensor import Labels, TensorBlock, TensorMap + +from . import _cg_cache, _dispatch + + +# ====================================================================== +# ===== Public API functions +# ====================================================================== + + +def correlate_density( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, +) -> Union[TensorMap, List[TensorMap]]: + """ + Takes iterative Clebsch-Gordan (CG) tensor products of a density descriptor + with itself up to the desired correlation order. Returns + :py:class:`TensorMap`(s) corresponding to the density correlations output + from the specified iteration(s). + + A density descriptor necessarily is body order 2 (i.e. correlation order 1), + but can be single- or multi-center. The output is a :py:class:`list` of + density correlations for each iteration specified in `output_selection`, up + to the target order passed in `correlation_order`. By default only the last + correlation (i.e. the correlation of order ``correlation_order``) is + returned. + + This function is an iterative special case of the more general + :py:func:`correlate_tensors`. As a density is being correlated with itself, + some redundant CG tensor products can be skipped with the `skip_redundant` + keyword. + + Selections on the angular and parity channels at each iteration can also be + controlled with arguments `angular_cutoff`, `angular_selection` and + `parity_selection`. + + :param density: A density descriptor of body order 2 (correlation order 1), + in :py:class:`TensorMap` format. This may be, for example, a rascaline + :py:class:`SphericalExpansion` or :py:class:`LodeSphericalExpansion`. + Alternatively, this could be multi-center descriptor, such as a pair + density. + :param correlation_order: The desired correlation order of the output + descriptor. Must be >= 1. + :param angular_cutoff: The maximum angular channel to compute at any given + CG iteration, applied globally to all iterations until the target + correlation order is reached. + :param selected_keys: :py:class:`Labels` or `List[:py:class:`Labels`]` + specifying the angular and/or parity channels to output at each + iteration. All :py:class:`Labels` objects passed here must only contain + key names "spherical_harmonics_l" and "inversion_sigma". If a single + :py:class:`Labels` object is passed, this is applied to the final + iteration only. If a :py:class:`list` of :py:class:`Labels` objects is + passed, each is applied to its corresponding iteration. If None is + passed, all angular and parity channels are output at each iteration, + with the global `angular_cutoff` applied if specified. + :param skip_redundant: Whether to skip redundant CG combinations. Defaults + to False, which means all combinations are performed. If a + :py:class:`list` of :py:class:`bool` is passed, this is applied to each + iteration. If a single :py:class:`bool` is passed, this is applied to + all iterations. + :param output_selection: A :py:class:`list` of :py:class:`bool` specifying + whether to output a :py:class:`TensorMap` for each iteration. If a + single :py:class:`bool` is passed as True, outputs from all iterations + will be returned. If a :py:class:`list` of :py:class:`bool` is passed, + this controls the output at each corresponding iteration. If None is + passed, only the final iteration is output. + + :return: A :py:class:`list` of :py:class:`TensorMap` corresponding to the + density correlations output from the specified iterations. If the output + from a single iteration is requested, a :py:class:`TensorMap` is + returned instead. + """ + return _correlate_density( + density, + correlation_order, + angular_cutoff, + selected_keys, + skip_redundant, + output_selection, + compute_metadata_only=False, + sparse=True, # sparse CG cache by default + ) + + +def correlate_density_metadata( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, +) -> Union[TensorMap, List[TensorMap]]: + """ + Returns the metadata-only :py:class:`TensorMap`(s) that would be output by + the function :py:func:`correlate_density` under the same settings, without + perfoming the actual Clebsch-Gordan tensor products. See this function for + full documentation. + """ + + return _correlate_density( + density, + correlation_order, + angular_cutoff, + selected_keys, + skip_redundant, + output_selection, + compute_metadata_only=True, + ) + + +# ==================================================================== +# ===== Private functions that do the work on the TensorMap level +# ==================================================================== + + +def _correlate_density( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, + compute_metadata_only: bool = False, + sparse: bool = True, +) -> Union[TensorMap, List[TensorMap]]: + """ + Performs the density correlations for public functions + :py:func:`correlate_density` and :py:func:`correlate_density_metadata`. + """ + # Check inputs + if correlation_order <= 1: + raise ValueError("`correlation_order` must be > 1") + # TODO: implement combinations of gradients too + if _dispatch.any([len(list(block.gradients())) > 0 for block in density]): + raise NotImplementedError( + "Clebsch Gordan combinations with gradients not yet implemented." + " Use metatensor.remove_gradients to remove gradients from the input." + ) + # Check metadata + if not ( + _dispatch.all(density.keys.names == ["spherical_harmonics_l", "species_center"]) + or _dispatch.all( + density.keys.names + == ["spherical_harmonics_l", "species_center", "species_neighbor"] + ) + ): + raise ValueError( + "input `density` must have key names" + ' ["spherical_harmonics_l", "species_center"] or' + ' ["spherical_harmonics_l", "species_center", "species_neighbor"]' + ) + if not _dispatch.all(density.component_names == ["spherical_harmonics_m"]): + raise ValueError( + "input `density` must have a single component" + " axis with name `spherical_harmonics_m`" + ) + n_iterations = correlation_order - 1 # num iterations + density = _standardize_keys(density) # standardize metadata + density_correlation = density # create a copy to combine with itself + + # Parse the selected keys + selected_keys = _parse_selected_keys( + n_iterations=n_iterations, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + like=density.keys.values, + ) + # Parse the bool flags that control skipping of redundant CG combinations + # and TensorMap output from each iteration + skip_redundant, output_selection = _parse_bool_iteration_filters( + n_iterations, + skip_redundant=skip_redundant, + output_selection=output_selection, + ) + + # Pre-compute the keys needed to perform each CG iteration + key_metadata = _precompute_keys( + density.keys, + density.keys, + n_iterations=n_iterations, + selected_keys=selected_keys, + skip_redundant=skip_redundant, + ) + # Compute CG coefficient cache + if compute_metadata_only: + cg_cache = None + else: + angular_max = max( + _dispatch.concatenate( + [density.keys.column("spherical_harmonics_l")] + + [mdata[2].column("spherical_harmonics_l") for mdata in key_metadata] + ) + ) + # TODO: keys have been precomputed, so perhaps we don't need to + # compute all CG coefficients up to angular_max here. + # TODO: use sparse cache by default until we understand under which + # circumstances (and if) dense is faster. + cg_cache = _cg_cache.ClebschGordanReal(angular_max, sparse=sparse) + + # Perform iterative CG tensor products + density_correlations = [] + for iteration in range(n_iterations): + # Define the correlation order of the current iteration + correlation_order_it = iteration + 2 + + # Combine block pairs + blocks_out = [] + for key_1, key_2, key_out in zip(*key_metadata[iteration]): + block_out = _combine_blocks_same_samples( + density_correlation[key_1], + density[key_2], + key_out["spherical_harmonics_l"], + cg_cache, + compute_metadata_only=compute_metadata_only, + ) + blocks_out.append(block_out) + keys_out = key_metadata[iteration][2] + density_correlation = TensorMap(keys=keys_out, blocks=blocks_out) + + # If this tensor is to be included in the output, move the [l1, l2, ...] + # keys to properties and store + if output_selection[iteration]: + density_correlations.append( + density_correlation.keys_to_properties( + [f"l{i}" for i in range(1, correlation_order_it + 1)] + + [f"k{i}" for i in range(2, correlation_order_it)] + ) + ) + + # Drop redundant key names. TODO: these should be part of the global + # matadata associated with the TensorMap. Awaiting this functionality in + # metatensor. + for i, tensor in enumerate(density_correlations): + keys = tensor.keys + if len(_dispatch.unique(tensor.keys.column("order_nu"))) == 1: + keys = keys.remove(name="order_nu") + if len(_dispatch.unique(tensor.keys.column("inversion_sigma"))) == 1: + keys = keys.remove(name="inversion_sigma") + density_correlations[i] = TensorMap( + keys=keys, blocks=[b.copy() for b in tensor.blocks()] + ) + + # Return a single TensorMap in the simple case + if len(density_correlations) == 1: + return density_correlations[0] + + # Otherwise return a list of TensorMaps + return density_correlations + + +# ================================================================== +# ===== Functions to handle metadata +# ================================================================== + + +def _standardize_keys(tensor: TensorMap) -> TensorMap: + """ + Takes a nu=1 tensor and standardizes its metadata. This involves: 1) moving + the "species_neighbor" key to properties, if present as a dimension in the + keys, and 2) adding dimensions in the keys for tracking the body order + ("order_nu") and parity ("inversion_sigma") of the blocks. + + Checking for the presence of the "species_neighbor" key in the keys allows + the option of the user pre-moving this key to the properties before calling + `n_body_iteration_single_center`, allowing sparsity in a set of global + neighbors to be created if desired. + + Assumes that the input `tensor` is nu=1, and has only even parity blocks. + """ + if "species_neighbor" in tensor.keys.names: + tensor = tensor.keys_to_properties(keys_to_move="species_neighbor") + keys = tensor.keys.insert( + name="order_nu", + values=_dispatch.int_array_like([1], like=tensor.keys.values), + index=0, + ) + keys = keys.insert( + name="inversion_sigma", + values=_dispatch.int_array_like([1], like=tensor.keys.values), + index=1, + ) + return TensorMap(keys=keys, blocks=[b.copy() for b in tensor.blocks()]) + + +def _parse_selected_keys( + n_iterations: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + like=None, +) -> List[Union[None, Labels]]: + """ + Parses the `selected_keys` argument passed to public functions. Checks the + values and returns a :py:class:`list` of :py:class:`Labels` objects, one for + each iteration of CG combination. + + `like` is required if a new :py:class:`Labels` object is to be created by + :py:mod:`_dispatch`. + """ + # Check angular_cutoff arg + if angular_cutoff is not None: + if not isinstance(angular_cutoff, int): + raise TypeError("`angular_cutoff` must be passed as an int") + if angular_cutoff < 1: + raise ValueError("`angular_cutoff` must be >= 1") + + if selected_keys is None: + if angular_cutoff is None: # no selections at all + selected_keys = [None] * n_iterations + else: + # Create a key selection with all angular channels <= the specified + # angular cutoff + selected_keys = [ + Labels( + names=["spherical_harmonics_l"], + values=_dispatch.int_range_like( + 0, angular_cutoff, like=like + ).reshape(-1, 1), + ) + ] * n_iterations + + if isinstance(selected_keys, Labels): + # Create a list, but only apply a key selection at the final iteration + selected_keys = [None] * (n_iterations - 1) + [selected_keys] + + # Check the selected_keys + if not isinstance(selected_keys, List): + raise TypeError( + "`selected_keys` must be a `Labels` or List[Union[None, `Labels`]]" + ) + if not len(selected_keys) == n_iterations: + raise ValueError( + "`selected_keys` must be a List[Union[None, Labels]] of length" + " `correlation_order` - 1" + ) + if not _dispatch.all( + [isinstance(val, (Labels, type(None))) for val in selected_keys] + ): + raise TypeError("`selected_keys` must be a Labels or List[Union[None, Labels]]") + + # Now iterate over each of the Labels (or None) in the list and check + for slct in selected_keys: + if slct is None: + continue + assert isinstance(slct, Labels) + if not _dispatch.all( + [ + name in ["spherical_harmonics_l", "inversion_sigma"] + for name in slct.names + ] + ): + raise ValueError( + "specified key names in `selected_keys` must be either" + " 'spherical_harmonics_l' or 'inversion_sigma'" + ) + if "spherical_harmonics_l" in slct.names: + if angular_cutoff is not None: + if not _dispatch.all( + slct.column("spherical_harmonics_l") <= angular_cutoff + ): + raise ValueError( + "specified angular channels in `selected_keys` must be <= the" + " specified `angular_cutoff`" + ) + if not _dispatch.all( + [angular_l >= 0 for angular_l in slct.column("spherical_harmonics_l")] + ): + raise ValueError( + "specified angular channels in `selected_keys` must be >= 0" + ) + if "inversion_sigma" in slct.names: + if not _dispatch.all( + [parity_s in [-1, +1] for parity_s in slct.column("inversion_sigma")] + ): + raise ValueError( + "specified parities in `selected_keys` must be -1 or +1" + ) + + return selected_keys + + +def _parse_bool_iteration_filters( + n_iterations: int, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, +) -> List[List[bool]]: + """ + Parses the `skip_redundant` and `output_selection` arguments passed to + public functions. + """ + if isinstance(skip_redundant, bool): + skip_redundant = [skip_redundant] * n_iterations + if not _dispatch.all([isinstance(val, bool) for val in skip_redundant]): + raise TypeError("`skip_redundant` must be a `bool` or `list` of `bool`") + if not len(skip_redundant) == n_iterations: + raise ValueError( + "`skip_redundant` must be a bool or `list` of `bool` of length" + " `correlation_order` - 1" + ) + if output_selection is None: + output_selection = [False] * (n_iterations - 1) + [True] + else: + if isinstance(output_selection, bool): + output_selection = [output_selection] * n_iterations + if not isinstance(output_selection, List): + raise TypeError("`output_selection` must be passed as `list` of `bool`") + + if not len(output_selection) == n_iterations: + raise ValueError( + "`output_selection` must be a ``list`` of ``bool`` of length" + " corresponding to the number of CG iterations" + ) + if not _dispatch.all([isinstance(v, bool) for v in output_selection]): + raise TypeError("`output_selection` must be passed as a `list` of `bool`") + if not _dispatch.all([isinstance(v, bool) for v in output_selection]): + raise TypeError("`output_selection` must be passed as a `list` of `bool`") + + return skip_redundant, output_selection + + +def _precompute_keys( + keys_1: Labels, + keys_2: Labels, + n_iterations: int, + selected_keys: List[Union[None, Labels]], + skip_redundant: List[bool], +) -> List[Tuple[Labels, List[List[int]]]]: + """ + Computes all the keys metadata needed to perform `n_iterations` of CG + combination steps. + + At each iteration, a full product of the keys of two tensors, i.e. `keys_1` + and `keys_2` is computed. Then, key selections are applied according to the + user-defined settings: the maximum angular channel cutoff + (`angular_cutoff`), and angular and/or parity selections specified in + `selected_keys`. + + If `skip_redundant` is True, then keys that represent redundant CG + operations are not included in the output keys at each step. + """ + keys_metadata = [] + keys_out = keys_1 + for iteration in range(n_iterations): + # Get the keys metadata for the combination of the 2 tensors + keys_1_entries, keys_2_entries, keys_out = _precompute_keys_full_product( + keys_1=keys_out, + keys_2=keys_2, + ) + if selected_keys[iteration] is not None: + keys_1_entries, keys_2_entries, keys_out = _apply_key_selection( + keys_1_entries, + keys_2_entries, + keys_out, + selected_keys=selected_keys[iteration], + ) + + if skip_redundant[iteration]: + keys_1_entries, keys_2_entries, keys_out = _remove_redundant_keys( + keys_1_entries, keys_2_entries, keys_out + ) + + # Check that some keys are produced as a result of the combination + if len(keys_out) == 0: + raise ValueError( + f"invalid selections: iteration {iteration + 1} produces no" + " valid combinations. Check the `angular_cutoff` and" + " `selected_keys` args and try again." + ) + + keys_metadata.append((keys_1_entries, keys_2_entries, keys_out)) + + return keys_metadata + + +def _precompute_keys_full_product( + keys_1: Labels, keys_2: Labels +) -> Tuple[List, List, Labels]: + """ + Given the keys of 2 TensorMaps, returns the keys that would be present after + a full CG product of these TensorMaps. + + Assumes that `keys_1` corresponds to a TensorMap with arbitrary body order, + while `keys_2` corresponds to a TensorMap with body order 1. `keys_1` must + follow the key name convention: + + ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center", + "l1", "l2", ..., f"l{`nu`}", "k2", ..., f"k{`nu`-1}"]. The "lx" columns + track the l values of the nu=1 blocks that were previously combined. The + "kx" columns tracks the intermediate lambda values of nu > 1 blocks that + have been combined. + + For instance, a TensorMap of body order nu=4 will have key names + ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center", + "l1", "l2", "l3", "l4", "k2", "k3"]. Two nu=1 TensorMaps with blocks of + order "l1" and "l2" were combined to form a nu=2 TensorMap with blocks of + order "k2". This was combined with a nu=1 TensorMap with blocks of order + "l3" to form a nu=3 TensorMap with blocks of order "k3". Finally, this was + combined with a nu=1 TensorMap with blocks of order "l4" to form a nu=4. + + .. math :: + + \\bra{ + n_1 l_1 ; n_2 l_2 k_2 ; ... ; + n_{\nu-1} l_{\\nu-1} k_{\\nu-1} ; + n_{\\nu} l_{\\nu} k_{\\nu}; \\lambda + } + \\ket{ \\rho^{\\otimes \\nu}; \\lambda M } + + `keys_2` must follow the key name convention: ["order_nu", + "inversion_sigma", "spherical_harmonics_l", "species_center"] + + Returned is Tuple[List, List, Labels]. The first two lists correspond to the + LabelsEntry objects of the keys being combined. The third element is a + Labels object corresponding to the keys of the output TensorMap. Each entry + in this Labels object corresponds to the keys is formed by combination of + the pair of blocks indexed by correspoding key pairs in the first two lists. + """ + # Get the correlation order of the first TensorMap. + unique_nu = _dispatch.unique(keys_1.column("order_nu")) + if len(unique_nu) > 1: + raise ValueError( + "keys_1 must correspond to a tensor of a single correlation order." + f" Found {len(unique_nu)} body orders: {unique_nu}" + ) + nu1 = unique_nu[0] + + # Define new correlation order of output TensorMap + nu = nu1 + 1 + + # The correlation order of the second TensorMap should be nu = 1. + assert _dispatch.all(keys_2.column("order_nu") == 1) + + # If nu1 = 1, the key names don't yet have any "lx" columns + if nu1 == 1: + l_list_names = [] + new_l_list_names = ["l1", "l2"] + else: + l_list_names = [f"l{angular_l}" for angular_l in range(1, nu1 + 1)] + new_l_list_names = l_list_names + [f"l{nu}"] + + # Check key names + assert _dispatch.all( + keys_1.names + == ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center"] + + l_list_names + + [f"k{k}" for k in range(2, nu1)] + ) + assert _dispatch.all( + keys_2.names + == ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center"] + ) + + # Define key names of output Labels (i.e. for combined TensorMap) + new_names = ( + ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center"] + + new_l_list_names + + [f"k{k}" for k in range(2, nu)] + ) + + new_key_values = [] + keys_1_entries = [] + keys_2_entries = [] + for key_1, key_2 in itertools.product(keys_1, keys_2): + # Unpack relevant key values + sig1, lam1, a = key_1.values[1:4] + sig2, lam2, a2 = key_2.values[1:4] + + # Only combine blocks of the same chemical species + if a != a2: + continue + + # First calculate the possible non-zero angular channels that can be + # formed from combination of blocks of order `lam1` and `lam2`. This + # corresponds to values in the inclusive range { |lam1 - lam2|, ..., + # |lam1 + lam2| } + nonzero_lams = _dispatch.int_range_like( + abs(lam1 - lam2), abs(lam1 + lam2) + 1, like=key_1.values + ) + + # Now iterate over the non-zero angular channels and apply the custom + # selections + for lambda_ in nonzero_lams: + # Calculate new sigma + sig = sig1 * sig2 * (-1) ** (lam1 + lam2 + lambda_) + + # Extract the l and k lists from keys_1 + l_list = key_1.values[4 : 4 + nu1].tolist() + k_list = key_1.values[4 + nu1 :].tolist() + + # Build the new keys values. l{nu} is `lam2`` (i.e. + # "spherical_harmonics_l" of the key from `keys_2`. k{nu-1} is + # `lam1` (i.e. "spherical_harmonics_l" of the key from `keys_1`). + new_vals = [nu, sig, lambda_, a] + l_list + [lam2] + k_list + [lam1] + new_key_values.append(new_vals) + keys_1_entries.append(key_1) + keys_2_entries.append(key_2) + + # Define new keys as the full product of keys_1 and keys_2 + keys_out = Labels( + names=new_names, + values=_dispatch.int_array_like(new_key_values, like=keys_1.values), + ) + + return keys_1_entries, keys_2_entries, keys_out + + +def _apply_key_selection( + keys_1_entries: List, keys_2_entries: List, keys_out: Labels, selected_keys: Labels +) -> Tuple[List, List, Labels]: + """ + Applies a selection according to `selected_keys` to the keys of an output + TensorMap `keys_out` produced by combination of blocks indexed by keys + entries in `keys_1_entries` and `keys_2_entries` lists. + + After application of the selections, returned is a reduced set of keys and + set of corresponding parents key entries. + + If a selection in `selected_keys` is not valid based on the keys in + `keys_out`, an error is raised. + """ + # Extract the relevant columns from `selected_keys` that the selection will + # be performed on + keys_out_vals = [[k[name] for name in selected_keys.names] for k in keys_out] + + # First check that all of the selected keys exist in the output keys + for slct in selected_keys.values: + if not _dispatch.any([_dispatch.all(slct == k) for k in keys_out_vals]): + raise ValueError( + f"selected key {selected_keys.names} = {slct} not found" + " in the output keys. Check the `selected_keys` argument." + ) + + # Build a mask of the selected keys + mask = [ + _dispatch.any([_dispatch.all(i == j) for j in selected_keys.values]) + for i in keys_out_vals + ] + + # Apply the mask to key entries and keys and return + keys_1_entries = [k for k, isin in zip(keys_1_entries, mask) if isin] + keys_2_entries = [k for k, isin in zip(keys_2_entries, mask) if isin] + keys_out = Labels(names=keys_out.names, values=keys_out.values[mask]) + + return keys_1_entries, keys_2_entries, keys_out + + +def _remove_redundant_keys( + keys_1_entries: List, keys_2_entries: List, keys_out: Labels +) -> Tuple[List, List, Labels]: + """ + For a Labels object `keys_out` that corresponds to the keys of a TensorMap + formed by combined of the blocks described by the entries in the lists + `keys_1_entries` and `keys_2_entries`, removes redundant keys. + + These are the keys that correspond to blocks that have the same sorted l + list. The block where the l values are already sorted (i.e. l1 <= l2 <= ... + <= ln) is kept. + """ + # Get and check the correlation order of the input keys + nu1 = keys_1_entries[0]["order_nu"] + nu2 = keys_2_entries[0]["order_nu"] + assert nu2 == 1 + + # Get the correlation order of the output TensorMap + nu = nu1 + 1 + + # Identify keys of redundant blocks and remove them + key_idxs_to_keep = [] + for key_idx, key in enumerate(keys_out): + # Get the important key values. This is all of the keys, excpet the k + # list + key_vals_slice = key.values[: 4 + (nu + 1)].tolist() + first_part, l_list = key_vals_slice[:4], key_vals_slice[4:] + + # Sort the l list + l_list_sorted = sorted(l_list) + + # Compare the sliced key with the one recreated when the l list is + # sorted. If they are identical, this is the key of the block that we + # want to compute a CG combination for. + key_slice_tuple = tuple(first_part + l_list) + key_slice_sorted_tuple = tuple(first_part + l_list_sorted) + if _dispatch.all(key_slice_tuple == key_slice_sorted_tuple): + key_idxs_to_keep.append(key_idx) + + # Build a reduced Labels object for the combined keys, with redundancies removed + keys_out_red = Labels( + names=keys_out.names, + values=_dispatch.int_array_like( + [keys_out[idx].values for idx in key_idxs_to_keep], + like=keys_1_entries[0].values, + ), + ) + + # Store the list of reduced entries that combine to form the reduced output keys + keys_1_entries_red = [keys_1_entries[idx] for idx in key_idxs_to_keep] + keys_2_entries_red = [keys_2_entries[idx] for idx in key_idxs_to_keep] + + return keys_1_entries_red, keys_2_entries_red, keys_out_red + + +# ================================================================== +# ===== Functions to perform the CG combinations of blocks +# ================================================================== + + +def _combine_blocks_same_samples( + block_1: TensorBlock, + block_2: TensorBlock, + lambda_: int, + cg_cache, + compute_metadata_only: bool = False, +) -> TensorBlock: + """ + For a given pair of TensorBlocks and desired angular channel, combines the + values arrays and returns a new TensorBlock. + """ + + # Do the CG combination - single center so no shape pre-processing required + if compute_metadata_only: + combined_values = _cg_cache.combine_arrays( + block_1.values, block_2.values, lambda_, cg_cache, return_empty_array=True + ) + else: + combined_values = _cg_cache.combine_arrays( + block_1.values, block_2.values, lambda_, cg_cache, return_empty_array=False + ) + + # Infer the new nu value: block 1's properties are nu pairs of + # "species_neighbor_x" and "nx". + combined_nu = int((len(block_1.properties.names) / 2) + 1) + + # Define the new property names for "nx" and "species_neighbor_x" + n_names = [f"n{i}" for i in range(1, combined_nu + 1)] + neighbor_names = [f"species_neighbor_{i}" for i in range(1, combined_nu + 1)] + prop_names = [item for i in zip(neighbor_names, n_names) for item in i] + + # Create a TensorBlock + combined_block = TensorBlock( + values=combined_values, + samples=block_1.samples, + components=[ + Labels( + names=["spherical_harmonics_m"], + values=_dispatch.int_range_like( + min_val=-lambda_, max_val=lambda_ + 1, like=block_1.values + ).reshape(-1, 1), + ), + ], + properties=Labels( + names=prop_names, + values=_dispatch.int_array_like( + [ + _dispatch.concatenate((b2, b1)) + for b2 in block_2.properties.values + for b1 in block_1.properties.values + ], + like=block_1.values, + ), + ), + ) + + return combined_block diff --git a/python/rascaline/tests/utils/clebsch_gordan.py b/python/rascaline/tests/utils/clebsch_gordan.py new file mode 100644 index 000000000..38b7fab32 --- /dev/null +++ b/python/rascaline/tests/utils/clebsch_gordan.py @@ -0,0 +1,540 @@ +# -*- coding: utf-8 -*- +import os +from typing import List + +import metatensor +import numpy as np +import pytest +from metatensor import Labels, TensorBlock, TensorMap + +import rascaline +from rascaline.utils import PowerSpectrum +from rascaline.utils.clebsch_gordan._cg_cache import ClebschGordanReal +from rascaline.utils.clebsch_gordan.clebsch_gordan import ( + _correlate_density, + _standardize_keys, + correlate_density, + correlate_density_metadata, +) + + +# Try to import some modules +ase = pytest.importorskip("ase") +import ase.io # noqa: E402 + + +try: + import metatensor.operations + + HAS_METATENSOR_OPERATIONS = True +except ImportError: + HAS_METATENSOR_OPERATIONS = False +try: + import sympy # noqa: F401 + + HAS_SYMPY = True +except ImportError: + HAS_SYMPY = False + +if HAS_SYMPY: + from .rotations import WignerDReal, transform_frame_o3, transform_frame_so3 + + +DATA_ROOT = os.path.join(os.path.dirname(__file__), "data") + +SPHEX_HYPERS = { + "cutoff": 3.0, # Angstrom + "max_radial": 6, # Exclusive + "max_angular": 4, # Inclusive + "atomic_gaussian_width": 0.2, + "radial_basis": {"Gto": {}}, + "cutoff_function": {"ShiftedCosine": {"width": 0.5}}, + "center_atom_weight": 1.0, +} + +SPHEX_HYPERS_SMALL = { + "cutoff": 3.0, # Angstrom + "max_radial": 1, # Exclusive + "max_angular": 2, # Inclusive + "atomic_gaussian_width": 0.2, + "radial_basis": {"Gto": {}}, + "cutoff_function": {"ShiftedCosine": {"width": 0.5}}, + "center_atom_weight": 1.0, +} + + +# ============ Pytest fixtures ============ + + +@pytest.fixture() +def cg_cache_sparse(): + return ClebschGordanReal(lambda_max=5, sparse=True) + + +@pytest.fixture() +def cg_cache_dense(): + return ClebschGordanReal(lambda_max=5, sparse=False) + + +# ============ Helper functions ============ + + +def h2_isolated(): + return ase.io.read(os.path.join(DATA_ROOT, "h2_isolated.xyz"), ":") + + +def h2o_isolated(): + return ase.io.read(os.path.join(DATA_ROOT, "h2o_isolated.xyz"), ":") + + +def h2o_periodic(): + return ase.io.read(os.path.join(DATA_ROOT, "h2o_periodic.xyz"), ":") + + +def wigner_d_matrices(lmax: int): + return WignerDReal(lmax=lmax) + + +def spherical_expansion(frames: List[ase.Atoms]): + """Returns a rascaline SphericalExpansion""" + calculator = rascaline.SphericalExpansion(**SPHEX_HYPERS) + return calculator.compute(frames) + + +def spherical_expansion_small(frames: List[ase.Atoms]): + """Returns a rascaline SphericalExpansion""" + calculator = rascaline.SphericalExpansion(**SPHEX_HYPERS_SMALL) + return calculator.compute(frames) + + +def power_spectrum(frames: List[ase.Atoms]): + """Returns a rascaline PowerSpectrum constructed from a + SphericalExpansion""" + return PowerSpectrum(rascaline.SphericalExpansion(**SPHEX_HYPERS)).compute(frames) + + +def power_spectrum_small(frames: List[ase.Atoms]): + """Returns a rascaline PowerSpectrum constructed from a + SphericalExpansion""" + return PowerSpectrum(rascaline.SphericalExpansion(**SPHEX_HYPERS_SMALL)).compute( + frames + ) + + +def get_norm(tensor: TensorMap): + """ + Calculates the norm used in CG iteration tests. Assumes that the TensorMap + is sliced to a single sample. + + For a given atomic sample, the norm is calculated for each feature vector, + as a sum over lambda, sigma, and m. + """ + # Check that there is only one sample + assert ( + len( + metatensor.unique_metadata( + tensor, "samples", ["structure", "center", "species_center"] + ).values + ) + == 1 + ) + norm = 0.0 + for key, block in tensor.items(): # Sum over lambda and sigma + angular_l = key["spherical_harmonics_l"] + norm += np.sum( + [ + np.linalg.norm(block.values[0, m, :]) ** 2 + for m in range(-angular_l, angular_l + 1) + ] + ) + + return norm + + +# ============ Test equivariance ============ + + +@pytest.mark.skipif( + not HAS_SYMPY or not HAS_METATENSOR_OPERATIONS, + reason="SymPy or metatensor-operations are not installed", +) +@pytest.mark.parametrize( + "frames, nu_target, angular_cutoff, selected_keys", + [ + ( + h2_isolated(), + 3, + None, + Labels( + names=["spherical_harmonics_l", "inversion_sigma"], + values=np.array([[0, 1], [4, 1], [5, 1]]), + ), + ), + (h2o_isolated(), 2, 5, None), + (h2o_periodic(), 2, 5, None), + ], +) +def test_so3_equivariance(frames, nu_target, angular_cutoff, selected_keys): + """ + Tests that the output of :py:func:`correlate_density` is equivariant under + SO(3) transformations. + """ + wig = wigner_d_matrices(nu_target * SPHEX_HYPERS["max_angular"]) + frames_so3 = [transform_frame_so3(frame, wig.angles) for frame in frames] + + nu_1 = spherical_expansion(frames) + nu_1_so3 = spherical_expansion(frames_so3) + + nu_3 = correlate_density( + density=nu_1, + correlation_order=nu_target, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + ) + nu_3_so3 = correlate_density( + density=nu_1_so3, + correlation_order=nu_target, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + ) + + nu_3_transf = wig.transform_tensormap_so3(nu_3) + assert metatensor.allclose(nu_3_transf, nu_3_so3) + + +@pytest.mark.skipif( + not HAS_SYMPY or not HAS_METATENSOR_OPERATIONS, + reason="SymPy or metatensor-operations are not installed", +) +@pytest.mark.parametrize( + "frames, nu_target, angular_cutoff, selected_keys", + [ + ( + h2_isolated(), + 3, + None, + Labels( + names=["spherical_harmonics_l"], + values=np.array([0, 4, 5]).reshape(-1, 1), + ), + ), + (h2o_isolated(), 2, 5, None), + (h2o_periodic(), 2, 5, None), + ], +) +def test_o3_equivariance(frames, nu_target, angular_cutoff, selected_keys): + """ + Tests that the output of :py:func:`correlate_density` is equivariant under + O(3) transformations. + """ + wig = wigner_d_matrices(nu_target * SPHEX_HYPERS["max_angular"]) + frames_o3 = [transform_frame_o3(frame, wig.angles) for frame in frames] + + nu_1 = spherical_expansion(frames) + nu_1_o3 = spherical_expansion(frames_o3) + + nu_3 = correlate_density( + density=nu_1, + correlation_order=nu_target, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + ) + nu_3_o3 = correlate_density( + density=nu_1_o3, + correlation_order=nu_target, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + ) + + nu_3_transf = wig.transform_tensormap_o3(nu_3) + assert metatensor.allclose(nu_3_transf, nu_3_o3) + + +# ============ Test lambda-SOAP vs PowerSpectrum ============ + + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2_isolated()]) +@pytest.mark.parametrize( + "sphex_powspec", + [ + (spherical_expansion, power_spectrum), + (spherical_expansion_small, power_spectrum_small), + ], +) +def test_lambda_soap_vs_powerspectrum(frames, sphex_powspec): + """ + Tests for exact equivalence between the invariant block of a generated + lambda-SOAP equivariant and the Python implementation of PowerSpectrum in + rascaline utils. + """ + # Build a PowerSpectrum + ps = sphex_powspec[1](frames) + + # Build a lambda-SOAP + density = sphex_powspec[0](frames) + lsoap = correlate_density( + density=density, + correlation_order=2, + selected_keys=Labels( + names=["spherical_harmonics_l"], values=np.array([0]).reshape(-1, 1) + ), + ) + keys = lsoap.keys.remove(name="spherical_harmonics_l") + lsoap = TensorMap(keys=keys, blocks=[b.copy() for b in lsoap.blocks()]) + + # Manipulate metadata to match that of PowerSpectrum: + # 1) remove components axis + # 2) change "l1" and "l2" properties dimensions to just "l" (as l1 == l2) + blocks = [] + for block in lsoap.blocks(): + n_samples, n_props = block.values.shape[0], block.values.shape[2] + new_props = block.properties + new_props = new_props.remove(name="l1") + new_props = new_props.rename(old="l2", new="l") + blocks.append( + TensorBlock( + values=block.values.reshape((n_samples, n_props)), + samples=block.samples, + components=[], + properties=new_props, + ) + ) + lsoap = TensorMap(keys=lsoap.keys, blocks=blocks) + + # Compare metadata + assert metatensor.equal_metadata(lsoap, ps) + + # allclose on values + assert metatensor.allclose(lsoap, ps) + + +# ============ Test norm preservation ============ + + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2_isolated(), h2o_periodic()]) +@pytest.mark.parametrize("correlation_order", [2, 3, 4]) +def test_correlate_density_norm(frames, correlation_order): + """ + Checks \\|ρ^\\nu\\| = \\|ρ\\|^\\nu in the case where l lists are not + sorted. If l lists are sorted, thus saving computation of redundant block + combinations, the norm check will not hold for target body order greater + than 2. + """ + + # Build nu=1 SphericalExpansion + nu1 = spherical_expansion_small(frames) + + # Build higher body order tensor without sorting the l lists + nux = correlate_density( + nu1, + correlation_order=correlation_order, + angular_cutoff=None, + selected_keys=None, + skip_redundant=False, + ) + # Build higher body order tensor *with* sorting the l lists + nux_sorted_l = correlate_density( + nu1, + correlation_order=correlation_order, + angular_cutoff=None, + selected_keys=None, + skip_redundant=True, + ) + + # Standardize the features by passing through the CG combination code but with + # no iterations (i.e. body order 1 -> 1) + nu1 = _standardize_keys(nu1) + + # Make only lambda and sigma part of keys + nu1 = nu1.keys_to_samples(["species_center"]) + nux = nux.keys_to_samples(["species_center"]) + nux_sorted_l = nux_sorted_l.keys_to_samples(["species_center"]) + + # The norm shoudl be calculated for each sample. First find the unqiue + # samples + uniq_samples = metatensor.unique_metadata( + nux, "samples", names=["structure", "center", "species_center"] + ) + grouped_labels = [ + Labels(names=nux.sample_names, values=uniq_samples.values[i].reshape(1, 3)) + for i in range(len(uniq_samples)) + ] + + # Calculate norms + norm_nu1 = 0.0 + norm_nux = 0.0 + norm_nux_sorted_l = 0.0 + for sample in grouped_labels: + # Slice the TensorMaps + nu1_sliced = metatensor.slice(nu1, "samples", labels=sample) + nux_sliced = metatensor.slice(nux, "samples", labels=sample) + nux_sorted_sliced = metatensor.slice(nux_sorted_l, "samples", labels=sample) + + # Calculate norms + norm_nu1 += get_norm(nu1_sliced) ** correlation_order + norm_nux += get_norm(nux_sliced) + norm_nux_sorted_l += get_norm(nux_sorted_sliced) + + # Without sorting the l list we should get the same norm + assert np.allclose(norm_nu1, norm_nux) + + # But with sorting the l list we should get a different norm + assert not np.allclose(norm_nu1, norm_nux_sorted_l) + + +# ============ Test CG cache ============ + + +@pytest.mark.parametrize("l1, l2", [(1, 2), (2, 3), (0, 5)]) +def test_clebsch_gordan_orthogonality(cg_cache_dense, l1, l2): + """ + Test orthogonality relationships of cached dense CG coefficients. + + See + https://en.wikipedia.org/wiki/Clebsch%E2%80%93Gordan_coefficients#Orthogonality_relations + for details. + """ + lam_min = abs(l1 - l2) + lam_max = l1 + l2 + + # We test lam dimension + # \sum_{-m1 \leq l1 \leq m1, -m2 \leq l2 \leq m2} + # <λμ|l1m1,l2m2> = δ_μμ' + for lam in range(lam_min, lam_max): + cg_mat = cg_cache_dense.coeffs[(l1, l2, lam)].reshape(-1, 2 * lam + 1) + dot_product = cg_mat.T @ cg_mat + diag_mask = np.zeros(dot_product.shape, dtype=np.bool_) + diag_mask[np.diag_indices(len(dot_product))] = True + assert np.allclose( + dot_product[~diag_mask], np.zeros(dot_product.shape)[~diag_mask] + ) + assert np.allclose(dot_product[diag_mask], dot_product[diag_mask][0]) + + # We test l1 l2 dimension + # \sum_{|l1-l2| \leq λ \leq l1+l2} \sum_{-μ \leq λ \leq μ} + # <λμ|l1m1,l2m2> = δ_m1m1' δ_m2m2' + l1l2_dim = (2 * l1 + 1) * (2 * l2 + 1) + dot_product = np.zeros((l1l2_dim, l1l2_dim)) + for lam in range(lam_min, lam_max + 1): + cg_mat = cg_cache_dense.coeffs[(l1, l2, lam)].reshape(-1, 2 * lam + 1) + dot_product += cg_mat @ cg_mat.T + diag_mask = np.zeros(dot_product.shape, dtype=np.bool_) + diag_mask[np.diag_indices(len(dot_product))] = True + + assert np.allclose(dot_product[~diag_mask], np.zeros(dot_product.shape)[~diag_mask]) + assert np.allclose(dot_product[diag_mask], dot_product[diag_mask][0]) + + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2_isolated(), h2o_isolated()]) +def test_correlate_density_dense_sparse_agree(frames): + """ + Tests for agreement between nu=3 tensors built using both sparse and dense + CG coefficient caches. + """ + density = spherical_expansion_small(frames) + + # NOTE: testing the private function here so we can control the use of + # sparse v dense CG cache + n_body_sparse = _correlate_density( + density, + correlation_order=3, + compute_metadata_only=False, + sparse=True, + ) + n_body_dense = _correlate_density( + density, + correlation_order=3, + compute_metadata_only=False, + sparse=False, + ) + + assert metatensor.allclose(n_body_sparse, n_body_dense, atol=1e-8, rtol=1e-8) + + +# ============ Test metadata ============ + + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2o_isolated()]) +@pytest.mark.parametrize("correlation_order", [2, 3]) +@pytest.mark.parametrize("skip_redundant", [True, False]) +def test_correlate_density_metadata_agree(frames, correlation_order, skip_redundant): + """ + Tests that the metadata of outputs from :py:func:`correlate_density` and + :py:func:`correlate_density_metadata` agree. + """ + for nu1 in [spherical_expansion_small(frames), spherical_expansion(frames)]: + # Build higher body order tensor with CG computation + nux = correlate_density( + nu1, + correlation_order=correlation_order, + angular_cutoff=None, + selected_keys=None, + skip_redundant=skip_redundant, + ) + # Build higher body order tensor without CG computation - i.e. metadata + # only + nux_metadata_only = correlate_density_metadata( + nu1, + correlation_order=correlation_order, + angular_cutoff=None, + selected_keys=None, + skip_redundant=skip_redundant, + ) + assert metatensor.equal_metadata(nux, nux_metadata_only) + + +@pytest.mark.parametrize("frames", [h2o_isolated()]) +@pytest.mark.parametrize( + "selected_keys", + [ + None, + Labels( + names=["spherical_harmonics_l"], values=np.array([1, 2, 4]).reshape(-1, 1) + ), + ], +) +@pytest.mark.parametrize("skip_redundant", [True, False]) +def test_correlate_density_angular_selection( + frames: List[ase.Atoms], + selected_keys: Labels, + skip_redundant: bool, +): + """ + Tests that the correct angular channels are outputted based on the + specified ``selected_keys``. + """ + nu_1 = spherical_expansion(frames) + + nu_2 = correlate_density( + density=nu_1, + correlation_order=2, + angular_cutoff=None, + selected_keys=selected_keys, + skip_redundant=skip_redundant, + ) + + if selected_keys is None: + assert np.all( + [ + angular in np.arange(SPHEX_HYPERS["max_angular"] * 2 + 1) + for angular in np.unique(nu_2.keys.column("spherical_harmonics_l")) + ] + ) + + else: + assert np.all( + np.sort(np.unique(nu_2.keys.column("spherical_harmonics_l"))) + == np.sort(selected_keys.column("spherical_harmonics_l")) + ) diff --git a/python/rascaline/tests/utils/data/h2_isolated.xyz b/python/rascaline/tests/utils/data/h2_isolated.xyz new file mode 100644 index 000000000..ec5f59680 --- /dev/null +++ b/python/rascaline/tests/utils/data/h2_isolated.xyz @@ -0,0 +1,4 @@ +2 +pbc="F F F" +H 1.97361700 1.73067300 2.47063400 +H 1.97361700 3.26932700 2.47063400 diff --git a/python/rascaline/tests/utils/data/h2o_isolated.xyz b/python/rascaline/tests/utils/data/h2o_isolated.xyz new file mode 100644 index 000000000..fc876d2ba --- /dev/null +++ b/python/rascaline/tests/utils/data/h2o_isolated.xyz @@ -0,0 +1,5 @@ +3 +pbc="F F F" +O 2.56633400 2.50000000 2.50370100 +H 1.97361700 1.73067300 2.47063400 +H 1.97361700 3.26932700 2.47063400 diff --git a/python/rascaline/tests/utils/data/h2o_periodic.xyz b/python/rascaline/tests/utils/data/h2o_periodic.xyz new file mode 100644 index 000000000..3374566e6 --- /dev/null +++ b/python/rascaline/tests/utils/data/h2o_periodic.xyz @@ -0,0 +1,5 @@ +3 +Lattice="5.0 0.0 0.0 0.0 5.0 0.0 0.0 0.0 5.0" pbc="T T T" +O 2.56633400 2.50000000 2.50370100 +H 1.97361700 1.73067300 2.47063400 +H 1.97361700 3.26932700 2.47063400 diff --git a/python/rascaline/tests/utils/rotations.py b/python/rascaline/tests/utils/rotations.py new file mode 100644 index 000000000..2abe22714 --- /dev/null +++ b/python/rascaline/tests/utils/rotations.py @@ -0,0 +1,330 @@ +""" +Class for generating real Wigner-D matrices, and using them to rotate ASE frames +and TensorMaps of density coefficients in the spherical basis. +""" +from typing import Sequence + +import pytest + + +ase = pytest.importorskip("ase") + +import numpy as np # noqa: E402 +from metatensor import TensorBlock, TensorMap # noqa: E402 +from scipy.spatial.transform import Rotation # noqa: E402 + + +try: + import torch # noqa: E402 + from torch import Tensor as TorchTensor # noqa: E402 +except ImportError: + + class TorchTensor: + pass + + +# ===== Functions for transformations in the Cartesian basis ===== + + +def cartesian_rotation(angles: Sequence[float]): + """ + Returns a Cartesian rotation matrix in the appropriate convention (ZYZ, + implicit rotations) to be consistent with the common Wigner D definition. + + `angles` correspond to the alpha, beta, gamma Euler angles in the ZYZ + convention, in radians. + """ + return Rotation.from_euler("ZYZ", angles).as_matrix() + + +def transform_frame_so3(frame: ase.Atoms, angles: Sequence[float]) -> ase.Atoms: + """ + Transforms the positions and cell coordinates of an ASE frame by a SO(3) + rigid rotation. + """ + new_frame = frame.copy() + + # Build cartesian rotation matrix + R = cartesian_rotation(angles) + + # Rotate its positions and cell + new_frame.positions = new_frame.positions @ R.T + new_frame.cell = new_frame.cell @ R.T + + return new_frame + + +def transform_frame_o3(frame: ase.Atoms, angles: Sequence[float]) -> ase.Atoms: + """ + Transforms the positions and cell coordinates of an ASE frame by an O(3) + rotation. This involves a rigid SO(3) rotation of the positions and cell + according to the Euler `angles`, then an inversion by multiplying just the + positions by -1. + """ + new_frame = frame.copy() + + # Build cartesian rotation matrix + R = cartesian_rotation(angles) + + # Rotate its positions and cell + new_frame.positions = new_frame.positions @ R.T + new_frame.cell = new_frame.cell @ R.T + + # Invert the atom positions + new_frame.positions *= -1 + + return new_frame + + +# ===== WignerDReal for transformations in the spherical basis ===== + + +class WignerDReal: + """ + A helper class to compute Wigner D matrices given the Euler angles of a rotation, + and apply them to spherical harmonics (or coefficients). Built to function with + real-valued coefficients. + """ + + def __init__(self, lmax: int, angles: Sequence[float] = None): + """ + Initialize the WignerDReal class. + + :param lmax: int, the maximum angular momentum channel for which the + Wigner D matrices are computed + :param angles: Sequence[float], the alpha, beta, gamma Euler angles, in + radians. + """ + self.lmax = lmax + # Randomly generate Euler angles between 0 and 2 pi if none are provided + if angles is None: + angles = np.random.uniform(size=(3)) * 2 * np.pi + self.angles = angles + self.rotation = cartesian_rotation(angles) + + r2c_mats = {} + c2r_mats = {} + for L in range(0, self.lmax + 1): + r2c_mats[L] = np.hstack( + [_r2c(np.eye(2 * L + 1)[i])[:, np.newaxis] for i in range(2 * L + 1)] + ) + c2r_mats[L] = np.conjugate(r2c_mats[L]).T + self.matrices = {} + for L in range(0, self.lmax + 1): + wig = _wigner_d(L, self.angles) + self.matrices[L] = np.real(c2r_mats[L] @ np.conjugate(wig) @ r2c_mats[L]) + + def rotate_coeff_vector( + self, + frame: ase.Atoms, + coeffs: np.ndarray, + lmax: dict, + nmax: dict, + ) -> np.ndarray: + """ + Rotates the irreducible spherical components (ISCs) of basis set + coefficients in the spherical basis passed in as a flat vector. + + Required is the basis set definition specified by ``lmax`` and ``nmax``. + This are dicts of the form: + + lmax = {symbol: lmax_value, ...} + nmax = {(symbol, l): nmax_value, ...} + + where ``symbol`` is the chemical symbol of the atom, ``lmax_value`` is + its corresponding max l channel value. For each combination of species + symbol and lmax, there exists a max radial channel value ``nmax_value``. + + Then, the assumed ordering of basis function coefficients follows a + hierarchy, which can be read as nested loops over the various indices. + Be mindful that some indices range are from 0 to x (exclusive) and + others from 0 to x + 1 (exclusive). The ranges reported below are + ordered. + + 1. Loop over atoms (index ``i``, of chemical species ``a``) in the + structure. ``i`` takes values 0 to N (** exclusive **), where N is the + number of atoms in the structure. + + 2. Loop over spherical harmonics channel (index ``l``) for each atom. + ``l`` takes values from 0 to ``lmax[a] + 1`` (** exclusive **), where + ``a`` is the chemical species of atom ``i``, given by the chemical + symbol at the ``i``th position of ``symbol_list``. + + 3. Loop over radial channel (index ``n``) for each atom ``i`` and + spherical harmonics channel ``l`` combination. ``n`` takes values from 0 + to ``nmax[(a, l)]`` (** exclusive **). + + 4. Loop over spherical harmonics component (index ``m``) for each atom. + ``m`` takes values from ``-l`` to ``l`` (** inclusive **). + + :param frame: the atomic structure in ASE format for which the + coefficients are defined. + :param coeffs: the coefficients in the spherical basis, as a flat + vector. + :param lmax: dict containing the maximum spherical harmonics (l) value + for each atom type. + :param nmax: dict containing the maximum radial channel (n) value for + each combination of atom type and l. + + :return: the rotated coefficients in the spherical basis, as a flat + vector with the same order as the input vector. + """ + # Initialize empty vector for storing the rotated ISCs + rot_vect = np.empty_like(coeffs) + + # Iterate over atomic species of the atoms in the frame + curr_idx = 0 + for symbol in frame.get_chemical_symbols(): + # Get the basis set lmax value for this species + sym_lmax = lmax[symbol] + for angular_l in range(sym_lmax + 1): + # Get the number of radial functions for this species and l value + sym_l_nmax = nmax[(symbol, angular_l)] + # Get the Wigner D Matrix for this l value + wig_mat = self.matrices[angular_l].T + for _n in range(sym_l_nmax): + # Retrieve the irreducible spherical component + isc = coeffs[curr_idx : curr_idx + (2 * angular_l + 1)] + # Rotate the ISC and store + rot_isc = isc @ wig_mat + rot_vect[curr_idx : curr_idx + (2 * angular_l + 1)][:] = rot_isc[:] + # Update the start index for the next ISC + curr_idx += 2 * angular_l + 1 + + return rot_vect + + def rotate_tensorblock(self, angular_l: int, block: TensorBlock) -> TensorBlock: + """ + Rotates a TensorBlock ``block``, represented in the spherical basis, + according to the Wigner D Real matrices for the given ``l`` value. + Assumes the components of the block are [("spherical_harmonics_m",),]. + """ + # Get the Wigner matrix for this l value + wig = self.matrices[angular_l].T + + # Copy the block + block_rotated = block.copy() + vals = block_rotated.values + + # Perform the rotation, either with numpy or torch, by taking the + # tensordot product of the irreducible spherical components. Modify + # in-place the values of the copied TensorBlock. + if isinstance(vals, TorchTensor): + wig = torch.tensor(wig) + block_rotated.values[:] = torch.tensordot( + vals.swapaxes(1, 2), wig, dims=1 + ).swapaxes(1, 2) + elif isinstance(block.values, np.ndarray): + block_rotated.values[:] = np.tensordot( + vals.swapaxes(1, 2), wig, axes=1 + ).swapaxes(1, 2) + else: + raise TypeError("TensorBlock values must be a numpy array or torch tensor.") + + return block_rotated + + def transform_tensormap_so3(self, tensor: TensorMap) -> TensorMap: + """ + Transforms a TensorMap by a by an SO(3) rigid rotation using Wigner-D + matrices. + + Assumes the input tensor follows the metadata structure consistent with + those produce by rascaline. + """ + # Retrieve the key and the position of the l value in the key names + keys = tensor.keys + idx_l_value = keys.names.index("spherical_harmonics_l") + + # Iterate over the blocks and rotate + rotated_blocks = [] + for key in keys: + # Retrieve the l value + angular_l = key[idx_l_value] + + # Rotate the block and store + rotated_blocks.append(self.rotate_tensorblock(angular_l, tensor[key])) + + return TensorMap(keys, rotated_blocks) + + def transform_tensormap_o3(self, tensor: TensorMap) -> TensorMap: + """ + Transforms a TensorMap by a by an O(3) transformation: this involves an + SO(3) rigid rotation using Wigner-D Matrices followed by an inversion. + + Assumes the input tensor follows the metadata structure consistent with + those produce by rascaline. + """ + # Retrieve the key and the position of the l value in the key names + keys = tensor.keys + idx_l_value = keys.names.index("spherical_harmonics_l") + + # Iterate over the blocks and rotate + new_blocks = [] + for key in keys: + # Retrieve the l value + angular_l = key[idx_l_value] + + # Rotate the block + new_block = self.rotate_tensorblock(angular_l, tensor[key]) + + # Work out the inversion multiplier according to the convention + inversion_multiplier = 1 + if key["spherical_harmonics_l"] % 2 == 1: + inversion_multiplier *= -1 + + # "inversion_sigma" may not be present if CG iterations haven't been + # performed (i.e. nu=1 rascaline SphericalExpansion) + if "inversion_sigma" in keys.names: + if key["inversion_sigma"] == -1: + inversion_multiplier *= -1 + + # Invert the block by applying the inversion multiplier + new_block = TensorBlock( + values=new_block.values * inversion_multiplier, + samples=new_block.samples, + components=new_block.components, + properties=new_block.properties, + ) + new_blocks.append(new_block) + + return TensorMap(keys, new_blocks) + + +# ===== Helper functions for WignerDReal + + +def _wigner_d(angular_l: int, angles: Sequence[float]) -> np.ndarray: + """ + Computes the Wigner D matrix: + D^l_{mm'}(alpha, beta, gamma) + from sympy and converts it to numerical values. + + `angles` are the alpha, beta, gamma Euler angles (radians, ZYZ convention) + and l the irrep. + """ + try: + from sympy.physics.wigner import wigner_d + except ModuleNotFoundError: + raise ModuleNotFoundError( + "Calculation of Wigner D matrices requires a sympy installation" + ) + return np.complex128(wigner_d(angular_l, *angles)) + + +def _r2c(sp): + """ + Real to complex SPH. Assumes a block with 2l+1 reals corresponding + to real SPH with m indices from -l to +l + """ + + i_sqrt_2 = 1.0 / np.sqrt(2) + + angular_l = (len(sp) - 1) // 2 # infers l from the vector size + rc = np.zeros(len(sp), dtype=np.complex128) + rc[angular_l] = sp[angular_l] + for m in range(1, angular_l + 1): + rc[angular_l + m] = ( + (sp[angular_l + m] + 1j * sp[angular_l - m]) * i_sqrt_2 * (-1) ** m + ) + rc[angular_l - m] = (sp[angular_l + m] - 1j * sp[angular_l - m]) * i_sqrt_2 + return rc diff --git a/tox.ini b/tox.ini index 2af417108..ed7ce7872 100644 --- a/tox.ini +++ b/tox.ini @@ -56,9 +56,14 @@ deps = {[testenv]metatensor-core-requirement} ase chemfiles + metatensor-operations pytest pytest-cov scipy + sympy + wigners + # TODO: add mops once support for windows available + # git+https://github.com/lab-cosmo/mops commands = pytest {[testenv]test_options} {posargs} From 10c663ea61faf69961680045afa426539d17937a Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Wed, 6 Dec 2023 18:11:18 +0100 Subject: [PATCH 22/30] Bump min Python version 3.7 --- .github/workflows/python-tests.yml | 4 ++-- .github/workflows/torch-tests.yml | 2 +- pyproject.toml | 2 +- python/rascaline-torch/pyproject.toml | 2 +- python/rascaline-torch/rascaline/torch/__init__.py | 12 ++---------- python/rascaline/rascaline/version.py | 12 ++---------- 6 files changed, 9 insertions(+), 25 deletions(-) diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 55773e416..889b38b42 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -19,7 +19,7 @@ jobs: matrix: include: - os: ubuntu-20.04 - python-version: "3.7" + python-version: "3.8" - os: ubuntu-20.04 python-version: "3.11" - os: macos-11 @@ -75,7 +75,7 @@ jobs: name: Python ${{ matrix.python-version }} / check build strategy: matrix: - python-version: ['3.7', '3.11'] + python-version: ['3.8', '3.11'] os: [ubuntu-20.04] steps: - uses: actions/checkout@v3 diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index 4fd2d77cb..c9557e22e 100644 --- a/.github/workflows/torch-tests.yml +++ b/.github/workflows/torch-tests.yml @@ -19,7 +19,7 @@ jobs: include: - os: ubuntu-20.04 torch-version: 1.11.* - python-version: "3.7" + python-version: "3.8" cargo-test-flags: --release - os: ubuntu-20.04 diff --git a/pyproject.toml b/pyproject.toml index 36feada37..2cce98f8f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rascaline" dynamic = ["version", "authors", "optional-dependencies"] -requires-python = ">=3.7" +requires-python = ">=3.8" readme = "README.rst" license = {text = "BSD-3-Clause"} diff --git a/python/rascaline-torch/pyproject.toml b/python/rascaline-torch/pyproject.toml index 4c4daacf4..098a8ba90 100644 --- a/python/rascaline-torch/pyproject.toml +++ b/python/rascaline-torch/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rascaline-torch" dynamic = ["version", "authors", "dependencies"] -requires-python = ">=3.7" +requires-python = ">=3.8" readme = "README.rst" license = {text = "BSD-3-Clause"} diff --git a/python/rascaline-torch/rascaline/torch/__init__.py b/python/rascaline-torch/rascaline/torch/__init__.py index 177d0c688..971e0a3a1 100644 --- a/python/rascaline-torch/rascaline/torch/__init__.py +++ b/python/rascaline-torch/rascaline/torch/__init__.py @@ -1,15 +1,7 @@ -import sys +import importlib.metadata -if (sys.version_info.major >= 3) and (sys.version_info.minor >= 8): - import importlib.metadata - - __version__ = importlib.metadata.version("rascaline-torch") - -else: - from pkg_resources import get_distribution - - __version__ = get_distribution("rascaline-torch").version +__version__ = importlib.metadata.version("rascaline-torch") from ._c_lib import _load_library diff --git a/python/rascaline/rascaline/version.py b/python/rascaline/rascaline/version.py index a7248880b..e69652e02 100644 --- a/python/rascaline/rascaline/version.py +++ b/python/rascaline/rascaline/version.py @@ -1,12 +1,4 @@ -import sys +import importlib.metadata -if (sys.version_info.major >= 3) and (sys.version_info.minor >= 8): - import importlib.metadata - - __version__ = importlib.metadata.version("rascaline") - -else: - from pkg_resources import get_distribution - - __version__ = get_distribution("rascaline").version +__version__ = importlib.metadata.version("rascaline") From 8f1de18ba17d63304c392ace66cd0c65db2620ab Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Wed, 6 Dec 2023 18:15:58 +0100 Subject: [PATCH 23/30] Various tox.ini improvments - Use default tox error for non existing environments - Use `description` for more info on what the env is doing available via `tox list` --- tox.ini | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/tox.ini b/tox.ini index ed7ce7872..dcd68c59f 100644 --- a/tox.ini +++ b/tox.ini @@ -1,5 +1,5 @@ [tox] -min_version = 4.0 +min_version = 4.9.0 # these are the default environments, i.e. the list of tests running when you # execute `tox` in the command-line without anything else envlist = @@ -31,14 +31,11 @@ test_options = --cov-report= \ --import-mode=append -commands = - # error if the user gives a wrong testenv name in `tox -e` - python -c "import sys; print('environement {env_name} does not exist'); sys.exit(1)" - [testenv:build-rascaline] -# This environment is only used to build the wheels which are then re-used by -# all other environments requiring rascaline to be installed +description = + This environment is only used to build the wheels which are then re-used by + all other environments requiring rascaline to be installed passenv = * deps = setuptools @@ -50,8 +47,9 @@ commands = [testenv:all-deps] -# Run Python unit tests with all dependencies installed (ase & chemfiles are -# optional dependencies) +description = + Run Python unit tests with all dependencies installed (ase & chemfiles are optional + dependencies) deps = {[testenv]metatensor-core-requirement} ase @@ -70,7 +68,7 @@ commands = [testenv:min-deps] -# Run Python unit tests with the minimal dependencies installed +description = Run Python unit tests with the minimal dependencies installed deps = {[testenv]metatensor-core-requirement} pytest @@ -81,6 +79,7 @@ commands = [testenv:torch-tests] +description = Run Python unit tests using torch deps = {[testenv]metatensor-torch-requirement} @@ -101,6 +100,7 @@ commands = pytest {[testenv]test_options} --assert=plain {posargs} [testenv:docs] +description = Build the package documentation. deps = -r docs/requirements.txt cmake @@ -118,7 +118,7 @@ commands = [testenv:docs-tests] -# this environement runs the doctests defined in any metatensor package +description = Runs the doctests defined in any metatensor package deps = {[testenv]metatensor-core-requirement} ase @@ -129,8 +129,9 @@ commands = [testenv:lint] -# lint the Python code with flake8 (code linter), black (code formatter), and -# isort (sorting of imports) +description = + lint the Python code with flake8 (code linter), black (code formatter), and isort + (sorting of imports) package = skip deps = black @@ -148,9 +149,8 @@ commands = [testenv:format] +description = Abuse tox to do actual formatting on all files. package = skip -# Abuse tox to do actual formatting. Users can call `tox -e format` to run -# formatting on all files deps = black blackdoc From 6c76b778bf6289be977b32e5314121be4b46834d Mon Sep 17 00:00:00 2001 From: Liam Marsh Date: Thu, 7 Dec 2023 16:59:21 +0100 Subject: [PATCH 24/30] added pyscf system --- .../rascaline/rascaline/systems/__init__.py | 13 +++ python/rascaline/rascaline/systems/pyscf.py | 102 ++++++++++++++++++ 2 files changed, 115 insertions(+) create mode 100644 python/rascaline/rascaline/systems/pyscf.py diff --git a/python/rascaline/rascaline/systems/__init__.py b/python/rascaline/rascaline/systems/__init__.py index 0f0f8d522..2000b006a 100644 --- a/python/rascaline/rascaline/systems/__init__.py +++ b/python/rascaline/rascaline/systems/__init__.py @@ -1,6 +1,7 @@ from .ase import HAVE_ASE, AseSystem from .base import SystemBase from .chemfiles import HAVE_CHEMFILES, ChemfilesSystem +from .pyscf import HAVE_PYSCF, PyscfSystem class IntoSystem: @@ -32,6 +33,15 @@ def __init__(self): .. _chemfiles.Frame: http://chemfiles.org/chemfiles.py/latest/reference/frame.html """ +if HAVE_PYSCF: + IntoSystem.__doc__ += """ + - `pyscf.gto.mole.Mole`_ or `pyscf.pbc.gto.cell.Cell`_: pyscf' Frame type. There is no associated neighbor + list implementation, the system will only be usable with + ``use_native_system=True`` + + .. _pyscf.gto.mole.Mole: https://pyscf.org/user/gto.html + .. _pyscf.pbc.gto.cell.Cell: https://pyscf.org/user/pbc/gto.html + """ def wrap_system(system: IntoSystem) -> SystemBase: """Wrap different systems implementation into the right class. @@ -53,6 +63,9 @@ def wrap_system(system: IntoSystem) -> SystemBase: if HAVE_ASE and AseSystem.can_wrap(system): return AseSystem(system) + + if HAVE_PYSCF and PyscfSystem.can_wrap(system): + return PyscfSystem(system) if HAVE_CHEMFILES and ChemfilesSystem.can_wrap(system): return ChemfilesSystem(system) diff --git a/python/rascaline/rascaline/systems/pyscf.py b/python/rascaline/rascaline/systems/pyscf.py new file mode 100644 index 000000000..4e27f14f6 --- /dev/null +++ b/python/rascaline/rascaline/systems/pyscf.py @@ -0,0 +1,102 @@ +import warnings + +import numpy as np + +from .._c_api import c_uintptr_t +from .base import SystemBase + + +try: + import pyscf + HAVE_PYSCF = True +except ImportError: + HAVE_PYSCF = False + + +# copied from pyscf/data/elements.py +def _std_symbol_without_ghost(symb_or_chg): + '''For a given atom symbol (lower case or upper case) or charge, return the + standardized atom symbol + ''' + if isinstance(symb_or_chg, (str, unicode)): + symb_or_chg = str(symb_or_chg.upper()) + rawsymb = pyscf.data.elements._rm_digit(symb_or_chg) + if rawsymb in _ELEMENTS_UPPER: + return _ELEMENTS_UPPER[rawsymb] + elif len(rawsymb) > 1 and symb_or_chg[0] == 'X' and symb_or_chg[:2] != 'XE': + rawsymb = rawsymb[1:] # Remove the prefix X + return _ELEMENTS_UPPER[rawsymb] + elif len(rawsymb) > 5 and rawsymb[:5] == 'GHOST': + rawsymb = rawsymb[5:] # Remove the prefix GHOST + return _ELEMENTS_UPPER[rawsymb] + else: + raise RuntimeError('Unsupported atom symbol %s' % symb_or_chg) + else: + return pyscf.data.elements.ELEMENTS[symb_or_chg] + +class PyscfSystem(SystemBase): + """Implements :py:class:`rascaline.SystemBase` wrapping a `pyscf.gto.mole.Mole`_ or `pyscf.pbc.gto.cell.Cell`_. + + Since pyscf does not offer a neighbors list, this + implementation of system can only be used with ``use_native_system=True`` in + :py:func:`rascaline.calculators.CalculatorBase.compute`. + + Atomic species are assigned as the atomic number if the atom ``type`` is one + of the periodic table elements; or their opposite if they are ghost atoms. + (Pyscf does not seem to support anything else) + + .. _pyscf.gto.mole.Mole: https://pyscf.org/user/gto.html + .. _pyscf.pbc.gto.cell.Cell: https://pyscf.org/user/pbc/gto.html + """ + + @staticmethod + def can_wrap(o): + return isinstance(o, (pyscf.gto.mole.Mole,pyscf.pbc.gto.cell.Cell)) + + def __init__(self, frame): + """ + :param frame : `chemfiles.Frame`_ object object to be wrapped + in this ``ChemfilesSystem`` + """ + super().__init__() + if not isinstance(frame, (pyscf.gto.mole.Mole,pyscf.pbc.gto.cell.Cell)): + raise Exception("this class expects pyscf.gto.mole.Mole or pyscf.pbc.gto.cell.Cell objects") + + self._frame = frame + self._species = self._frame.atom_charges().copy() # dtype=int32 + for atm_i, species in enumerate(self._species): + if species == 0: + symb = self._frame.atom_symbol(atm_i) + chg = pyscf.data.elements.index(symb) + self._species[atm_i] = -chg + self.is_periodic = isinstance(self._frame, pyscf.pbc.gto.cell.Cell) + + def size(self): + return self._frame.natm + + def species(self): + return self._species + + def positions(self): + return self._frame.atom_coords() + + def cell(self): + if self.is_periodic: + return self._frame.a + else: + return np.zeros((3,3),float) + + def compute_neighbors(self, cutoff): + raise Exception( + "pyscf systems can only be used with 'use_native_system=True'" + ) + + def pairs(self): + raise Exception( + "pyscf systems can only be used with 'use_native_system=True'" + ) + + def pairs_containing(self, center): + raise Exception( + "pyscf systems can only be used with 'use_native_system=True'" + ) From 46364c7fe8935d186febbdbd419163c7939b021b Mon Sep 17 00:00:00 2001 From: Liam Marsh Date: Thu, 7 Dec 2023 17:26:06 +0100 Subject: [PATCH 25/30] fix flake8/black linting and tox integration --- .../rascaline/rascaline/systems/__init__.py | 9 ++-- python/rascaline/rascaline/systems/pyscf.py | 51 +++++++++---------- tox.ini | 2 + 3 files changed, 31 insertions(+), 31 deletions(-) diff --git a/python/rascaline/rascaline/systems/__init__.py b/python/rascaline/rascaline/systems/__init__.py index 2000b006a..6c3d08496 100644 --- a/python/rascaline/rascaline/systems/__init__.py +++ b/python/rascaline/rascaline/systems/__init__.py @@ -35,14 +35,15 @@ def __init__(self): if HAVE_PYSCF: IntoSystem.__doc__ += """ - - `pyscf.gto.mole.Mole`_ or `pyscf.pbc.gto.cell.Cell`_: pyscf' Frame type. There is no associated neighbor - list implementation, the system will only be usable with - ``use_native_system=True`` + - `pyscf.gto.mole.Mole`_ or `pyscf.pbc.gto.cell.Cell`_: pyscf' Frame type. + There is no associated neighbor list implementation, the system will only + be usable with ``use_native_system=True`` .. _pyscf.gto.mole.Mole: https://pyscf.org/user/gto.html .. _pyscf.pbc.gto.cell.Cell: https://pyscf.org/user/pbc/gto.html """ + def wrap_system(system: IntoSystem) -> SystemBase: """Wrap different systems implementation into the right class. @@ -63,7 +64,7 @@ def wrap_system(system: IntoSystem) -> SystemBase: if HAVE_ASE and AseSystem.can_wrap(system): return AseSystem(system) - + if HAVE_PYSCF and PyscfSystem.can_wrap(system): return PyscfSystem(system) diff --git a/python/rascaline/rascaline/systems/pyscf.py b/python/rascaline/rascaline/systems/pyscf.py index 4e27f14f6..2f0ac6d85 100644 --- a/python/rascaline/rascaline/systems/pyscf.py +++ b/python/rascaline/rascaline/systems/pyscf.py @@ -1,13 +1,11 @@ -import warnings - import numpy as np -from .._c_api import c_uintptr_t from .base import SystemBase try: import pyscf + HAVE_PYSCF = True except ImportError: HAVE_PYSCF = False @@ -15,27 +13,29 @@ # copied from pyscf/data/elements.py def _std_symbol_without_ghost(symb_or_chg): - '''For a given atom symbol (lower case or upper case) or charge, return the + """For a given atom symbol (lower case or upper case) or charge, return the standardized atom symbol - ''' - if isinstance(symb_or_chg, (str, unicode)): + """ + if isinstance(symb_or_chg, str): symb_or_chg = str(symb_or_chg.upper()) rawsymb = pyscf.data.elements._rm_digit(symb_or_chg) - if rawsymb in _ELEMENTS_UPPER: - return _ELEMENTS_UPPER[rawsymb] - elif len(rawsymb) > 1 and symb_or_chg[0] == 'X' and symb_or_chg[:2] != 'XE': + if rawsymb in pyscf.data.elements._ELEMENTS_UPPER: + return pyscf.data.elements._ELEMENTS_UPPER[rawsymb] + elif len(rawsymb) > 1 and symb_or_chg[0] == "X" and symb_or_chg[:2] != "XE": rawsymb = rawsymb[1:] # Remove the prefix X - return _ELEMENTS_UPPER[rawsymb] - elif len(rawsymb) > 5 and rawsymb[:5] == 'GHOST': + return pyscf.data.elements._ELEMENTS_UPPER[rawsymb] + elif len(rawsymb) > 5 and rawsymb[:5] == "GHOST": rawsymb = rawsymb[5:] # Remove the prefix GHOST - return _ELEMENTS_UPPER[rawsymb] + return pyscf.data.elements._ELEMENTS_UPPER[rawsymb] else: - raise RuntimeError('Unsupported atom symbol %s' % symb_or_chg) + raise RuntimeError("Unsupported atom symbol %s" % symb_or_chg) else: return pyscf.data.elements.ELEMENTS[symb_or_chg] + class PyscfSystem(SystemBase): - """Implements :py:class:`rascaline.SystemBase` wrapping a `pyscf.gto.mole.Mole`_ or `pyscf.pbc.gto.cell.Cell`_. + """Implements :py:class:`rascaline.SystemBase` wrapping a `pyscf.gto.mole.Mole`_ + or `pyscf.pbc.gto.cell.Cell`_. Since pyscf does not offer a neighbors list, this implementation of system can only be used with ``use_native_system=True`` in @@ -51,7 +51,7 @@ class PyscfSystem(SystemBase): @staticmethod def can_wrap(o): - return isinstance(o, (pyscf.gto.mole.Mole,pyscf.pbc.gto.cell.Cell)) + return isinstance(o, (pyscf.gto.mole.Mole, pyscf.pbc.gto.cell.Cell)) def __init__(self, frame): """ @@ -59,8 +59,11 @@ def __init__(self, frame): in this ``ChemfilesSystem`` """ super().__init__() - if not isinstance(frame, (pyscf.gto.mole.Mole,pyscf.pbc.gto.cell.Cell)): - raise Exception("this class expects pyscf.gto.mole.Mole or pyscf.pbc.gto.cell.Cell objects") + if not isinstance(frame, (pyscf.gto.mole.Mole, pyscf.pbc.gto.cell.Cell)): + raise Exception( + "this class expects pyscf.gto.mole.Mole" + + "or pyscf.pbc.gto.cell.Cell objects" + ) self._frame = frame self._species = self._frame.atom_charges().copy() # dtype=int32 @@ -84,19 +87,13 @@ def cell(self): if self.is_periodic: return self._frame.a else: - return np.zeros((3,3),float) + return np.zeros((3, 3), float) def compute_neighbors(self, cutoff): - raise Exception( - "pyscf systems can only be used with 'use_native_system=True'" - ) + raise Exception("pyscf systems can only be used with 'use_native_system=True'") def pairs(self): - raise Exception( - "pyscf systems can only be used with 'use_native_system=True'" - ) + raise Exception("pyscf systems can only be used with 'use_native_system=True'") def pairs_containing(self, center): - raise Exception( - "pyscf systems can only be used with 'use_native_system=True'" - ) + raise Exception("pyscf systems can only be used with 'use_native_system=True'") diff --git a/tox.ini b/tox.ini index dcd68c59f..b8bbb1dbb 100644 --- a/tox.ini +++ b/tox.ini @@ -59,6 +59,7 @@ deps = pytest-cov scipy sympy + pyscf wigners # TODO: add mops once support for windows available # git+https://github.com/lab-cosmo/mops @@ -122,6 +123,7 @@ description = Runs the doctests defined in any metatensor package deps = {[testenv]metatensor-core-requirement} ase + pyscf pytest commands = From 0bd5ab49f0e43b2f7185cc961efa3d2ee9a085b9 Mon Sep 17 00:00:00 2001 From: Liam Marsh Date: Fri, 8 Dec 2023 09:30:16 +0100 Subject: [PATCH 26/30] pyscf systems: changed management of length units, added tests [WIP] --- python/rascaline/rascaline/systems/pyscf.py | 16 ++- python/rascaline/tests/systems/pyscf.py | 105 ++++++++++++++++++++ 2 files changed, 119 insertions(+), 2 deletions(-) create mode 100644 python/rascaline/tests/systems/pyscf.py diff --git a/python/rascaline/rascaline/systems/pyscf.py b/python/rascaline/rascaline/systems/pyscf.py index 2f0ac6d85..f32e4d3df 100644 --- a/python/rascaline/rascaline/systems/pyscf.py +++ b/python/rascaline/rascaline/systems/pyscf.py @@ -44,6 +44,11 @@ class PyscfSystem(SystemBase): Atomic species are assigned as the atomic number if the atom ``type`` is one of the periodic table elements; or their opposite if they are ghost atoms. (Pyscf does not seem to support anything else) + + Please note that while pyscf uses Bohrs as length units internally, + we convert those back into Angströms for rascaline. + A pyscf object's "unit" attribute determines the units of the coordinates + given *to pyscf*, which are by default angströms. .. _pyscf.gto.mole.Mole: https://pyscf.org/user/gto.html .. _pyscf.pbc.gto.cell.Cell: https://pyscf.org/user/pbc/gto.html @@ -81,11 +86,18 @@ def species(self): return self._species def positions(self): - return self._frame.atom_coords() + return pyscf.data.nist.BOHR * self._frame.atom_coords() def cell(self): if self.is_periodic: - return self._frame.a + cell = self._frame.a + if self._frame.unit[0].lower() == 'a': + # assume angströms, we are good to go + return cell + else: + # assume bohrs, correct this + return pyscf.data.nist.BOHS * np.asarray(cell) + return else: return np.zeros((3, 3), float) diff --git a/python/rascaline/tests/systems/pyscf.py b/python/rascaline/tests/systems/pyscf.py new file mode 100644 index 000000000..0d19f72a0 --- /dev/null +++ b/python/rascaline/tests/systems/pyscf.py @@ -0,0 +1,105 @@ +import numpy as np +import pytest + +from rascaline.systems import AseSystem + + +pyscf = pytest.importorskip("pyscf") + + +@pytest.fixture +def system(): + atoms = pyscf.M( + atom = "C 0 0 0; O 0 0 1.4; O 0 0 -1.6", + ) + #atoms.pbc = [False, False, False] + return PyscfSystem(atoms) + + +def test_system_implementation(system): + assert system.size() == 3 + assert np.all(system.species() == [6, 8, 8]) + + positions = [ + (0, 0, 0), + (0, 0, 1.4), + (0, 0, -1.6), + ] + assert np.all(system.positions() == positions) + assert np.all(system.cell() == [[0, 0, 0], [0, 0, 0], [0, 0, 0]]) + + +def test_pairs(system): + system.compute_neighbors(1.5) + pairs = system.pairs() + + assert len(pairs) == 1 + assert pairs[0][:2] == (0, 1) + assert pairs[0][2] == 1.4 + assert np.all(pairs[0][3] == [0, 0, 1.4]) + + system.compute_neighbors(2.5) + pairs = system.pairs() + assert len(pairs) == 2 + assert pairs[0][:2] == (0, 1) + assert pairs[0][2] == 1.4 + assert np.all(pairs[0][3] == [0, 0, 1.4]) + + assert pairs[1][:2] == (0, 2) + assert pairs[1][2] == 1.6 + assert np.all(pairs[1][3] == [0, 0, -1.6]) + + system.compute_neighbors(3.5) + pairs = system.pairs() + assert len(pairs) == 3 + assert pairs[0][:2] == (0, 1) + assert pairs[0][2] == 1.4 + assert np.all(pairs[0][3] == [0, 0, 1.4]) + + assert pairs[1][:2] == (0, 2) + assert pairs[1][2] == 1.6 + assert np.all(pairs[1][3] == [0, 0, -1.6]) + + assert pairs[2][:2] == (1, 2) + assert pairs[2][2] == 3.0 + assert np.all(pairs[2][3] == [0, 0, -3.0]) + + +def test_pairs_containing(system): + system.compute_neighbors(1.5) + pairs = system.pairs_containing(0) + assert len(pairs) == 1 + assert pairs[0][:2] == (0, 1) + + pairs = system.pairs_containing(1) + assert len(pairs) == 1 + assert pairs[0][:2] == (0, 1) + + pairs = system.pairs_containing(2) + assert len(pairs) == 0 + + system.compute_neighbors(3.5) + pairs = system.pairs_containing(0) + assert len(pairs) == 2 + assert pairs[0][:2] == (0, 1) + assert pairs[1][:2] == (0, 2) + + pairs = system.pairs_containing(1) + assert len(pairs) == 2 + assert pairs[0][:2] == (0, 1) + assert pairs[1][:2] == (1, 2) + + pairs = system.pairs_containing(2) + assert len(pairs) == 2 + assert pairs[0][:2] == (0, 2) + assert pairs[1][:2] == (1, 2) + + +def test_pbc_data(): + #import pyscf.pbo + atoms = pyscf.pbo.gto.Cell( + atom = "H 0 0 0; H 1 1 1", + a = np.array([[2,0,0],[0,2,1],[0,0,2]],dtype=float) + ) + ras_sys = PyscfSystem(atoms) + assert ras_sys.positions() == np.array([[0,0,0],[1,1,1]],dtype=float) From dba772d696c3f8cd90f7f9cfa786f470f561e428 Mon Sep 17 00:00:00 2001 From: Liam Marsh Date: Fri, 8 Dec 2023 16:58:31 +0100 Subject: [PATCH 27/30] fixed tests --- python/rascaline/rascaline/systems/pyscf.py | 27 ++--- python/rascaline/tests/systems/pyscf.py | 104 +++++--------------- 2 files changed, 39 insertions(+), 92 deletions(-) diff --git a/python/rascaline/rascaline/systems/pyscf.py b/python/rascaline/rascaline/systems/pyscf.py index f32e4d3df..b4e8b8bb7 100644 --- a/python/rascaline/rascaline/systems/pyscf.py +++ b/python/rascaline/rascaline/systems/pyscf.py @@ -21,7 +21,7 @@ def _std_symbol_without_ghost(symb_or_chg): rawsymb = pyscf.data.elements._rm_digit(symb_or_chg) if rawsymb in pyscf.data.elements._ELEMENTS_UPPER: return pyscf.data.elements._ELEMENTS_UPPER[rawsymb] - elif len(rawsymb) > 1 and symb_or_chg[0] == "X" and symb_or_chg[:2] != "XE": + elif len(rawsymb) > 1 and (symb_or_chg[0] == "X" and symb_or_chg[:2] != "XE"): rawsymb = rawsymb[1:] # Remove the prefix X return pyscf.data.elements._ELEMENTS_UPPER[rawsymb] elif len(rawsymb) > 5 and rawsymb[:5] == "GHOST": @@ -34,17 +34,19 @@ def _std_symbol_without_ghost(symb_or_chg): class PyscfSystem(SystemBase): - """Implements :py:class:`rascaline.SystemBase` wrapping a `pyscf.gto.mole.Mole`_ - or `pyscf.pbc.gto.cell.Cell`_. + """Implements :py:class:`rascaline.SystemBase` wrapping a + `pyscf.gto.mole.Mole`_ or `pyscf.pbc.gto.cell.Cell`_. Since pyscf does not offer a neighbors list, this - implementation of system can only be used with ``use_native_system=True`` in + implementation of system can only be used with + ``use_native_system=True`` in :py:func:`rascaline.calculators.CalculatorBase.compute`. - Atomic species are assigned as the atomic number if the atom ``type`` is one - of the periodic table elements; or their opposite if they are ghost atoms. + Atomic species are assigned as the atomic number if the atom ``type`` is + one of the periodic table elements; or their opposite if they are + ghost atoms. (Pyscf does not seem to support anything else) - + Please note that while pyscf uses Bohrs as length units internally, we convert those back into Angströms for rascaline. A pyscf object's "unit" attribute determines the units of the coordinates @@ -56,7 +58,11 @@ class PyscfSystem(SystemBase): @staticmethod def can_wrap(o): - return isinstance(o, (pyscf.gto.mole.Mole, pyscf.pbc.gto.cell.Cell)) + # assumption: if we have a periodic system, then pyscf.pbc is defined + if hasattr(pyscf, "pbc"): + return isinstance(o, (pyscf.gto.mole.Mole, pyscf.pbc.gto.cell.Cell)) + else: + return isinstance(o, pyscf.gto.mole.Mole) def __init__(self, frame): """ @@ -64,7 +70,7 @@ def __init__(self, frame): in this ``ChemfilesSystem`` """ super().__init__() - if not isinstance(frame, (pyscf.gto.mole.Mole, pyscf.pbc.gto.cell.Cell)): + if not self.can_wrap(frame): raise Exception( "this class expects pyscf.gto.mole.Mole" + "or pyscf.pbc.gto.cell.Cell objects" @@ -91,13 +97,12 @@ def positions(self): def cell(self): if self.is_periodic: cell = self._frame.a - if self._frame.unit[0].lower() == 'a': + if self._frame.unit[0].lower() == "a": # assume angströms, we are good to go return cell else: # assume bohrs, correct this return pyscf.data.nist.BOHS * np.asarray(cell) - return else: return np.zeros((3, 3), float) diff --git a/python/rascaline/tests/systems/pyscf.py b/python/rascaline/tests/systems/pyscf.py index 0d19f72a0..875159f60 100644 --- a/python/rascaline/tests/systems/pyscf.py +++ b/python/rascaline/tests/systems/pyscf.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from rascaline.systems import AseSystem +from rascaline.systems import PyscfSystem pyscf = pytest.importorskip("pyscf") @@ -10,9 +10,9 @@ @pytest.fixture def system(): atoms = pyscf.M( - atom = "C 0 0 0; O 0 0 1.4; O 0 0 -1.6", + atom="C 0 0 0; O 0 0 1.4; O 0 0 -1.6", ) - #atoms.pbc = [False, False, False] + # atoms.pbc = [False, False, False] return PyscfSystem(atoms) @@ -20,86 +20,28 @@ def test_system_implementation(system): assert system.size() == 3 assert np.all(system.species() == [6, 8, 8]) - positions = [ - (0, 0, 0), - (0, 0, 1.4), - (0, 0, -1.6), - ] - assert np.all(system.positions() == positions) + positions = np.array( + [ + (0, 0, 0), + (0, 0, 1.4), + (0, 0, -1.6), + ] + ) + print(system.positions(), positions) + assert np.allclose(system.positions(), positions, rtol=1e-14) assert np.all(system.cell() == [[0, 0, 0], [0, 0, 0], [0, 0, 0]]) -def test_pairs(system): - system.compute_neighbors(1.5) - pairs = system.pairs() - - assert len(pairs) == 1 - assert pairs[0][:2] == (0, 1) - assert pairs[0][2] == 1.4 - assert np.all(pairs[0][3] == [0, 0, 1.4]) - - system.compute_neighbors(2.5) - pairs = system.pairs() - assert len(pairs) == 2 - assert pairs[0][:2] == (0, 1) - assert pairs[0][2] == 1.4 - assert np.all(pairs[0][3] == [0, 0, 1.4]) - - assert pairs[1][:2] == (0, 2) - assert pairs[1][2] == 1.6 - assert np.all(pairs[1][3] == [0, 0, -1.6]) - - system.compute_neighbors(3.5) - pairs = system.pairs() - assert len(pairs) == 3 - assert pairs[0][:2] == (0, 1) - assert pairs[0][2] == 1.4 - assert np.all(pairs[0][3] == [0, 0, 1.4]) - - assert pairs[1][:2] == (0, 2) - assert pairs[1][2] == 1.6 - assert np.all(pairs[1][3] == [0, 0, -1.6]) - - assert pairs[2][:2] == (1, 2) - assert pairs[2][2] == 3.0 - assert np.all(pairs[2][3] == [0, 0, -3.0]) - - -def test_pairs_containing(system): - system.compute_neighbors(1.5) - pairs = system.pairs_containing(0) - assert len(pairs) == 1 - assert pairs[0][:2] == (0, 1) - - pairs = system.pairs_containing(1) - assert len(pairs) == 1 - assert pairs[0][:2] == (0, 1) - - pairs = system.pairs_containing(2) - assert len(pairs) == 0 - - system.compute_neighbors(3.5) - pairs = system.pairs_containing(0) - assert len(pairs) == 2 - assert pairs[0][:2] == (0, 1) - assert pairs[1][:2] == (0, 2) - - pairs = system.pairs_containing(1) - assert len(pairs) == 2 - assert pairs[0][:2] == (0, 1) - assert pairs[1][:2] == (1, 2) - - pairs = system.pairs_containing(2) - assert len(pairs) == 2 - assert pairs[0][:2] == (0, 2) - assert pairs[1][:2] == (1, 2) - - def test_pbc_data(): - #import pyscf.pbo - atoms = pyscf.pbo.gto.Cell( - atom = "H 0 0 0; H 1 1 1", - a = np.array([[2,0,0],[0,2,1],[0,0,2]],dtype=float) - ) + import pyscf.pbc + + atoms = pyscf.pbc.gto.Cell( + atom="H 0 0 0; H 1 1 1", + a=np.array([[2, 0, 0], [0, 2, 1], [0, 0, 2]], dtype=float), + ).build() ras_sys = PyscfSystem(atoms) - assert ras_sys.positions() == np.array([[0,0,0],[1,1,1]],dtype=float) + assert np.allclose( + ras_sys.positions(), + np.array([[0, 0, 0], [1, 1, 1]], dtype=float), + rtol=1e-14, + ) From 77739718e3d53c0f4286fb15a9ecf4777455cbe5 Mon Sep 17 00:00:00 2001 From: Liam Marsh Date: Mon, 11 Dec 2023 10:46:43 +0100 Subject: [PATCH 28/30] fix dependency management for tox --- tox.ini | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tox.ini b/tox.ini index b8bbb1dbb..0e8594950 100644 --- a/tox.ini +++ b/tox.ini @@ -47,9 +47,10 @@ commands = [testenv:all-deps] + description = - Run Python unit tests with all dependencies installed (ase & chemfiles are optional - dependencies) + Run Python unit tests with all dependencies installed (ase, pyscf, + and chemfiles are optional dependencies) deps = {[testenv]metatensor-core-requirement} ase @@ -59,11 +60,10 @@ deps = pytest-cov scipy sympy - pyscf + pyscf;platform_system!="Windows" wigners - # TODO: add mops once support for windows available - # git+https://github.com/lab-cosmo/mops - + # 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} From 24363ad5d9ad9790afcf1600f793931d64a567c0 Mon Sep 17 00:00:00 2001 From: Liam Marsh Date: Mon, 11 Dec 2023 14:03:36 +0100 Subject: [PATCH 29/30] fix tox dependencies again --- tox.ini | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tox.ini b/tox.ini index 0e8594950..0eec76844 100644 --- a/tox.ini +++ b/tox.ini @@ -47,7 +47,7 @@ commands = [testenv:all-deps] - +# note: platform_system can be "Linux","Darwin", or "Windows". description = Run Python unit tests with all dependencies installed (ase, pyscf, and chemfiles are optional dependencies) @@ -123,7 +123,7 @@ description = Runs the doctests defined in any metatensor package deps = {[testenv]metatensor-core-requirement} ase - pyscf + pyscf;platform_system!="Windows" pytest commands = From f597fbb862437294c00b8cb75825463c84ed0489 Mon Sep 17 00:00:00 2001 From: Liam Marsh Date: Tue, 12 Dec 2023 12:19:45 +0100 Subject: [PATCH 30/30] pyscf systems: add explicit test for unit management --- python/rascaline/rascaline/systems/pyscf.py | 2 +- python/rascaline/tests/systems/pyscf.py | 24 +++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/python/rascaline/rascaline/systems/pyscf.py b/python/rascaline/rascaline/systems/pyscf.py index b4e8b8bb7..868a2b314 100644 --- a/python/rascaline/rascaline/systems/pyscf.py +++ b/python/rascaline/rascaline/systems/pyscf.py @@ -102,7 +102,7 @@ def cell(self): return cell else: # assume bohrs, correct this - return pyscf.data.nist.BOHS * np.asarray(cell) + return pyscf.data.nist.BOHR * np.asarray(cell) else: return np.zeros((3, 3), float) diff --git a/python/rascaline/tests/systems/pyscf.py b/python/rascaline/tests/systems/pyscf.py index 875159f60..35d39d99e 100644 --- a/python/rascaline/tests/systems/pyscf.py +++ b/python/rascaline/tests/systems/pyscf.py @@ -45,3 +45,27 @@ def test_pbc_data(): np.array([[0, 0, 0], [1, 1, 1]], dtype=float), rtol=1e-14, ) + + +def test_explicit_units(): + import pyscf.pbc + + cell = np.array([[2, 0, 0], [0, 2, 1], [0, 0, 2]], dtype=float) + + at1 = pyscf.pbc.gto.Cell( + atom="H 0 0 0; H 1 1 1", + a=cell, + unit="Angstrom", + ).build() + at2 = pyscf.pbc.gto.Cell( + atom=[("H", at1.atom_coord(0)), ("H", at1.atom_coord(1))], + a=cell / pyscf.data.nist.BOHR, + unit="Bohr", + ).build() + at1 = PyscfSystem(at1) + at2 = PyscfSystem(at2) + + assert np.allclose(at1.positions(), at2.positions()) + assert np.allclose(at1.positions(), np.array([[0, 0, 0], [1, 1, 1]], dtype=float)) + assert np.allclose(at1.cell(), at2.cell()) + assert np.allclose(at1.cell(), cell)