diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml index e77afa378..f24c30ed2 100644 --- a/.github/workflows/build-wheels.yml +++ b/.github/workflows/build-wheels.yml @@ -24,6 +24,79 @@ concurrency: jobs: + build-core-wheels: + runs-on: ${{ matrix.os }} + name: ${{ matrix.name }} + strategy: + matrix: + include: + - name: x86_64 Linux + os: ubuntu-24.04 + rust-target: x86_64-unknown-linux-gnu + cibw-arch: x86_64 + - name: arm64 Linux + os: ubuntu-24.04-arm + rust-target: aarch64-unknown-linux-gnu + cibw-arch: aarch64 + - name: arm64 macOS + os: macos-15 + rust-target: aarch64-apple-darwin + cibw-arch: arm64 + - name: x86_64 Windows + os: windows-2022 + rust-target: x86_64-pc-windows-msvc + cibw-arch: AMD64 + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + + - name: setup rust + uses: dtolnay/rust-toolchain@master + with: + toolchain: stable + target: ${{ matrix.rust-target }} + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.14" + + - name: install dependencies + run: python -m pip install cibuildwheel twine + + - name: build custom manylinux docker images + if: matrix.os == 'ubuntu-24.04' || matrix.os == 'ubuntu-24.04-arm' + run: | + docker buildx build \ + -t rustc-manylinux_2_28_${{ matrix.cibw-arch }} \ + python/scripts/rustc-manylinux_2_28_${{ matrix.cibw-arch }} + + - name: build metatomic-core wheel + run: python -m cibuildwheel python/metatomic_core + env: + CIBW_BUILD: cp311-* + CIBW_SKIP: "*musllinux*" + CIBW_ARCHS: ${{ matrix.cibw-arch }} + CIBW_BUILD_VERBOSITY: 1 + CIBW_MANYLINUX_X86_64_IMAGE: rustc-manylinux_2_28_x86_64 + CIBW_MANYLINUX_AARCH64_IMAGE: rustc-manylinux_2_28_aarch64 + CIBW_ENVIRONMENT: > + MACOSX_DEPLOYMENT_TARGET=11 + # do not complain for missing shared libraries provided by other packages + CIBW_REPAIR_WHEEL_COMMAND_MACOS: | + delocate-wheel --ignore-missing-dependencies --require-archs {delocate_archs} -w {dest_dir} -v {wheel} + CIBW_REPAIR_WHEEL_COMMAND_LINUX: | + auditwheel repair --exclude libmetatensor.so -w {dest_dir} {wheel} + + - name: check wheels with twine + run: twine check wheelhouse/*.whl + + - uses: actions/upload-artifact@v7 + with: + name: core-wheel-${{ matrix.os }}-${{ matrix.cibw-arch }} + path: ./wheelhouse/*.whl + build-torch-wheels: runs-on: ${{ matrix.os }} name: ${{ matrix.name }} (torch v${{ matrix.torch-version }}) @@ -89,8 +162,8 @@ jobs: if: matrix.os == 'ubuntu-24.04' || matrix.os == 'ubuntu-24.04-arm' run: | docker buildx build \ - -t gcc11-manylinux_2_28_${{ matrix.cibw-arch }} \ - python/scripts/gcc11-manylinux_2_28_${{ matrix.cibw-arch }} + -t rustc-manylinux_2_28_${{ matrix.cibw-arch }} \ + python/scripts/rustc-manylinux_2_28_${{ matrix.cibw-arch }} - name: build metatomic-torch wheel run: python -m cibuildwheel python/metatomic_torch @@ -99,10 +172,19 @@ jobs: CIBW_SKIP: "*musllinux*" CIBW_ARCHS: ${{ matrix.cibw-arch }} CIBW_BUILD_VERBOSITY: 1 - CIBW_MANYLINUX_X86_64_IMAGE: gcc11-manylinux_2_28_x86_64 - CIBW_MANYLINUX_AARCH64_IMAGE: gcc11-manylinux_2_28_aarch64 + CIBW_MANYLINUX_X86_64_IMAGE: rustc-manylinux_2_28_x86_64 + CIBW_MANYLINUX_AARCH64_IMAGE: rustc-manylinux_2_28_aarch64 + # METATOMIC_NO_LOCAL_DEPS is set to 1 when building a tag of + # metatomic-torch, which will force to use the version of + # metatomic-core already released on PyPI. Otherwise, this will use + # the version of metatomic-core from git checkout (in case there are + # unreleased breaking changes). + # + # This means that when releasing a breaking change in metatomic-core, + # the full release should be available on PyPI before pushing the new + # metatomic-torch tag. CIBW_ENVIRONMENT: > - METATOMIC_NO_LOCAL_DEPS=1 + METATOMIC_NO_LOCAL_DEPS=${{ startsWith(github.ref, 'refs/tags/metatomic-torch-v') && '1' || '0' }} METATOMIC_TORCH_BUILD_WITH_TORCH_VERSION=${{ matrix.torch-version }}.* PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu MACOSX_DEPLOYMENT_TARGET=11 @@ -205,6 +287,9 @@ jobs: - name: install dependencies run: python -m pip install wheel build twine + - name: build metatomic-core sdist + run: python -m build python/metatomic_core --sdist --outdir=dist/ + - name: build metatomic-torch sdist run: python -m build python/metatomic_torch --sdist --outdir=dist/ @@ -222,6 +307,7 @@ jobs: - name: create C++ tarballs run: | + ./scripts/package-core.sh dist/cxx/ ./scripts/package-torch.sh dist/cxx/ - uses: actions/upload-artifact@v7 @@ -234,12 +320,19 @@ jobs: merge-and-release: name: Merge and release wheels/sdists - needs: [merge-torch-wheels, build-others] + needs: [build-core-wheels, merge-torch-wheels, build-others] runs-on: ubuntu-24.04 permissions: contents: write pull-requests: write steps: + - name: Download metatomic-core wheels + uses: actions/download-artifact@v8 + with: + path: wheels + pattern: core-wheel-* + merge-multiple: true + - name: Download metatomic-torch wheels uses: actions/download-artifact@v8 with: @@ -268,6 +361,17 @@ jobs: name: wheels description: ⚙️ Download Python wheels for this pull-request (you can install these with pip) + - name: upload to GitHub release (metatomic-core) + if: startsWith(github.ref, 'refs/tags/metatomic-core-v') + uses: softprops/action-gh-release@v3 + with: + files: | + wheels/cxx/metatomic-core-cxx-*.tar.gz + wheels/metatomic_core-* + prerelease: ${{ contains(github.ref, '-rc') }} + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: upload to GitHub release (metatomic) if: startsWith(github.ref, 'refs/tags/metatomic-v') uses: softprops/action-gh-release@v3 @@ -311,8 +415,8 @@ jobs: test-build-external: # This checks building the wheels with external libraries. This setup is - # mainly used for the conda packages metatensor-*-python, which use the - # libmetatensor-* conda packages to provide the native code. + # mainly used for the conda packages metatomic-*-python, which use the + # libmetatomic-* conda packages to provide the native code. runs-on: ${{ matrix.os }} name: External libraries / ${{ matrix.os }} defaults: @@ -338,6 +442,12 @@ jobs: with: fetch-depth: 0 + - name: setup rust + uses: dtolnay/rust-toolchain@master + with: + toolchain: stable + target: ${{ matrix.rust-target }} + - name: setup Python uses: actions/setup-python@v7 with: @@ -356,9 +466,9 @@ jobs: - name: setup libmetatensor run: | - curl --location -O https://github.com/metatensor/metatensor/releases/download/metatensor-core-v0.2.3/metatensor-core-cxx-0.2.3.tar.gz - tar xf metatensor-core-cxx-0.2.3.tar.gz - cmake -B build-metatensor -S metatensor-core-cxx-0.2.3 \ + curl --location -O https://github.com/metatensor/metatensor/releases/download/metatensor-core-v0.2.4/metatensor-core-cxx-0.2.4.tar.gz + tar xf metatensor-core-cxx-0.2.4.tar.gz + cmake -B build-metatensor -S metatensor-core-cxx-0.2.4 \ -DMETATENSOR_INSTALL_BOTH_STATIC_SHARED=OFF \ -DCMAKE_INSTALL_PREFIX=$CMAKE_PREFIX_PATH \ -DCMAKE_BUILD_TYPE=Debug @@ -375,6 +485,15 @@ jobs: cmake --build build-metatensor-torch --config Debug cmake --install build-metatensor-torch --config Debug + - name: build libmetatomic + run: | + cmake -B build-metatomic-core -S metatomic-core \ + -DMETATOMIC_INSTALL_BOTH_STATIC_SHARED=OFF \ + -DCMAKE_INSTALL_PREFIX=$CMAKE_PREFIX_PATH \ + -DCMAKE_BUILD_TYPE=Debug + cmake --build build-metatomic-core --config Debug + cmake --install build-metatomic-core --config Debug + - name: build libmetatomic-torch run: | cmake -B build-metatomic-torch -S metatomic-torch \ @@ -383,6 +502,14 @@ jobs: cmake --build build-metatomic-torch --config Debug cmake --install build-metatomic-torch --config Debug + - name: build metatomic-core wheels + run: | + python -m build python/metatomic_core --wheel --outdir=dist/ + # check that the wheel is using an external library + unzip -l dist/metatomic_core*.whl | grep "_external.py" + env: + METATOMIC_CORE_PYTHON_USE_EXTERNAL_LIB: "ON" + - name: build metatomic-torch wheels run: | pip install build diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml new file mode 100644 index 000000000..3984f879a --- /dev/null +++ b/.github/workflows/python-tests.yml @@ -0,0 +1,95 @@ +name: Python tests + +on: + push: + branches: [main] + pull_request: + # Check all PR + +concurrency: + group: python-tests-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +jobs: + python-tests: + runs-on: ${{ matrix.os }} + name: ${{ matrix.os }} / Python ${{ matrix.python-version }} / Torch ${{ matrix.torch-version }} + strategy: + matrix: + include: + - os: ubuntu-24.04 + python-version: "3.11" + torch-version: "2.3" + numpy-version-pin: "<2.0" + - os: ubuntu-24.04 + python-version: "3.11" + torch-version: "2.13" + - os: ubuntu-24.04 + # TorchScript is no longer supported in Python 3.14 + # so we keep a test with 3.13 to make sure this doesn't break + python-version: "3.13" + torch-version: "2.13" + - os: ubuntu-24.04 + python-version: "3.14" + torch-version: "2.13" + - os: macos-15 + python-version: "3.14" + torch-version: "2.13" + - os: windows-2022 + python-version: "3.14" + torch-version: "2.13" + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + + - name: setup Python + uses: actions/setup-python@v7 + with: + python-version: ${{ matrix.python-version }} + + - name: setup rust + uses: dtolnay/rust-toolchain@master + with: + toolchain: stable + + - name: Setup sccache + if: ${{ !env.ACT }} + uses: mozilla-actions/sccache-action@v0.0.10 + with: + version: "v0.10.0" + + - name: setup MSVC command prompt + uses: ilammy/msvc-dev-cmd@v1 + + - name: Setup sccache environnement variables + if: ${{ !env.ACT }} + run: | + echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV + echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV + echo "CMAKE_C_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV + echo "CMAKE_CXX_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV + + - name: install tests dependencies + run: | + python -m pip install --upgrade pip + python -m pip install tox coverage + + - name: run tests + run: tox -e lint,core-tests,torch-tests,docs-tests + env: + PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu + METATOMIC_TESTS_TORCH_VERSION: ${{ matrix.torch-version }} + + - name: combine Python coverage files + shell: bash + run: | + coverage combine .tox/*/.coverage + coverage xml + + - name: upload to codecov.io + uses: codecov/codecov-action@v7 + with: + fail_ci_if_error: true + files: coverage.xml + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/.github/workflows/rust-tests.yml b/.github/workflows/rust-tests.yml new file mode 100644 index 000000000..5951f68bc --- /dev/null +++ b/.github/workflows/rust-tests.yml @@ -0,0 +1,189 @@ +name: Rust tests + +on: + push: + branches: [main] + pull_request: + # Check all PR + +concurrency: + group: rust-tests-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +jobs: + rust-tests: + name: ${{ matrix.os }} / Rust ${{ matrix.rust-version }}${{ matrix.extra-name }} + runs-on: ${{ matrix.os }} + container: ${{ matrix.container }} + defaults: + run: + shell: "bash" + env: + CMAKE_CXX_COMPILER: ${{ matrix.cxx }} + CMAKE_C_COMPILER: ${{ matrix.cc }} + CMAKE_GENERATOR: ${{ matrix.cmake-generator }} + strategy: + matrix: + include: + # test our MSRV + - os: ubuntu-24.04 + rust-version: 1.88 + rust-target: x86_64-unknown-linux-gnu + cxx: g++ + cc: gcc + cargo: cargo + cmake-generator: Unix Makefiles + + # check the build on a stock Ubuntu 22.04, which uses cmake 3.22, and + # using cargo/rustc from APT + - os: ubuntu-24.04 + rust-version: from APT + container: ubuntu:22.04 + rust-target: x86_64-unknown-linux-gnu + extra-name: ", cmake 3.22" + cxx: g++ + cc: gcc + cargo: cargo-1.89 + cmake-generator: Unix Makefiles + + - os: macos-15 + rust-version: stable + rust-target: aarch64-apple-darwin + cxx: clang++ + cc: clang + cargo: cargo + cmake-generator: Unix Makefiles + + - os: windows-2022 + rust-version: stable + rust-target: x86_64-pc-windows-msvc + extra-name: " / MSVC" + cxx: cl.exe + cc: cl.exe + cargo: cargo + cmake-generator: Visual Studio 17 2022 + + - os: windows-2022 + rust-version: stable + rust-target: x86_64-pc-windows-gnu + extra-name: " / MinGW" + cxx: g++.exe + cc: gcc.exe + cargo: cargo + cmake-generator: MinGW Makefiles + steps: + - name: install dependencies in container + if: matrix.container == 'ubuntu:22.04' + run: | + apt update + apt install -y software-properties-common + apt install -y cmake make gcc g++ git curl python3-venv cargo-1.89 + + # for some reason, cargo-1.89 from APT tries to find `rustdoc` and + # not `rustdoc-1.89`, so we force it to use the correct one + echo "RUSTDOC=rustdoc-1.89" >> "$GITHUB_ENV" + + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: Configure git safe directory + if: matrix.container == 'ubuntu:22.04' + run: git config --global --add safe.directory /__w/metatomic/metatomic + + - name: setup rust + uses: dtolnay/rust-toolchain@master + if: matrix.container == null + with: + toolchain: ${{ matrix.rust-version }} + target: ${{ matrix.rust-target }} + + - name: setup Python + uses: actions/setup-python@v6 + if: matrix.container == null + with: + python-version: "3.14" + + - name: install valgrind + if: matrix.do-valgrind + run: | + sudo apt-get update + sudo apt-get install -y valgrind + + - name: Setup sccache + if: ${{ !env.ACT }} + uses: mozilla-actions/sccache-action@v0.0.10 + with: + version: "v0.15.0" + + - name: Setup sccache environnement variables + if: ${{ !env.ACT }} + run: | + echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV + echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV + echo "CMAKE_C_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV + echo "CMAKE_CXX_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV + + - name: run tests + env: + RUST_BACKTRACE: full + run: | + ${{ matrix.cargo }} test --package metatomic-core --target ${{ matrix.rust-target }} + + - name: check that the header was already up to date + run: | + git diff --exit-code + + # check that the C API declarations are correctly documented and used + prevent-bitrot: + runs-on: ubuntu-24.04 + name: check C API declarations + steps: + - uses: actions/checkout@v6 + + - name: setup Python + uses: actions/setup-python@v6 + with: + python-version: "3.14" + + - name: install python dependencies + run: | + pip install pycparser + + - name: check that C API functions are all documented + run: | + python scripts/check-c-api-docs.py + + # make sure no debug print stays in the code + check-debug-prints: + runs-on: ubuntu-24.04 + name: check leftover debug print + + steps: + - uses: actions/checkout@v6 + + - name: install ripgrep + run: | + wget https://github.com/BurntSushi/ripgrep/releases/download/13.0.0/ripgrep-13.0.0-x86_64-unknown-linux-musl.tar.gz + tar xf ripgrep-13.0.0-x86_64-unknown-linux-musl.tar.gz + echo "$(pwd)/ripgrep-13.0.0-x86_64-unknown-linux-musl" >> $GITHUB_PATH + + - name: check for leftover dbg! + run: | + # use ripgrep (rg) to check for instances of `dbg!` in rust files. + # rg will return 1 if it fails to find a match, so we invert it again + # with the `!` builtin to get the error/success in CI + + ! rg "dbg!" --type=rust --quiet + + - name: check for leftover \#include + run: | + ! rg "" --iglob "\!metatomic-core/tests/cpp/external/catch/catch.hpp" --quiet + + - name: check for leftover std::cout + run: | + ! rg "cout" --iglob "\!metatomic-core/tests/cpp/external/catch/catch.hpp" --quiet + + - name: check for leftover std::cerr + run: | + ! rg "cerr" --iglob "\!metatomic-core/tests/cpp/external/catch/catch.hpp" --quiet diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index 3c4a90973..9290638e1 100644 --- a/.github/workflows/torch-tests.yml +++ b/.github/workflows/torch-tests.yml @@ -13,81 +13,91 @@ concurrency: jobs: tests: runs-on: ${{ matrix.os }} - name: ${{ matrix.os }} / Python ${{ matrix.python-version }} / Torch ${{ matrix.torch-version }} + name: ${{ matrix.os }} / Torch ${{ matrix.torch-version }}${{ matrix.extra-name }} + container: ${{ matrix.container }} strategy: matrix: include: - os: ubuntu-24.04 - python-version: "3.10" - torch-version: "2.3" - - os: ubuntu-24.04 - python-version: "3.10" torch-version: "2.13" - - os: ubuntu-24.04 - # Keep a building with Python 3.13 since TorchScript is deprecated - # in Python 3.14 - python-version: "3.13" - torch-version: "2.13" - - os: ubuntu-24.04 python-version: "3.14" - torch-version: "2.13" + cargo-test-flags: --release + do-valgrind: true + + # check the build on a stock Ubuntu 22.04, which uses cmake 3.22 + - os: ubuntu-24.04 + container: ubuntu:22.04 + extra-name: ", cmake 3.22" + torch-version: "2.3" + cargo-test-flags: "" + - os: macos-15 - python-version: "3.14" torch-version: "2.13" - - os: windows-2022 python-version: "3.14" + cargo-test-flags: --release + + - os: windows-2022 torch-version: "2.13" + python-version: "3.14" + cargo-test-flags: --release steps: + - name: install dependencies in container + if: matrix.container == 'ubuntu:22.04' + env: + # python3.11 pulls in tzdata, which otherwise asks for a timezone + DEBIAN_FRONTEND: noninteractive + run: | + apt update + apt install -y software-properties-common + add-apt-repository ppa:deadsnakes/ppa + apt install -y cmake make gcc g++ git curl python3.11 python3.11-venv + + update-alternatives --install /usr/local/bin/python python /usr/bin/python3.11 1 + - uses: actions/checkout@v7 with: fetch-depth: 0 - name: setup Python - uses: actions/setup-python@v7 + uses: actions/setup-python@v6 + if: matrix.container == null with: python-version: ${{ matrix.python-version }} + - name: Configure git safe directory + if: matrix.container == 'ubuntu:22.04' + run: git config --global --add safe.directory /__w/metatomic/metatomic + + - name: setup rust + uses: dtolnay/rust-toolchain@master + with: + toolchain: stable + + - name: install valgrind + if: matrix.do-valgrind + run: | + sudo apt-get update + sudo apt-get install -y valgrind + - name: Setup sccache + if: ${{ !env.ACT }} uses: mozilla-actions/sccache-action@v0.0.11 with: version: "v0.10.0" - - name: setup MSVC command prompt - uses: ilammy/msvc-dev-cmd@v1 - - name: Setup sccache environnement variables + if: ${{ !env.ACT }} run: | echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV echo "CMAKE_C_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV echo "CMAKE_CXX_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV - - name: install tests dependencies - run: | - python -m pip install --upgrade pip - python -m pip install tox coverage - - - name: run Python tests - run: tox -e lint,torch-tests,docs-tests + - name: run TorchScript C++ tests + run: cargo test --package metatomic-torch ${{ matrix.cargo-test-flags }} env: + # Use the CPU only version of torch when building/running the code PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu METATOMIC_TESTS_TORCH_VERSION: ${{ matrix.torch-version }} - - - name: run C++ tests - run: tox -e torch-tests-cxx,torch-install-tests-cxx - env: - PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu - METATOMIC_TESTS_TORCH_VERSION: ${{ matrix.torch-version }} - - - name: combine Python coverage files - shell: bash - run: | - coverage combine .tox/*/.coverage - coverage xml - - - name: upload to codecov.io - uses: codecov/codecov-action@v7 - with: - fail_ci_if_error: true - files: coverage.xml - token: ${{ secrets.CODECOV_TOKEN }} + CXXFLAGS: ${{ matrix.cxx-flags }} + RUST_BACKTRACE: full diff --git a/.gitignore b/.gitignore index ab865aa23..2b7ae69db 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,8 @@ build/ htmlcov/ .coverage* coverage.xml + +Cargo.lock +target/ + +.cache/ diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 50c8dc986..c42cbff81 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -16,6 +16,10 @@ on metatomic: - **git**: the software we use for version control of the source code. See https://git-scm.com/downloads for installation instructions. +- **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.88 to build metatomic. - **Python**: you can install ``Python`` and ``pip`` on your operating system. We require a Python version of at least 3.9. - **tox**: a Python test runner, see https://tox.readthedocs.io/en/latest/. You @@ -28,17 +32,21 @@ not have to interact with them directly: - **a C++ compiler** we need a compiler supporting C++11. GCC >= 7, clang >= 5 and MSVC >= 19 should all work, although MSVC is not yet tested continuously. +.. _rustup: https://rustup.rs +.. _`cargo` : https://doc.rust-lang.org/cargo/ +.. _tox: https://tox.readthedocs.io/en/latest + .. admonition:: Optional tools Depending on which part of the code you are working on, you might experience a - lot of time spent re-compiling code, even if you did not directly change them. - For faster builds (and in turn faster tests), you can use compiler cache, like - `sccache`_ or the classic `ccache`_ to reduce the recompilation of unchanged - source code. To do this, you should install and configure one of these tools - (we suggest ``sccache`` since it also supports Rust), and then configure - ``cmake`` and ``cargo`` to use them by setting environnement variables. On - Linux and macOS, you should set the following (look up how to do set - environment variable with your shell): + lot of time spend re-compiling Rust or C++ code, even if you did not change + them. If you'd like faster builds (and in turn faster tests), you can use + `sccache`_ or the classic `ccache`_ to only re-run the compiler if the + corresponding source code changed. To do this, you should install and configure + one of these tools (we suggest sccache since it also supports Rust), and then + configure cmake and cargo to use them by setting environnement variables. On + Linux and macOS, you should set the following (look up how to do set environment + variable with your shell): .. code-block:: bash @@ -88,36 +96,68 @@ changes: Running tests ------------- -The continuous integration pipeline is based on `tox`_. You can run all tests +The continuous integration pipeline is based on `cargo`_. You can run all tests with: .. code-block:: bash cd - tox + cargo test # or cargo test --release to run tests in release mode -These are exactly the same tests that will be performed online in our Github CI +These are exactly the same tests that will be performed online in our GitHub CI workflows. You can also run only a subset of tests with one of these commands: +- ``cargo test`` runs everything + +- ``cargo test --package=metatomic-core`` to run the C++ tests only; + + - ``cargo test --test=run-cxx-tests`` will run the unit tests C and C++ API; + - ``cargo test --test=check-cxx-install`` will try to build a basic project + depending on metatomic-core with cmake; + +- ``cargo test --package=metatomic-torch`` to run the C++ TorchScript tests only; + + - ``cargo test --test=run-torch-tests`` will run the unit tests for the + TorchScript C++ extension; + - ``cargo test --test=check-torch-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=metatomic-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 + they still work. + +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 -- system`` + +Also, you can run individual Python tests using `tox`_ if you wish to run a +subset of Python tests, for example: + .. code-block:: bash tox -e lint # check files for formatting errors + tox -e core-tests # unit tests for metatomic-core, in Python tox -e torch-tests # unit tests for metatomic-torch, in Python - tox -e torch-tests-cxx # unit tests for metatomic-torch, in C++ - tox -e torch-install-tests-cxx # testing that the C++ code is a valid CMake package + tox -e ase-tests # unit tests for metatomic-ase, in Python + tox -e torchsim-tests # unit tests for metatomic-torchsim, in Python tox -e docs-tests # doctests (checking inline examples) for all packages - tox -e lint # code style tox -e format # format all files -The last command ``tox -e format`` will use ``tox`` to do actual formatting -instead of just checking it, you can use this to automatically fix some of the -issues detected by ``tox -e lint``. +The last command ``tox -e format`` will use tox to do actual formatting instead +of just checking it, you can use to automatically fix some of the issues +detected by ``tox -e lint``. -You can run only a subset of the tests with ``tox -e torch-tests -- +You can run only a subset of the tests with ``tox -e core-tests -- ``, replacing ```` with the path to the files you -want to test, e.g. ``tox -e tests -- tests/system.py``. +want to test, e.g. ``tox -e core-tests -- tests/utils.py``. Controlling test behavior with environment variables ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -126,11 +166,15 @@ There are a handful of environment variables that you can set to control the behavior of tests: - ``METATOMIC_DISABLE_VALGRIND=1`` will disable the use of `valgrind`_ for the - C++ tests. Valgrind is a tool that check for memory errors in native code, but it makes the tests run quite a bit slower; + C++ tests. Valgrind is a tool that check for memory errors in native code, but + it makes the tests run quite a bit slower; - ``METATOMIC_TESTS_TORCH_VERSION`` allow you to run the tests against a specific PyTorch version instead of the latest one. For example, setting ``METATOMIC_TESTS_TORCH_VERSION=2.4`` will run the tests against PyTorch 2.4; +- ``METATOMIC_BUILD_TYPE`` can be set to ``release`` or ``debug`` to force one + of the build option. Debug builds have more internal consistency checks, while + release builds are faster. - ``PIP_EXTRA_INDEX_URL`` can be used to pull PyTorch (or other dependencies) from a different index. This can be useful on Linux if you have issues with CUDA, since the default PyTorch version expects CUDA to be available. A diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 000000000..1a233774c --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,8 @@ +[workspace] +resolver = "2" + +members = [ + "metatomic-core", + "metatomic-torch", + "python", +] diff --git a/docs/Doxyfile b/docs/Doxyfile index f48f15ed9..5cf71fe6f 100644 --- a/docs/Doxyfile +++ b/docs/Doxyfile @@ -991,7 +991,9 @@ WARN_LOGFILE = # spaces. See also FILE_PATTERNS and EXTENSION_MAPPING # Note: If this tag is empty the current directory is searched. -INPUT = ../metatomic-torch/include/metatomic \ +INPUT = ../metatomic-core/include/ \ + ../metatomic-core/include/metatomic \ + ../metatomic-torch/include/metatomic \ ../metatomic-torch/include/metatomic/torch # This tag can be used to specify the character encoding of the source files diff --git a/docs/generate_examples/conf.py b/docs/generate_examples/conf.py index d6289e679..586df5208 100644 --- a/docs/generate_examples/conf.py +++ b/docs/generate_examples/conf.py @@ -13,16 +13,16 @@ HERE = os.path.dirname(__file__) ROOT = os.path.realpath(os.path.join(HERE, "..", "..")) +EXAMPLES = ["c", "torch", "ase", "torchsim"] sphinx_gallery_conf = { - "filename_pattern": ".*", + "filename_pattern": r"\.py", "copyfile_regex": r".*\.(example|mts|xyz)", - "examples_dirs": [ - os.path.join(ROOT, "python", "examples"), - ], + "examples_dirs": [os.path.join(ROOT, "examples", e) for e in EXAMPLES], "gallery_dirs": [ - os.path.join(ROOT, "docs", "src", "examples"), + os.path.join(ROOT, "docs", "src", "examples", e) for e in EXAMPLES ], + "example_extensions": {".py", ".c", ".cpp"}, "matplotlib_animations": False, "image_scrapers": ("matplotlib", ChemiscopeScraper()), "remove_config_comments": True, diff --git a/docs/src/.gitignore b/docs/src/.gitignore index b46f9144a..9288ffdc7 100644 --- a/docs/src/.gitignore +++ b/docs/src/.gitignore @@ -1,2 +1,3 @@ -examples/ sg_execution_times.rst +examples/* +!examples/index.rst diff --git a/docs/src/conf.py b/docs/src/conf.py index cd9d303b3..598749091 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -1,6 +1,8 @@ +import glob import os import subprocess import sys +import zipfile from datetime import datetime from sphinx.domains.c import CObject @@ -20,6 +22,7 @@ import lammps_lexer # noqa: E402 from sphinx.highlighting import lexers # noqa: E402 + lexers["LAMMPS"] = lammps_lexer.LAMMPSLexer(startinline=True) @@ -102,10 +105,37 @@ def generate_examples(): os.environ["METATENSOR_IMPORT_FOR_SPHINX"] = "1" +# Extra files (globs relative to the example directory) to add to the zip files +# sphinx-gallery generates for the corresponding gallery. +EXTRA_ZIP_FILES = { + "c": ["utils/*.h"], +} + + +def add_extra_files_to_zips(app): + """Add the files from ``EXTRA_ZIP_FILES`` to the corresponding gallery zip files.""" + for example, patterns in EXTRA_ZIP_FILES.items(): + examples_dir = os.path.join(ROOT, "examples", example) + gallery_dir = os.path.join(ROOT, "docs", "src", "examples", example) + + files = [] + for pattern in patterns: + files += glob.glob(pattern, root_dir=examples_dir, recursive=True) + names = sorted(file.replace(os.sep, "/") for file in files) + + for zip_path in glob.glob(os.path.join(gallery_dir, "*.zip")): + with zipfile.ZipFile(zip_path, mode="a") as archive: + included = set(archive.namelist()) + for name in names: + if name not in included: + archive.write(os.path.join(examples_dir, name), name) + + def setup(app): build_doxygen_docs() generate_examples() + app.connect("builder-inited", add_extra_files_to_zips, priority=600) app.add_css_file("css/metatomic.css") @@ -191,6 +221,7 @@ def setup(app): # URL redirects redirects = { + # outputs renamed to quantities "outputs/charges.html": "/quantities/charge.html", "outputs/energy.html": "/quantities/energy.html", "outputs/features.html": "/quantities/feature.html", @@ -202,6 +233,13 @@ def setup(app): "outputs/positions.html": "/quantities/position.html", "outputs/variants.html": "/quantities/variants.html", "outputs/velocities.html": "/quantities/velocity.html", + # example re-organization + "examples/1-export-atomistic-model.html": "/examples/torch/1-export-atomistic-model.html", # noqa: E501 + "examples/2-running-ase-md.html": "/examples/ase/1-md.html", + "examples/3-atomistic-model-with-nl.html": "/examples/torch/2-atomistic-model-with-nl.html", # noqa: E501 + "examples/4-profiling.html": "/examples/torch/3-profiling.html", + "examples/5-torchsim-getting-started.html": "/examples/torchsim/1-getting-started.html", # noqa: E501 + "examples/6-torchsim-batched.html": "/examples/torchsim/2-batched-md.html", } # -- Options for HTML output ------------------------------------------------- diff --git a/docs/src/core/CHANGELOG.md b/docs/src/core/CHANGELOG.md new file mode 120000 index 000000000..a344bc46b --- /dev/null +++ b/docs/src/core/CHANGELOG.md @@ -0,0 +1 @@ +../../../metatomic-core/CHANGELOG.md \ No newline at end of file diff --git a/docs/src/core/index.rst b/docs/src/core/index.rst new file mode 100644 index 000000000..405bfce59 --- /dev/null +++ b/docs/src/core/index.rst @@ -0,0 +1,21 @@ +Core Classes +============ + +WIP + + +.. toctree:: + :maxdepth: 2 + + reference/c/index + reference/cxx/index + reference/python/index + reference/json-formats + units + + +.. toctree:: + :maxdepth: 1 + :hidden: + + CHANGELOG.md diff --git a/docs/src/core/reference/c/index.rst b/docs/src/core/reference/c/index.rst new file mode 100644 index 000000000..f190a5e74 --- /dev/null +++ b/docs/src/core/reference/c/index.rst @@ -0,0 +1,17 @@ +.. _c-api-core: + +C API reference +=============== + +WIP + +The functions and types provided in ``metatomic.h`` can be grouped in four +main groups: + +.. toctree:: + :maxdepth: 1 + + system + model + plugin + misc diff --git a/docs/src/core/reference/c/misc.rst b/docs/src/core/reference/c/misc.rst new file mode 100644 index 000000000..6aec886bc --- /dev/null +++ b/docs/src/core/reference/c/misc.rst @@ -0,0 +1,56 @@ +Miscellaneous +============= + +Version number +^^^^^^^^^^^^^^ + +.. doxygenfunction:: mta_version + +.. c:macro:: METATOMIC_VERSION + + Macro containing the compile-time version of metatomic, as a string + +.. c:macro:: METATOMIC_VERSION_MAJOR + + Macro containing the compile-time **major** version number of metatomic, as + an integer + +.. c:macro:: METATOMIC_VERSION_MINOR + + Macro containing the compile-time **minor** version number of metatomic, as + an integer + +.. c:macro:: METATOMIC_VERSION_PATCH + + Macro containing the compile-time **patch** version number of metatomic, as + an integer + + +Error handling +^^^^^^^^^^^^^^ + +.. doxygenfunction:: mta_last_error + +.. doxygenfunction:: mta_set_last_error + +.. doxygenenum:: mta_status_t + + +String manipulation +^^^^^^^^^^^^^^^^^^^ + +.. doxygentypedef:: mta_string_t + +.. doxygenfunction:: mta_string_create + +.. doxygenfunction:: mta_string_free + +.. doxygenfunction:: mta_string_view + +.. doxygenfunction:: mta_format_metadata + + +Unit conversion +^^^^^^^^^^^^^^^ + +.. doxygenfunction:: mta_unit_conversion_factor diff --git a/docs/src/core/reference/c/model.rst b/docs/src/core/reference/c/model.rst new file mode 100644 index 000000000..6a3d9ee38 --- /dev/null +++ b/docs/src/core/reference/c/model.rst @@ -0,0 +1,16 @@ +Model +===== + +.. doxygenstruct:: mta_model_t + :members: + +The following functions operate on :c:type:`mta_model_t`: + +- :c:func:`mta_load_model`: TODO summary +- :c:func:`mta_execute_model`: TODO summary + +-------------------------------------------------------------------------------- + +.. doxygenfunction:: mta_load_model + +.. doxygenfunction:: mta_execute_model diff --git a/docs/src/core/reference/c/plugin.rst b/docs/src/core/reference/c/plugin.rst new file mode 100644 index 000000000..952650f4c --- /dev/null +++ b/docs/src/core/reference/c/plugin.rst @@ -0,0 +1,16 @@ +Plugin system +============= + +.. doxygenstruct:: mta_plugin_t + :members: + +The following functions operate on :c:type:`mta_plugin_t`: + +- :c:func:`mta_register_plugin`: TODO summary +- :c:func:`mta_load_plugin`: TODO summary + +-------------------------------------------------------------------------------- + +.. doxygenfunction:: mta_register_plugin + +.. doxygenfunction:: mta_load_plugin diff --git a/docs/src/core/reference/c/system.rst b/docs/src/core/reference/c/system.rst new file mode 100644 index 000000000..0141a2c66 --- /dev/null +++ b/docs/src/core/reference/c/system.rst @@ -0,0 +1,57 @@ +System +====== + +.. doxygentypedef:: mta_system_t + +The following functions operate on :c:type:`mta_system_t`: + +- :c:func:`mta_system_create`: create a new system from types, positions, cell, and PBC data +- :c:func:`mta_system_free`: free a system handle +- :c:func:`mta_system_size`: get the number of atoms in a system +- :c:func:`mta_system_get_data`: get a borrowed DLPack tensor for some system data +- :c:func:`mta_system_get_length_unit`: get the length unit of a system +- :c:func:`mta_system_add_pairs`: add a pair list to a system +- :c:func:`mta_system_get_pairs`: get a borrowed view of a pair list from a system +- :c:func:`mta_system_known_pairs`: get all pair list options known by a system +- :c:func:`mta_system_add_custom_data`: add custom data to a system +- :c:func:`mta_system_get_custom_data`: get a borrowed view of custom data by name +- :c:func:`mta_system_known_custom_data`: get all custom data names known by a system + +- :c:func:`mta_save`: save a system to a file +- :c:func:`mta_save_buffer`: save a system to a buffer +- :c:func:`mta_load`: load a system from a file +- :c:func:`mta_load_buffer`: load a system from a buffer + +-------------------------------------------------------------------------------- + +.. doxygenfunction:: mta_system_create + +.. doxygenfunction:: mta_system_free + +.. doxygenfunction:: mta_system_size + +.. doxygenfunction:: mta_system_get_data + +.. doxygenfunction:: mta_system_get_length_unit + +.. doxygenfunction:: mta_system_add_pairs + +.. doxygenfunction:: mta_system_get_pairs + +.. doxygenfunction:: mta_system_known_pairs + +.. doxygenfunction:: mta_system_add_custom_data + +.. doxygenfunction:: mta_system_get_custom_data + +.. doxygenfunction:: mta_system_known_custom_data + +.. doxygenfunction:: mta_save + +.. doxygenfunction:: mta_save_buffer + +.. doxygenfunction:: mta_load + +.. doxygenfunction:: mta_load_buffer + +.. doxygenenum:: mta_system_data_kind diff --git a/docs/src/core/reference/cxx/index.rst b/docs/src/core/reference/cxx/index.rst new file mode 100644 index 000000000..9a4a7add3 --- /dev/null +++ b/docs/src/core/reference/cxx/index.rst @@ -0,0 +1,17 @@ +.. _cxx-api-core: + +C++ API reference +================= + +WIP + +The functions and types provided in ``metatomic.hpp`` can be grouped in four +main groups: + +.. toctree:: + :maxdepth: 1 + + system + model + plugin + misc diff --git a/docs/src/core/reference/cxx/misc.rst b/docs/src/core/reference/cxx/misc.rst new file mode 100644 index 000000000..26ba29607 --- /dev/null +++ b/docs/src/core/reference/cxx/misc.rst @@ -0,0 +1,14 @@ +Miscellaneous +============= + + +Error handling +^^^^^^^^^^^^^^ + +.. doxygenclass:: metatomic::Error + + +Unit conversion +^^^^^^^^^^^^^^^ + +.. doxygenfunction:: metatomic::unit_conversion_factor diff --git a/docs/src/core/reference/cxx/model.rst b/docs/src/core/reference/cxx/model.rst new file mode 100644 index 000000000..5411064fa --- /dev/null +++ b/docs/src/core/reference/cxx/model.rst @@ -0,0 +1,13 @@ +Model +===== + +.. TODO: Model classes + +.. doxygenclass:: metatomic::Quantity + :members: + +.. doxygenclass:: metatomic::ModelMetadata + :members: + +.. doxygenclass:: metatomic::ModelCapabilities + :members: diff --git a/docs/src/core/reference/cxx/plugin.rst b/docs/src/core/reference/cxx/plugin.rst new file mode 100644 index 000000000..67cd50b04 --- /dev/null +++ b/docs/src/core/reference/cxx/plugin.rst @@ -0,0 +1,2 @@ +Plugin system +============= diff --git a/docs/src/core/reference/cxx/system.rst b/docs/src/core/reference/cxx/system.rst new file mode 100644 index 000000000..54cad2391 --- /dev/null +++ b/docs/src/core/reference/cxx/system.rst @@ -0,0 +1,25 @@ +System +====== + +.. doxygenclass:: metatomic::System + :members: + +.. doxygenclass:: metatomic::PairListOptions + :members: + +Serialization +------------- + +Systems can be saved to a file or serialized into an in-memory byte buffer. +Loading a system requires an array-creation callback, which allocates the +arrays of the reconstructed system. + +.. doxygenfunction:: metatomic::io::save + +.. doxygenfunction:: metatomic::io::save_buffer + +.. doxygenfunction:: metatomic::io::load + +.. doxygenfunction:: metatomic::io::load_buffer(const uint8_t* buffer, uintptr_t buffer_count, mts_create_array_callback_t create_array) + +.. doxygenfunction:: metatomic::io::load_buffer(const Buffer& buffer, mts_create_array_callback_t create_array) diff --git a/docs/src/core/reference/json-formats.rst b/docs/src/core/reference/json-formats.rst new file mode 100644 index 000000000..4abd8d127 --- /dev/null +++ b/docs/src/core/reference/json-formats.rst @@ -0,0 +1,226 @@ +.. _core-json-formats: + +JSON data formats +================= + +Some metatomic data structures are exchanged across the C API as JSON-encoded +strings rather than dedicated C types. This page documents the exact JSON +representation of each such structure, so that engines and models written in any +language can produce and consume them. + +.. _core-json-pair-options: + +Pair list options +----------------- + +The JSON representation of a requested pair list (also known as a neighbor +list). This is used for example by :c:func:`mta_system_add_pairs`, +:c:func:`mta_system_get_pairs` and :c:func:`mta_system_known_pairs`. + +.. code-block:: json + + { + "type": "metatomic_pair_list_options", + "cutoff": "0x400c000000000000", + "full_list": false, + "strict": false, + "requestors": ["my-model"] + } + +``type`` + Must be the string ``"metatomic_pair_list_options"``. + +``cutoff`` + Cutoff radius for the pair list in the length unit of the model. Must be a + positive finite number. + + It is stored as a string containing the hexadecimal representation of the + 64-bit integer with the same bit pattern as the ``cutoff`` floating-point + value (i.e. reinterpreting the ``double`` as a ``uint64_t``). + +``full_list`` + Boolean. If ``true``, the list is a full list containing both ``i -> j`` + and ``j -> i`` for each pair, if ``false``, it is a half list containing + only ``i -> j``. + +``strict`` + Boolean. If ``true``, the list is guaranteed to contain only atoms within + the cutoff, if ``false``, it may also include some pairs slightly beyond the + cutoff. + +``requestors`` + Optional array of strings identifying who requested this pair list. May be + omitted, in which case it is treated as an empty list. + + +.. _core-json-quantity: + +Quantities +---------- + +The JSON representation of a physical quantity, used to represent custom models +inputs and outputs. This is used for example in +:c:member:`mta_model_t.requested_inputs` and the ``outputs`` field of the +model's capabilities. + +.. code-block:: json + + { + "type": "metatomic_quantity", + "name": "energy", + "unit": "eV", + "sample_kind": "system" + "gradients": ["positions"] + "description": "Potential energy of the system", + } + +``type`` + Must be the string ``"metatomic_quantity"``. + +``name`` + Name of the quantity, this this can be a standard name from the list of + :ref:`standard-quantities`, or a custom name of the form + ``::[/]`` + +``unit`` + Unit of the quantity. + +``gradients`` + Array of strings identifying the gradients for this quantity. This can be an + empty array if the quantity has no gradients. Valid values for the gradients + are ``"positions"``, and ``"strain"``. + +``sample_kind`` + Kind of sample for which this quantity is defined. This can be one of the + following: ``"atom"``, ``"system"`` or ``"atom_pair"``. + + +.. _core-json-model-metadata: + +Model metadata +-------------- + +The JSON representation of a model's metadata. This is used for example by +:c:member:`mta_model_t.metadata`. + +.. code-block:: json + + { + "type": "metatomic_model_metadata", + "name": "MyCoolModel v1.2", + "authors": ["Alice Smith", "Bob Johnson "], + "description": "A machine learning potential for water", + "references": { + "model": ["doi:10.1234/model-paper"], + "architecture": ["doi:10.1234/arch-paper"], + "implementation": ["https://github.com/example/mycoolmodel"] + }, + "extra": { + "training_set": "QM9", + "cutoff": "4.5" + } + } + +``type`` + Must be the string ``"metatomic_model_metadata"``. + +``name`` + Name of the model, e.g. ``"MyCoolModel v1.2"``. + +``authors`` + Array of strings identifying the authors of the model. Each string can be a + name or a name with an email address, e.g. ``"Alice Smith"`` or + ``"Bob Johnson "``. + +``description`` + A free-text description of the model. + +``references`` + An object with three keys, each containing an array of strings (DOIs, URLs, + or any other format): + + ``model`` + References about the model as a whole, e.g. a paper describing the model + or a website presenting it. + + ``architecture`` + References about the architecture of the model, e.g. papers describing + the mathematical form of the model. + + ``implementation`` + References about the implementation of the model, e.g. a link to the + source code repository or a paper describing the software. + +``extra`` + An object with string values, providing any additional key-value pairs the + model author wishes to include. This can be used for any purpose. + +.. _core-json-model-capabilities: + +Model capabilities +------------------ + +The JSON representation of a model's capabilities, describing which outputs it +provides, which atomic types it supports, and other constraints. This is used +for example by :c:member:`mta_model_t.capabilities`. + +.. code-block:: json + + { + "type": "metatomic_model_capabilities", + "outputs": [ + { + "type": "metatomic_quantity", + "name": "energy", + "unit": "eV", + "sample_kind": "system", + "gradients": ["positions"], + "description": "Potential energy of the system" + }, + { + "type": "metatomic_quantity", + "name": "energy/pbe0", + "unit": "eV", + "sample_kind": "system", + "gradients": ["positions", "strain"], + "description": "Potential energy of the system" + }, + ], + "atomic_types": [1, 6, 8], + "interaction_range": 5.0, + "length_unit": "angstrom", + "supported_devices": ["cpu", "cuda"], + "dtype": "float32" + } + +``type`` + Must be the string ``"metatomic_model_capabilities"``. + +``outputs`` + Array of :ref:`quantity objects ` describing the + outputs this model can provide. + +``atomic_types`` + Array of integers listing the atomic types this model supports. The meaning + of these integers is up to the model, and is not required to be the atomic + numbers. + +``interaction_range`` + The interaction range of the model in the length unit of the model. This is + the maximum distance between two atoms for which the model's output can + depend on their relative position. Must be a non-negative number. + +``length_unit`` + String identifying the length unit used by the model, e.g. ``"angstrom"`` or + ``"nanometer"``. This must be a valid :ref:`unit expression ` with + dimensions compatible with length. + +``supported_devices`` + Array of strings listing the devices on which the model can run. Valid + values are ``"cpu"``, ``"cuda"``, ``"rocm"``, and ``"metal"``. + +``dtype`` + The data type of the model, used for all inputs and outputs. Must be either + ``"float32"`` or ``"float64"``. The model is free to use different data + types for internal computations, but all inputs and outputs must be in this + data type. diff --git a/docs/src/core/reference/python/index.rst b/docs/src/core/reference/python/index.rst new file mode 100644 index 000000000..7eae04901 --- /dev/null +++ b/docs/src/core/reference/python/index.rst @@ -0,0 +1,13 @@ +.. _python-api-core: + +Python API reference +==================== + +WIP + +.. toctree:: + :maxdepth: 1 + + system + metadata + misc diff --git a/docs/src/core/reference/python/metadata.rst b/docs/src/core/reference/python/metadata.rst new file mode 100644 index 000000000..8296ca81b --- /dev/null +++ b/docs/src/core/reference/python/metadata.rst @@ -0,0 +1,22 @@ +Metadata +======== + +.. currentmodule:: metatomic + +The following classes are used to store metadata about a model. +:py:class:`Quantity` represent a physical quantity used as input or output of a +model; :py:class:`ModelCapabilities` stores information about what a model can +do; and :py:class:`ModelMetadata` stores meta-information about a model: list of +authors, :py:class:`references `, *etc.* + +.. autoclass:: Quantity + :members: + +.. autoclass:: ModelCapabilities + :members: + +.. autoclass:: ModelMetadata + :members: + +.. autoclass:: References + :members: diff --git a/docs/src/core/reference/python/misc.rst b/docs/src/core/reference/python/misc.rst new file mode 100644 index 000000000..547d8900b --- /dev/null +++ b/docs/src/core/reference/python/misc.rst @@ -0,0 +1,9 @@ +Miscellaneous +============= + +.. currentmodule:: metatomic + +Error handling +^^^^^^^^^^^^^^ + +.. autoclass:: MetatomicError diff --git a/docs/src/core/reference/python/system.rst b/docs/src/core/reference/python/system.rst new file mode 100644 index 000000000..f1275421a --- /dev/null +++ b/docs/src/core/reference/python/system.rst @@ -0,0 +1,7 @@ +System +====== + +.. currentmodule:: metatomic + +.. autoclass:: PairListOptions + :members: diff --git a/docs/src/core/units.rst b/docs/src/core/units.rst new file mode 100644 index 000000000..6c50603ca --- /dev/null +++ b/docs/src/core/units.rst @@ -0,0 +1,101 @@ +.. _units: + +Units +^^^^^ + +Models in metatensor can use arbitrary units for their inputs and outputs. The +unit conversion system allows models to specify the units they expect and +receive data in any compatible unit, with automatic conversion handled by +during model execution. + +Unit parsing is handled by one of the following functions: + +- :c:func:`mta_unit_conversion_factor` in C +- :cpp:func:`metatomic::unit_conversion_factor` in C++ + +These functions parses two unit expressions, checks that they have compatible +physical dimensions, and returns the multiplicative conversion factor. For +example, in C++: + +.. code-block:: C++ + + // How many eV are in one kJ/mol? + double factor = metatomic::unit_conversion_factor("kJ/mol", "eV"); + // factor ≈ 0.01036 + + // How many GPa are in one eV/A^3? + factor = metatomic::unit_conversion_factor("eV/A^3", "GPa"); + // factor ≈ 160.22 + +If either (or both) unit strings are empty, the conversion returns ``1.0`` +without checking dimensions. This makes it safe to pass optional/unknown units. + +.. _known-base-units: + +Base units +~~~~~~~~~~ + +Unit expressions are built from the following base units. Matching is +case-insensitive, and whitespace is ignored. + +**Temperature**: + ``Kelvin`` (``K``) + +**Length**: + ``angstrom`` (``A``), ``Bohr``, ``meter`` (``m``), ``centimeter`` (``cm``), + ``millimeter`` (``mm``), ``micrometer`` (``um``, ``µm``), ``nanometer`` (``nm``) + +**Energy**: + ``eV``, ``meV``, ``Hartree``, ``kcal``, ``kJ``, ``Joule`` (``J``), ``Rydberg`` (``Ry``) + +**Time**: + ``second`` (``s``), ``millisecond`` (``ms``), ``microsecond`` (``us``, ``µs``), + ``nanosecond`` (``ns``), ``picosecond`` (``ps``), ``femtosecond`` (``fs``) + +**Mass**: + ``Dalton`` (``u``), ``kilogram`` (``kg``), ``gram`` (``g``), ``electron_mass`` (``m_e``) + +**Charge**: + ``e``, ``Coulomb`` (``C``) + +**Pressure**: + ``Pascal`` (``Pa``), ``kiloPascal`` (``kPa``), ``MegaPascal`` (``MPa``), + ``GigaPascal`` (``GPa``), ``bar``, ``atm`` + +**Electric Dipole Moment**: + ``Debye`` (``D``) + +**Dimensionless**: + ``mol`` + +**Derived constants**: + ``hbar`` + +Expression syntax +~~~~~~~~~~~~~~~~~ + +Base units can be combined using the following operators: + +- Multiplication: ``*`` or whitespace (``kJ mol``, ``kJ*mol``) +- Division: ``/`` (``kJ/mol``) +- Exponentiation: ``^`` (``A^3``, ``m^2``) +- Parentheses: ``()`` for grouping (``(eV*u)^(1/2)``) + +Fractional powers + Exponents can be integers (``A^3``) or fractions enclosed in parentheses + (``^(1/2)``, ``^(2/3)``). Fractional powers are supported only when the + result has integer physical dimensions — for example ``(eV*u)^(1/2)`` + computes momentum with dimensions :math:`[L T^{-1} M]`. + +Numeric literals + Bare numbers can be used as dimensionless quantity expressions, e.g. + ``"2"`` evaluates to the conversion factor ``2.0``. This is useful when a + model needs to define a unit that is simply a scalar multiple of another. + +Examples of valid compound expressions: + +- ``kJ/mol`` --- energy per mole +- ``eV/Angstrom^3`` or ``eV/A^3`` --- pressure +- ``(eV*u)^(1/2)`` --- momentum (fractional powers) +- ``Hartree/Bohr`` --- force in atomic units +- ``nm/fs`` --- velocity diff --git a/docs/src/devdoc/get-started.rst b/docs/src/devdoc/get-started.rst new file mode 100644 index 000000000..4c19e4ef6 --- /dev/null +++ b/docs/src/devdoc/get-started.rst @@ -0,0 +1,6 @@ +.. _devdoc-get-started: + +Getting started +=============== + +.. include:: ../../../CONTRIBUTING.rst diff --git a/docs/src/devdoc/index.rst b/docs/src/devdoc/index.rst new file mode 100644 index 000000000..43755fdf6 --- /dev/null +++ b/docs/src/devdoc/index.rst @@ -0,0 +1,26 @@ +.. _devdoc: + +Developer documentation +####################### + +This developer documentation contains the following sections: + +1. :ref:`devdoc-get-started` explains how you can start developing code and + documentation; + +.. toctree:: + :maxdepth: 2 + + get-started + +Development team +---------------- + +Metatensor is developed in the `COSMO laboratory`_ at `EPFL`_, and made +available under the `BSD 3-clauses license `_. We welcome +contributions from anyone, feel free to contact us if you need some help working +with the code! + +.. _COSMO laboratory: https://www.epfl.ch/labs/cosmo/ +.. _EPFL: https://www.epfl.ch/ +.. _LICENSE: https://github.com/metatensor/metatensor/blob/main/LICENSE diff --git a/docs/src/examples/index.rst b/docs/src/examples/index.rst new file mode 100644 index 000000000..06a75e558 --- /dev/null +++ b/docs/src/examples/index.rst @@ -0,0 +1,80 @@ +.. _atomistic-tutorials: + +Tutorials +========= + +The first set of tutorials on this page is about existing integrations between +``metatomic`` and simulation engines. You can also find more examples in the +:ref:`engines` section. These tutorials are intended for users who want to use +existing metatomic models with existing simulation engines. + +.. toctree:: + :maxdepth: 1 + :hidden: + + ase/index + torchsim/index + + +.. grid:: + + .. grid-item-card:: ASE tutorials + :link: ase-tutorials + :link-type: ref + :columns: 12 12 6 6 + :margin: 0 3 0 0 + :img-top: /../static/images/logo-ase.* + :class-img-top: mta-card-img-top + + How to use ``metatomic`` with the Atomic Simulation Environment (ASE). + + .. grid-item-card:: Torch-Sim tutorials + :link: torchsim-tutorials + :link-type: ref + :columns: 12 12 6 6 + :margin: 0 3 0 0 + :img-top: /../static/images/logo-radical-ai.* + :class-img-top: mta-card-img-top + + How to use existing ``metatomic`` models with Torch-Sim, a simulation + engine for batched molecular dynamics simulations, based on PyTorch. + + +-------------------------------------------------------------------------------- + +The second set of tutorials on this page is intended for developers who want to +use ``metatomic`` to either create new models, or run existing models in a new +simulation engine. These tutorials require existing knowledge of the +corresponding programming languages and machine learning frameworks. + +.. toctree:: + :maxdepth: 1 + :hidden: + + c/index + torch/index + +.. grid:: + + .. grid-item-card:: C API tutorials + :link: c-tutorials + :link-type: ref + :columns: 12 12 6 6 + :margin: 0 3 0 0 + :img-top: /../static/images/logo-c.* + :class-img-top: mta-card-img-top + + How to use the C API of ``metatomic`` both to create custom atomistic + models; and to load and run existing atomistic models from simulation + engines. + + .. grid-item-card:: PyTorch tutorials + :link: torch-tutorials + :link-type: ref + :columns: 12 12 6 6 + :margin: 0 3 0 0 + :img-top: /../static/images/logo-torch.* + :class-img-top: mta-card-img-top + + How to use the PyTorch API of ``metatomic`` to define custom atomistic + models. diff --git a/docs/src/index.rst b/docs/src/index.rst index d94c6ded2..441356c29 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -92,8 +92,10 @@ existing trained models, look into the metatrain_ project instead. overview installation + core/index torch/index quantities/index engines/index examples/index + devdoc/index cite diff --git a/docs/src/quantities/mass.rst b/docs/src/quantities/mass.rst index b2a46b7ab..ed5ab88d8 100644 --- a/docs/src/quantities/mass.rst +++ b/docs/src/quantities/mass.rst @@ -37,7 +37,7 @@ following metadata: - the ``"mass"`` quantity must not have any components * - properties - - ``"mass`` + - ``"mass"`` - The ``"mass"`` quantity must have a single property dimension named ``"mass"``, with a single entry set to ``0``. diff --git a/docs/src/quantities/non_conservative.rst b/docs/src/quantities/non_conservative.rst index 29e5ee24c..9085553d2 100644 --- a/docs/src/quantities/non_conservative.rst +++ b/docs/src/quantities/non_conservative.rst @@ -133,7 +133,7 @@ and must have the following metadata: * - keys - ``"_"`` - the keys must have a single dimension named ``"_"``, with a single entry - set to ``0``. The ``"non_conservative_force"`` quantity is always + set to ``0``. The ``"non_conservative_stress"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. diff --git a/docs/src/quantities/velocity.rst b/docs/src/quantities/velocity.rst index 4ad546662..f868ce0dc 100644 --- a/docs/src/quantities/velocity.rst +++ b/docs/src/quantities/velocity.rst @@ -36,7 +36,7 @@ following metadata: * - components - ``"xyz"`` - The ``"velocity"`` quantity must have a single component dimension named - ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The position + ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The velocity is always a 3D vector, and the order of the components is ``x, y, z``. * - properties diff --git a/docs/src/torch/reference/index.rst b/docs/src/torch/reference/index.rst index f3d4cb61d..31360d3f2 100644 --- a/docs/src/torch/reference/index.rst +++ b/docs/src/torch/reference/index.rst @@ -8,7 +8,6 @@ API reference systems models/index - units wrappers o3 symmetrized-model diff --git a/docs/src/torch/reference/misc.rst b/docs/src/torch/reference/misc.rst index 00bf79f06..10d3e636d 100644 --- a/docs/src/torch/reference/misc.rst +++ b/docs/src/torch/reference/misc.rst @@ -7,3 +7,5 @@ simulation engine to use metatomic models. .. autofunction:: metatomic.torch.pick_device .. autofunction:: metatomic.torch.pick_output + +.. autofunction:: metatomic.torch.unit_conversion_factor diff --git a/docs/src/torch/reference/units.rst b/docs/src/torch/reference/units.rst deleted file mode 100644 index cba2a7397..000000000 --- a/docs/src/torch/reference/units.rst +++ /dev/null @@ -1,71 +0,0 @@ -Unit conversions -================ - -.. autofunction:: metatomic.torch.unit_conversion_factor - -The :py:func:`unit_conversion_factor` function accepts any valid unit expression -built from base units combined with operators. There is no need to specify a -physical quantity --- the parser automatically verifies dimensional -compatibility between the source and target units. - -.. _known-base-units: - -Supported base units -~~~~~~~~~~~~~~~~~~~~ - -Unit expressions are built from the following base units. Matching is -case-insensitive, and whitespace is ignored. - - -**Temperature**: - ``Kelvin`` (``K``) - -**Length**: - ``angstrom`` (``A``), ``Bohr``, ``meter`` (``m``), ``centimeter`` (``cm``), - ``millimeter`` (``mm``), ``micrometer`` (``um``, ``µm``), ``nanometer`` (``nm``) - -**Energy**: - ``eV``, ``meV``, ``Hartree``, ``kcal``, ``kJ``, ``Joule`` (``J``), ``Rydberg`` (``Ry``) - -**Time**: - ``second`` (``s``), ``millisecond`` (``ms``), ``microsecond`` (``us``, ``µs``), - ``nanosecond`` (``ns``), ``picosecond`` (``ps``), ``femtosecond`` (``fs``) - -**Mass**: - ``Dalton`` (``u``), ``kilogram`` (``kg``), ``gram`` (``g``), ``electron_mass`` (``m_e``) - -**Charge**: - ``e``, ``Coulomb`` (``C``) - -**Pressure**: - ``Pascal`` (``Pa``), ``kiloPascal`` (``kPa``), ``MegaPascal`` (``MPa``), ``GigaPascal`` (``GPa``), ``bar``, ``atm`` - -**Electric Dipole Moment**: - ``Debye`` (``D``) - -**Dimensionless**: - ``mol`` - -**Derived constants**: - ``hbar`` - -Expression syntax -~~~~~~~~~~~~~~~~~~~ - -Base units can be combined using the following operators: - -- Multiplication: ``*`` or whitespace (``kJ mol``, ``kJ*mol``) -- Division: ``/`` (``kJ/mol``) -- Exponentiation: ``^`` (``A^3``, ``m^2``) -- Parentheses: ``()`` for grouping (``(eV*u)^(1/2)``) - -Examples of valid compound expressions: - -- ``kJ/mol`` --- energy per mole -- ``eV/Angstrom^3`` or ``eV/A^3`` --- pressure -- ``(eV*u)^(1/2)`` --- momentum (fractional powers) -- ``Hartree/Bohr`` --- force in atomic units -- ``nm/fs`` --- velocity - -The parser automatically checks that both unit expressions have matching -physical dimensions before computing the conversion factor. diff --git a/docs/static/css/metatomic.css b/docs/static/css/metatomic.css index 41c6d74fe..106ffe394 100644 --- a/docs/static/css/metatomic.css +++ b/docs/static/css/metatomic.css @@ -68,3 +68,10 @@ body[data-theme="auto"] { .font-size-small { font-size: small !important; } + +.mta-card-img-top { + width: 25% !important; + max-height: 25cqw; + margin: auto; + margin-top: 0.5em; +} diff --git a/docs/static/images/logo-torch.png b/docs/static/images/logo-torch.png new file mode 100644 index 000000000..dfd55b9ab Binary files /dev/null and b/docs/static/images/logo-torch.png differ diff --git a/python/examples/.gitignore b/examples/ase/.gitignore similarity index 100% rename from python/examples/.gitignore rename to examples/ase/.gitignore diff --git a/python/examples/2-running-ase-md.py b/examples/ase/1-md.py similarity index 100% rename from python/examples/2-running-ase-md.py rename to examples/ase/1-md.py diff --git a/examples/ase/README.rst b/examples/ase/README.rst new file mode 100644 index 000000000..e682cd02d --- /dev/null +++ b/examples/ase/README.rst @@ -0,0 +1,4 @@ +.. _ase-tutorials: + +ASE tutorials +============= diff --git a/examples/c/1-create-system.c b/examples/c/1-create-system.c new file mode 100644 index 000000000..3fc4fcff1 --- /dev/null +++ b/examples/c/1-create-system.c @@ -0,0 +1,278 @@ +// .. _c-tutorial-create-system: +// +// Creating ``mta_system_t`` +// ========================= +// +// When integrating metatomic into an existing simulation code, the atomic +// types, positions, cell, and periodic boundary conditions are usually +// already stored in memory as plain arrays — for example as ``double**`` +// pointers. This example shows how to wrap such existing data into DLPack +// tensors, and use them to create a :c:type:`mta_system_t`. +// +// The same approach works for any data layout, as long as you can describe +// it with a DLPack tensor. + +#include +#include +#include +#include + +#include + + +// %% +// +// DLPack tensors +// -------------- +// +// This tutorial shows a basic way to create DLPack tensors from existing data. +// You should also explore the corresponding documentation in the DLPack header +// file, which describes the full `DLPack API`_ and options. +// +// .. _DLPack API: https://dmlc.github.io/dlpack/latest/ +// +// We get the dlpack header from the vendored version in the metatensor package, +// which is the same that metatomic uses internally. You can also bring your own +// copy of the DLPack header, or use the one from your framework (PyTorch, +// TensorFlow, …) as long at it is at least version 1.0. + +#include + +// %% +// +// We'll need a context to store the shape and strides of the DLPack tensor. The +// context is owned by the DLPack tensor, and will be freed when the tensor is +// freed. The whole DLManagedTensorVersioned is passed to a custom deleter +// function when the tensor is no longer needed, which can free the context and +// the tensor itself. + +typedef struct CustomDLPackContext { + int64_t* shape; + int64_t* strides; +} CustomDLPackContext; + +void dlpack_deleter(DLManagedTensorVersioned *self) { + if (!self) { + return; + } + + CustomDLPackContext* ctx = (CustomDLPackContext*)self->manager_ctx; + if (ctx) { + free(ctx->shape); + free(ctx->strides); + free(ctx); + } + free(self); +} + +// %% +// +// We then define a helper function to create a DLPack tensor from a flat data +// buffer. The tensor is created as a row-major, contiguous tensor on CPU, with +// the specified shape and data type. The caller owns the data buffer, and is +// responsible for freeing it after the tensor is no longer needed. + +static DLManagedTensorVersioned* tensor_from_data( + void *data, + const int64_t *shape, + int32_t ndim, + DLDataType dtype +) { + CustomDLPackContext* ctx = malloc(sizeof(CustomDLPackContext)); + if (!ctx) { + return NULL; + } + + // copy the shape into a new buffer owned by the DLPack tensor. + ctx->shape = malloc(ndim * sizeof(int64_t)); + ctx->strides = malloc(ndim * sizeof(int64_t)); + if (!ctx->shape || !ctx->strides) { + free(ctx->shape); + free(ctx->strides); + free(ctx); + return NULL; + } + memcpy(ctx->shape, shape, ndim * sizeof(int64_t)); + + // set the strides to indicate a contiguous row-major tensor + int64_t stride = 1; + for (int32_t i = ndim - 1; i >= 0; i--) { + ctx->strides[i] = stride; + stride *= shape[i]; + } + + // Create the DLPack tensor + DLManagedTensorVersioned* tensor = calloc(1, sizeof(*tensor)); + if (!tensor) { + free(ctx->shape); + free(ctx->strides); + free(ctx); + return NULL; + } + + tensor->version.major = DLPACK_MAJOR_VERSION; + tensor->version.minor = DLPACK_MINOR_VERSION; + tensor->manager_ctx = ctx; + tensor->deleter = dlpack_deleter; + + // Set the flags to indicate that the tensor is read-only. + tensor->flags = DLPACK_FLAG_BITMASK_READ_ONLY; + + tensor->dl_tensor.data = data; + // offset in bytes from the beginning of the data buffer to the first + // element of the tensor. + tensor->dl_tensor.byte_offset = 0; + + // device the tensor is on. Here we use CPU, device 0 (the only CPU device). + tensor->dl_tensor.device.device_type = kDLCPU; + tensor->dl_tensor.device.device_id = 0; + + // data type of the tensor + tensor->dl_tensor.dtype = dtype; + + // number of dimensions, shape, and strides, re-using the buffers we + // allocated above. + tensor->dl_tensor.ndim = ndim; + tensor->dl_tensor.shape = ctx->shape; + tensor->dl_tensor.strides = ctx->strides; + + return tensor; +} + +// %% + +int main(void) { + +// %% +// +// Build the ``positions`` and ``cell`` tensor +// ------------------------------------------- +// +// The positions and cell tensors can be either ``float32`` or ``float64``. They +// would typically wrap existing data from the simulation code, here we create +// them inline for demonstration purposes. + +const int64_t n_atoms = 4; +double positions_data[] = { + 0.0, 0.0, 0.0, + 0.5, 0.5, 0.0, + 0.5, 0.0, 0.5, + 0.0, 0.5, 0.5, +}; + +double cell_data[] = { + 1.0, 0.0, 0.0, + 0.0, 1.0, 0.0, + 0.0, 0.0, 1.0, +}; + +DLManagedTensorVersioned* positions = tensor_from_data( + /*data=*/ positions_data, + /*shape=*/(int64_t[]){n_atoms, 3}, + /*ndim=*/ 2, + /*dtype=*/(DLDataType){.code = kDLFloat, .bits = 64, .lanes = 1} +); + +DLManagedTensorVersioned* cell = tensor_from_data( + /*data=*/ cell_data, + /*shape=*/(int64_t[]){3, 3}, + /*ndim=*/ 2, + /*dtype=*/(DLDataType){.code = kDLFloat, .bits = 64, .lanes = 1} +); + +// %% +// +// Build the ``types`` tensor +// -------------------------- +// +// Atomic types must be an ``int32`` tensor of shape ``(n_atoms,)``. Note that +// the atomic types are not necessarily the same as the atomic numbers, and can +// be any integer values that the model understands. It can be useful to let +// users provide a mapping from the atomic tags used in the simulation code to +// the atomic types used by the model. + +int32_t types_data[] = { + 1, 1, 6, 6 +}; + +DLManagedTensorVersioned *types = tensor_from_data( + /*data=*/ types_data, + /*shape=*/(int64_t[]){n_atoms}, + /*ndim=*/ 1, + /*dtype=*/(DLDataType){.code = kDLInt, .bits = 32, .lanes = 1} +); + +// %% +// +// Build the ``pbc`` tensor +// ------------------------ +// +// Periodic boundary conditions are a ``bool`` tensor of shape ``(3,)``, +// one entry per axis. For a fully periodic system all three are ``true``. + +bool pbc_data[] = {true, true, true}; +DLManagedTensorVersioned *pbc = tensor_from_data( + /*data=*/ pbc_data, + /*shape=*/(int64_t[]){3}, + /*ndim=*/ 1, + /*dtype=*/(DLDataType){.code = kDLBool, .bits = 8, .lanes = 1} +); + +// %% +// +// Create the system +// ----------------- +// +// :c:func:`mta_system_create` takes ownership of the four DLPack tensors; +// they must not be used afterwards. The returned :c:type:`mta_system_t` +// must be freed with :c:func:`mta_system_free` once you are done with it. + +mta_system_t* system = NULL; +mta_status_t status = mta_system_create( + "Angstrom", types, positions, cell, pbc, &system +); + +if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, /*origin=*/NULL, /*data=*/NULL); + fprintf(stderr, "failed to create system: %s\n", error_message); + return EXIT_FAILURE; +} + +// %% +// +// Use the system +// -------------- +// +// Now that we have a :c:type:`mta_system_t`, we can use it with the rest of the +// metatomic API, pass it to a model, etc. Here we just query its size and print +// it. + +uintptr_t size = 0; +status = mta_system_size(system, &size); +if (status == MTA_SUCCESS) { + printf("created system with %lu atoms\n", (unsigned long)size); +} else { + printf("failed to get system size\n"); + mta_system_free(system); + return EXIT_FAILURE; +} + + +// %% +// +// Cleanup +// ------- +// +// Free the system once it is no longer needed. The DLPack tensors have already +// been consumed by ``mta_system_create`` and must not be freed again. + +status = mta_system_free(system); +if (status != MTA_SUCCESS) { + fprintf(stderr, "failed to free system memory\n"); + return EXIT_FAILURE; +}; + +// %% + +return EXIT_SUCCESS; } diff --git a/examples/c/2-using-system.c b/examples/c/2-using-system.c new file mode 100644 index 000000000..b9e821f1f --- /dev/null +++ b/examples/c/2-using-system.c @@ -0,0 +1,442 @@ +// .. _c-tutorial-use-system: +// +// Using ``mta_system_t`` +// ====================== +// +// This tutorial explores how to access data stored inside a +// :c:type:`mta_system_t`. We look at retrieving the basic tensors (positions, +// cell, types, pbc), working with pair lists (neighbor lists), and storing +// custom per-system data. + +#include +#include +#include +#include +#include + +#include + +// Function to create a system that we will use in this tutorial +static mta_system_t* create_system_for_tutorial(); + +// %% +// +// This tutorial uses the same code as the :ref:`previous one +// ` to create a system. In practice, the +// :c:type:`mta_system_t` is created by a simulation engine, and then passed to +// the model, which can acess the data inside in a similar way regardless of +// wether the system was created from C, Python, or any other supported +// language. +// + +#include "utils/dlpack.h" // tensor_from_data +#include "utils/mts_array.h" // make_mts_array + + +// %% +// +// In this tutorial, the function ``create_system_for_tutorial()`` plays the +// role of a simulation engine: it creates the basic system data, then creates a +// pair list and some custom data, attaches both to the system, and returns the +// fully-populated system. +// +// .. raw:: html +// +//
Implementation of create_system_for_tutorial() + +static mta_system_t* create_system_for_tutorial() { + // The basic tensor data (positions, cell, types, pbc) is referenced + // directly by the DLPack tensors, so it must stay alive as long as the + // system. We use ``static`` arrays to keep it alive without polluting the + // global scope. + static double POSITIONS_DATA[] = { + 0.0, 0.0, 0.0, + 0.5, 0.5, 0.0, + 0.5, 0.0, 0.5, + 0.0, 0.5, 0.5, + }; + + static double CELL_DATA[] = { + 1.0, 0.0, 0.0, + 0.0, 1.0, 0.0, + 0.0, 0.0, 1.0, + }; + + static int32_t TYPES_DATA[] = {1, 1, 6, 6}; + + static bool PBC_DATA[] = {true, true, true}; + + DLDataType f64_dtype = {.code = kDLFloat, .bits = 64, .lanes = 1}; + DLDataType i32_dtype = {.code = kDLInt, .bits = 32, .lanes = 1}; + DLDataType bool_dtype = {.code = kDLBool, .bits = 8, .lanes = 1}; + + const int64_t n_atoms = 4; + DLManagedTensorVersioned* positions = tensor_from_data( + POSITIONS_DATA, (int64_t[]){n_atoms, 3}, 2, f64_dtype + ); + + DLManagedTensorVersioned* cell = tensor_from_data( + CELL_DATA, (int64_t[]){3, 3}, 2, f64_dtype + ); + + DLManagedTensorVersioned* types = tensor_from_data( + TYPES_DATA, (int64_t[]){n_atoms}, 1, i32_dtype + ); + + DLManagedTensorVersioned* pbc = tensor_from_data( + PBC_DATA, (int64_t[]){3}, 1, bool_dtype + ); + + mta_system_t* system = NULL; + mta_status_t status = mta_system_create( + "Angstrom", types, positions, cell, pbc, &system + ); + + if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to create system: %s\n", error_message); + return NULL; + } + + int32_t pair_samples[] = { + 0, 1, 0, 0, 0, + 0, 2, 0, 0, 0, + 1, 3, 0, 0, 0, + }; + + double pair_distances[] = { + 0.5, 0.5, 0.0, + 0.5, 0.0, 0.5, + 0.0, 0.5, 0.5, + }; + + int32_t xyz_values[] = {0, 1, 2}; + int32_t distance_values[] = {0}; + + const char* sample_dimensions[] = { + "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c" + }; + struct mts_array_t samples_array = make_mts_array( + pair_samples, (uintptr_t[]){3, 5}, 2, i32_dtype + ); + const mts_labels_t* samples = mts_labels(sample_dimensions, 5, samples_array); + + const char* component_dimensions[] = {"xyz"}; + struct mts_array_t comp_array = make_mts_array( + xyz_values, (uintptr_t[]){3, 1}, 2, i32_dtype + ); + const mts_labels_t* component = mts_labels(component_dimensions, 1, comp_array); + const mts_labels_t* components[] = {component}; + + const char* properties_dimensions[] = {"distance"}; + struct mts_array_t prop_array = make_mts_array( + distance_values, (uintptr_t[]){1, 1}, 2, i32_dtype + ); + const mts_labels_t* properties = mts_labels(properties_dimensions, 1, prop_array); + + struct mts_array_t values_array = make_mts_array( + pair_distances, (uintptr_t[]){3, 3, 1}, 3, f64_dtype + ); + mts_block_t* pairs = mts_block( + values_array, samples, components, 1, properties + ); + + const char* options = + "{\"type\": \"metatomic_pair_list_options\"," + " \"cutoff\": \"0x4008000000000000\"," + " \"full_list\": true," + " \"strict\": false," + " \"requestors\": [\"tutorial\"]}"; + + status = mta_system_add_pairs(system, options, pairs); + if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to add pairs to system: %s\n", error_message); + mts_labels_free(samples); + mts_labels_free(component); + mts_labels_free(properties); + mta_system_free(system); + return NULL; + } + + mts_labels_free(samples); + mts_labels_free(component); + mts_labels_free(properties); + + // ------------------------------------------------------------------ // + // Custom data: a TensorMap with a single block of per-atom values + + double custom_values[] = {0.42, -0.31, 0.15, -0.08}; + int32_t custom_keys[] = {0}; + int32_t custom_samples[] = {0, 1, 2, 3}; + int32_t custom_properties[] = {0}; + + const char* key_dims[] = {"_"}; + struct mts_array_t key_array = make_mts_array( + custom_keys, (uintptr_t[]){1, 1}, 2, i32_dtype + ); + const mts_labels_t* keys = mts_labels(key_dims, 1, key_array); + + const char* custom_sample_dims[] = {"atom"}; + struct mts_array_t custom_samples_array = make_mts_array( + custom_samples, (uintptr_t[]){4, 1}, 2, i32_dtype + ); + const mts_labels_t* custom_labels = mts_labels( + custom_sample_dims, 1, custom_samples_array + ); + + const char* custom_prop_dims[] = {"property"}; + struct mts_array_t custom_prop_array = make_mts_array( + custom_properties, (uintptr_t[]){1, 1}, 2, i32_dtype + ); + const mts_labels_t* custom_labels_props = mts_labels( + custom_prop_dims, 1, custom_prop_array + ); + + struct mts_array_t values_array_custom = make_mts_array( + custom_values, (uintptr_t[]){4, 1}, 2, f64_dtype + ); + mts_block_t* block = mts_block( + values_array_custom, custom_labels, NULL, 0, custom_labels_props + ); + assert(block != NULL); + + mts_block_t* blocks[] = {block}; + mts_tensormap_t* custom = mts_tensormap(keys, blocks, 1); + assert(custom != NULL); + + status = mta_system_add_custom_data(system, "tutorial::charges", custom); + if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to add custom data: %s\n", error_message); + mts_labels_free(custom_labels); + mts_labels_free(custom_labels_props); + mta_system_free(system); + return NULL; + } + + mts_labels_free(custom_labels); + mts_labels_free(custom_labels_props); + + return system; +} + +// %% +// +// .. raw:: html +// +//
+ +int main(void) { + +// %% +// +// Let's get a system from somewhere, and have a look as what's inside + +mta_system_t* system = create_system_for_tutorial(); + +if (system == NULL) { + fprintf(stderr, "Failed to create tutorial system\n"); + return EXIT_FAILURE; +} + +// %% +// +// Global information +// ------------------ +// +// First, we can access some global information about a system, such as the +// number of atoms making up the system: + +uintptr_t size = 0; +mta_status_t status = mta_system_size(system, &size); + +if (status != MTA_SUCCESS) { + fprintf(stderr, "failed to get system size\n"); + mta_system_free(system); + return EXIT_FAILURE; +} +assert(size == 4); +printf("this system contains %lu atoms\n", (unsigned long)size); + +// %% +// +// We can also access the unit used for all length data in this system. +// +// This is returned as a newly allocated :c:type:`mta_string_t`. Callers of +// functions that return :c:type:`mta_string_t` can get a view (i.e. a ``const +// char*`` to a null-terminated string) and are responsible for freeing the +// memory. + +mta_string_t length_unit = NULL; +status = mta_system_get_length_unit(system, &length_unit); +if (status != MTA_SUCCESS) { + fprintf(stderr, "failed to get length unit\n"); + mta_system_free(system); + return EXIT_FAILURE; +} +printf("length unit: %s\n", mta_string_view(length_unit)); +assert(strcmp(mta_string_view(length_unit), "Angstrom") == 0); + +mta_string_free(length_unit); + +// %% +// +// Tensor data +// ----------- +// +// The main tensors in the system (positions, cell, atomic types and pbc) can +// all be accessed with :c:func:`mta_system_get_data`, passing a different +// :c:enum:`mta_system_data_kind` for each tensor. + +DLManagedTensorVersioned* positions = NULL; +status = mta_system_get_data(system, MTA_SYSTEM_DATA_POSITIONS, &positions); + +// this tensor contains 64-bit floats +assert(positions->dl_tensor.dtype.code == kDLFloat); +assert(positions->dl_tensor.dtype.bits == 64); +assert(positions->dl_tensor.dtype.lanes == 1); + +double* pos_ptr = (double*)((uint8_t*)positions->dl_tensor.data + positions->dl_tensor.byte_offset); + +assert(positions->dl_tensor.ndim == 2); +assert(positions->dl_tensor.shape[0] == 4); +assert(positions->dl_tensor.shape[1] == 3); + +// check that the code is contiguous and row-major +if (positions->dl_tensor.strides != NULL) { + assert(positions->dl_tensor.strides[0] == 3); + assert(positions->dl_tensor.strides[1] == 1); +} + +// Acces the data in a linear fashion +assert(pos_ptr[0] == 0.0 && pos_ptr[1] == 0.0 && pos_ptr[2] == 0.0); +assert(pos_ptr[3] == 0.5 && pos_ptr[4] == 0.5 && pos_ptr[5] == 0.0); +assert(pos_ptr[6] == 0.5 && pos_ptr[7] == 0.0 && pos_ptr[8] == 0.5); +assert(pos_ptr[9] == 0.0 && pos_ptr[10] == 0.5 && pos_ptr[11] == 0.5); + +// %% +// +// When done with a dlpack tensor, one must release it +if (positions->deleter) { + positions->deleter(positions); +} + +// %% +// +// We can do something similar with the pbc data: + +DLManagedTensorVersioned* pbc = NULL; +status = mta_system_get_data(system, MTA_SYSTEM_DATA_PBC, &pbc); + +// this tensor contains booleans +assert(pbc->dl_tensor.dtype.code == kDLBool); +assert(pbc->dl_tensor.dtype.bits == 8); +assert(pbc->dl_tensor.dtype.lanes == 1); + +bool* pbc_ptr = (bool*)((uint8_t*)pbc->dl_tensor.data + pbc->dl_tensor.byte_offset); + +assert(pbc->dl_tensor.ndim == 1); +assert(pbc->dl_tensor.shape[0] == 3); + +// check that the code is contiguous and row-major +if (pbc->dl_tensor.strides != NULL) { + assert(pbc->dl_tensor.strides[0] == 1); +} + +// Acces the data in a linear fashion +assert(pbc_ptr[0] == true && pbc_ptr[1] == true && pbc_ptr[2] == true); + +// release the dlpack tensor +if (pbc->deleter) { + pbc->deleter(pbc); +} + +// %% +// +// Pair lists +// ---------- +// +// A pair list (neighbor list) is a :c:type:`mts_block_t` where each sample +// represents a pair of atoms, and the values contains the distance vector +// between them. +// +// The pair list is identified by its JSON-serialized pair list options, which +// are typically declared by a model through +// :c:func:`mta_model_t.requested_pair_lists`. + +const char* pair_options = + "{\"type\": \"metatomic_pair_list_options\"," + " \"cutoff\": \"0x4008000000000000\"," + " \"full_list\": true," + " \"strict\": false," + " \"requestors\": [\"tutorial\"]}"; + +// The block returned by this function is a borrowed view, do not free it. +const mts_block_t* pairs = NULL; +status = mta_system_get_pairs(system, pair_options, &pairs); +if (status != MTA_SUCCESS) { + fprintf(stderr, "failed to get pairs from system\n"); + mta_system_free(system); + return EXIT_FAILURE; +} +assert(pairs != NULL); +printf("successfully retrieved pair list\n"); + +// %% +// +// We can also list all known pair lists in the system: + +mta_string_t known = NULL; +status = mta_system_known_pairs(system, &known); +if (status == MTA_SUCCESS && known != NULL) { + printf("known pair lists: %s\n", mta_string_view(known)); + mta_string_free(known); +} + +// %% +// +// Custom data +// ----------- +// +// Custom data allows models to attach arbitrary per-system data to a system, +// stored as a named :c:type:`mts_tensormap_t`. The name must follow the usual +// quantity naming convention. + + +// The returned tensor map is a borrowed view, do not free it. +const mts_tensormap_t* retrieved = NULL; +status = mta_system_get_custom_data(system, "tutorial::charges", &retrieved); +if (status != MTA_SUCCESS) { + fprintf(stderr, "failed to get custom data from system\n"); + mta_system_free(system); + return EXIT_FAILURE; +} +assert(retrieved != NULL); +printf("successfully retrieved custom data 'tutorial::charges'\n"); + +// %% +// +// And we can list all known custom data names: + +mta_string_t names = NULL; +status = mta_system_known_custom_data(system, &names); +if (status == MTA_SUCCESS && names != NULL) { + printf("known custom data: %s\n", mta_string_view(names)); + mta_string_free(names); +} + +// %% +// +// Once done, we can cleanup the system. + +status = mta_system_free(system); +if (status != MTA_SUCCESS) { + fprintf(stderr, "failed to free system memory\n"); + return EXIT_FAILURE; +}; + +return EXIT_SUCCESS; } diff --git a/examples/c/3-model.c b/examples/c/3-model.c new file mode 100644 index 000000000..40dd3e734 --- /dev/null +++ b/examples/c/3-model.c @@ -0,0 +1,882 @@ +// .. _c-tutorial-model: +// +// Defining a custom model +// ======================= +// +// We will now explore how to implement a metatomic model purely in C. While +// most users will never need to do this, it is useful to understand the +// underlying mechanics, and when writing new bindings to the metatomic C API. +// +// In the C API, a metatomic model is represented by the :c:type:`mta_model_t` +// struct. This struct contains a ``void*`` data pointer, and multiple function +// pointers, making what's sometimes called a "vtable" or `virtual table +// `_. Each of the function +// pointer take the data pointer as a first argument, allowing the model to +// store its own private data and state. The vtable contains functions to query +// the model's metadata, capabilities, requested inputs, pair-list options, and +// requested outputs, as well as a function to execute the model. +// +// The model in this tutorial will provide a single ``"energy"`` output, +// computing a shifted Lennard-Jones pair potential: +// +// .. math:: +// +// E = 4 \epsilon \left[ +// \left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^{6} +// \right] - E_{\mathrm{shift}}, +// +// with :math:`E_{\mathrm{shift}}` chosen so the energy is zero at the cutoff +// :math:`r_c`. + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +// %% +// +// Some helpers for the tutorial + +#include "utils/dlpack.h" // tensor_from_data +#include "utils/mts_array.h" // make_mts_array + +// Read a mts_block_t values as a DLPack tensor. The DLpack tensor must be +// released with `view->deleter(view)` when done. +static mts_status_t block_dlpack_data(const mts_block_t* block, DLManagedTensorVersioned** view) { + struct mts_array_t array = {0}; + mts_status_t status = mts_block_data((mts_block_t*)block, &array); + if (status != MTS_SUCCESS) { + return status; + } + + DLDevice cpu = {.device_type = kDLCPU, .device_id = 0}; + DLPackVersion version = {.major = DLPACK_MAJOR_VERSION, .minor = DLPACK_MINOR_VERSION}; + status = array.as_dlpack(array.ptr, view, cpu, NULL, version); + + if (array.destroy != NULL) { + array.destroy(array.ptr); + } + + return status; +} + +// Read a mts_labels_t values as a DLPack tensor. The DLpack tensor must be +// released with `view->deleter(view)` when done. +static mts_status_t labels_dlpack_data(const mts_labels_t* labels, DLManagedTensorVersioned** view) { + struct mts_array_t array = {0}; + mts_status_t status = mts_labels_values(labels, &array); + if (status != MTS_SUCCESS) { + return status; + } + + DLDevice cpu = {.device_type = kDLCPU, .device_id = 0}; + DLPackVersion version = {.major = DLPACK_MAJOR_VERSION, .minor = DLPACK_MINOR_VERSION}; + status = array.as_dlpack(array.ptr, view, cpu, NULL, version); + + if (array.destroy != NULL) { + array.destroy(array.ptr); + } + + return status; +} + +// Get a CPU pointer to the float64 data in a DLPack tensor. The tensor must be +// contiguous and have the right type. +static double* dlpack_double_data(DLManagedTensorVersioned* view) { + assert(view != NULL); + assert(view->dl_tensor.device.device_type == kDLCPU); + assert(view->dl_tensor.dtype.code == kDLFloat); + assert(view->dl_tensor.dtype.bits == 64); + assert(view->dl_tensor.dtype.lanes == 1); + assert((view->flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0); + + assert(view->dl_tensor.ndim >= 1); + assert(view->dl_tensor.shape != NULL); + assert(view->dl_tensor.strides == NULL || view->dl_tensor.strides[view->dl_tensor.ndim - 1] == 1); + + return (double*)((uint8_t*)view->dl_tensor.data + view->dl_tensor.byte_offset); +} + +// Get a CPU pointer to the int32 data in a DLPack tensor. The tensor must be +// contiguous and have the right type. +static int32_t* dlpack_int32_data(DLManagedTensorVersioned* view) { + assert(view != NULL); + assert(view->dl_tensor.device.device_type == kDLCPU); + assert(view->dl_tensor.dtype.code == kDLInt); + assert(view->dl_tensor.dtype.bits == 32); + assert(view->dl_tensor.dtype.lanes == 1); + assert((view->flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0); + + assert(view->dl_tensor.ndim >= 1); + assert(view->dl_tensor.shape != NULL); + assert(view->dl_tensor.strides == NULL || view->dl_tensor.strides[view->dl_tensor.ndim - 1] == 1); + + return (int32_t*)((uint8_t*)view->dl_tensor.data + view->dl_tensor.byte_offset); +} + +// %% +// +// The model's physics +// ------------------- +// +// Let's start by implementing the Lennard-Jones potential, nothing too fancy +// here. The model's private data will be a struct storing the parameters +// required by the computation. + +typedef struct LennardJones { + double sigma; + double epsilon; + double cutoff; + double shift; +} LennardJones; + +// Compute the shift so the potential is zero at the cutoff. +static double lj_shift(double cutoff, double sigma, double epsilon) { + double x = sigma / cutoff; + double x6 = x * x * x * x * x * x; + return 4.0 * epsilon * (x6 * x6 - x6); +} + +// Compute the energy and force contribution from one pair of atoms. The force +// is returned on the first atom; the second atom's force is the same vector, +// negated. +static void lj_pair( + const LennardJones* parameters, + const double distance[3], + double* energy, + double force_on_first[3] +) { + double r2 = distance[0] * distance[0] + + distance[1] * distance[1] + + distance[2] * distance[2]; + + double cutoff2 = parameters->cutoff * parameters->cutoff; + if (r2 <= 0.0 || r2 >= cutoff2) { + *energy = 0.0; + force_on_first[0] = force_on_first[1] = force_on_first[2] = 0.0; + return; + } + + double sigma_r_2 = (parameters->sigma * parameters->sigma) / r2; + double sigma_r_6 = sigma_r_2 * sigma_r_2 * sigma_r_2; + double sigma_r_12 = sigma_r_6 * sigma_r_6; + *energy = 4.0 * parameters->epsilon * (sigma_r_12 - sigma_r_6) - parameters->shift; + + // dE/d(r^2) = (12 epsilon / r^2) (inv6 - 2 inv12) + // r^2 = |pos_2 - pos_1|^2, so d(r^2)/d(pos_1) = -2 d + // force on atom 1 is -dE/d(pos_1) = 2 dE/d(r^2) d + double dedr2 = (12.0 * parameters->epsilon / r2) * (sigma_r_6 - 2.0 * sigma_r_12); + force_on_first[0] = 2.0 * dedr2 * distance[0]; + force_on_first[1] = 2.0 * dedr2 * distance[1]; + force_on_first[2] = 2.0 * dedr2 * distance[2]; +} + +// %% +// +// Creating the ``mta_model_t`` callbacks +// -------------------------------------- +// +// We need to implement the following callbacks for the model: +// +// - :c:func:`mta_model_t.unload`: free the model's private data +// - :c:func:`mta_model_t.metadata`: return a JSON string describing the model +// - :c:func:`mta_model_t.capabilities`: return a JSON string describing the +// model's capabilities +// - :c:func:`mta_model_t.requested_pair_lists`: return a JSON string describing +// the pair-list options requested by the model +// - :c:func:`mta_model_t.requested_inputs`: return a JSON string describing any +// extra inputs +// - :c:func:`mta_model_t.execute_inner`: compute the model's outputs for the +// given systems +// +// All the callbacks take the model's private data pointer as their first +// argument, and return a status code. :c:func:`mta_model_t.unload` is +// straightforward: it just frees the private data. + +static mta_status_t lj_unload(void* model_data) { + free(model_data); + return MTA_SUCCESS; +} + +// %% +// +// :c:func:`mta_model_t.metadata` returns the human-facing description of the +// model: name, authors, references, etc. All the information is passed around +// as JSON strings, which the caller takes ownership of. The expected JSON +// structure is documented in :ref:`this page `. +// +// :c:type:`mta_string_t` is used to allocate and pass strings around instead of +// plain NULL-terminated ``const char*`` so that metatomic can manage the memory +// in a way that is compatible with all supported languages. + + +static mta_status_t lj_metadata(const void* model_data, mta_string_t* out) { + (void)model_data; + *out = mta_string_create("{" + "\"type\": \"metatomic_model_metadata\"," + "\"name\": \"Lennard-Jones model for tutorials\"," + "\"authors\": [\"metatomic authors\"]," + "\"references\": {\"model\": [], \"architecture\": [], \"implementation\": []}" + "}"); + return (*out != NULL) ? MTA_SUCCESS : MTA_INTERNAL_ERROR; +} + +// %% +// +// :c:func:`mta_model_t.capabilities` tells an engine what the model can +// compute. Here the energy output advertises a ``"positions"`` gradient, +// meaning callers can request forces from this model. For a short range model +// like this one, the ``"interaction_range"`` is set to the cutoff. +// +// Other notable information in the capabilities is the list of supported atomic +// types, the supported devices, and the data type used for all inputs and +// outputs. +// +// If any of the callback in :c:type:`mta_model_t` fail, they should return a +// non-success status code and set an error message with +// :c:func:`mta_set_last_error`. The error message is a UTF-8 string, which can +// be retrieved by the caller with :c:func:`mta_last_error`. + +static mta_status_t error(mta_status_t status, const char* message) { + assert(status != MTA_SUCCESS); + mta_set_last_error( + /*message=*/message, + /*origin=*/"model tutorial", + /*data=*/NULL, + /*data_deleter=*/NULL + ); + return status; +} + +static mta_status_t lj_capabilities(const void* model_data, mta_string_t* out) { + const LennardJones* parameters = (const LennardJones*)model_data; + char json[512]; + int printed = snprintf(json, sizeof(json), + "{" + "\"type\": \"metatomic_model_capabilities\"," + "\"outputs\": [{" + "\"type\": \"metatomic_quantity\"," + "\"name\": \"energy\"," + "\"unit\": \"eV\"," + "\"gradients\": [\"positions\"]," + "\"sample_kind\": \"system\"" + "}]," + "\"atomic_types\": [1]," + "\"interaction_range\": %.17g," + "\"length_unit\": \"Angstrom\"," + "\"supported_devices\": [\"cpu\"]," + "\"dtype\": \"float64\"" + "}", + parameters->cutoff + ); + + if (printed < 0 || (size_t)printed >= sizeof(json)) { + return error(MTA_MODEL_ERROR, "failed to format capabilities JSON"); + } + + *out = mta_string_create(json); + + if (*out == NULL) { + return error(MTA_MEMORY_ERROR, "failed to allocate capabilities JSON"); + } + + return MTA_SUCCESS; +} + +// %% +// +// The model can request pair lists (neighor lists) from the engine, which will +// be computed by the engine and attached to the :c:type:`mta_system_t` before +// calling the model. Multiple pair lists with different options can all be +// requested simulataneously, and the engine will compute them all. +// +// These requests are also exchanged as JSON. Here we request a half list (each +// pair once) with the model's cutoff. The cutoff is represented by its IEEE-754 +// bit pattern to make sure the value can not be changed by serialization. + +// format the pair options for the given cutoff +static mta_status_t format_pair_options(double cutoff, char* buffer, size_t size) { + // extract the bit pattern of the cutoff + uint64_t bits; + memcpy(&bits, &cutoff, sizeof(bits)); + // format the bits as hex string + char cutoff_hex[32]; + snprintf(cutoff_hex, sizeof(cutoff_hex), "0x%" PRIx64, bits); + + int printed = snprintf(buffer, size, "{" + "\"type\": \"metatomic_pair_list_options\"," + " \"cutoff\": \"%s\"," + " \"full_list\": false," + " \"strict\": true," + "\"requestors\": [\"lj-tutorial\"]" + "}", + cutoff_hex + ); + + if (printed < 0 || (size_t)printed >= size) { + return error(MTA_MODEL_ERROR, "failed to format pair list options JSON"); + } + + return MTA_SUCCESS; +} + +static mta_status_t lj_requested_pair_lists(const void* model_data, mta_string_t* out) { + const LennardJones* parameters = (const LennardJones*)model_data; + + char json[512] = {0}; + mta_status_t status = format_pair_options(parameters->cutoff, json, sizeof(json)); + if (status != MTA_SUCCESS) { + return status; + } + + char json_array[512] = {0}; + int printed = snprintf(json_array, sizeof(json_array), "[%s]", json); + + if (printed < 0 || (size_t)printed >= sizeof(json_array)) { + return error(MTA_MODEL_ERROR, "failed to format pair list options JSON"); + } + + *out = mta_string_create(json_array); + + if (*out == NULL) { + return error(MTA_MEMORY_ERROR, "failed to allocate pair list options JSON"); + } + + return MTA_SUCCESS; +} + +// %% +// +// Finallty, the model can also request additional inputs from the engine, which +// are passed as a JSON-formatted list of :ref:`core-json-quantity` objects. +// Here we don't need any extra inputs, so we return an empty list. + +static mta_status_t lj_requested_inputs(const void* model_data, mta_string_t* out) { + (void)model_data; + *out = mta_string_create("[]"); // no extra inputs beyond the system + if (*out == NULL) { + return error(MTA_MEMORY_ERROR, "failed to allocate requested inputs JSON"); + } + return MTA_SUCCESS; +} + +// %% +// +// ``mta_model_t.execute_inner`` +// ----------------------------- +// +// The most important callback in :c:type:`mta_model_t` is +// :c:func:`mta_model_t.execute_inner`, which actually executes the model. +// +// It is named ``execute_inner`` to indicate that it should not be directly +// called by the engine, but rather through the free function +// :c:func:`mta_execute_model`, which handles unit conversion and consistency +// checks. +// +// This function takes the following parameters: +// +// - ``model_data`` is a pointer to the model data, here ``LennardJones*``; +// - ``systems`` is an array of systems on which the model shoudl be executed; +// - ``systems_count`` is the size of the ``systems`` array; +// - ``selected_atoms`` is an optional ``mts_labels_t`` object with "system" and +// "atom" dimensions, indicating which atoms should be included in the +// outputs. If ``NULL``, all atoms participate in the output; +// - ``requested_outputs_json`` is a JSON-encoded list of +// :ref:`core-json-quantity` that the model should compute as outputs; +// - ``outputs`` is an array (of the same size as ``requested_outputs_json``) +// with space for all requested outputs; +// - ``outputs_count`` is the size of the ``output`` array. + +// helper to create the energy tensormap output from the energy and forces +static mta_status_t create_energy_tensormap( + double energy, + const double* energy_gradient, + uintptr_t n_atoms, + mts_tensormap_t** tensor +); + +static mta_status_t lj_execute_inner( + void* model_data, + const mta_system_t* const* systems, + uintptr_t systems_count, + const mts_labels_t* selected_atoms, + const char* requested_outputs_json, + mts_tensormap_t** outputs, + uintptr_t outputs_count +) { + if (selected_atoms != NULL) { + return error(MTA_INVALID_PARAMETER_ERROR, "this model does not support selected_atoms"); + } + + if (outputs_count == 0) { + // no output requested + return MTA_SUCCESS; + } + + // if some output was requested, we assume it is the total energy, this + // should be checked properly in an actual model + (void)requested_outputs_json; + + if (systems_count != 1) { + return error(MTA_INVALID_PARAMETER_ERROR, "this model only supports a single system"); + } + + const LennardJones* parameters = (const LennardJones*)model_data; + const mta_system_t* system = systems[0]; + + uintptr_t n_atoms = 0; + assert(mta_system_size(system, &n_atoms) == MTA_SUCCESS); + + + // get the pair list the engine computed, using the pair options JSON as a key + char options_json[512] = {0}; + mta_status_t status = format_pair_options(parameters->cutoff, options_json, sizeof(options_json)); + if (status != MTA_SUCCESS) { + return status; + } + + const mts_block_t* pairs = NULL; + status = mta_system_get_pairs(system, options_json, &pairs); + if (status != MTA_SUCCESS) { + return status; + } + + // extract the array from the block, and get a CPU pointer to the values. + // The array is a 2D array of shape (n_pairs, 3, 1), with the 3D distance + // vector for each pair. + DLManagedTensorVersioned* distances_dlpack = NULL; + mts_status_t mts_status = block_dlpack_data(pairs, &distances_dlpack); + if (mts_status != MTS_SUCCESS) { + return error(MTA_METATENSOR_ERROR, "failed to get pair list distances as DLPack tensor"); + } + assert(distances_dlpack != NULL); + assert(distances_dlpack->dl_tensor.ndim == 3); + assert(distances_dlpack->dl_tensor.shape[1] == 3); + assert(distances_dlpack->dl_tensor.shape[2] == 1); + + double* distances = dlpack_double_data(distances_dlpack); + int64_t n_pairs = distances_dlpack->dl_tensor.shape[0]; + + // get the samples associated with the pairs `mts_block_t`. These contain + // the indices of the two atoms in each pair, which we will need to compute + // the forces. + const mts_labels_t* pairs_samples = mts_block_labels(pairs, 0); + DLManagedTensorVersioned* samples_dlpack = NULL; + mts_status = labels_dlpack_data(pairs_samples, &samples_dlpack); + if (mts_status != MTS_SUCCESS) { + return error(MTA_METATENSOR_ERROR, "failed to get pair list samples as DLPack tensor"); + } + assert(samples_dlpack != NULL); + assert(samples_dlpack->dl_tensor.ndim == 2); + // the dimensions of the samples are [i, j, shift_a, shift_b, shift_c] + assert(samples_dlpack->dl_tensor.shape[1] == 5); + + int32_t* samples = dlpack_int32_data(samples_dlpack); + + double* energy_gradient = calloc(n_atoms * 3, sizeof(double)); + double energy = 0.0; + for (int64_t pair_i = 0; pair_i < n_pairs; pair_i++) { + double vector[3] = { + distances[3 * pair_i + 0], + distances[3 * pair_i + 1], + distances[3 * pair_i + 2] + }; + + double pair_energy = 0.0; + double force_on_first[3] = {0.0, 0.0, 0.0}; + lj_pair(parameters, vector, &pair_energy, force_on_first); + energy += pair_energy; + + int32_t i = samples[pair_i * 5 + 0]; + int32_t j = samples[pair_i * 5 + 1]; + for (int d = 0; d < 3; d++) { + energy_gradient[i * 3 + d] -= force_on_first[d]; + energy_gradient[j * 3 + d] += force_on_first[d]; // Newton's third law + } + } + + // fill the output + status = create_energy_tensormap(energy, energy_gradient, n_atoms, outputs); + + // cleanup + free(energy_gradient); + + if (distances_dlpack->deleter != NULL) { + distances_dlpack->deleter(distances_dlpack); + } + if (samples_dlpack->deleter != NULL) { + samples_dlpack->deleter(samples_dlpack); + } + + mts_labels_free(pairs_samples); + + return status; +} + +// %% +// +// .. raw:: html +// +//
Implementation of create_energy_tensormap() + +static mta_status_t create_energy_tensormap( + double energy, + const double* energy_gradient, + uintptr_t n_atoms, + mts_tensormap_t** tensor +) { + DLDataType f64_dtype = {.code = kDLFloat, .bits = 64, .lanes = 1}; + DLDataType i32_dtype = {.code = kDLInt, .bits = 32, .lanes = 1}; + + // create all the labels and arrays for the energy block + mts_array_t values = make_mts_array(&energy, (uintptr_t[]){1, 1}, 2, f64_dtype); + + // the samples are a single-element array with a single "system" sample + int32_t zero = 0; + mts_array_t samples_values = make_mts_array(&zero, (uintptr_t[]){1, 1}, 2, i32_dtype); + + const char* system_dims[] = {"system"}; + const mts_labels_t* samples = mts_labels(system_dims, 1, samples_values); + if (samples == NULL) { + return error(MTA_METATENSOR_ERROR, "failed to create samples labels"); + } + + // properties are a single-element array with a single "energy" property + struct mts_array_t properties_values = make_mts_array(&zero, (uintptr_t[]){1, 1}, 2, i32_dtype); + const char* energy_dims[] = {"energy"}; + const mts_labels_t* properties = mts_labels(energy_dims, 1, properties_values); + if (properties == NULL) { + return error(MTA_METATENSOR_ERROR, "failed to create properties labels"); + } + + // create the energy block with no components + mts_block_t* block = mts_block(values, samples, NULL, 0, properties); + if (block == NULL) { + return error(MTA_METATENSOR_ERROR, "failed to create energy block"); + } + + // samples for the gradients of the energy + int32_t* gradient_samples_scratch = malloc(sizeof(int32_t) * n_atoms * 3); + for (uintptr_t i = 0; i < n_atoms; i++) { + gradient_samples_scratch[3 * i + 0] = 0; // parent energy sample (row 0) + gradient_samples_scratch[3 * i + 1] = 0; // system index + gradient_samples_scratch[3 * i + 2] = (int32_t)i; // atom i + } + const char* grad_sample_dims[] = {"sample", "system", "atom"}; + mts_array_t grad_sample_values = make_mts_array( + gradient_samples_scratch, (uintptr_t[]){n_atoms, 3}, 2, i32_dtype + ); + free(gradient_samples_scratch); + + const mts_labels_t* gradient_samples = mts_labels(grad_sample_dims, 3, grad_sample_values); + if (gradient_samples == NULL) { + return error(MTA_METATENSOR_ERROR, "failed to create gradient samples labels"); + } + + // components for the gradients of the energy + int32_t xyz_values[] = {0, 1, 2}; + struct mts_array_t xyz_array = make_mts_array( + xyz_values, (uintptr_t[]){3, 1}, 2, i32_dtype + ); + const char* xyz_dims[] = {"xyz"}; + const mts_labels_t* xyz = mts_labels(xyz_dims, 1, xyz_array); + if (xyz == NULL) { + return error(MTA_METATENSOR_ERROR, "failed to create xyz components labels"); + } + + // create the gradient block and attach it to the energy block + mts_array_t gradient_values = make_mts_array( + energy_gradient, (uintptr_t[]){n_atoms, 3, 1}, 3, f64_dtype + ); + + mts_block_t* gradient_block = mts_block(gradient_values, gradient_samples, &xyz, 1, properties); + if (gradient_block == NULL) { + return error(MTA_METATENSOR_ERROR, "failed to create energy gradient block"); + } + + mts_status_t status = mts_block_add_gradient(block, "positions", gradient_block); + if (status != MTS_SUCCESS) { + return error(MTA_METATENSOR_ERROR, "failed to attach energy gradient block"); + } + + // create the keys for the energy tensormap + struct mts_array_t key_array = make_mts_array(&zero, (uintptr_t[]){1, 1}, 2, i32_dtype); + const char* key_dims[] = {"_"}; // energy is always a single-block map + const mts_labels_t* keys = mts_labels(key_dims, 1, key_array); + if (keys == NULL) { + return error(MTA_METATENSOR_ERROR, "failed to create keys labels"); + } + + mts_block_t* blocks[] = {block}; + *tensor = mts_tensormap(keys, blocks, 1); + if (*tensor == NULL) { + return error(MTA_METATENSOR_ERROR, "failed to create energy tensormap"); + } + + mts_labels_free(samples); + mts_labels_free(properties); + mts_labels_free(xyz); + mts_labels_free(keys); + return MTA_SUCCESS; +} + +// %% +// +// .. raw:: html +// +//
+// +// +// Running the model +// ----------------- +// +// We now have all the building blocks to create a model. This would typically +// be done through a :c:type:`mta_plugin_t`, which we will explore in more +// details in the :ref:`next tutorial `. + +static mta_model_t make_lennard_jones_model(double cutoff) { + LennardJones* data = malloc(sizeof(LennardJones)); + assert(data != NULL); + data->sigma = 1.0; + data->epsilon = 1.0; + data->cutoff = cutoff; + data->shift = lj_shift(data->cutoff, data->sigma, data->epsilon); + + mta_model_t model = { + .data = data, + .unload = lj_unload, + .metadata = lj_metadata, + .capabilities = lj_capabilities, + .requested_pair_lists = lj_requested_pair_lists, + .requested_inputs = lj_requested_inputs, + .execute_inner = lj_execute_inner, + }; + return model; +} + +// %% +// +// +// We'll use a system containing two atoms ``distance`` apart along *z*, with a +// single pair between them. See the previous tutorials +// (:ref:`c-tutorial-create-system` and :ref:`c-tutorial-use-system`) for more +// details. + + +static mta_system_t* make_two_atom_system(double distance, double cutoff); + +// %% +// +// .. raw:: html +// +//
Implementation of make_two_atom_system() + +static mta_system_t* make_two_atom_system(double distance, double cutoff) { + DLDataType i32_dtype = {.code = kDLInt, .bits = 32, .lanes = 1}; + DLDataType bool_dtype = {.code = kDLBool, .bits = 8, .lanes = 1}; + DLDataType f64_dtype = {.code = kDLFloat, .bits = 64, .lanes = 1}; + + static double positions_data[6]; + positions_data[0] = 0.0; positions_data[1] = 0.0; positions_data[2] = 0.0; + positions_data[3] = 0.0; positions_data[4] = 0.0; positions_data[5] = distance; + + // non-periodic: cell must be all zeros + static double cell_data[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; + static int32_t types_data[] = {1, 1}; + static bool pbc_data[] = {false, false, false}; + + DLManagedTensorVersioned* positions = tensor_from_data( + positions_data, (int64_t[]){2, 3}, 2, f64_dtype + ); + DLManagedTensorVersioned* cell = tensor_from_data( + cell_data, (int64_t[]){3, 3}, 2, f64_dtype + ); + DLManagedTensorVersioned* types = tensor_from_data( + types_data, (int64_t[]){2}, 1, i32_dtype + ); + DLManagedTensorVersioned* pbc = tensor_from_data( + pbc_data, (int64_t[]){3}, 1, bool_dtype + ); + + mta_system_t* system = NULL; + mta_status_t create_status = mta_system_create("Angstrom", types, positions, cell, pbc, &system); + if (create_status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to create system: %s\n", error_message); + return NULL; + } + + int32_t pair_samples[] = {0, 1, 0, 0, 0}; + struct mts_array_t samples_array = make_mts_array( + pair_samples, (uintptr_t[]){1, 5}, 2, i32_dtype + ); + const char* sample_dims[] = {"first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"}; + const mts_labels_t* samples = mts_labels(sample_dims, 5, samples_array); + + int32_t xyz_values[] = {0, 1, 2}; + struct mts_array_t xyz_array = make_mts_array( + xyz_values, (uintptr_t[]){3, 1}, 2, i32_dtype + ); + const char* xyz_dims[] = {"xyz"}; + const mts_labels_t* xyz = mts_labels(xyz_dims, 1, xyz_array); + const mts_labels_t* components[] = {xyz}; + + int32_t zero = 0; + struct mts_array_t prop_array = make_mts_array( + &zero, (uintptr_t[]){1, 1}, 2, i32_dtype + ); + const char* distance_dims[] = {"distance"}; + const mts_labels_t* properties = mts_labels(distance_dims, 1, prop_array); + + double disp_data[] = {0.0, 0.0, distance}; + struct mts_array_t values = make_mts_array( + disp_data, (uintptr_t[]){1, 3, 1}, 3, f64_dtype + ); + mts_block_t* pairs = mts_block(values, samples, components, 1, properties); + + char options[512] = {0}; + format_pair_options(cutoff, options, sizeof(options)); + mta_system_add_pairs(system, options, pairs); + + mts_labels_free(samples); + mts_labels_free(xyz); + mts_labels_free(properties); + return system; +} + +// %% +// +// .. raw:: html +// +//
+ +int main(void) { + +// %% +// +// Now, we create the model and a system with two atoms 1.3 σ apart, just past +// the potential minimum, so they attract. We then use +// :c:func:`mta_execute_model` to run the model on the system, and request the +// energy and its gradient with respect to the atomic positions. + +double cutoff = 3.0; + +mta_model_t model = make_lennard_jones_model(cutoff); +mta_system_t* system = make_two_atom_system(/*distance=*/1.3, cutoff); +const mta_system_t* systems[] = {system}; + +const char* requested_outputs = "[{" + "\"type\": \"metatomic_quantity\"," + "\"name\": \"energy\"," + "\"unit\": \"eV\"," + "\"gradients\": [\"positions\"]," + "\"sample_kind\": \"system\"" +"}]"; + +mts_tensormap_t* output = NULL; +mta_status_t status; +status = mta_execute_model( + /*model=*/model, + /*systems=*/systems, + /*systems_count=*/1, + /*selected_atoms=*/NULL, + /*requested_outputs_json=*/requested_outputs, + /*check_consistency=*/true, + /*outputs=*/&output, + /*outputs_count=*/1 +); + +bool failed = false; +DLManagedTensorVersioned* energy_dlpack = NULL; +DLManagedTensorVersioned* gradient_dlpack = NULL; + +if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to run model: %s\n", error_message); + + failed = true; + goto cleanup; +} + +// %% +// +// Finally, we can look into the output and extract the energy and forces. The +// energy is a single scalar, while the forces are a 2×3 array (2 atoms, 3 +// dimensions). The forces are the negative of the gradient of the energy with +// respect to the atomic positions. + +mts_block_t* block = NULL; +mts_status_t mts_status = mts_tensormap_block_by_id(output, &block, 0); +if (mts_status != MTS_SUCCESS) { + fprintf(stderr, "failed to get energy block from output\n"); + failed = true; + goto cleanup; +} + +mts_status = block_dlpack_data(block, &energy_dlpack); +if (mts_status != MTS_SUCCESS) { + fprintf(stderr, "failed to get energy block as DLPack tensor\n"); + failed = true; + goto cleanup; +} +double energy = dlpack_double_data(energy_dlpack)[0]; +assert(fabs(energy - (-0.651537)) < 1e-6); + +mts_block_t* gradient_block = NULL; +mts_status = mts_block_gradient(block, "positions", &gradient_block); +if (mts_status != MTS_SUCCESS) { + fprintf(stderr, "failed to get gradient block\n"); + failed = true; + goto cleanup; +} + +mts_status = block_dlpack_data(gradient_block, &gradient_dlpack); +if (mts_status != MTS_SUCCESS) { + fprintf(stderr, "failed to get gradient block as DLPack tensor\n"); + failed = true; + goto cleanup; +} +double* gradient = dlpack_double_data(gradient_dlpack); + +// -forces on the first atom +assert(gradient[0] == 0.0); +assert(gradient[1] == 0.0); +assert(fabs(gradient[2] - (-2.239980)) < 1e-6); + +// -forces on the second atom +assert(gradient[3] == 0.0); +assert(gradient[4] == 0.0); +assert(fabs(gradient[5] - (2.239980)) < 1e-6); + +// %% +// + +cleanup: + +if (energy_dlpack != NULL && energy_dlpack->deleter != NULL) { + energy_dlpack->deleter(energy_dlpack); +} + +if (gradient_dlpack != NULL && gradient_dlpack->deleter != NULL) { + gradient_dlpack->deleter(gradient_dlpack); +} + +mts_tensormap_free(output); +mta_system_free(system); +model.unload(model.data); + +if (failed) { + return EXIT_FAILURE; +} else { + return EXIT_SUCCESS; +} + +} diff --git a/examples/c/4-plugin.c b/examples/c/4-plugin.c new file mode 100644 index 000000000..955b3575a --- /dev/null +++ b/examples/c/4-plugin.c @@ -0,0 +1,273 @@ +// .. _c-tutorial-plugin: +// +// Creating a plugin +// ================= +// +// The :ref:`previous tutorial ` constructed a model directly, +// which requires its implementation to be linked with the program using it. +// +// Plugins allow models to be loaded at runtime using shared libraries. This +// way, simulation engines can use different models without containing any +// model-specific code. The engine will register plugins with +// :c:func:`mta_load_plugin` and then load models with :c:func:`mta_load_model`. +// The latter will then query all loaded plugin to find one that can load the +// requested model. +// +// In the C API, plugins are represented by the :c:type:`mta_plugin_t` struct, +// which contains a :c:func:`mta_plugin_t.load_model` callback. + +#include +#include +#include + +#include + +// %% +// +// A minimal model +// --------------- +// +// Let's start by defining a stub model that does not actually compute anything, +// but implements the required callbacks. For more details on the model +// callbacks, see the :ref:`previous tutorial `. + +static mta_status_t error(mta_status_t status, const char* message) { + assert(status != MTA_SUCCESS); + mta_set_last_error( + /*message=*/message, + /*origin=*/"plugin tutorial", + /*data=*/NULL, + /*data_deleter=*/NULL + ); + return status; +} + +// mta_model_t.unload implementation +static mta_status_t stub_unload(void* model_data) { + (void)model_data; + return MTA_SUCCESS; +} + +// mta_model_t.metadata implementation +static mta_status_t stub_metadata(const void* model_data, mta_string_t* out) { + (void)model_data; + *out = mta_string_create("{" + "\"type\": \"metatomic_model_metadata\"," + "\"name\": \"Stub model for tutorials\"," + "\"authors\": [\"metatomic authors\"]," + "\"references\": {\"model\": [], \"architecture\": [], \"implementation\": []}" + "}"); + + if (*out == NULL) { + return error(MTA_MEMORY_ERROR, "failed to allocate metadata JSON"); + } + return MTA_SUCCESS; +} + +// mta_model_t.capabilities implementation +static mta_status_t stub_capabilities(const void* model_data, mta_string_t* out) { + (void)model_data; + *out = mta_string_create("{" + "\"type\": \"metatomic_model_capabilities\"," + "\"outputs\": []," + "\"atomic_types\": []," + "\"interaction_range\": 0.0," + "\"length_unit\": \"Angstrom\"," + "\"supported_devices\": [\"cpu\"]," + "\"dtype\": \"float64\"" + "}"); + + if (*out == NULL) { + return error(MTA_MEMORY_ERROR, "failed to allocate capabilities JSON"); + } + return MTA_SUCCESS; +} + +// mta_model_t.requested_pair_lists implementation +static mta_status_t stub_requested_pair_lists(const void* model_data, mta_string_t* out) { + (void)model_data; + *out = mta_string_create("[]"); + if (*out == NULL) { + return error(MTA_MEMORY_ERROR, "failed to allocate requested pair lists JSON"); + } + return MTA_SUCCESS; +} +// mta_model_t.requested_inputs implementation +static mta_status_t stub_requested_inputs(const void* model_data, mta_string_t* out) { + (void)model_data; + *out = mta_string_create("[]"); + if (*out == NULL) { + return error(MTA_MEMORY_ERROR, "failed to allocate requested inputs JSON"); + } + return MTA_SUCCESS; +} + +static mta_status_t stub_execute_inner( + void* model_data, + const mta_system_t* const* systems, + uintptr_t systems_count, + const mts_labels_t* selected_atoms, + const char* requested_outputs_json, + mts_tensormap_t** outputs, + uintptr_t outputs_count +) { + (void)model_data; + (void)systems; + (void)systems_count; + (void)selected_atoms; + (void)requested_outputs_json; + (void)outputs; + (void)outputs_count; + return error(MTA_MODEL_ERROR, "this model is a stub and does not compute anything"); +} + +// %% +// +// The plugin +// ---------- +// +// Each plugin must implement the ``load_model`` callback. This function is +// called by :c:func:`mta_load_model` to load a model, and given a string +// indicating which model to load and a JSON string with options. The +// ``load_from`` string indicates which model to load, and typically is either a +// model name or the path to a model file. The JSON options contain any model +// and/or plugin parameters as a JSON object. +// +// It MUST either fills the :c:type:`mta_model_t` vtable or returns +// :c:enumerator:`MTA_UNSUPPORTED_MODEL_ERROR` to indicate that the plugin does +// not know how to load the requested model. + +static mta_status_t demo_load_model( + const char* load_from, + const char* options_json, + mta_model_t* model +) { + (void)options_json; + if (strcmp(load_from, "plugin-tutorial-stub-model") != 0) { + return MTA_UNSUPPORTED_MODEL_ERROR; + } + + model->unload = stub_unload; + model->metadata = stub_metadata; + model->capabilities = stub_capabilities; + model->requested_pair_lists = stub_requested_pair_lists; + model->requested_inputs = stub_requested_inputs; + model->execute_inner = stub_execute_inner; + return MTA_SUCCESS; +} + +// %% +// +// We then need to make the plugin accessible to metatomic and +// :c:func:`mta_load_plugin`. This is done by the :c:macro:`MTA_REGISTER_PLUGIN` +// macro. +// +// This macro takes the name of a function that will be called to register the +// plugin, and a block that creates the :c:type:`mta_plugin_t` struct, and use +// them to define a registration function. The block must return a status code +// indicating sucess or failure. + +MTA_REGISTER_PLUGIN(register_plugin, { + fprintf(stderr, "registering plugin 'tutorial-plugin'\n"); + mta_plugin_t plugin = { + .abi_version = MTA_ABI_VERSION, + .name = "tutorial-plugin", + .load_model = demo_load_model, + }; + return register_plugin(plugin); +}); + +// %% +// + +int main(void) { + + +// %% +// +// We can load the plugin from the current binary by using ``NULL`` as the path +// in :c:func:`mta_load_plugin`. +// +// .. warning:: +// +// To be able to load a plugin directly from an executable, all symbols +// should be exported. This is the default on macOS and Windows, but on +// linux we need to link the executable with ``-Wl,--export-dynamic``, or +// set the ``ENABLE_EXPORTS`` property on the target in CMake. + +mta_status_t status = mta_load_plugin(NULL); + +if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to load plugin: %s\n", error_message); + return EXIT_FAILURE; +} + +// %% +// +// Once all relevant plugins have been loaded, we can load a model with +// :c:func:`mta_load_model`. The easiest way to do this is to leave +// ``plugin_name=NULL`` and let metatomic find a plugin that can load the model. + +mta_model_t model = {0}; +status = mta_load_model( + /*load_from=*/"plugin-tutorial-stub-model", + /*options_json=*/NULL, + /*plugin_name=*/NULL, + &model +); + +if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to load model: %s\n", error_message); + return EXIT_FAILURE; +} + +// %% +// +// Once loaded, we can use the model as usual. For example, we can query the +// metadata + +mta_string_t metadata = NULL; +status = model.metadata(model.data, &metadata); +if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to get model metadata: %s\n", error_message); + return EXIT_FAILURE; +} + +assert(metadata != NULL); +assert(strstr(mta_string_view(metadata), "Stub model for tutorials") != NULL); +mta_string_free(metadata); + +model.unload(model.data); + +// %% +// +// We can also specify the plugin to use when loading a model, preventing +// metatomic to try to use a different plugin to load this model: + +memset(&model, 0, sizeof(model)); +status = mta_load_model( + /*load_from=*/"plugin-tutorial-stub-model", + /*options_json=*/NULL, + /*plugin_name=*/"tutorial-plugin", + &model +); + +if (status != MTA_SUCCESS) { + const char* error_message = NULL; + mta_last_error(&error_message, NULL, NULL); + fprintf(stderr, "failed to load model: %s\n", error_message); + return EXIT_FAILURE; +} +model.unload(model.data); + +// %% +// + +return EXIT_SUCCESS; +} diff --git a/examples/c/README.rst b/examples/c/README.rst new file mode 100644 index 000000000..b56164cd3 --- /dev/null +++ b/examples/c/README.rst @@ -0,0 +1,4 @@ +.. _c-tutorials: + +C API tutorials +=============== diff --git a/examples/c/utils/dlpack.h b/examples/c/utils/dlpack.h new file mode 100644 index 000000000..1a9088f2a --- /dev/null +++ b/examples/c/utils/dlpack.h @@ -0,0 +1,94 @@ +#ifndef MTA_EXAMPLE_UTILS_DLPACK_H +#define MTA_EXAMPLE_UTILS_DLPACK_H + +#include +#include +#include + +#include + +// Context for a DLPack tensor, which keeps the allocations for shape and +// strides. +typedef struct CustomDLPackContext { + int64_t* shape; + int64_t* strides; +} CustomDLPackContext; + +// Deleter for a DLPack tensor, which frees the context and the tensor itself. +// We do not free the data buffer, as it is owned by the caller. +static inline void dlpack_deleter(DLManagedTensorVersioned *self) { + if (!self) { + return; + } + + CustomDLPackContext* ctx = (CustomDLPackContext*)self->manager_ctx; + if (ctx) { + free(ctx->shape); + free(ctx->strides); + free(ctx); + } + free(self); +} + +// Create a DLPack tensor from a flat data buffer. The tensor is created as a +// row-major, contiguous tensor on CPU, with the specified shape and data type. +// The caller owns the data buffer, and is responsible for freeing it after the +// tensor is no longer needed. +static inline DLManagedTensorVersioned* tensor_from_data( + void *data, + const int64_t *shape, + int32_t ndim, + DLDataType dtype +) { + CustomDLPackContext* ctx = malloc(sizeof(CustomDLPackContext)); + if (!ctx) { + return NULL; + } + + ctx->shape = malloc(ndim * sizeof(int64_t)); + ctx->strides = malloc(ndim * sizeof(int64_t)); + if (!ctx->shape || !ctx->strides) { + free(ctx->shape); + free(ctx->strides); + free(ctx); + return NULL; + } + memcpy(ctx->shape, shape, ndim * sizeof(int64_t)); + + int64_t stride = 1; + for (int32_t i = ndim - 1; i >= 0; i--) { + ctx->strides[i] = stride; + stride *= shape[i]; + } + + DLManagedTensorVersioned* tensor = calloc(1, sizeof(*tensor)); + if (!tensor) { + free(ctx->shape); + free(ctx->strides); + free(ctx); + return NULL; + } + + tensor->version.major = DLPACK_MAJOR_VERSION; + tensor->version.minor = DLPACK_MINOR_VERSION; + tensor->manager_ctx = ctx; + tensor->deleter = dlpack_deleter; + + tensor->flags = DLPACK_FLAG_BITMASK_READ_ONLY; + + tensor->dl_tensor.data = data; + tensor->dl_tensor.byte_offset = 0; + + tensor->dl_tensor.device.device_type = kDLCPU; + tensor->dl_tensor.device.device_id = 0; + + tensor->dl_tensor.dtype = dtype; + + tensor->dl_tensor.ndim = ndim; + tensor->dl_tensor.shape = ctx->shape; + tensor->dl_tensor.strides = ctx->strides; + + return tensor; +} + +#endif // MTA_EXAMPLE_UTILS_DLPACK_H diff --git a/examples/c/utils/mts_array.h b/examples/c/utils/mts_array.h new file mode 100644 index 000000000..a32677c79 --- /dev/null +++ b/examples/c/utils/mts_array.h @@ -0,0 +1,133 @@ +#ifndef MTA_EXAMPLE_UTILS_ARRAY_H +#define MTA_EXAMPLE_UTILS_ARRAY_H + +#include +#include +#include + +#include + +#include "./dlpack.h" + +// To work with metatensor's ``mts_block_t`` and ``mts_tensormap_t`` in C, we +// need an ``mts_array_t`` - a vtable-based abstraction over n-dimensional +// arrays. Below is a minimal implementation backed by a flat data buffer that +// the array owns: the data is copied into a heap allocation when the array is +// created, and released when the array is destroyed. + +// Context data for the array, which is stored in the ``ptr`` field of the +// ``mts_array_t``. The context is owned by the array, and will be freed when +// the array is destroyed. +typedef struct BasicMtsArray { + void* data; + uintptr_t ndim; + uintptr_t shape[4]; + DLDataType dtype; +} BasicMtsArray; + +// Destroy the array, freeing the data buffer and the context itself. +static inline void array_destroy(void* array) { + BasicMtsArray* a = (BasicMtsArray*)array; + free(a->data); + free(a); +} + +// Return the origin of the array. +static inline mts_status_t array_origin(const void* array, mts_data_origin_t* origin) { + static mts_data_origin_t BASIC_MTS_ARRAY_ORIGIN = 0; + + (void)array; + if (BASIC_MTS_ARRAY_ORIGIN == 0) { + mts_register_data_origin("tutorial-mts-array", &BASIC_MTS_ARRAY_ORIGIN); + } + *origin = BASIC_MTS_ARRAY_ORIGIN; + return MTS_SUCCESS; +} + +// Return the device of the array. +static inline mts_status_t array_device(const void* array, DLDevice* device) { + (void)array; + device->device_type = kDLCPU; + device->device_id = 0; + return MTS_SUCCESS; +} + +// Return the data type of the array. +static inline mts_status_t array_dtype(const void* array, DLDataType* dtype) { + *dtype = ((const BasicMtsArray*)array)->dtype; + return MTS_SUCCESS; +} + +// Return the array as a DLPack tensor. The DLpack tensor must be released with +// `tensor->deleter(tensor)` when done. +static inline mts_status_t array_as_dlpack( + void* array, + DLManagedTensorVersioned** tensor, + DLDevice device, + const int64_t* stream, + DLPackVersion max_version +) { + (void)stream; + (void)max_version; + BasicMtsArray* a = (BasicMtsArray*)array; + if (device.device_type != kDLCPU) { + return MTS_CALLBACK_ERROR; + } + static_assert(sizeof(uintptr_t) == sizeof(int64_t), "int64_t and uintptr_t must be the same size"); + *tensor = tensor_from_data(a->data, (const int64_t*)a->shape, (int32_t)a->ndim, a->dtype); + return MTS_SUCCESS; +} + +// Get the shape of the array. +static inline mts_status_t array_shape( + const void* array, + const uintptr_t** shape, + uintptr_t* shape_count +) { + const BasicMtsArray* a = (const BasicMtsArray*)array; + *shape = a->shape; + *shape_count = a->ndim; + return MTS_SUCCESS; +} + +// Create a new ``mts_array_t`` backed by a flat data buffer. The array will +// copy the data into a heap allocation, and will free it when the array is +// destroyed. The array will have the specified shape and data type. +static inline struct mts_array_t make_mts_array( + const void* data, const uintptr_t* shape, uintptr_t ndim, DLDataType dtype +) { + BasicMtsArray* raw = malloc(sizeof(BasicMtsArray)); + + // copy the data into a buffer owned by the array + size_t data_size = 1; + for (uintptr_t i = 0; i < ndim; i++) { + data_size *= shape[i]; + } + data_size *= (dtype.bits / 8); + raw->data = malloc(data_size); + memcpy(raw->data, data, data_size); + + raw->ndim = ndim; + for (uintptr_t i = 0; i < ndim; i++) { + raw->shape[i] = shape[i]; + } + raw->dtype = dtype; + + struct mts_array_t result = {0}; + result.ptr = raw; + result.destroy = array_destroy; + result.origin = array_origin; + result.device = array_device; + result.dtype = array_dtype; + result.as_dlpack = array_as_dlpack; + result.shape = array_shape; + result.from_dlpack = NULL; + result.reshape = NULL; + result.swap_axes = NULL; + result.create = NULL; + result.copy = NULL; + result.move_data = NULL; + return result; +} + +#endif // MTA_EXAMPLE_UTILS_ARRAY_H diff --git a/examples/torch/.gitignore b/examples/torch/.gitignore new file mode 100644 index 000000000..b0096484e --- /dev/null +++ b/examples/torch/.gitignore @@ -0,0 +1,2 @@ +*.pt +*.json.gz diff --git a/python/examples/1-export-atomistic-model.py b/examples/torch/1-export-atomistic-model.py similarity index 100% rename from python/examples/1-export-atomistic-model.py rename to examples/torch/1-export-atomistic-model.py diff --git a/python/examples/3-atomistic-model-with-nl.py b/examples/torch/2-atomistic-model-with-nl.py similarity index 99% rename from python/examples/3-atomistic-model-with-nl.py rename to examples/torch/2-atomistic-model-with-nl.py index 2072df0c9..4cd5f420f 100644 --- a/python/examples/3-atomistic-model-with-nl.py +++ b/examples/torch/2-atomistic-model-with-nl.py @@ -16,7 +16,7 @@ by the simulation engine and attached to the :py:class:`Systems`. The :py:class:`Systems` with the neighbor list is then passed to the model. -.. figure:: ../../static/images/nl-dataflow.* +.. figure:: ../../../static/images/nl-dataflow.* :width: 600px :align: center diff --git a/python/examples/4-profiling.py b/examples/torch/3-profiling.py similarity index 100% rename from python/examples/4-profiling.py rename to examples/torch/3-profiling.py diff --git a/examples/torch/README.rst b/examples/torch/README.rst new file mode 100644 index 000000000..3fa1a8830 --- /dev/null +++ b/examples/torch/README.rst @@ -0,0 +1,4 @@ +.. _torch-tutorials: + +PyTorch tutorials +================= diff --git a/python/examples/liquid-argon.xyz b/examples/torch/liquid-argon.xyz similarity index 100% rename from python/examples/liquid-argon.xyz rename to examples/torch/liquid-argon.xyz diff --git a/examples/torchsim/.gitignore b/examples/torchsim/.gitignore new file mode 100644 index 000000000..b0096484e --- /dev/null +++ b/examples/torchsim/.gitignore @@ -0,0 +1,2 @@ +*.pt +*.json.gz diff --git a/python/examples/5-torchsim-getting-started.py b/examples/torchsim/1-getting-started.py similarity index 96% rename from python/examples/5-torchsim-getting-started.py rename to examples/torchsim/1-getting-started.py index 3e9d6e091..f82734d0d 100644 --- a/python/examples/5-torchsim-getting-started.py +++ b/examples/torchsim/1-getting-started.py @@ -25,8 +25,11 @@ from typing import Dict, List, Optional import ase.build +import matplotlib.pyplot as plt import torch from metatensor.torch import Labels, TensorBlock, TensorMap +from torch_sim.integrators import nve_init, nve_step +from torch_sim.units import MetalUnits import metatomic.torch as mta from metatomic_torchsim import MetatomicModel @@ -151,10 +154,6 @@ def forward( # ``nve_init`` samples momenta from a Maxwell-Boltzmann distribution at the # given temperature, and ``nve_step`` advances by one timestep: -import matplotlib.pyplot as plt # noqa: E402 -from torch_sim.integrators import nve_init, nve_step # noqa: E402 -from torch_sim.units import MetalUnits # noqa: E402 - sim_state = ts.initialize_state(atoms, device=model.device, dtype=model.dtype) diff --git a/python/examples/6-torchsim-batched.py b/examples/torchsim/2-batched-md.py similarity index 100% rename from python/examples/6-torchsim-batched.py rename to examples/torchsim/2-batched-md.py diff --git a/examples/torchsim/README.rst b/examples/torchsim/README.rst new file mode 100644 index 000000000..0ef6aa268 --- /dev/null +++ b/examples/torchsim/README.rst @@ -0,0 +1,4 @@ +.. _torchsim-tutorials: + +Torch-Sim tutorials +=================== diff --git a/metatomic-core/CHANGELOG.md b/metatomic-core/CHANGELOG.md new file mode 100644 index 000000000..160995db2 --- /dev/null +++ b/metatomic-core/CHANGELOG.md @@ -0,0 +1,18 @@ +# Changelog + +All notable changes to metatomic-core are documented here, following the [keep +a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows +[Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased](https://github.com/metatensor/metatensor/) + + diff --git a/metatomic-core/CMakeLists.txt b/metatomic-core/CMakeLists.txt new file mode 100644 index 000000000..0eb97f20f --- /dev/null +++ b/metatomic-core/CMakeLists.txt @@ -0,0 +1,463 @@ +# This file defines the CMake build system for the C and C++ API of metatomic. +# +# This API is implemented in Rust, in the metatomic-core crate, but Rust users +# of the API should use the metatomic crate instead, wrapping metatomic-core in +# an easier to use, idiomatic Rust API. +cmake_minimum_required(VERSION 3.22) + +# Is metatomic the main project configured by the user? Or is this being used +# as a submodule/subdirectory? +if (${CMAKE_CURRENT_SOURCE_DIR} STREQUAL ${CMAKE_SOURCE_DIR}) + set(METATOMIC_MAIN_PROJECT ON) +else() + set(METATOMIC_MAIN_PROJECT OFF) +endif() + +if(${METATOMIC_MAIN_PROJECT} AND NOT "${CACHED_LAST_CMAKE_VERSION}" VERSION_EQUAL ${CMAKE_VERSION}) + # We use CACHED_LAST_CMAKE_VERSION to only print the cmake version + # once in the configuration log + set(CACHED_LAST_CMAKE_VERSION ${CMAKE_VERSION} CACHE INTERNAL "Last version of cmake used to configure") + message(STATUS "Running CMake version ${CMAKE_VERSION}") +endif() + +if (POLICY CMP0077) + # use variables to set OPTIONS + cmake_policy(SET CMP0077 NEW) +endif() + +file(STRINGS "Cargo.toml" CARGO_TOML_CONTENT) +foreach(line ${CARGO_TOML_CONTENT}) + string(REGEX REPLACE "^version = \"(.*)\"" "\\1" METATOMIC_VERSION ${line}) + if (NOT ${CMAKE_MATCH_COUNT} EQUAL 0) + # stop on the first regex match, this should be the right version + break() + endif() +endforeach() + +include(cmake/dev-versions.cmake) +create_development_version("${METATOMIC_VERSION}" METATOMIC_FULL_VERSION "metatomic-core-v") +message(STATUS "Building metatomic-core v${METATOMIC_FULL_VERSION}") + +# strip any -dev/-rc suffix on the version since project(VERSION) does not support it +string(REGEX REPLACE "([0-9]*)\\.([0-9]*)\\.([0-9]*).*" "\\1.\\2.\\3" METATOMIC_VERSION ${METATOMIC_FULL_VERSION}) +project(metatomic + VERSION ${METATOMIC_VERSION} + LANGUAGES C CXX # we need to declare a language to access CMAKE_SIZEOF_VOID_P later +) +set(PROJECT_VERSION ${METATOMIC_FULL_VERSION}) + + +# We follow the standard CMake convention of using BUILD_SHARED_LIBS to provide +# either a shared or static library as a default target. But since cargo always +# builds both versions by default, we also install both versions by default. +# `METATOMIC_INSTALL_BOTH_STATIC_SHARED=OFF` allow to disable this behavior, and +# only install the file corresponding to `BUILD_SHARED_LIBS=ON/OFF`. +# +# BUILD_SHARED_LIBS controls the `metatomic` cmake target, making it an alias of +# either `metatomic::static` or `metatomic::shared`. This is mainly relevant +# when using metatomic from another cmake project, either as a submodule or from +# an installed library (see cmake/metatomic-config.cmake) +option(BUILD_SHARED_LIBS "Use a shared library by default instead of a static one" ON) +option(METATOMIC_INSTALL_BOTH_STATIC_SHARED "Install both shared and static libraries" ON) + +set(RUST_BUILD_TARGET "${RUST_BUILD_TARGET}" CACHE STRING "Cross-compilation target for rust code. Leave empty to build for the host") +set(EXTRA_RUST_FLAGS "${EXTRA_RUST_FLAGS}" CACHE STRING "Flags used to build rust code") + +include(GNUInstallDirs) + +if("${CMAKE_BUILD_TYPE}" STREQUAL "" AND "${CMAKE_CONFIGURATION_TYPES}" STREQUAL "") + message(STATUS "Setting build type to 'release' as none was specified.") + set(CMAKE_BUILD_TYPE "release" + CACHE STRING + "Choose the type of build, options are: debug or release" + FORCE) + set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS release debug) +endif() + +if(${METATOMIC_MAIN_PROJECT} AND NOT "${CACHED_LAST_CMAKE_BUILD_TYPE}" STREQUAL "${CMAKE_BUILD_TYPE}") + set(CACHED_LAST_CMAKE_BUILD_TYPE ${CMAKE_BUILD_TYPE} CACHE INTERNAL "Last build type used in configuration") + message(STATUS "Building metatomic in ${CMAKE_BUILD_TYPE} mode") +endif() + + +function(check_compatible_versions _actual_ _requested_) + if(${_actual_} MATCHES "^([0-9]+)\\.([0-9]+)") + set(_actual_major_ "${CMAKE_MATCH_1}") + set(_actual_minor_ "${CMAKE_MATCH_2}") + else() + message(FATAL_ERROR "Failed to parse actual version: ${_actual_}") + endif() + + if(${_requested_} MATCHES "^([0-9]+)\\.([0-9]+)") + set(_requested_major_ "${CMAKE_MATCH_1}") + set(_requested_minor_ "${CMAKE_MATCH_2}") + else() + message(FATAL_ERROR "Failed to parse requested version: ${_requested_}") + endif() + + if (${_requested_major_} EQUAL 0 AND ${_actual_minor_} EQUAL ${_requested_minor_}) + # major version is 0 and same minor version, everything is fine + elseif (${_actual_major_} EQUAL ${_requested_major_}) + # same major version, everything is fine + else() + # not compatible + message(FATAL_ERROR "Incompatible versions: we need ${_requested_}, but we got ${_actual_}") + endif() +endfunction() + + +set(REQUIRED_METATENSOR_VERSION "0.2.4") +# Either metatensor is built as part of the same CMake project, or we try to +# find the corresponding CMake package +if (TARGET metatensor) + get_target_property(METATENSOR_BUILD_VERSION metatensor BUILD_VERSION) + check_compatible_versions(${METATENSOR_BUILD_VERSION} ${REQUIRED_METATENSOR_VERSION}) +else() + find_package(metatensor ${REQUIRED_METATENSOR_VERSION} CONFIG REQUIRED) +endif() + +include(cmake/nlohmann_json.cmake) + +include(cmake/detect_cargo.cmake) + +# ============================================================================ # +# determine Cargo flags + +set(CARGO_BUILD_ARG "") + +if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.lock) + set(CARGO_BUILD_ARG "${CARGO_BUILD_ARG};--locked") +endif() + +# TODO: support multiple configuration generators (MSVC, ...) +string(TOLOWER ${CMAKE_BUILD_TYPE} BUILD_TYPE) +if ("${BUILD_TYPE}" STREQUAL "debug") + set(CARGO_BUILD_TYPE "debug") +elseif("${BUILD_TYPE}" STREQUAL "release") + set(CARGO_BUILD_ARG "${CARGO_BUILD_ARG};--release") + set(CARGO_BUILD_TYPE "release") +elseif("${BUILD_TYPE}" STREQUAL "relwithdebinfo") + set(CARGO_BUILD_ARG "${CARGO_BUILD_ARG};--release") + set(CARGO_BUILD_TYPE "release") +else() + message(FATAL_ERROR "unsuported build type: ${CMAKE_BUILD_TYPE}") +endif() + +set(CARGO_TARGET_DIR ${CMAKE_CURRENT_BINARY_DIR}/target) +set(CARGO_BUILD_ARG "${CARGO_BUILD_ARG};--target-dir=${CARGO_TARGET_DIR}") + +if (WIN32) + # on Windows, we need to use the same ABI in both CMake and cargo. If the + # user did not explicitly request a target, we can try to set it ourself, + # otherwise we just check that it matches what we expect. + if (MSVC) + if ("${RUST_BUILD_TARGET}" STREQUAL "") + set(RUST_BUILD_TARGET "${RUST_HOST_ARCH}-pc-windows-msvc") + message(STATUS "Setting rust target to ${RUST_BUILD_TARGET}") + elseif(NOT "${RUST_BUILD_TARGET}" MATCHES "-pc-windows-msvc") + message(FATAL_ERROR "CMake is building with MSVC but the Rust target is ${RUST_BUILD_TARGET}") + endif() + endif() + + if (MINGW) + if ("${RUST_BUILD_TARGET}" STREQUAL "") + set(RUST_BUILD_TARGET "${RUST_HOST_ARCH}-pc-windows-gnu") + message(STATUS "Setting rust target to ${RUST_BUILD_TARGET}") + elseif(NOT "${RUST_BUILD_TARGET}" MATCHES "-pc-windows-gnu") + message(FATAL_ERROR "CMake is building with MinGW but the Rust target is ${RUST_BUILD_TARGET}") + endif() + endif() +endif() + +# Handle cross compilation with RUST_BUILD_TARGET +if ("${RUST_BUILD_TARGET}" STREQUAL "") + if (${METATOMIC_MAIN_PROJECT}) + message(STATUS "Compiling to host (${RUST_HOST_TARGET})") + endif() + + set(CARGO_OUTPUT_DIR "${CARGO_TARGET_DIR}/${CARGO_BUILD_TYPE}") + set(RUST_BUILD_TARGET ${RUST_HOST_TARGET}) +else() + if (${METATOMIC_MAIN_PROJECT}) + message(STATUS "Cross-compiling to ${RUST_BUILD_TARGET}") + endif() + + set(CARGO_BUILD_ARG "${CARGO_BUILD_ARG};--target=${RUST_BUILD_TARGET}") + set(CARGO_OUTPUT_DIR "${CARGO_TARGET_DIR}/${RUST_BUILD_TARGET}/${CARGO_BUILD_TYPE}") +endif() + +# Get the list of libraries linked by default by cargo/rustc to add when linking +# to metatomic::static +if (CARGO_VERSION_CHANGED) + include(cmake/tempdir.cmake) + get_tempdir(TMPDIR) + + # Adapted from https://github.com/corrosion-rs/corrosion/blob/dc1e4e5/cmake/FindRust.cmake + execute_process( + COMMAND "${CARGO_EXE}" new --lib _cargo_required_libs + WORKING_DIRECTORY "${TMPDIR}" + RESULT_VARIABLE cargo_new_result + ERROR_QUIET + ) + + if (cargo_new_result) + message(FATAL_ERROR "could not create empty project to find default static libs: ${cargo_new_result}") + endif() + + file(APPEND "${TMPDIR}/_cargo_required_libs/Cargo.toml" "[lib]\ncrate-type=[\"staticlib\"]") + + execute_process( + COMMAND ${CARGO_EXE} rustc --color never --target=${RUST_BUILD_TARGET} -- --print=native-static-libs + WORKING_DIRECTORY "${TMPDIR}/_cargo_required_libs" + RESULT_VARIABLE cargo_static_libs_result + ERROR_VARIABLE cargo_static_libs_stderr + ) + + # clean up the files + file(REMOVE_RECURSE "${TMPDIR}") + + if (cargo_static_libs_result) + message(FATAL_ERROR + "could not extract default static libs (status ${cargo_static_libs_result}), stderr:\n${cargo_static_libs_stderr}" + ) + endif() + + # The pattern starts with `native-static-libs:` and goes to the end of the line. + if (cargo_static_libs_stderr MATCHES "native-static-libs: ([^\r\n]+)\r?\n") + string(REPLACE " " ";" "libs_list" "${CMAKE_MATCH_1}") + set(stripped_lib_list "") + foreach(lib ${libs_list}) + # Strip leading `-l` (unix) and potential .lib suffix (windows) + string(REGEX REPLACE "^-l" "" "stripped_lib" "${lib}") + string(REGEX REPLACE "\.lib$" "" "stripped_lib" "${stripped_lib}") + list(APPEND stripped_lib_list "${stripped_lib}") + endforeach() + + # Special case `msvcrt` to link with the debug version in Debug mode. + list(TRANSFORM stripped_lib_list REPLACE "^msvcrt$" "\$<\$:msvcrtd>") + # Don't try to pass a linker *flag* where CMake expects libraries + list(REMOVE_ITEM stripped_lib_list "/defaultlib:msvcrt") + + if (APPLE) + # Prevent warnings about duplicated `System` in linked libraries + # from Apple's `ld` + list(REMOVE_ITEM stripped_lib_list "System") + endif() + + list(REMOVE_DUPLICATES stripped_lib_list) + set(CARGO_DEFAULT_LIBRARIES "${stripped_lib_list}" CACHE INTERNAL "list of implicitly linked libraries") + + if (${METATOMIC_MAIN_PROJECT}) + message(STATUS "Cargo default link libraries are: ${CARGO_DEFAULT_LIBRARIES}") + endif() + else() + message(FATAL_ERROR "could not find default static libs: `native-static-libs` not found in: `${cargo_static_libs_stderr}`") + endif() +endif() + +file(GLOB_RECURSE ALL_RUST_SOURCES + ${PROJECT_SOURCE_DIR}/Cargo.toml + ${PROJECT_SOURCE_DIR}/src/**.rs +) + +add_library(metatomic::shared SHARED IMPORTED GLOBAL) +set(METATOMIC_SHARED_LOCATION "${CARGO_OUTPUT_DIR}/${CMAKE_SHARED_LIBRARY_PREFIX}metatomic${CMAKE_SHARED_LIBRARY_SUFFIX}") +set(METATOMIC_IMPLIB_LOCATION "${METATOMIC_SHARED_LOCATION}.lib") + +if (MINGW) + # `rustc` does not follow the usual naming scheme for DLL with mingw (it + # would typically be 'libmetatomic.dll') + set(METATOMIC_SHARED_LOCATION "${CARGO_OUTPUT_DIR}/metatomic.dll") + set(METATOMIC_IMPLIB_LOCATION "${CARGO_OUTPUT_DIR}/libmetatomic.dll.a") +endif() + +add_library(metatomic::static STATIC IMPORTED GLOBAL) +set(METATOMIC_STATIC_LOCATION "${CARGO_OUTPUT_DIR}/${CMAKE_STATIC_LIBRARY_PREFIX}metatomic${CMAKE_STATIC_LIBRARY_SUFFIX}") + +get_filename_component(METATOMIC_SHARED_LIB_NAME ${METATOMIC_SHARED_LOCATION} NAME) +get_filename_component(METATOMIC_IMPLIB_NAME ${METATOMIC_IMPLIB_LOCATION} NAME) +get_filename_component(METATOMIC_STATIC_LIB_NAME ${METATOMIC_STATIC_LOCATION} NAME) + +# We need to add some metadata to the shared library to enable linking to it +# without using an absolute path. +if (UNIX) + if (APPLE) + # set the install name to `@rpath/libmetatomic.dylib` + set(CARGO_RUSTC_ARGS "-Clink-arg=-Wl,-install_name,@rpath/${METATOMIC_SHARED_LIB_NAME}") + set_target_properties(metatomic::shared PROPERTIES + IMPORTED_SONAME @rpath/${METATOMIC_SHARED_LIB_NAME} + ) + else() # LINUX + # set the SONAME to libmetatomic.so + set(CARGO_RUSTC_ARGS "-Clink-arg=-Wl,-soname,${METATOMIC_SHARED_LIB_NAME}") + set_target_properties(metatomic::shared PROPERTIES + IMPORTED_SONAME ${METATOMIC_SHARED_LIB_NAME} + ) + endif() +else() + set(CARGO_RUSTC_ARGS "") +endif() + +if (NOT "${EXTRA_RUST_FLAGS}" STREQUAL "") + set(CARGO_RUSTC_ARGS "${CARGO_RUSTC_ARGS};${EXTRA_RUST_FLAGS}") +endif() + +# Set environment variables for cargo build +set(CARGO_ENV "METATOMIC_FULL_VERSION=${METATOMIC_FULL_VERSION}") +if (NOT "${CMAKE_OSX_DEPLOYMENT_TARGET}" STREQUAL "") + list(APPEND CARGO_ENV "MACOSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}") +endif() + +if (METATOMIC_INSTALL_BOTH_STATIC_SHARED) + set(CARGO_BUILD_ARG "${CARGO_BUILD_ARG};--crate-type=cdylib;--crate-type=staticlib") + set(CARGO_OUTPUTS ${METATOMIC_SHARED_LOCATION} ${METATOMIC_STATIC_LOCATION}) + if (WIN32) + list(APPEND CARGO_OUTPUTS ${METATOMIC_IMPLIB_LOCATION}) + set(FILE_CREATED_MESSAGE "${METATOMIC_SHARED_LIB_NAME}, ${METATOMIC_STATIC_LIB_NAME}, and ${METATOMIC_IMPLIB_NAME}") + else() + set(FILE_CREATED_MESSAGE "${METATOMIC_SHARED_LIB_NAME} and ${METATOMIC_STATIC_LIB_NAME}") + endif() +else() + if (BUILD_SHARED_LIBS) + set(CARGO_BUILD_ARG "${CARGO_BUILD_ARG};--crate-type=cdylib") + set(CARGO_OUTPUTS ${METATOMIC_SHARED_LOCATION}) + if (WIN32) + list(APPEND CARGO_OUTPUTS ${METATOMIC_IMPLIB_LOCATION}) + set(FILE_CREATED_MESSAGE "${METATOMIC_SHARED_LIB_NAME} and ${METATOMIC_IMPLIB_NAME}") + else() + set(FILE_CREATED_MESSAGE "${METATOMIC_SHARED_LIB_NAME}") + endif() + else() + set(CARGO_BUILD_ARG "${CARGO_BUILD_ARG};--crate-type=staticlib") + set(CARGO_OUTPUTS ${METATOMIC_STATIC_LOCATION}) + set(FILE_CREATED_MESSAGE "${METATOMIC_STATIC_LIB_NAME}") + endif() +endif() + +add_custom_command( + OUTPUT ${CARGO_OUTPUTS} + COMMAND ${CMAKE_COMMAND} -E env ${CARGO_ENV} + ${CARGO_EXE} rustc ${CARGO_BUILD_ARG} -- ${CARGO_RUSTC_ARGS} + WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} + DEPENDS ${ALL_RUST_SOURCES} + COMMENT "Building ${FILE_CREATED_MESSAGE} with cargo" + VERBATIM +) +add_custom_target(cargo-build-metatomic ALL DEPENDS ${CARGO_OUTPUTS}) + +# Auto-generate a header containing the version number as #define +set(_path_ "${CMAKE_CURRENT_BINARY_DIR}/generated-version.h") +file(WRITE ${_path_} "#pragma once\n\n") +file(APPEND ${_path_} "/** Full version of metatomic as a string */\n") +file(APPEND ${_path_} "#define METATOMIC_VERSION \"${METATOMIC_FULL_VERSION}\"\n\n") +file(APPEND ${_path_} "/** Major version number of metatomic as an integer */\n") +file(APPEND ${_path_} "#define METATOMIC_VERSION_MAJOR ${PROJECT_VERSION_MAJOR}\n\n") +file(APPEND ${_path_} "/** Minor version number of metatomic as an integer */\n") +file(APPEND ${_path_} "#define METATOMIC_VERSION_MINOR ${PROJECT_VERSION_MINOR}\n\n") +file(APPEND ${_path_} "/** Patch version number of metatomic as an integer */\n") +file(APPEND ${_path_} "#define METATOMIC_VERSION_PATCH ${PROJECT_VERSION_PATCH}\n") + +file(MAKE_DIRECTORY ${PROJECT_BINARY_DIR}/include/metatomic) +set(_destination_ "${CMAKE_CURRENT_BINARY_DIR}/include/metatomic/version.h") +file(COPY_FILE ${_path_} ${_destination_} ONLY_IF_DIFFERENT) + +add_dependencies(metatomic::shared cargo-build-metatomic) +add_dependencies(metatomic::static cargo-build-metatomic) + +set_target_properties(metatomic::shared PROPERTIES + IMPORTED_LOCATION ${METATOMIC_SHARED_LOCATION} + INTERFACE_INCLUDE_DIRECTORIES "${CMAKE_CURRENT_SOURCE_DIR}/include;${CMAKE_CURRENT_BINARY_DIR}/include" + BUILD_VERSION "${METATOMIC_FULL_VERSION}" +) +target_compile_features(metatomic::shared INTERFACE cxx_std_17) + +if (WIN32) + set_target_properties(metatomic::shared PROPERTIES + IMPORTED_IMPLIB ${METATOMIC_IMPLIB_LOCATION} + ) +endif() + +set_target_properties(metatomic::static PROPERTIES + IMPORTED_LOCATION ${METATOMIC_STATIC_LOCATION} + INTERFACE_INCLUDE_DIRECTORIES "${CMAKE_CURRENT_SOURCE_DIR}/include;${CMAKE_CURRENT_BINARY_DIR}/include" + INTERFACE_LINK_LIBRARIES "${CARGO_DEFAULT_LIBRARIES}" + BUILD_VERSION "${METATOMIC_FULL_VERSION}" +) +target_compile_features(metatomic::static INTERFACE cxx_std_17) + +if (TARGET metatensor::static) + target_link_libraries(metatomic::static INTERFACE metatensor::static) +else() + target_link_libraries(metatomic::static INTERFACE metatensor) +endif() + +if (TARGET metatensor::shared) + target_link_libraries(metatomic::shared INTERFACE metatensor::shared) +else() + target_link_libraries(metatomic::shared INTERFACE metatensor) +endif() + +target_link_libraries(metatomic::static INTERFACE nlohmann_json::nlohmann_json) +target_link_libraries(metatomic::shared INTERFACE nlohmann_json::nlohmann_json) + +if(APPLE) + target_link_libraries(metatomic::static INTERFACE + "-framework Metal" "-framework CoreGraphics" "-framework CoreFoundation" "-framework Foundation" objc + ) +endif() + + +if (BUILD_SHARED_LIBS) + add_library(metatomic ALIAS metatomic::shared) +else() + add_library(metatomic ALIAS metatomic::static) +endif() + +#------------------------------------------------------------------------------# +# Installation configuration +#------------------------------------------------------------------------------# +include(CMakePackageConfigHelpers) +configure_package_config_file( + ${PROJECT_SOURCE_DIR}/cmake/metatomic-config.in.cmake + ${PROJECT_BINARY_DIR}/metatomic-config.cmake + INSTALL_DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/metatomic +) +write_basic_package_version_file( + metatomic-config-version.cmake + VERSION ${METATOMIC_FULL_VERSION} + COMPATIBILITY SameMinorVersion +) + +install(FILES "include/metatomic.h" DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +install(FILES "include/metatomic.hpp" DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +install(DIRECTORY "include/metatomic" DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +install(FILES "${CMAKE_CURRENT_BINARY_DIR}/include/metatomic/version.h" DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/metatomic) + +if (METATOMIC_INSTALL_BOTH_STATIC_SHARED OR BUILD_SHARED_LIBS) + if (WIN32) + # DLL files should go in /bin + install( + FILES ${METATOMIC_SHARED_LOCATION} + DESTINATION ${CMAKE_INSTALL_BINDIR} + PERMISSIONS OWNER_EXECUTE OWNER_WRITE OWNER_READ GROUP_EXECUTE GROUP_READ WORLD_READ WORLD_EXECUTE + ) + # .lib files should go in /lib + install(FILES ${METATOMIC_IMPLIB_LOCATION} DESTINATION ${CMAKE_INSTALL_LIBDIR}) + else() + install( + FILES ${METATOMIC_SHARED_LOCATION} + DESTINATION ${CMAKE_INSTALL_LIBDIR} + PERMISSIONS OWNER_EXECUTE OWNER_WRITE OWNER_READ GROUP_EXECUTE GROUP_READ WORLD_READ WORLD_EXECUTE + ) + endif() +endif() + +if (METATOMIC_INSTALL_BOTH_STATIC_SHARED OR NOT BUILD_SHARED_LIBS) + install(FILES ${METATOMIC_STATIC_LOCATION} DESTINATION ${CMAKE_INSTALL_LIBDIR}) +endif() + +install(FILES + ${PROJECT_BINARY_DIR}/metatomic-config-version.cmake + ${PROJECT_BINARY_DIR}/metatomic-config.cmake + DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/metatomic +) diff --git a/metatomic-core/Cargo.toml b/metatomic-core/Cargo.toml new file mode 100644 index 000000000..6c51a185d --- /dev/null +++ b/metatomic-core/Cargo.toml @@ -0,0 +1,43 @@ +[package] +name = "metatomic-core" +version = "0.1.0" +edition = "2024" +publish = false +rust-version = "1.88" +exclude = [ + "tests" +] + +[lib] +crate-type = ["cdylib", "staticlib"] +name = "metatomic" +bench = false + +[dependencies] +metatensor = { version = "0.5.1" } +dlpk = { version = "0.4.1", features = ["ndarray"]} +json = "0.12" +libloading = "0.9" +ndarray = "0.17" + +# For serialization of the systems +zip = { version = "8.6.0", default-features = false } +byteorder = {version = "1"} + +# For custom kernels +cudarc = {version = "0.19", default-features = false, features=["std", "cuda-13030", "driver", "nvrtc", "dynamic-loading"]} +lru = "0.18" + +[target.'cfg(target_os = "macos")'.dependencies] +objc2-metal = { version = "0.3", features = ["block2"] } +objc2 = "0.6" +objc2-foundation = "0.3" + +[build-dependencies] +cbindgen = { version = "0.29", default-features = false } + +[dev-dependencies] +lazy_static = "1" +which = "8" +approx = "0.5" +ndarray = { version = "0.17", features = ["approx"] } diff --git a/metatomic-core/Clippy.toml b/metatomic-core/Clippy.toml new file mode 100644 index 000000000..49c5aa7b9 --- /dev/null +++ b/metatomic-core/Clippy.toml @@ -0,0 +1 @@ +doc-valid-idents = ["DLPack", "ROCm", ".."] diff --git a/metatomic-core/build.rs b/metatomic-core/build.rs new file mode 100644 index 000000000..9e7b4e3e3 --- /dev/null +++ b/metatomic-core/build.rs @@ -0,0 +1,108 @@ +#![allow(clippy::field_reassign_with_default)] + +use std::path::PathBuf; + +fn main() { + let crate_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap(); + + let generated_comment = "\ +/* ============ Automatically generated file, DO NOT EDIT. ============== * + * * + * This file is automatically generated from the metatomic sources, * + * using cbindgen. If you want to change this file (including documentation), * + * make the corresponding changes in the rust sources and regenerate it. * + * ============================================================================= */"; + + let mut config: cbindgen::Config = Default::default(); + config.language = cbindgen::Language::C; + config.cpp_compat = true; + config.include_guard = Some("METATOMIC_H".into()); + config.include_version = false; + config.documentation = true; + config.documentation_style = cbindgen::DocumentationStyle::Doxy; + config.line_endings = cbindgen::LineEndingStyle::LF; + config.autogen_warning = Some(generated_comment.into()); + config.sys_includes.push("stdio.h".into()); + config.sys_includes.push("metatensor.h".into()); + config.includes.push("metatomic/version.h".into()); + + config.export = cbindgen::ExportConfig { + include: vec!["mta_.*".into()], + // These are done manually below + exclude: vec!["mta_opaque_string_t".into(), "mta_system_t".into()], + ..Default::default() + }; + config.after_includes = Some(" + +#ifndef MTA_EXPORT + #if defined(_WIN32) || defined(__CYGWIN__) + #define MTA_EXPORT __declspec(dllexport) + #else + #define MTA_EXPORT __attribute__((visibility(\"default\"))) + #endif +#endif + +#ifndef MTA_EXTERN_C + #ifdef __cplusplus + #define MTA_EXTERN_C extern \"C\" + #else + #define MTA_EXTERN_C + #endif +#endif + +/** + * Define the exported plugin entry points. + * + * This macro should be used once in each plugin shared library with a + * `mta_plugin_t` expression. It exports the plugin ABI version and a + * registration function used by `mta_load_plugin`. + */ +#define MTA_REGISTER_PLUGIN(register_fn_name, ...) \\ + MTA_EXTERN_C MTA_EXPORT mta_status_t mta_plugin_init(int abi, void *data) { \\ + if (abi != MTA_ABI_VERSION) { \\ + char message[256]; \\ + snprintf(message, sizeof(message), \\ + \"Metatomic plugin ABI version mismatch: expected %d, got %d\", \\ + MTA_ABI_VERSION, abi \\ + ); \\ + mta_set_last_error(message, \"MTA_REGISTER_PLUGIN\", NULL, NULL); \\ + return MTA_INVALID_PARAMETER_ERROR; \\ + } \\ + mta_status_t (*register_fn_name)(mta_plugin_t) = (mta_status_t (*)(mta_plugin_t))data; \\ + __VA_ARGS__; \\ + return MTA_SUCCESS; \\ + } + +/** Heap allocated storage for mta_string_t */ +typedef struct mta_opaque_string_t mta_opaque_string_t; + +/** + * Opaque handle to an atomistic system. + * + * The system owns DLPack tensors for types, positions, cell, and PBC, as well + * as metatensor blocks for pair lists and tensor maps for custom data. + */ +typedef struct mta_system_t mta_system_t;".into()); + + let result = cbindgen::Builder::new() + .with_crate(crate_dir) + .with_config(config) + .generate() + .map(|data| { + let mut path = PathBuf::from("include"); + path.push("metatomic.h"); + data.write_to_file(&path); + }); + + // if not ok, rerun the build script unconditionally + if result.is_ok() { + println!("cargo:rerun-if-changed=src"); + println!("cargo:rerun-if-changed=build.rs"); + } + + if std::env::var("METATOMIC_FULL_VERSION").is_err() { + let version = std::env::var("CARGO_PKG_VERSION").expect("missing CARGO_PKG_VERSION"); + println!("cargo:rustc-env=METATOMIC_FULL_VERSION={}+rust", version); + } + println!("cargo:rerun-if-env-changed=METATOMIC_FULL_VERSION"); +} diff --git a/metatomic-core/cmake/detect_cargo.cmake b/metatomic-core/cmake/detect_cargo.cmake new file mode 100644 index 000000000..0af2293e7 --- /dev/null +++ b/metatomic-core/cmake/detect_cargo.cmake @@ -0,0 +1,181 @@ +# This module finds a suitable cargo binary. It tries plain "cargo" first, then +# searches for versioned cargo binaries (e.g. cargo-1.82) commonly installed on +# Ubuntu. If a binary is found but too old, it continues searching for a newer +# one. +# +# Sets: +# CARGO_EXE - path to the chosen cargo binary +# CARGO_VERSION - parsed version string (e.g. 1.74.0) +# RUST_HOST_TARGET - host target triple (e.g. x86_64-unknown-linux-gnu) +# RUST_HOST_ARCH - host CPU architecture (e.g. x86_64) +# CACHED_LAST_CARGO_VERSION - cache variable for change detection +# CARGO_VERSION_CHANGED - true if the version differs from the last run + +set(REQUIRED_RUST_VERSION "1.88.0") + +# --------------------------------------------------------------------------- +# Helper: run cargo --version --verbose, extract version & host target +# --------------------------------------------------------------------------- +function(_try_cargo _exe _ok_var _version_var _host_target_var _host_arch_var) + execute_process( + COMMAND "${_exe}" "--version" "--verbose" + RESULT_VARIABLE _status + OUTPUT_VARIABLE _raw + ERROR_QUIET + ) + + if (NOT _status EQUAL 0) + set(${_ok_var} FALSE PARENT_SCOPE) + return() + endif() + + set(_ok TRUE) + set(_version "") + set(_host_target "") + + if (_raw MATCHES "cargo ([0-9]+\\.[0-9]+\\.[0-9]+)") + set(_version "${CMAKE_MATCH_1}") + else() + set(_ok FALSE) + endif() + + if (_raw MATCHES "host: ([a-zA-Z0-9_\\-]*)\n") + set(_host_target "${CMAKE_MATCH_1}") + else() + set(_ok FALSE) + endif() + + set(${_ok_var} ${_ok} PARENT_SCOPE) + set(${_version_var} "${_version}" PARENT_SCOPE) + set(${_host_target_var} "${_host_target}" PARENT_SCOPE) + + if (_host_target MATCHES "([a-zA-Z0-9_]*)\\-") + set(${_host_arch_var} "${CMAKE_MATCH_1}" PARENT_SCOPE) + else() + set(${_host_arch_var} "" PARENT_SCOPE) + endif() +endfunction() + +# --------------------------------------------------------------------------- +# Step 1: try plain "cargo" (or respect a pre-defined CARGO_EXE) +# --------------------------------------------------------------------------- +set(_cargo_found FALSE) +if (DEFINED CARGO_EXE AND NOT CARGO_EXE STREQUAL "CARGO_EXE-NOTFOUND") + _try_cargo("${CARGO_EXE}" _ok _ver _target _arch) + if (_ok AND ${_ver} VERSION_GREATER_EQUAL ${REQUIRED_RUST_VERSION}) + set(_cargo_found TRUE) + set(CARGO_VERSION "${_ver}") + set(RUST_HOST_TARGET "${_target}") + set(RUST_HOST_ARCH "${_arch}") + else() + # Cache is stale or binary changed; re-search below + message(STATUS "cargo at ${CARGO_EXE} is not usable, searching for alternatives...") + unset(CARGO_EXE) + unset(CARGO_EXE CACHE) + endif() +endif() + +if (NOT _cargo_found) + find_program(_cargo_vanilla "cargo") + if (_cargo_vanilla) + _try_cargo("${_cargo_vanilla}" _ok _ver _target _arch) + if (_ok AND ${_ver} VERSION_GREATER_EQUAL ${REQUIRED_RUST_VERSION}) + set(_cargo_found TRUE) + set(CARGO_EXE "${_cargo_vanilla}") + set(CARGO_VERSION "${_ver}") + set(RUST_HOST_TARGET "${_target}") + set(RUST_HOST_ARCH "${_arch}") + endif() + endif() +endif() + +# --------------------------------------------------------------------------- +# Step 2: search for versioned cargo-* binaries across PATH +# --------------------------------------------------------------------------- +if (NOT _cargo_found) + # Collect all directories to search + set(_search_dirs ${CMAKE_PROGRAM_PATH}) + + if (WIN32) + foreach(_dir IN LISTS $ENV{PATH}) + list(APPEND _search_dirs "${_dir}") + endforeach() + else() + string(REPLACE ":" ";" _sys_path "$ENV{PATH}") + list(APPEND _search_dirs ${_sys_path}) + endif() + + if (NOT "$ENV{HOME}" STREQUAL "") + list(APPEND _search_dirs "$ENV{HOME}/.cargo/bin") + endif() + + set(_cargo_candidates "") + foreach(_dir IN LISTS _search_dirs) + if (IS_DIRECTORY "${_dir}") + file(GLOB _bins "${_dir}/cargo-*") + list(APPEND _cargo_candidates ${_bins}) + endif() + endforeach() + + if (_cargo_candidates) + list(REMOVE_DUPLICATES _cargo_candidates) + endif() + + set(_best_exe "") + set(_best_version "0.0.0") + set(_best_target "") + set(_best_arch "") + + foreach(_bin IN LISTS _cargo_candidates) + _try_cargo("${_bin}" _ok _ver _target _arch) + if (_ok AND ${_ver} VERSION_GREATER_EQUAL ${REQUIRED_RUST_VERSION} + AND ${_ver} VERSION_GREATER ${_best_version}) + set(_best_exe "${_bin}") + set(_best_version "${_ver}") + set(_best_target "${_target}") + set(_best_arch "${_arch}") + endif() + endforeach() + + if (_best_exe) + set(_cargo_found TRUE) + set(CARGO_EXE "${_best_exe}") + set(CARGO_VERSION "${_best_version}") + set(RUST_HOST_TARGET "${_best_target}") + set(RUST_HOST_ARCH "${_best_arch}") + endif() +endif() + +# --------------------------------------------------------------------------- +# Final validation +# --------------------------------------------------------------------------- +if (NOT _cargo_found) + message(FATAL_ERROR + "could not find a suitable cargo binary (version >= ${REQUIRED_RUST_VERSION}).\n" + "Please install Rust from https://www.rust-lang.org/tools/install\n" + "or set CARGO_EXE to point to your cargo binary before calling CMake." + ) +endif() + +if (NOT RUST_HOST_TARGET) + message(FATAL_ERROR + "failed to determine host target from cargo --version --verbose" + ) +endif() + +if (NOT RUST_HOST_ARCH) + message(FATAL_ERROR + "failed to determine host CPU arch from target: ${RUST_HOST_TARGET}" + ) +endif() + +# --------------------------------------------------------------------------- +# Cache for change detection across CMake re-configures +# --------------------------------------------------------------------------- +if (NOT "${CACHED_LAST_CARGO_VERSION}" STREQUAL "${CARGO_VERSION}") + set(CACHED_LAST_CARGO_VERSION "${CARGO_VERSION}" + CACHE INTERNAL "Last version of cargo used in configuration" + ) + message(STATUS "Using cargo version ${CARGO_VERSION} at ${CARGO_EXE}") + set(CARGO_VERSION_CHANGED TRUE) +endif() diff --git a/metatomic-core/cmake/dev-versions.cmake b/metatomic-core/cmake/dev-versions.cmake new file mode 100644 index 000000000..543296493 --- /dev/null +++ b/metatomic-core/cmake/dev-versions.cmake @@ -0,0 +1,91 @@ +# Parse a `_version_` number, and store its components in `_major_` `_minor_` +# `_patch_` and `_rc_` +function(parse_version _version_ _major_ _minor_ _patch_ _rc_) + string(REGEX MATCH "([0-9]+)\\.([0-9]+)\\.([0-9]+)(-rc)?([0-9]+)?" _ "${_version_}") + + if(${CMAKE_MATCH_COUNT} EQUAL 3) + set(${_rc_} "" PARENT_SCOPE) + elseif(${CMAKE_MATCH_COUNT} EQUAL 5) + set(${_rc_} ${CMAKE_MATCH_5} PARENT_SCOPE) + else() + message(FATAL_ERROR "invalid version string ${_version_}") + endif() + + set(${_major_} ${CMAKE_MATCH_1} PARENT_SCOPE) + set(${_minor_} ${CMAKE_MATCH_2} PARENT_SCOPE) + set(${_patch_} ${CMAKE_MATCH_3} PARENT_SCOPE) +endfunction() + +# Get the time of the last modification since the last tag/release, and a hash +# of the latest commit/full state of a dirty repository +function(git_version_info _tag_prefix_ _output_n_commits_ _output_git_hash_) + set(_script_ "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../../scripts/git-version-info.py") + + if (EXISTS "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/git_version_info") + # When building from a tarball, the script is executed and the result + # put in this file + file(STRINGS "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/git_version_info" _file_content_) + list(GET _file_content_ 0 _n_commits_) + list(GET _file_content_ 1 _git_hash_) + + elseif (EXISTS "${_script_}") + # When building from a checkout, we'll need to run the script + find_package(Python COMPONENTS Interpreter REQUIRED) + execute_process( + COMMAND "${Python_EXECUTABLE}" "${_script_}" "${_tag_prefix_}" + RESULT_VARIABLE _status_ + OUTPUT_VARIABLE _stdout_ + ERROR_VARIABLE _stderr_ + WORKING_DIRECTORY ${CMAKE_CURRENT_FUNCTION_LIST_DIR} + ) + + if (NOT ${_status_} EQUAL 0) + message(WARNING + "git-version-info.py failed, version number might be wrong:\nstdout: ${_stdout_}\nstderr: ${_stderr_}") + set(${_output_} 0 PARENT_SCOPE) + return() + endif() + + if (NOT "${_stderr_}" STREQUAL "") + message(WARNING "git-version-info.py gave some errors, version number might be wrong:\nstdout: ${_stdout_}\nstderr: ${_stderr_}") + endif() + + string(REPLACE "\n" ";" _lines_ ${_stdout_}) + list(GET _lines_ 0 _n_commits_) + list(GET _lines_ 1 _git_hash_) + else() + message(FATAL_ERROR "could not update git version information") + endif() + + string(STRIP ${_n_commits_} _n_commits_) + set(${_output_n_commits_} ${_n_commits_} PARENT_SCOPE) + + string(STRIP ${_git_hash_} _git_hash_) + set(${_output_git_hash_} ${_git_hash_} PARENT_SCOPE) +endfunction() + + +# Take the version declared in the package, and increase the right number if we +# are actually installing a developement version from after the latest git tag +function(create_development_version _version_ _output_ _tag_prefix_) + git_version_info("${_tag_prefix_}" _n_commits_ _git_hash_) + + parse_version(${_version_} _major_ _minor_ _patch_ _rc_) + if(${_n_commits_} STREQUAL "0") + # we are building a release, leave the version number as-is + if("${_rc_}" STREQUAL "") + set(${_output_} "${_major_}.${_minor_}.${_patch_}" PARENT_SCOPE) + else() + set(${_output_} "${_major_}.${_minor_}.${_patch_}-rc${_rc_}" PARENT_SCOPE) + endif() + else() + # we are building a development version, increase the right part of the version + if("${_rc_}" STREQUAL "") + math(EXPR _minor_ "${_minor_} + 1") + set(${_output_} "${_major_}.${_minor_}.0-dev${_n_commits_}+${_git_hash_}" PARENT_SCOPE) + else() + math(EXPR _rc_ "${_rc_} + 1") + set(${_output_} "${_major_}.${_minor_}.${_patch_}-rc${_rc_}-dev${_n_commits_}+${_git_hash_}" PARENT_SCOPE) + endif() + endif() +endfunction() diff --git a/metatomic-core/cmake/metatomic-config.in.cmake b/metatomic-core/cmake/metatomic-config.in.cmake new file mode 100644 index 000000000..b40085dd1 --- /dev/null +++ b/metatomic-core/cmake/metatomic-config.in.cmake @@ -0,0 +1,114 @@ +@PACKAGE_INIT@ + +cmake_minimum_required(VERSION 3.22) + +include(CMakeFindDependencyMacro) +include(FindPackageHandleStandardArgs) + +if(metatomic_FOUND) + return() +endif() + +enable_language(CXX) + +# use the same version for metatensor-core as the main CMakeLists.txt +set(REQUIRED_METATENSOR_VERSION @REQUIRED_METATENSOR_VERSION@) +find_package(metatensor ${REQUIRED_METATENSOR_VERSION} CONFIG REQUIRED) + +get_filename_component(METATOMIC_PREFIX_DIR "${CMAKE_CURRENT_LIST_DIR}/@PACKAGE_RELATIVE_PATH@" ABSOLUTE) + +# nlohmann_json is either vendored by us (its headers are installed alongside +# metatomic's own, in the same include directory, but not its CMake package +# config) or was found as a system package when building metatomic; only look +# for the system package in the latter case. +if (@nlohmann_json_FOUND@) + find_dependency(nlohmann_json 3.11.0) + set(METATOMIC_NLOHMANN_JSON_LIBRARY "nlohmann_json::nlohmann_json") +else() + add_library(metatomic_nlohmann_json INTERFACE) + target_include_directories(metatomic_nlohmann_json INTERFACE "${METATOMIC_PREFIX_DIR}/@CMAKE_INSTALL_INCLUDEDIR@/metatomic/third-party") + set(METATOMIC_NLOHMANN_JSON_LIBRARY "metatomic_nlohmann_json") +endif() + +if (WIN32) + set(METATOMIC_SHARED_LOCATION ${METATOMIC_PREFIX_DIR}/@CMAKE_INSTALL_BINDIR@/@METATOMIC_SHARED_LIB_NAME@) + set(METATOMIC_IMPLIB_LOCATION ${METATOMIC_PREFIX_DIR}/@CMAKE_INSTALL_LIBDIR@/@METATOMIC_IMPLIB_NAME@) +else() + set(METATOMIC_SHARED_LOCATION ${METATOMIC_PREFIX_DIR}/@CMAKE_INSTALL_LIBDIR@/@METATOMIC_SHARED_LIB_NAME@) +endif() + +set(METATOMIC_STATIC_LOCATION ${METATOMIC_PREFIX_DIR}/@CMAKE_INSTALL_LIBDIR@/@METATOMIC_STATIC_LIB_NAME@) +set(METATOMIC_INCLUDE ${METATOMIC_PREFIX_DIR}/@CMAKE_INSTALL_INCLUDEDIR@/) + +if (NOT EXISTS ${METATOMIC_INCLUDE}/metatomic.h OR NOT EXISTS ${METATOMIC_INCLUDE}/metatomic.hpp) + message(FATAL_ERROR "could not find metatomic headers in '${METATOMIC_INCLUDE}', please re-install metatomic") +endif() + + +# Shared library target +if (@METATOMIC_INSTALL_BOTH_STATIC_SHARED@ OR @BUILD_SHARED_LIBS@) + if (NOT EXISTS ${METATOMIC_SHARED_LOCATION}) + message(FATAL_ERROR "could not find metatomic library at '${METATOMIC_SHARED_LOCATION}', please re-install metatomic") + endif() + + add_library(metatomic::shared SHARED IMPORTED) + set_target_properties(metatomic::shared PROPERTIES + IMPORTED_LOCATION ${METATOMIC_SHARED_LOCATION} + INTERFACE_INCLUDE_DIRECTORIES ${METATOMIC_INCLUDE} + BUILD_VERSION "@METATOMIC_FULL_VERSION@" + ) + + target_compile_features(metatomic::shared INTERFACE cxx_std_17) + target_link_libraries(metatomic::shared INTERFACE metatensor ${METATOMIC_NLOHMANN_JSON_LIBRARY}) + + if (WIN32) + if (NOT EXISTS ${METATOMIC_IMPLIB_LOCATION}) + message(FATAL_ERROR "could not find metatomic library at '${METATOMIC_IMPLIB_LOCATION}', please re-install metatomic") + endif() + + set_target_properties(metatomic::shared PROPERTIES + IMPORTED_IMPLIB ${METATOMIC_IMPLIB_LOCATION} + ) + endif() +endif() + + +# Static library target +if (@METATOMIC_INSTALL_BOTH_STATIC_SHARED@ OR NOT @BUILD_SHARED_LIBS@) + if (NOT EXISTS ${METATOMIC_STATIC_LOCATION}) + message(FATAL_ERROR "could not find metatomic library at '${METATOMIC_STATIC_LOCATION}', please re-install metatomic") + endif() + + add_library(metatomic::static STATIC IMPORTED) + set_target_properties(metatomic::static PROPERTIES + IMPORTED_LOCATION ${METATOMIC_STATIC_LOCATION} + INTERFACE_INCLUDE_DIRECTORIES ${METATOMIC_INCLUDE} + INTERFACE_LINK_LIBRARIES "@CARGO_DEFAULT_LIBRARIES@" + BUILD_VERSION "@METATOMIC_FULL_VERSION@" + ) + + target_compile_features(metatomic::static INTERFACE cxx_std_17) + + target_link_libraries(metatomic::static INTERFACE metatensor) + target_link_libraries(metatomic::static INTERFACE ${METATOMIC_NLOHMANN_JSON_LIBRARY}) + + if(APPLE) + target_link_libraries(metatomic::static INTERFACE + "-framework Metal" "-framework CoreGraphics" "-framework CoreFoundation" "-framework Foundation" objc + ) + endif() +endif() + +# Export either the shared or static library as the metatomic target +if (@BUILD_SHARED_LIBS@) + add_library(metatomic ALIAS metatomic::shared) +else() + add_library(metatomic ALIAS metatomic::static) +endif() + + +if (@BUILD_SHARED_LIBS@) + find_package_handle_standard_args(metatomic DEFAULT_MSG METATOMIC_SHARED_LOCATION METATOMIC_INCLUDE) +else() + find_package_handle_standard_args(metatomic DEFAULT_MSG METATOMIC_STATIC_LOCATION METATOMIC_INCLUDE) +endif() diff --git a/metatomic-core/cmake/nlohmann_json.cmake b/metatomic-core/cmake/nlohmann_json.cmake new file mode 100644 index 000000000..3a2203e15 --- /dev/null +++ b/metatomic-core/cmake/nlohmann_json.cmake @@ -0,0 +1,47 @@ +# Find or fetch nlohmann JSON library +# +# This module first tries to find nlohmann_json via find_package. +# If that fails, it falls back to fetching it via FetchContent. +# +# After including this module, you can link against nlohmann_json::nlohmann_json + +# Guard against multiple inclusion +if(TARGET nlohmann_json::nlohmann_json) + return() +endif() + +if (POLICY CMP0135) + cmake_policy(SET CMP0135 NEW) # DOWNLOAD_EXTRACT_TIMESTAMP TRUE in FetchContent_Declare +endif() + +include(FetchContent) + +find_package(nlohmann_json 3.11.0 QUIET) + +if(nlohmann_json_FOUND) + message(STATUS "Found nlohmann_json via find_package: ${nlohmann_json_VERSION}") +else() + message(STATUS "nlohmann_json not found via find_package, fetching from GitHub") + + # Fetch the release tarball, which contains the CMake build files and headers + # but not the benchmark reports with very long filenames that break Windows. + FetchContent_Declare( + nlohmann_json + URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz + ) + + set(JSON_BuildTests OFF CACHE INTERNAL "") + # Don't use nlohmann_json's own install rules, they would also install its + # CMake package config and pkg-config files, and we don't want to advertise + # a system-wide nlohmann_json package to external users. + set(JSON_Install OFF CACHE INTERNAL "") + + FetchContent_MakeAvailable(nlohmann_json) + + # nlohmann_json is header-only, so we can install the headers ourselves, + # alongside metatomic's own (they are used in our public headers). + install( + DIRECTORY "${nlohmann_json_SOURCE_DIR}/include/nlohmann" + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/metatomic/third-party + ) +endif() diff --git a/metatomic-core/cmake/tempdir.cmake b/metatomic-core/cmake/tempdir.cmake new file mode 100644 index 000000000..52e4805fc --- /dev/null +++ b/metatomic-core/cmake/tempdir.cmake @@ -0,0 +1,51 @@ +# Create a temporary directory using mktemp on *nix and powershell on windows +function(get_tempdir _outvar_) + # special case for github actions, where $TEMP might + # exist but point to nowhere/a non writable location + # https://docs.github.com/en/actions/learn-github-actions/variables + if (DEFINED ENV{RUNNER_TEMP}) + string(RANDOM LENGTH 12 _dirname_) + set(_output_ $ENV{RUNNER_TEMP}/${_dirname_}) + file(TO_NATIVE_PATH "${_output_}" _output_) + file(MAKE_DIRECTORY ${_output_}) + set(${_outvar_} ${_output_} PARENT_SCOPE) + return() + endif() + + find_program(MKTEMP_EXE NAMES mktemp) + if(MKTEMP_EXE) + execute_process( + COMMAND ${MKTEMP_EXE} -d + OUTPUT_VARIABLE _output_ + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _status_ + ) + + if(_status_ EQUAL 0) + file(MAKE_DIRECTORY ${_output_}) + set(${_outvar_} ${_output_} PARENT_SCOPE) + return() + endif() + endif() + + + find_program(POWERSHELL_EXE NAMES pwsh powershell) + if(POWERSHELL_EXE) + execute_process( + COMMAND ${POWERSHELL_EXE} -c "[System.IO.Path]::GetTempPath()" + OUTPUT_VARIABLE _output_ + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _status_ + ) + + if(_status_ EQUAL 0) + string(RANDOM LENGTH 12 _dirname_) + set(_output_ ${_output_}${_dirname_}) + file(MAKE_DIRECTORY ${_output_}) + set(${_outvar_} ${_output_} PARENT_SCOPE) + return() + endif() + endif() + + message(FATAL_ERROR "Could not find mktemp or PowerShell to make temporary directory") +endfunction() diff --git a/metatomic-core/include/metatomic.h b/metatomic-core/include/metatomic.h new file mode 100644 index 000000000..040e5c19c --- /dev/null +++ b/metatomic-core/include/metatomic.h @@ -0,0 +1,781 @@ +#ifndef METATOMIC_H +#define METATOMIC_H + +/* ============ Automatically generated file, DO NOT EDIT. ============== * + * * + * This file is automatically generated from the metatomic sources, * + * using cbindgen. If you want to change this file (including documentation), * + * make the corresponding changes in the rust sources and regenerate it. * + * ============================================================================= */ + +#include +#include +#include +#include +#include +#include +#include "metatomic/version.h" + + +#ifndef MTA_EXPORT + #if defined(_WIN32) || defined(__CYGWIN__) + #define MTA_EXPORT __declspec(dllexport) + #else + #define MTA_EXPORT __attribute__((visibility("default"))) + #endif +#endif + +#ifndef MTA_EXTERN_C + #ifdef __cplusplus + #define MTA_EXTERN_C extern "C" + #else + #define MTA_EXTERN_C + #endif +#endif + +/** + * Define the exported plugin entry points. + * + * This macro should be used once in each plugin shared library with a + * `mta_plugin_t` expression. It exports the plugin ABI version and a + * registration function used by `mta_load_plugin`. + */ +#define MTA_REGISTER_PLUGIN(register_fn_name, ...) \ + MTA_EXTERN_C MTA_EXPORT mta_status_t mta_plugin_init(int abi, void *data) { \ + if (abi != MTA_ABI_VERSION) { \ + char message[256]; \ + snprintf(message, sizeof(message), \ + "Metatomic plugin ABI version mismatch: expected %d, got %d", \ + MTA_ABI_VERSION, abi \ + ); \ + mta_set_last_error(message, "MTA_REGISTER_PLUGIN", NULL, NULL); \ + return MTA_INVALID_PARAMETER_ERROR; \ + } \ + mta_status_t (*register_fn_name)(mta_plugin_t) = (mta_status_t (*)(mta_plugin_t))data; \ + __VA_ARGS__; \ + return MTA_SUCCESS; \ + } + +/** Heap allocated storage for mta_string_t */ +typedef struct mta_opaque_string_t mta_opaque_string_t; + +/** + * Opaque handle to an atomistic system. + * + * The system owns DLPack tensors for types, positions, cell, and PBC, as well + * as metatensor blocks for pair lists and tensor maps for custom data. + */ +typedef struct mta_system_t mta_system_t; + +/** + * ABI version of the metatomic plugin interface. + * + * This increases anytime the plugin or model C API changes in a non backward + * compatible way. Plugins compiled with an incompatible ABI version will be + * rejected at registration time. + */ +#define MTA_ABI_VERSION 1 + +/** + * Status type returned by all functions in the C API. + * + * The value 0 (`MTA_SUCCESS`) indicates success, while any non-zero value indicates an error. + */ +typedef enum mta_status_t { + /** + * Status code indicating success + */ + MTA_SUCCESS = 0, + /** + * Status code indicating invalid function parameters + */ + MTA_INVALID_PARAMETER_ERROR = 1, + /** + * Status code indicating I/O errors + */ + MTA_IO_ERROR = 2, + /** + * Status code indicating memory allocation errors + */ + MTA_MEMORY_ERROR = 3, + /** + * Status code indicating serialization/deserialization errors + */ + MTA_SERIALIZATION_ERROR = 4, + /** + * Status code indicating dlpack errors + */ + MTA_DLPACK_ERROR = 5, + /** + * Status code indicating metatensor errors + */ + MTA_METATENSOR_ERROR = 6, + /** + * Status code used by plugins when a model is not supported by the + * current plugin + */ + MTA_UNSUPPORTED_MODEL_ERROR = 7, + /** + * Status code used by model for any error that does not fit the cases above + */ + MTA_MODEL_ERROR = 8, + /** + * Status code used when there is an internal error + */ + MTA_INTERNAL_ERROR = 255, +} mta_status_t; + +/** + * Kind of data always stored in a system. + * + * Other kinds of data can be stored with `mta_system_add_custom_data` and + * retrieved with `mta_system_get_custom_data`. + */ +typedef enum mta_system_data_kind { + MTA_SYSTEM_DATA_TYPES = 0, + MTA_SYSTEM_DATA_POSITIONS = 1, + MTA_SYSTEM_DATA_CELL = 2, + MTA_SYSTEM_DATA_PBC = 3, +} mta_system_data_kind; + +/** + * An heap-allocated UTF-8 string passed across the C API boundary. + * + * This is used whenever a C API function or callback needs to return a string. + * + * A null pointer represents an absent or empty string. Use `mta_string_create` + * to allocate, `mta_string_free` to release, and `mta_string_view` to get a + * pointer to the inner C string. + */ +typedef mta_opaque_string_t *mta_string_t; + +/** + * A model that computes physical properties of atomistic systems. + * + * `mta_model_t` is a small virtual table: `data` holds the model's own state, + * and the function pointers describe what the model can do. A model is usually + * produced by a plugin's `load_model` callback (see `mta_load_model`) and then + * executed with `mta_execute_model`. + * + * Every callback receives `data` as its first argument. metatomic treats + * `data` as opaque and only hands it back to the callbacks. Callbacks should + * report any error by saving it with `mta_set_last_error` and returning a + * non-success `mta_status_t`. + */ +typedef struct mta_model_t { + /** + * Opaque pointer to the model's internal state + * + * Its layout and meaning are private to the model implementation. It is + * initialized by whoever creates the model (e.g. a plugin's `load_model`) + * and released by `unload`. + */ + void *data; + /** + * Release the resources owned by `model_data` + * + * Called exactly once when the model is no longer needed. May be `NULL` if + * the model owns no resources. + * + * @param model_data the model's `data` pointer + * @return `MTA_SUCCESS` on success, another status code on error + */ + enum mta_status_t (*unload)(void *model_data); + /** + * Get the capabilities of the model as a JSON string. + * + * @verbatim embed:rst:leading-asterisk + * The expected JSON structure is documented in :ref:`core-json-model-capabilities`. + * @endverbatim + * + * @param model_data the model's `data` pointer + * @param capabilities_json output string, set to a JSON-serialized + * `ModelCapabilities` object. The caller takes ownership and must + * free it with `mta_string_free`. + * @return `MTA_SUCCESS` on success, another status code on error + */ + enum mta_status_t (*capabilities)(const void *model_data, mta_string_t *capabilities_json); + /** + * Get metadata describing the model (name, authors, references, ...) as a + * JSON string. + * + * @verbatim embed:rst:leading-asterisk + * The expected JSON structure is documented in :ref:`core-json-model-metadata`. + * @endverbatim + * + * @param model_data the model's `data` pointer + * @param metadata_json output string, set to a JSON-serialized + * `ModelMetadata` object. The caller takes ownership and must + * free it with `mta_string_free`. + * @return `MTA_SUCCESS` on success, another status code on error + */ + enum mta_status_t (*metadata)(const void *model_data, mta_string_t *metadata_json); + /** + * List the pair lists (neighbor lists) the model needs as input as a JSON + * string. + * + * @verbatim embed:rst:leading-asterisk + * + * The engine is expected to compute these and attach them to every system + * with :c:func:`mta_system_add_pairs` before calling + * :c:func:`mta_execute_model`. + * + * The expected JSON structure for each pair list is documented in :ref:`core-json-pair-options`. + * + * @endverbatim + * + * @param model_data the model's `data` pointer + * @param pair_options_json output string, set to a JSON array of + * `PairListOptions` objects. The caller takes ownership and must + * free it with `mta_string_free`. + * @return `MTA_SUCCESS` on success, another status code on error + */ + enum mta_status_t (*requested_pair_lists)(const void *model_data, mta_string_t *pair_options_json); + /** + * List the additional per-system inputs the model needs as a JSON string. + * + * @verbatim embed:rst:leading-asterisk + * + * These correspond to custom data the engine should attach to every system + * with :c:func:`mta_system_add_custom_data` before execution. + * + * The expected JSON structure for each input is documented in :ref:`core-json-quantity`. + * + * @endverbatim + * + * @param model_data the model's `data` pointer + * @param inputs_json output string, set to a JSON array of `Quantity` + * objects, one per requested input. The caller takes ownership and + * must free it with `mta_string_free`. + * @return `MTA_SUCCESS` on success, another status code on error + */ + enum mta_status_t (*requested_inputs)(const void *model_data, mta_string_t *inputs_json); + /** + * Run the model and compute the requested outputs + * + * @verbatim embed:rst:leading-asterisk + * + * This performs the model's actual computation. This should not be called + * directly, but rather through :c:func:`mta_execute_model`, which handles + * unit conversion and can check inputs and output data for consistency. + * + * @endverbatim + * + * @param model_data the model's `data` pointer + * @param systems array of `systems_count` systems to run the model on + * @param systems_count number of entries in `systems` + * @param selected_atoms optional labels selecting the subset of atoms to + * compute outputs for, or `NULL` to use all atoms. When set, it has the + * dimensions `"system"` and `"atom"` holding 0-based indices. + * @param requested_outputs_json JSON string containing an array of + * `Quantity`, one for each output the model should produce + * @param outputs array of `outputs_count` tensor maps to fill, one per + * requested output and in the same order + * @param outputs_count number of entries in `outputs`, must equal + * `requested_outputs_count` + * @return `MTA_SUCCESS` on success, another status code on error + */ + enum mta_status_t (*execute_inner)(void *model_data, + const mta_system_t *const *systems, + uintptr_t systems_count, + const mts_labels_t *selected_atoms, + const char *requested_outputs_json, + mts_tensormap_t **outputs, + uintptr_t outputs_count); +} mta_model_t; + +/** + * A metatomic plugin definition. + */ +typedef struct mta_plugin_t { + /** + * ABI version this plugin was compiled against, this should be set to + * `MTA_ABI_VERSION` when creating the plugin struct. + */ + int32_t abi_version; + /** + * Name of the plugin, as a null-terminated UTF-8 string. This is the name + * specified in `mta_load_model` when trying to load a model with a + * specific plugin. The name must be unique among all registered plugins. + */ + const char *name; + /** + * Callback function to load a model. This function should try to load a + * model from `load_from` (which can be a file path, a model name, etc.) + * and a set of key/values options passed as a JSON string. + * + * If the plugin can load the model, it should fill `model` with a pointer + * to a valid `mta_model_t` struct and return `MTA_SUCCESS`. If the data in + * `load_from` does not correspond to a model supported by the plugin, it + * should return `MTA_UNSUPPORTED_MODEL_ERROR`. If an error occurs while + * loading the model, it should return another status code and save an + * error message with `mta_set_last_error`. + * + * @param load_from a null-terminated UTF-8 string describing where to load + * the model from (e.g. a file path, a model name, etc.). The + * interpretation of this string is up to the plugin. + * @param options_json a null-terminated UTF-8 string containing a set of + * string keys and string value options for loading the model. + * @param model output pointer to the loaded model. The caller takes + * ownership of the model and must unload it when the model is no + * longer needed. + * @return `MTA_SUCCESS` if the model was loaded successfully, + * `MTA_UNSUPPORTED_MODEL_ERROR` if the plugin can not load the model, + * or another status code if an error occurs. + */ + enum mta_status_t (*load_model)(const char *load_from, + const char *options_json, + struct mta_model_t *model); +} mta_plugin_t; + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +/** + * Get last error message that was created on the current thread. + */ +enum mta_status_t mta_last_error(const char **message, const char **origin, void **data); + +/** + * Set last error message for the current thread. + */ +enum mta_status_t mta_set_last_error(const char *message, + const char *origin, + void *data, + void (*data_deleter)(void*)); + +/** + * Get the runtime version of the metatomic library as a string. + * + * This version follows the `..[-]` format. + */ +const char *mta_version(void); + +/** + * Allocate a new `mta_string_t` by copying the null-terminated C string + * `string`. + * + * The returned string must be freed with `mta_string_free`. + * + * @param string A pointer to a null-terminated C string. Must not be null. + * @return A new `mta_string_t` containing a copy of `string`, or null if an + * error occurred. You can check the error with `mta_last_error`. + */ +mta_string_t mta_string_create(const char *string); + +/** + * Free a `mta_string_t` previously created by `mta_string_create`. + * + * @param string A `mta_string_t` to free. Can be null, in which case this function is a no-op. + */ +void mta_string_free(mta_string_t string); + +/** + * Return a pointer to the null-terminated string data inside `string`. + * + * The pointer is valid only for the lifetime of `string`. + * + * @param string A `mta_string_t` containing the string to view. Must not be null. + * @return A pointer to the null-terminated C string inside `string` + */ +const char *mta_string_view(mta_string_t string); + +/** + * Get the multiplicative conversion factor to use to convert from `from_unit` + * to `to_unit`. Both units are parsed as expressions (e.g. `kJ / mol / A^2`, + * `(eV * u)^(1/2)`) and their dimensions must match. + * + * @verbatim embed:rst:leading-asterisk + * + * .. seealso:: + * + * The general documentation for :ref:`units`, with the expression + * syntax and list of supported base units. + * + * @endverbatim + * + * @param from_unit A null-terminated C string containing the unit to convert from. + * @param to_unit A null-terminated C string containing the unit to convert to. + * @param conversion A pointer to a `double` where the conversion factor will be stored. + * @return The status code of the operation. If this code is not `MTA_SUCCESS`, + * you can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_unit_conversion_factor(const char *from_unit, + const char *to_unit, + double *conversion); + +/** + * Create a new system from raw DLPack tensors. + * + * This function **takes ownership** of `types`, `positions`, `cell`, and + * `pbc`. The caller must not use these tensors after calling this function. + * + * @param length_unit A null-terminated C string containing the length unit + * (e.g. "Angstrom", "nanometer"). Must not be null. + * @param types A DLPack managed tensor with shape `(n_atoms,)` and dtype + * `int32`. Ownership is transferred. + * @param positions A DLPack managed tensor with shape `(n_atoms, 3)` and + * dtype `float32` or `float64`. Ownership is transferred. + * @param cell A DLPack managed tensor with shape `(3, 3)` and the same dtype + * as `positions`. Ownership is transferred. + * @param pbc A DLPack managed tensor with shape `(3,)` and dtype `bool`. + * Ownership is transferred. + * @param system Output parameter, set to the newly created system handle. + * The caller takes ownership and must free it with `mta_system_free`. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_create(const char *length_unit, + DLManagedTensorVersioned *types, + DLManagedTensorVersioned *positions, + DLManagedTensorVersioned *cell, + DLManagedTensorVersioned *pbc, + mta_system_t **system); + +/** + * Free a system previously created by `mta_system_create`. + * + * If there are outstanding borrowed views (from `mta_system_get_data`), the + * system's data will remain alive until all views are released. + * + * @param system The system handle to free. Can be null, in which case this + * function is a no-op. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_free(mta_system_t *system); + +/** + * Get the number of atoms in a system. + * + * @param system The system handle. Must not be null. + * @param size Output parameter, set to the number of atoms. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_size(const mta_system_t *system, uintptr_t *size); + +/** + * Get a DLPack tensor from a system for the requested data. + * + * This function **returns a borrowed view** of the system's internal data. + * The returned `DLManagedTensorVersioned` has a custom deleter that decrements + * the system's reference count, keeping the system alive as long as the + * borrowed view exists. + * + * The caller is responsible for calling the deleter on the returned tensor + * when it is no longer needed. The tensor shares the data pointer with the + * system; do **not** modify it. + * + * @param system The system handle. Must not be null. + * @param request Which data to retrieve (types, positions, cell, or PBC). + * @param data Output parameter, set to a pointer to a newly allocated + * `DLManagedTensorVersioned` containing the requested data. The caller + * takes ownership and must call the deleter when done. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_get_data(const mta_system_t *system, + enum mta_system_data_kind request, + DLManagedTensorVersioned **data); + +/** + * Get the length unit of a system. + * + * This function returns a new `mta_string_t` that the caller must free with + * `mta_string_free`. + * + * @param system The system handle. Must not be null. + * @param length_unit Output parameter, set to the length unit string. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_get_length_unit(const mta_system_t *system, mta_string_t *length_unit); + +/** + * Add a pair list (neighbor list) to a system. + * + * This function **takes ownership** of `pairs`. The caller must not use the + * block after calling this function. + * + * @param system The system handle. Must not be null. + * @param options A JSON-serialized `PairListOptions` object. Must not be null. + * @param pairs A `mts_block_t` containing the pair data. Ownership is + * transferred. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_add_pairs(mta_system_t *system, + const char *options, + mts_block_t *pairs); + +/** + * Get a pair list from a system. + * + * **Returns a borrowed view** of the pair list. The system must outlive the + * returned pointer. Do **not** free the returned block. + * + * @param system The system handle. Must not be null. + * @param options A JSON-serialized `PairListOptions` object identifying which + * pair list to retrieve. Must not be null. + * @param pairs Output parameter, set to a pointer to the pair list block, or + * NULL if no pair list matches the options. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_get_pairs(const mta_system_t *system, + const char *options, + const mts_block_t **pairs); + +/** + * Get all pair list options known by a system. + * + * This function returns a new `mta_string_t` containing a JSON array of + * `PairListOptions` objects. The caller must free it with `mta_string_free`. + * + * @param system The system handle. Must not be null. + * @param pairs_options Output parameter, set to a JSON string containing an + * array of `PairListOptions` objects. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_known_pairs(const mta_system_t *system, mta_string_t *pairs_options); + +/** + * Add custom data to a system. + * + * This function **takes ownership** of `data`. The caller must not use the + * tensor map after calling this function. + * + * @param system The system handle. Must not be null. + * @param name A null-terminated C string containing the name of the custom + * data. Must not be null. + * @param data A `mts_tensormap_t` containing the custom data. Ownership is + * transferred. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_add_custom_data(mta_system_t *system, + const char *name, + mts_tensormap_t *data); + +/** + * Get custom data from a system by name. + * + * **Returns a borrowed view** of the custom data. The system must outlive the + * returned pointer. Do **not** free the returned tensor map. + * + * @param system The system handle. Must not be null. + * @param name A null-terminated C string containing the name of the custom + * data. Must not be null. + * @param data Output parameter, set to a pointer to the custom data tensor + * map, or an error if no data with the given name exists. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_get_custom_data(const mta_system_t *system, + const char *name, + const mts_tensormap_t **data); + +/** + * Get all custom data names known by a system. + * + * **Returns a new** `mta_string_t` containing a JSON array of strings. The + * caller must free it with `mta_string_free`. + * + * @param system The system handle. Must not be null. + * @param names Output parameter, set to a JSON string containing an array of + * custom data names. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_system_known_custom_data(const mta_system_t *system, mta_string_t *names); + +/** + * Execute a model to compute the requested outputs for a set of systems + * + * This is the main entry point to run a model loaded through the C API. It + * validates the arguments and delegates the computation to the model's + * `execute_inner` callback. + * + * @param model the model to execute + * @param systems array of `systems_count` systems to run the model on + * @param systems_count number of entries in `systems` + * @param selected_atoms optional labels selecting the subset of atoms to + * compute outputs for, or `NULL` to use all atoms + * @param requested_outputs_json JSON string containing an array of + * `Quantity`, one for each output the model should produce + * @param check_consistency if `true`, run additional checks on the + * inputs and on the data produced by the model + * @param outputs array of `outputs_count` tensor maps to fill, one per + * requested output and in the same order. The caller takes ownership of + * the returned tensor maps. + * @param outputs_count number of entries in `outputs`, must equal + * `requested_outputs_count` + * @return `MTA_SUCCESS` on success, another status code on error (the message + * is available through `mta_last_error`) + */ +enum mta_status_t mta_execute_model(struct mta_model_t model, + const mta_system_t *const *systems, + uintptr_t systems_count, + const mts_labels_t *selected_atoms, + const char *requested_outputs_json, + bool check_consistency, + mts_tensormap_t **outputs, + uintptr_t outputs_count); + +/** + * Render model metadata as a human-readable string + * + * @param metadata a JSON-serialized `ModelMetadata` object as produced by a + * model's `metadata` callback. Must not be null. + * @param printed output string, set to a human-readable rendering of the + * metadata. The caller takes ownership and must free it with + * `mta_string_free`. + * @return `MTA_SUCCESS` on success, another status code on error + */ +enum mta_status_t mta_format_metadata(const char *metadata, mta_string_t *printed); + +/** + * Register a plugin. This is passed as a callback to the `MTA_REGISTER_PLUGIN` + * macro, and should not be called directly by C or C++ plugin implementations. + * + * @param plugin the plugin to register + * @return `MTA_SUCCESS` if the plugin was registered successfully, or another + * status code if an error occurs. You can get more details about the error + * with `mta_last_error`. + */ +enum mta_status_t mta_register_plugin(struct mta_plugin_t plugin); + +/** + * Load the shared library at `path` and register the plugin contained within. + * + * The library must export the symbols generated by the `MTA_REGISTER_PLUGIN` + * macro. + * + * @param path a null-terminated UTF-8 string containing the path to the plugin + * shared library, or `NULL` to load the plugin from the current binary + * @return `MTA_SUCCESS` if the plugin was loaded successfully, or another + * status code if an error occurs. You can get more details about the + * error with `mta_last_error`. + */ +enum mta_status_t mta_load_plugin(const char *path); + +/** + * Load a model from `load_from` with the given options. + * + * If `plugin_name` is a NULL pointer, metatomic will try to determine the + * correct plugin to use by checking the `load_from` parameter. If we can not + * determine the correct plugin, we then try to load the model with each + * registered plugin until one succeeds. + * + * If `plugin_name` is given, then we only try to load the model with the + * specified plugin, and return an error if the plugin can not load the model. + * + * @param plugin_name optional null-terminated UTF-8 string containing the name + * of the plugin to use for loading the model, or `NULL` to let metatomic + * search for a correct plugin + * @param load_from a null-terminated UTF-8 string describing where to load the + * model from (e.g. a file path, a model name, etc.). The interpretation + * of this string is up to the plugin. + * @param options_json a null-terminated UTF-8 string containing a set of string + * keys and string value options for loading the model. The interpretation + * of these options is up to the plugin. + * @param model output pointer to the loaded model. The caller takes ownership of + * the model and must unload it when the model is no longer needed. + * @return `MTA_SUCCESS` if the model was loaded successfully, or another + * status code if an error occurs. You can get more details about the + * error with `mta_last_error`. + */ +enum mta_status_t mta_load_model(const char *load_from, + const char *options_json, + const char *plugin_name, + struct mta_model_t *model); + +/** + * Save a system to a file. + * + * The format consists of a zip archive containing NPY files for the system's + * data (types, positions, cell, pbc), a `info.json` file for metadata, and + * optional sub-directories for pair lists (`pairs//options.json` and + * `pairs//data.mts`) and custom data (`data/.mts`). + * + * @param path A null-terminated C string containing the file path. Must not be + * null. + * @param system The system to save. Must not be null. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_save(const char *path, const mta_system_t *system); + +/** + * Save a system to an in-memory buffer. + * + * The buffer is grown as needed using the provided `realloc` callback. On + * success, `*buffer` points to the serialized data and `*buffer_count` + * contains the number of bytes written. + * + * @param buffer Pointer to the buffer pointer. On input, `*buffer` may be NULL + * (in which case `*buffer_count` must be 0). On output, `*buffer` is + * updated to point to the serialized data. + * @param buffer_count Pointer to the buffer size. On input, `*buffer_count` + * must contain the current allocation size. On output, it is set to the + * number of bytes written. + * @param realloc_user_data User data passed as the first argument to + * `realloc`. + * @param realloc Callback to grow the buffer. Must not be NULL. + * @param system The system to save. Must not be null. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_save_buffer(uint8_t **buffer, + uintptr_t *buffer_count, + void *realloc_user_data, + mts_realloc_buffer_t realloc, + const mta_system_t *system); + +/** + * Load a system from a file. + * + * The file must have been written by `mta_save` and contain a valid metatomic + * system. + * + * @param path A null-terminated C string containing the file path. Must not be + * null. + * @param create_array Callback to allocate arrays for the system's data. Must + * not be NULL. + * @param system Output parameter, set to the newly created system handle. + * The caller takes ownership and must free it with `mta_system_free`. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_load(const char *path, + mts_create_array_callback_t create_array, + mta_system_t **system); + +/** + * Load a system from an in-memory buffer. + * + * The buffer must contain data serialized by `mta_save_buffer` (or the + * equivalent Rust function). + * + * @param buffer Pointer to the serialized data. Must not be NULL. + * @param buffer_size Number of bytes in `buffer`. + * @param create_array Callback to allocate arrays for the system's data. Must + * not be NULL. + * @param system Output parameter, set to the newly created system handle. + * The caller takes ownership and must free it with `mta_system_free`. + * @return `MTA_SUCCESS` on success, or another status code if an error occurs. + * You can get more details about the error with `mta_last_error`. + */ +enum mta_status_t mta_load_buffer(const uint8_t *buffer, + uintptr_t buffer_size, + mts_create_array_callback_t create_array, + mta_system_t **system); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif /* METATOMIC_H */ diff --git a/metatomic-core/include/metatomic.hpp b/metatomic-core/include/metatomic.hpp new file mode 100644 index 000000000..e4fa4c930 --- /dev/null +++ b/metatomic-core/include/metatomic.hpp @@ -0,0 +1,7 @@ +#include "metatomic/utils.hpp" // IWYU pragma: export +#include "metatomic/system.hpp" // IWYU pragma: export +#include "metatomic/model.hpp" // IWYU pragma: export +#include "metatomic/plugin.hpp" // IWYU pragma: export +#include "metatomic/errors.hpp" // IWYU pragma: export +#include "metatomic/metadata.hpp" // IWYU pragma: export +#include "metatomic/io.hpp" // IWYU pragma: export diff --git a/metatomic-core/include/metatomic/errors.hpp b/metatomic-core/include/metatomic/errors.hpp new file mode 100644 index 000000000..b275750a3 --- /dev/null +++ b/metatomic-core/include/metatomic/errors.hpp @@ -0,0 +1,100 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include + +namespace metatomic { + + /// Exception class used for all errors in metatomic + class Error: public std::runtime_error { + public: + /// Create a new MetatomicError with the given `message` + Error(const std::string& message): std::runtime_error(message) {} + }; + + namespace details { + /// Check if a return status from the C API indicates an error, and if it is + /// the case, throw an exception of type `metatomic::Error` with the last + /// error message from the library. + inline void check_status(mta_status_t status) { + if (status == MTA_SUCCESS) { + return; + } + + const char* message = nullptr; + const char* origin = nullptr; + void* data = nullptr; + mta_last_error(&message, &origin, &data); + if (origin != nullptr && std::strcmp(origin, "C++ exception") == 0 && data != nullptr) { + std::rethrow_exception(*static_cast(data)); + } else { + throw Error(message == nullptr ? "unknown error" : message); + } + } + + /// Call the given `function` with the given `args` (the function should + /// return an `mta_status_t`), catching any C++ exception, and translating + /// them to native metatomic error code. + /// + /// This is required to prevent callbacks unwinding through the C API. + template + inline mta_status_t catch_exceptions(Function function, Args ...args) { + try { + function(std::move(args)...); + return MTA_SUCCESS; + } catch (...) { + auto* exception_ptr = new std::exception_ptr(std::current_exception()); + + std::string message; + try { + std::rethrow_exception(*exception_ptr); + } catch (const std::exception& e) { + message = e.what(); + } catch (...) { + message = "C++ code threw an exception that was not a std::exception"; + } + + auto status = mta_set_last_error( + message.c_str(), + "C++ exception", + exception_ptr, + [](void *ptr) { delete static_cast(ptr); } + ); + + if (status != MTA_SUCCESS) { + // If we failed to set the error, we are in a very bad state, + // but we should still try to report the original error + // message if possible. + std::fprintf(stderr, "INTERNAL ERROR: unable to set last error after C++ callback failure (status: %d). ", status); + std::fprintf(stderr, "C++ error was: %s\n", message.c_str()); + delete exception_ptr; + } + + return MTA_MODEL_ERROR; + } + } + + /// Check if a pointer allocated by the C API is null, and if it is the + /// case, throw an exception of type `metatomic::Error` with the last + /// error message from the library. + inline void check_pointer(const void* pointer) { + if (pointer == nullptr) { + const char* message = nullptr; + const char* origin = nullptr; + void* data = nullptr; + mta_last_error(&message, &origin, &data); + if (origin != nullptr && std::strcmp(origin, "C++ exception") == 0 && data != nullptr) { + std::rethrow_exception(*static_cast(data)); + } else { + throw Error(message == nullptr ? "unknown error" : message); + } + } + } + } // namespace details + +} // namespace metatomic diff --git a/metatomic-core/include/metatomic/io.hpp b/metatomic-core/include/metatomic/io.hpp new file mode 100644 index 000000000..a314a8190 --- /dev/null +++ b/metatomic-core/include/metatomic/io.hpp @@ -0,0 +1,116 @@ +#pragma once + +#include +#include +#include + +#include + +#include +#include + +namespace metatomic { +namespace io { + +/// Save a system to a file. +/// +/// @param path path of the file to create or overwrite +/// @param system system to serialize +inline void save(const std::string& path, const System& system) { + details::check_status(mta_save(path.c_str(), system.as_mta_system_t())); +} + +/// Serialize a system into a byte container. +/// +/// `Buffer` must be constructible from a pair of iterators over bytes. The +/// serialization is performed using a `std::vector` and copied into +/// the requested container type. +/// +/// @tparam Buffer byte-container type, such as `std::vector` +/// @param system system to serialize +/// @return serialized system data +template +Buffer save_buffer(const System& system) { + auto buffer = metatomic::io::save_buffer>(system); + return Buffer(buffer.begin(), buffer.end()); +} + +/// Serialize a system into a `std::vector`. +/// +/// The C API grows the vector through a reallocation callback. The returned +/// vector contains exactly the number of bytes produced by the serializer. +/// +/// @param system system to serialize +/// @return serialized system data +template <> +inline std::vector save_buffer>(const System& system) { + std::vector buffer; + + auto* ptr = buffer.data(); + auto size = buffer.size(); + + auto realloc = [](void* user_data, uint8_t*, uintptr_t new_size) { + auto* buffer = reinterpret_cast*>(user_data); + buffer->resize(new_size, '\0'); + return buffer->data(); + }; + + details::check_status(mta_save_buffer(&ptr, &size, &buffer, realloc, system.as_mta_system_t())); + + buffer.resize(size, '\0'); + + return buffer; +} + +/// Load a system from a file. +/// +/// @param path path of the serialized system file +/// @param create_array callback used to create arrays during deserialization +/// @return reconstructed system +inline System load( + const std::string& path, + mts_create_array_callback_t create_array = metatensor::details::default_create_array +) { + mta_system_t* ptr = nullptr; + details::check_status(mta_load(path.c_str(), create_array, &ptr)); + details::check_pointer(ptr); + return System::unsafe_from_ptr(ptr); +} + +/// Load a system from a contiguous byte buffer. +/// +/// @param buffer serialized system data +/// @param buffer_count number of bytes available at `buffer` +/// @param create_array callback used to create arrays during deserialization +/// @return reconstructed system +inline System load_buffer( + const uint8_t* buffer, + uintptr_t buffer_count, + mts_create_array_callback_t create_array = metatensor::details::default_create_array +) { + mta_system_t* ptr = nullptr; + details::check_status(mta_load_buffer(buffer, buffer_count, create_array, &ptr)); + details::check_pointer(ptr); + return System::unsafe_from_ptr(ptr); +} + +/// Load a system from a byte container. +/// +/// The container must provide contiguous storage through `data()` and report +/// its size in bytes through `size()`. +/// +/// @tparam Buffer contiguous byte-container type +/// @param buffer serialized system data +/// @param create_array callback used to create arrays during deserialization +/// @return reconstructed system +template +System load_buffer( + const Buffer& buffer, + mts_create_array_callback_t create_array = metatensor::details::default_create_array +) { + static_assert(sizeof(typename Buffer::value_type) == sizeof(uint8_t), "`Buffer` must be a container of uint8_t or equivalent"); + return metatomic::io::load_buffer(reinterpret_cast(buffer.data()), buffer.size(), create_array); +} + +} // namespace io +} // namespace metatomic diff --git a/metatomic-core/include/metatomic/metadata.hpp b/metatomic-core/include/metatomic/metadata.hpp new file mode 100644 index 000000000..544d3b2e3 --- /dev/null +++ b/metatomic-core/include/metatomic/metadata.hpp @@ -0,0 +1,1352 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include // std::move +#include // std::isfinite +#include // std::memcpy +#include // std::uint64_t, std::int64_t +#include // std::isxdigit + +#include +#include + +namespace metatomic { + namespace detail { + + inline std::vector read_string_array( + const nlohmann::json& j, const std::string& key, const char* context + ) { + if (!j.contains(key) || !j[key].is_array()) { + throw metatomic::Error("'" + key + "' in " + context + " must be an array"); + } + + std::vector result; + for (const auto& item : j[key]) { + if (!item.is_string()) { + throw metatomic::Error("'" + key + "' in " + context + " must be an array of strings"); + } + result.push_back(item.get()); + } + return result; + } + + } // namespace detail + + /// Options for the calculation of a pair list (neighbor list) + class PairListOptions final { + private: + /// Cutoff radius for this pair list in the length unit of the model + double cutoff_; + /// Whether the list is a full list (contains both the pair `i -> j` and `j -> i`) + /// or a half list (contains only `i -> j`) + bool full_list_; + /// Whether the list guarantees that only atoms within the cutoff are + /// included (strict) or may also include pairs slightly beyond the cutoff + /// (non-strict) + bool strict_ = true; + /// List of strings describing who requested this pair list + std::vector requestors_; + + PairListOptions( + double cutoff, + bool full_list, + bool strict, + std::vector requestors + ) : cutoff_(cutoff), full_list_(full_list), + strict_(strict), requestors_(std::move(requestors)) {} + + public: + /// Get the cutoff radius for this pair list. + double cutoff() const { + return cutoff_; + } + + /// Get whether this pair list is a full list. + bool full_list() const { + return full_list_; + } + + /// Get whether this pair list is strict. + bool strict() const { + return strict_; + } + + /// Get the list of requestors for this pair list. + const std::vector& requestors() const { + return requestors_; + } + + /// Check if two `PairListOptions` are equal. + /// + /// The list of requestors is ignored when checking for equality. + bool operator==(const PairListOptions& other) const { + return cutoff_ == other.cutoff_ && + full_list_ == other.full_list_ && + strict_ == other.strict_; + } + + /// Check if two `PairListOptions` are different. + /// + /// The list of requestors is ignored when checking for equality. + bool operator!=(const PairListOptions& other) const { + return !(*this == other); + } + + /// Builder for `PairListOptions`. + /// + /// Use `PairListOptions::builder()` to create a new builder, set the + /// required fields via the setters, and call `build()` to obtain + /// a fully-initialized `PairListOptions`. + class Builder { + private: + std::optional cutoff_; + std::optional full_list_; + bool strict_ = true; + std::vector requestors_; + + /// Validate that `value` is a finite positive number + /// + /// @throw metatomic::Error if `value` is not a finite positive number + static void validate_cutoff(double value) { + if (!std::isfinite(value) || value <= 0.0) { + throw metatomic::Error("cutoff must be a finite positive number"); + } + } + + /// Add `requestor` to `requestors_`, ignoring empty strings and + /// duplicates, keeping first-seen order. + /// + /// @param requestors the list of requestors to add to + /// @param requestor the requestor to add + static void add_requestor_to(std::vector& requestors, const std::string& requestor) { + if (!requestor.empty() && std::find(requestors.begin(), requestors.end(), requestor) == requestors.end()) { + requestors.push_back(requestor); + } + } + + public: + /// Set the cutoff radius for this pair list. + /// + /// @throw metatomic::Error if the value is not a finite positive number. + Builder& cutoff(double value) { + validate_cutoff(value); + cutoff_ = value; + return *this; + } + + /// Set whether this pair list is a full list. + Builder& full_list(bool value) { + full_list_ = value; + return *this; + } + + /// Set whether this pair list is strict. Defaults to `true`. + Builder& strict(bool value) { + strict_ = value; + return *this; + } + + /// Set the list of requestors for this pair list. + Builder& requestors(std::vector value) { + requestors_ = std::move(value); + return *this; + } + + /// Add a requestor to the list. + /// + /// Empty strings and duplicates are ignored, keeping first-seen order. + Builder& add_requestor(const std::string& requestor) { + add_requestor_to(requestors_, requestor); + return *this; + } + + /// Build a fully-initialized `PairListOptions`. + /// + /// @throw metatomic::Error if `cutoff` or `full_list` has not been set. + /// + /// This moves the builder's fields; calling `build()` a second time + /// produces an object with moved-from values. Builders are intended + /// as one-shot temporaries. + [[nodiscard]] PairListOptions build() { + if (!cutoff_.has_value()) { + throw metatomic::Error("cutoff must be set before building PairListOptions"); + } + if (!full_list_.has_value()) { + throw metatomic::Error("full_list must be set before building PairListOptions"); + } + return PairListOptions( + cutoff_.value(), + full_list_.value(), + strict_, + std::move(requestors_) + ); + } + }; + + /// Create a new `Builder` for `PairListOptions`. + [[nodiscard]] static Builder builder() { + return Builder{}; + } + }; + + inline void to_json(nlohmann::json& j, const PairListOptions& p){ + // Store cutoff as hex-encoded bit pattern + // Floating-point round-trip conversions is exact + double cutoff = p.cutoff(); + uint64_t bits; + std::memcpy(&bits, &cutoff, sizeof(double)); + std::ostringstream oss; + oss << "0x" << std::hex << bits; + + j = nlohmann::json{ + {"type", "metatomic_pair_list_options"}, + {"cutoff", oss.str()}, + {"full_list", p.full_list()}, + {"strict", p.strict()}, + {"requestors", p.requestors()} + }; + } + + inline PairListOptions from_json( + const nlohmann::json& j, nlohmann::detail::identity_tag + ) { + if (!j.is_object()) { + throw metatomic::Error("invalid JSON data for PairListOptions, expected an object"); + } + + if (!j.contains("type") || !j["type"].is_string() || j["type"].get() != "metatomic_pair_list_options") { + throw metatomic::Error("'type' in JSON for PairListOptions must be 'metatomic_pair_list_options'"); + } + + // Cutoff is an hex-encoded string + if (!j.contains("cutoff") || !j["cutoff"].is_string()) { + throw metatomic::Error("'cutoff' in JSON for PairListOptions must be a hex-encoded string"); + } + std::string cutoff_str = j["cutoff"].get(); + + // Strip "0x" prefix if present + if (cutoff_str.size() >= 2 && cutoff_str[0] == '0' && cutoff_str[1] == 'x') { + cutoff_str = cutoff_str.substr(2); + } + + uint64_t bits; + try { + // std::isxdigit checks for hex digits + if (cutoff_str.empty() || !std::all_of(cutoff_str.begin(), cutoff_str.end(), [](unsigned char c) { return std::isxdigit(c); })) { + throw metatomic::Error("'cutoff' in JSON for PairListOptions must be a hex-encoded string"); + } + + std::size_t pos = 0; + bits = std::stoull(cutoff_str, &pos, 16); + if (pos != cutoff_str.size()) { + throw metatomic::Error("'cutoff' in JSON for PairListOptions must be a hex-encoded string"); + } + } catch (...) { + throw metatomic::Error("'cutoff' in JSON for PairListOptions must be a hex-encoded string"); + } + double cutoff; + std::memcpy(&cutoff, &bits, sizeof(double)); + + if (!std::isfinite(cutoff) || cutoff <= 0.0) { + throw metatomic::Error("'cutoff' in JSON for PairListOptions must be a finite positive number"); + } + + if (!j.contains("full_list") || !j["full_list"].is_boolean()) { + throw metatomic::Error("'full_list' in JSON for PairListOptions must be a boolean"); + } + bool full_list = j["full_list"].get(); + + if (!j.contains("strict") || !j["strict"].is_boolean()) { + throw metatomic::Error("'strict' in JSON for PairListOptions must be a boolean"); + } + bool strict = j["strict"].get(); + + auto p = PairListOptions::builder() + .cutoff(cutoff) + .full_list(full_list) + .strict(strict); + + if (j.contains("requestors")) { + if (!j["requestors"].is_array()) { + throw metatomic::Error("'requestors' in JSON for PairListOptions must be an array"); + } + + for (const auto& requestor : j["requestors"]) { + if (!requestor.is_string()) { + throw metatomic::Error("'requestors' in JSON for PairListOptions must be an array of strings"); + } + p.add_requestor(requestor.get()); + } + } + + return p.build(); + } + + // Forward declarations + // The ModelMetadata::print function uses to_json + class ModelMetadata; + void to_json(nlohmann::json&, const ModelMetadata&); + + class ModelMetadata final { + public: + /// References for a model, divided into three categories: references about + /// the model as a whole, references about the architecture of the model, + /// and references about the implementation of the model. + class References final { + private: + /// The references about the model as a whole, e.g. a paper describing the + /// model or a website presenting it. + std::vector model_; + /// The references about the architecture of the model, e.g. papers + /// describing the mathematical form of the model. + std::vector architecture_; + /// The references about the implementation of the model, e.g. a link to + /// the source code repository or a paper describing the software. + std::vector implementation_; + + // Private constructor + References( + std::vector model, + std::vector architecture, + std::vector implementation + ) : model_(std::move(model)), architecture_(std::move(architecture)), + implementation_(std::move(implementation)) {} + + public: + /// Get the references about the model as a whole. + const std::vector& model() const { + return model_; + } + + /// Get the references about the architecture of the model. + const std::vector& architecture() const { + return architecture_; + } + + /// Get the references about the implementation of the model. + const std::vector& implementation() const { + return implementation_; + } + + /// Builder for `References`. + /// + /// Use `References::builder()` to create a new builder, set the + /// fields via the setters, and call `build()` to obtain a + /// fully-initialized `References`. All fields default to empty + /// lists, so `build()` always succeeds. + class Builder { + private: + std::vector model_; + std::vector architecture_; + std::vector implementation_; + + public: + /// Set the references about the model as a whole. + Builder& model(std::vector value) { + model_ = std::move(value); + return *this; + } + + /// Add a reference about the model as a whole. + Builder& add_model(const std::string& reference) { + model_.push_back(reference); + return *this; + } + + /// Set the references about the architecture of the model. + Builder& architecture(std::vector value) { + architecture_ = std::move(value); + return *this; + } + + /// Add a reference about the architecture of the model. + Builder& add_architecture(const std::string& reference) { + architecture_.push_back(reference); + return *this; + } + + /// Set the references about the implementation of the model. + Builder& implementation(std::vector value) { + implementation_ = std::move(value); + return *this; + } + + /// Add a reference about the implementation of the model. + Builder& add_implementation(const std::string& reference) { + implementation_.push_back(reference); + return *this; + } + + /// Build a fully-initialized `References`. + /// + /// This moves the builder's fields; calling `build()` a second + /// time produces an object with moved-from values. Builders are + /// intended as one-shot temporaries. + [[nodiscard]] References build() { + return References( + std::move(model_), + std::move(architecture_), + std::move(implementation_) + ); + } + }; + + /// Create a new `Builder` for `References`. + [[nodiscard]] static Builder builder() { + return Builder{}; + } + }; + + private: + std::string name_; + std::vector authors_; + std::string description_; + References references_; + // BTreeMap in Rust is an ordered map + std::map extra_; + + /// Private constructor + ModelMetadata( + std::string name, + std::vector authors, + std::string description, + References references, + std::map extra + ) : name_(std::move(name)), authors_(std::move(authors)), + description_(std::move(description)), references_(std::move(references)), + extra_(std::move(extra)) {} + + public: + /// Get the name of the model. + const std::string& name() const { + return name_; + } + + /// Get the list of authors of the model. + const std::vector& authors() const { + return authors_; + } + + /// Get the description of the model. + const std::string& description() const { + return description_; + } + + /// Get the references for the model. + const References& references() const { + return references_; + } + + /// Get the extra metadata for the model. + const std::map& extra() const { + return extra_; + } + + /// Builder for `ModelMetadata`. + /// + /// Use `ModelMetadata::builder()` to create a new builder, set the + /// fields via the setters, and call `build()` to obtain a + /// fully-initialized `ModelMetadata`. All fields have defaults, so + /// `build()` always succeeds. + class Builder { + private: + std::string name_; + std::vector authors_; + std::string description_; + References::Builder references_builder_; + std::map extra_; + + public: + /// Set the name of the model. + Builder& name(std::string value) { + name_ = std::move(value); + return *this; + } + + /// Set the list of authors of the model. + Builder& authors(std::vector value) { + authors_ = std::move(value); + return *this; + } + + /// Add an author to the list of authors. + Builder& add_author(const std::string& author) { + authors_.push_back(author); + return *this; + } + + /// Set the description of the model. + Builder& description(std::string value) { + description_ = std::move(value); + return *this; + } + + /// Set the references for the model. + Builder& references(References value) { + references_builder_ + .model(value.model()) + .architecture(value.architecture()) + .implementation(value.implementation()); + return *this; + } + + /// Add a reference to the given section. + /// + /// @param section reference section, one of "model", "architecture", or + /// "implementation" + /// @param reference the reference to add + /// @throw metatomic::Error if `section` is not one of the allowed values + Builder& add_reference(const std::string& section, const std::string& reference) { + if (section == "model") { + references_builder_.add_model(reference); + } else if (section == "architecture") { + references_builder_.add_architecture(reference); + } else if (section == "implementation") { + references_builder_.add_implementation(reference); + } else { + throw metatomic::Error( + "reference section must be 'model', 'architecture', or 'implementation', got '" + section + "'" + ); + } + return *this; + } + + /// Set the extra metadata for the model. + Builder& extra(std::map value) { + extra_ = std::move(value); + return *this; + } + + /// Add a key/value pair to the extra metadata. + /// + /// If the key already exists, its value is overwritten. + /// + /// @param key key for the extra metadata entry + /// @param value value for the extra metadata entry + Builder& add_extra(const std::string& key, const std::string& value) { + extra_[key] = value; + return *this; + } + + /// Build a fully-initialized `ModelMetadata`. + /// + /// This moves the builder's fields; calling `build()` a second + /// time produces an object with moved-from values. Builders are + /// intended as one-shot temporaries. + [[nodiscard]] ModelMetadata build() { + return ModelMetadata( + std::move(name_), + std::move(authors_), + std::move(description_), + references_builder_.build(), + std::move(extra_) + ); + } + }; + + /// Create a new `Builder` for `ModelMetadata`. + [[nodiscard]] static Builder builder() { + return Builder{}; + } + + /// Print the metadata as a human-readable string. + std::string print() const { + // Re-use C API to avoid re-implementing 'normalize_withespace' and 'wrap_80_chars' + mta_string_t mta_string; + nlohmann::json j; + + to_json(j, *this); + auto status = mta_format_metadata(j.dump().c_str(), &mta_string); + details::check_status(status); + + std::string output = mta_string_view(mta_string); + mta_string_free(mta_string); + + return output; + } + }; + + inline void to_json(nlohmann::json& j, const ModelMetadata::References& r) { + j = nlohmann::json{ + {"model", r.model()}, + {"architecture", r.architecture()}, + {"implementation", r.implementation()} + }; + } + + inline ModelMetadata::References from_json( + const nlohmann::json& j, nlohmann::detail::identity_tag + ) { + if (!j.is_object()) { + throw metatomic::Error("invalid JSON data for references in ModelMetadata, expected an object"); + } + + return ModelMetadata::References::builder() + .model(detail::read_string_array(j, "model", "references of ModelMetadata")) + .architecture(detail::read_string_array(j, "architecture", "references of ModelMetadata")) + .implementation(detail::read_string_array(j, "implementation", "references of ModelMetadata")) + .build(); + } + + inline void to_json(nlohmann::json& j, const ModelMetadata& m) { + j = nlohmann::json{ + {"type", "metatomic_model_metadata"}, + {"name", m.name()}, + {"authors", m.authors()}, + {"description", m.description()}, + {"references", m.references()}, + {"extra", m.extra()} + }; + } + + inline ModelMetadata from_json( + const nlohmann::json& j, nlohmann::detail::identity_tag + ) { + if (!j.is_object()) { + throw metatomic::Error("invalid JSON data for ModelMetadata, expected an object"); + } + + if (!j.contains("type") || !j["type"].is_string() || j["type"].get() != "metatomic_model_metadata") { + throw metatomic::Error("'type' in JSON for ModelMetadata must be 'metatomic_model_metadata'"); + } + + if (!j.contains("name") || !j["name"].is_string()) { + throw metatomic::Error("'name' in JSON for ModelMetadata must be a string"); + } + std::string name = j["name"].get(); + + auto authors = metatomic::detail::read_string_array(j, "authors", "JSON for ModelMetadata"); + + if (!j.contains("description") || !j["description"].is_string()) { + throw metatomic::Error("'description' in JSON for ModelMetadata must be a string"); + } + std::string description = j["description"].get(); + + if (!j.contains("references") || !j["references"].is_object()) { + throw metatomic::Error("invalid JSON data for references in ModelMetadata, expected an object"); + } + auto references = j["references"].get(); + + if (!j.contains("extra") || !j["extra"].is_object()) { + throw metatomic::Error("'extra' in JSON for ModelMetadata must be an object"); + } + std::map extra; + for (const auto& item : j["extra"].items()) { + if (!item.value().is_string()) { + throw metatomic::Error("'extra' in JSON for ModelMetadata must be an object with string values"); + } + extra[item.key()] = item.value().get(); + } + + // Validate authors content + for (const auto& author : authors) { + if (author.empty()) { + throw metatomic::Error("author can not be empty string in ModelMetadata"); + } + } + + // Validate references content + for (const auto& ref : references.model()) { + if (ref.empty()) { + throw metatomic::Error("reference can not be empty string (in 'model' section)"); + } + } + + for (const auto& ref : references.architecture()) { + if (ref.empty()) { + throw metatomic::Error("reference can not be empty string (in 'architecture' section)"); + } + } + + for (const auto& ref : references.implementation()) { + if (ref.empty()) { + throw metatomic::Error("reference can not be empty string (in 'implementation' section)"); + } + } + + return ModelMetadata::builder() + .name(std::move(name)) + .authors(std::move(authors)) + .description(std::move(description)) + .references(std::move(references)) + .extra(std::move(extra)) + .build(); + } + + /// The kind of samples a quantity can be associated with + enum class SampleKind { + /// The quantity is defined for each atom (e.g. atomic energy, charge, ...) + Atom, + /// The quantity is defined for the whole system (e.g. total energy, ...) + System, + /// The quantity is defined for each pair of atoms (e.g. hamiltonian elements, ...) + AtomPair, + }; + + /// The gradients a quantity can have + enum class Gradients { + /// Gradients with respect to atomic positions + Positions, + /// Gradients with respect to the strain (typically used for stress) + Strain, + }; + + /// A quantity that a model can use as input or output + class Quantity final { + private: + /// Name of the quantity, this can be a standard name from + /// https://docs.metatensor.org/metatomic/latest/quantities/index.html, or + /// a custom name of the form `::[/]` + std::string name_; + /// Unit of the quantity + std::string unit_; + /// Description of the quantity, used to provide more details about the + /// quantity, especially when a model defines multiple variants of the same + /// quantity. An empty string is treated as no description. + std::string description_; + /// List of explicit gradients for this quantity + std::vector gradients_; + /// The kind of samples this quantity is associated with + SampleKind sample_kind_; + + // Private constructor + Quantity( + std::string name, + std::string unit, + SampleKind sample_kind, + std::string description, + std::vector gradients + ) : name_(std::move(name)), unit_(std::move(unit)), + description_(std::move(description)), gradients_(std::move(gradients)), + sample_kind_(sample_kind) {} + + public: + /// Get the name of this quantity. + const std::string& name() const { + return name_; + } + + /// Get the unit of this quantity. + const std::string& unit() const { + return unit_; + } + + /// Get the description of this quantity. + const std::string& description() const { + return description_; + } + + /// Get the list of explicit gradients for this quantity. + const std::vector& gradients() const { + return gradients_; + } + + /// Get the kind of samples this quantity is associated with. + SampleKind sample_kind() const { + return sample_kind_; + } + + /// Builder for `Quantity`. + /// + /// Use `Quantity::builder()` to create a new builder, set the required + /// fields via the setters, and call `build()` to obtain a fully + /// initialized `Quantity`. + /// + /// `name`, `unit`, and `sample_kind` are required and must be set before + /// calling `build()`, otherwise `build()` throws `metatomic::Error`. + /// `description` defaults to an empty string and `gradients` defaults to + /// an empty list. + class Builder { + private: + std::optional name_; + std::optional unit_; + std::string description_; + std::vector gradients_; + std::optional sample_kind_; + + public: + /// Set the name of this quantity. + Builder& name(std::string value) { + name_ = std::move(value); + return *this; + } + + /// Set the unit of this quantity. + Builder& unit(std::string value) { + unit_ = std::move(value); + return *this; + } + + /// Set the description of this quantity. + Builder& description(std::string value) { + description_ = std::move(value); + return *this; + } + + /// Set the list of explicit gradients for this quantity. + Builder& gradients(std::vector value) { + gradients_ = std::move(value); + return *this; + } + + /// Add an explicit gradient to this quantity. + Builder& add_gradient(Gradients gradient) { + gradients_.push_back(gradient); + return *this; + } + + /// Set the kind of samples this quantity is associated with. + Builder& sample_kind(SampleKind value) { + sample_kind_ = value; + return *this; + } + + /// Build a fully-initialized `Quantity`. + /// + /// @throw metatomic::Error if `name`, `unit`, or `sample_kind` has + /// not been set. + /// + /// This moves the builder's fields; calling `build()` a second time + /// produces an object with moved-from values. Builders are intended + /// as one-shot temporaries. + [[nodiscard]] Quantity build() { + if (!name_.has_value()) { + throw metatomic::Error("name must be set before building Quantity"); + } + if (!unit_.has_value()) { + throw metatomic::Error("unit must be set before building Quantity"); + } + if (!sample_kind_.has_value()) { + throw metatomic::Error("sample_kind must be set before building Quantity"); + } + return Quantity( + std::move(name_).value(), + std::move(unit_).value(), + sample_kind_.value(), + std::move(description_), + std::move(gradients_) + ); + } + }; + + /// Create a new `Builder` for `Quantity`. + [[nodiscard]] static Builder builder() { + return Builder{}; + } + }; + + /// Capabilities of a model: which outputs it provides, which atoms it + /// supports, etc. + class ModelCapabilities final { + public: + /// The data type of a model, used for all inputs and outputs. + enum class DType { + /// 32-bit floating point, following the IEEE 754 standard + Float32, + /// 64-bit floating point, following the IEEE 754 standard + Float64, + }; + + /// A device on which a model can run. + enum class Device { + CPU, + CUDA, + ROCM, + Metal, + }; + + using SampleKind = metatomic::SampleKind; ///< Alias for top-level `metatomic::SampleKind` + using Gradients = metatomic::Gradients; ///< Alias for top-level `metatomic::Gradients` + using Quantity = metatomic::Quantity; ///< Alias for top-level `metatomic::Quantity` + + private: + /// The outputs this model can provide + std::vector outputs_; + /// The atomic types this model supports. The meaning of the integers in + /// this list is up to the model, and is not required to be the atomic + /// numbers. + std::vector atomic_types_; + /// The interaction range of the model (in the length unit of the model), + /// i.e. the maximum distance between two atoms for which the model's output + /// can depend on their relative position. + double interaction_range_; + /// The length unit of the model, e.g. "angstrom" or "nanometer". This is + /// used to interpret the `interaction_range` and convert the inputs. + std::string length_unit_; + /// The devices on which the model can run, e.g. `["cpu", "cuda"]`. + std::vector supported_devices_; + /// The data type of the model, used for all inputs and outputs. + DType dtype_; + + /// Validate that `value` is non-negative, throwing `metatomic::Error` + /// otherwise. Used by both the class setters and the `Builder` setters so + /// validation lives in a single place. + static void validate_interaction_range(double value) { + if (value < 0.0) { + throw metatomic::Error("interaction_range must be non-negative"); + } + } + + /// Private constructor + ModelCapabilities( + std::vector atomic_types, + double interaction_range, + std::string length_unit, + std::vector supported_devices, + DType dtype, + std::vector outputs + ) : outputs_(std::move(outputs)), + atomic_types_(std::move(atomic_types)), + interaction_range_(interaction_range), + length_unit_(std::move(length_unit)), + supported_devices_(std::move(supported_devices)), + dtype_(dtype) {} + + public: + /// Get the list of outputs this model can provide. + const std::vector& outputs() const { + return outputs_; + } + + /// Get the atomic types this model supports. + const std::vector& atomic_types() const { + return atomic_types_; + } + + /// Get the interaction range of the model. + double interaction_range() const { + return interaction_range_; + } + + /// Get the length unit of the model. + const std::string& length_unit() const { + return length_unit_; + } + + /// Get the devices on which this model can run. + const std::vector& supported_devices() const { + return supported_devices_; + } + + /// Get the data type of the model. + DType dtype() const { + return dtype_; + } + + /// Builder for `ModelCapabilities`. + /// + /// Use `ModelCapabilities::builder()` to create a new builder, set the + /// required fields via the fluent setters, and call `build()` to obtain + /// a fully-initialized `ModelCapabilities`. + /// + /// `atomic_types`, `interaction_range`, `length_unit`, + /// `supported_devices`, and `dtype` are required and must be set before + /// calling `build()`, otherwise `build()` throws `metatomic::Error`. + /// `outputs` defaults to an empty list. + class Builder { + private: + std::vector outputs_; + std::optional> atomic_types_; + std::optional interaction_range_; + std::optional length_unit_; + std::optional> supported_devices_; + std::optional dtype_; + + public: + /// Set the list of outputs this model can provide. + Builder& outputs(std::vector value) { + outputs_ = std::move(value); + return *this; + } + + /// Add an output to the list of outputs this model can provide. + Builder& add_output(const Quantity& output) { + outputs_.push_back(output); + return *this; + } + + /// Set the atomic types this model supports. + Builder& atomic_types(std::vector value) { + atomic_types_ = std::move(value); + return *this; + } + + /// Add an atomic type to the list of atomic types this model supports. + Builder& add_atomic_type(int64_t atomic_type) { + if (!atomic_types_.has_value()) { + atomic_types_ = std::vector(); + } + atomic_types_->push_back(atomic_type); + return *this; + } + + /// Set the interaction range of the model. + /// + /// @throw metatomic::Error if the value is negative. + Builder& interaction_range(double value) { + ModelCapabilities::validate_interaction_range(value); + interaction_range_ = value; + return *this; + } + + /// Set the length unit of the model. + Builder& length_unit(std::string value) { + length_unit_ = std::move(value); + return *this; + } + + /// Set the devices on which this model can run. + Builder& supported_devices(std::vector value) { + supported_devices_ = std::move(value); + return *this; + } + + /// Add a device to the list of devices on which this model can run. + Builder& add_supported_device(Device device) { + if (!supported_devices_.has_value()) { + supported_devices_ = std::vector(); + } + supported_devices_->push_back(device); + return *this; + } + + /// Set the data type of the model. + Builder& dtype(DType value) { + dtype_ = value; + return *this; + } + + /// Build a fully-initialized `ModelCapabilities`. + /// + /// @throw metatomic::Error if `atomic_types`, `interaction_range`, + /// `length_unit`, `supported_devices`, or `dtype` has not been + /// set. + /// + /// This moves the builder's fields; calling `build()` a second time + /// produces an object with moved-from values. Builders are intended + /// as one-shot temporaries. + [[nodiscard]] ModelCapabilities build() { + if (!atomic_types_.has_value()) { + throw metatomic::Error("atomic_types must be set before building ModelCapabilities"); + } + if (!interaction_range_.has_value()) { + throw metatomic::Error("interaction_range must be set before building ModelCapabilities"); + } + if (!length_unit_.has_value()) { + throw metatomic::Error("length_unit must be set before building ModelCapabilities"); + } + if (!supported_devices_.has_value()) { + throw metatomic::Error("supported_devices must be set before building ModelCapabilities"); + } + if (!dtype_.has_value()) { + throw metatomic::Error("dtype must be set before building ModelCapabilities"); + } + return ModelCapabilities( + std::move(atomic_types_).value(), + interaction_range_.value(), + std::move(length_unit_).value(), + std::move(supported_devices_).value(), + dtype_.value(), + std::move(outputs_) + ); + } + }; + + /// Create a new `Builder` for `ModelCapabilities`. + [[nodiscard]] static Builder builder() { + return Builder{}; + } + }; + + inline void to_json(nlohmann::json& j, const ModelCapabilities::DType& dtype) { + switch (dtype) { + case ModelCapabilities::DType::Float32: + j = "float32"; + break; + case ModelCapabilities::DType::Float64: + j = "float64"; + break; + default: + throw metatomic::Error("invalid dtype in ModelCapabilities"); + } + } + + inline void from_json(const nlohmann::json& j, ModelCapabilities::DType& dtype) { + if (!j.is_string()) { + throw metatomic::Error("dtype in JSON for ModelCapabilities must be a string"); + } + + std::string s = j.get(); + if (s == "float32") { + dtype = ModelCapabilities::DType::Float32; + } else if (s == "float64") { + dtype = ModelCapabilities::DType::Float64; + } else { + throw metatomic::Error( + "invalid string for dtype in JSON for ModelCapabilities, expected 'float32' or 'float64'" + ); + } + } + + inline void to_json(nlohmann::json& j, const ModelCapabilities::Device& device) { + switch (device) { + case ModelCapabilities::Device::CPU: + j = "cpu"; + break; + case ModelCapabilities::Device::CUDA: + j = "cuda"; + break; + case ModelCapabilities::Device::ROCM: + j = "rocm"; + break; + case ModelCapabilities::Device::Metal: + j = "metal"; + break; + default: + throw metatomic::Error("invalid device in ModelCapabilities"); + } + } + + inline void from_json(const nlohmann::json& j, ModelCapabilities::Device& device) { + if (!j.is_string()) { + throw metatomic::Error("device in JSON for ModelCapabilities must be a string"); + } + + std::string s = j.get(); + if (s == "cpu") { + device = ModelCapabilities::Device::CPU; + } else if (s == "cuda") { + device = ModelCapabilities::Device::CUDA; + } else if (s == "rocm") { + device = ModelCapabilities::Device::ROCM; + } else if (s == "metal") { + device = ModelCapabilities::Device::Metal; + } else { + throw metatomic::Error( + "invalid string for device in JSON for ModelCapabilities, expected 'cpu', 'cuda', 'rocm', or 'metal'" + ); + } + } + + inline void to_json(nlohmann::json& j, const SampleKind& kind) { + switch (kind) { + case SampleKind::Atom: + j = "atom"; + break; + case SampleKind::System: + j = "system"; + break; + case SampleKind::AtomPair: + j = "atom_pair"; + break; + default: + throw metatomic::Error("invalid sample_kind in Quantity"); + } + } + + inline void from_json(const nlohmann::json& j, SampleKind& kind) { + if (!j.is_string()) { + throw metatomic::Error("'sample_kind' in JSON for Quantity must be a string"); + } + + std::string s = j.get(); + if (s == "atom") { + kind = SampleKind::Atom; + } else if (s == "system") { + kind = SampleKind::System; + } else if (s == "atom_pair") { + kind = SampleKind::AtomPair; + } else { + throw metatomic::Error( + "'sample_kind' in JSON for Quantity must be 'atom', 'system' or 'atom_pair', got '" + s + "'" + ); + } + } + + inline void to_json(nlohmann::json& j, const Gradients& gradients) { + switch (gradients) { + case Gradients::Positions: + j = "positions"; + break; + case Gradients::Strain: + j = "strain"; + break; + default: + throw metatomic::Error("invalid gradients in Quantity"); + } + } + + inline void from_json(const nlohmann::json& j, Gradients& gradients) { + if (!j.is_string()) { + throw metatomic::Error("'gradients' in JSON for Quantity must be a string"); + } + + std::string s = j.get(); + if (s == "positions") { + gradients = Gradients::Positions; + } else if (s == "strain") { + gradients = Gradients::Strain; + } else { + throw metatomic::Error( + "'gradients' in JSON for Quantity must be 'positions' or 'strain', got '" + s + "'" + ); + } + } + + inline void to_json(nlohmann::json& j, const Quantity& q) { + j = nlohmann::json{ + {"type", "metatomic_quantity"}, + {"name", q.name()}, + {"unit", q.unit()}, + {"gradients", q.gradients()}, + {"sample_kind", q.sample_kind()} + }; + + if (!q.description().empty()) { + j["description"] = q.description(); + } + } + + inline Quantity from_json( + const nlohmann::json& j, nlohmann::detail::identity_tag + ) { + if (!j.is_object()) { + throw metatomic::Error("invalid JSON data for Quantity, expected an object"); + } + + if (!j.contains("type") || !j["type"].is_string() || j["type"].get() != "metatomic_quantity") { + throw metatomic::Error("'type' in JSON for Quantity must be 'metatomic_quantity'"); + } + + if (!j.contains("name") || !j["name"].is_string()) { + throw metatomic::Error("'name' in JSON for Quantity must be a string"); + } + std::string name = j["name"].get(); + + if (!j.contains("unit") || !j["unit"].is_string()) { + throw metatomic::Error("'unit' in JSON for Quantity must be a string"); + } + std::string unit = j["unit"].get(); + + std::string description; + if (j.contains("description")) { + if (!j["description"].is_string()) { + throw metatomic::Error("'description' in JSON for Quantity must be a string"); + } + description = j["description"].get(); + } + + if (!j.contains("gradients") || !j["gradients"].is_array()) { + throw metatomic::Error("'gradients' in JSON for Quantity must be an array"); + } + std::vector gradients; + for (const auto& gradient : j["gradients"]) { + gradients.push_back(gradient.get()); + } + + if (!j.contains("sample_kind") || !j["sample_kind"].is_string()) { + throw metatomic::Error("'sample_kind' in JSON for Quantity must be a string"); + } + auto sample_kind = j["sample_kind"].get(); + + return Quantity::builder() + .name(std::move(name)) + .unit(std::move(unit)) + .sample_kind(sample_kind) + .description(std::move(description)) + .gradients(std::move(gradients)) + .build(); + } + + inline void to_json(nlohmann::json& j, const ModelCapabilities& c) { + j = nlohmann::json{ + {"type", "metatomic_model_capabilities"}, + {"outputs", c.outputs()}, + {"atomic_types", c.atomic_types()}, + {"interaction_range", c.interaction_range()}, + {"length_unit", c.length_unit()}, + {"supported_devices", c.supported_devices()}, + {"dtype", c.dtype()} + }; + } + + inline ModelCapabilities from_json( + const nlohmann::json& j, nlohmann::detail::identity_tag + ) { + if (!j.is_object()) { + throw metatomic::Error("invalid JSON data for ModelCapabilities, expected an object"); + } + + if (!j.contains("type") || !j["type"].is_string() || j["type"].get() != "metatomic_model_capabilities") { + throw metatomic::Error("'type' in JSON for ModelCapabilities must be 'metatomic_model_capabilities'"); + } + + if (!j.contains("outputs") || !j["outputs"].is_array()) { + throw metatomic::Error("'outputs' in JSON for ModelCapabilities must be an array"); + } + std::vector outputs; + for (const auto& output : j["outputs"]) { + outputs.push_back(output.get()); + } + + if (!j.contains("atomic_types") || !j["atomic_types"].is_array()) { + throw metatomic::Error("'atomic_types' in JSON for ModelCapabilities must be an array"); + } + std::vector atomic_types; + for (const auto& atomic_type : j["atomic_types"]) { + if (!atomic_type.is_number_integer()) { + throw metatomic::Error("'atomic_types' in JSON for ModelCapabilities must be an array of integers"); + } + atomic_types.push_back(atomic_type.get()); + } + + if (!j.contains("interaction_range") || !j["interaction_range"].is_number()) { + throw metatomic::Error("'interaction_range' in JSON for ModelCapabilities must be a number"); + } + double interaction_range = j["interaction_range"].get(); + if (interaction_range < 0.0) { + throw metatomic::Error("'interaction_range' in JSON for ModelCapabilities must be non-negative"); + } + + if (!j.contains("length_unit") || !j["length_unit"].is_string()) { + throw metatomic::Error("'length_unit' in JSON for ModelCapabilities must be a string"); + } + std::string length_unit = j["length_unit"].get(); + + // Validate that `length_unit` has the dimension of length by asking the + // C API for a conversion factor to meters. The call only succeeds when + // the dimensions match; otherwise `check_status` throws with the C API's + // dimension-mismatch message. + double conversion_factor = 0.0; + auto status = mta_unit_conversion_factor(length_unit.c_str(), "m", &conversion_factor); + metatomic::details::check_status(status); + + if (!j.contains("supported_devices") || !j["supported_devices"].is_array()) { + throw metatomic::Error("'supported_devices' in JSON for ModelCapabilities must be an array"); + } + std::vector supported_devices; + for (const auto& device : j["supported_devices"]) { + supported_devices.push_back(device.get()); + } + + if (!j.contains("dtype") || !j["dtype"].is_string()) { + throw metatomic::Error("dtype in JSON for ModelCapabilities must be a string"); + } + auto dtype = j["dtype"].get(); + + return ModelCapabilities::builder() + .atomic_types(std::move(atomic_types)) + .interaction_range(interaction_range) + .length_unit(std::move(length_unit)) + .supported_devices(std::move(supported_devices)) + .dtype(dtype) + .outputs(std::move(outputs)) + .build(); + } + +} // namespace metatomic diff --git a/metatomic-core/include/metatomic/model.hpp b/metatomic-core/include/metatomic/model.hpp new file mode 100644 index 000000000..b6aad0187 --- /dev/null +++ b/metatomic-core/include/metatomic/model.hpp @@ -0,0 +1,454 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +namespace metatomic { + /// Render model metadata as a human-readable string. + /// + /// @param metadata a JSON-serialized `ModelMetadata` object as produced by a + /// model's `metadata` callback + /// @return a human-readable rendering of the metadata + inline std::string format_metadata(const std::string& metadata) { + mta_string_t printed = nullptr; + auto status = mta_format_metadata(metadata.c_str(), &printed); + details::check_status(status); + + return details::string_from_mta(printed); + } + + /// Abstract base class for atomistic models. + /// + /// This class provides a C++ interface for implementing custom models. Users + /// can inherit from this class, override the virtual methods, and then + /// convert the model to a `mta_model_t` with `BaseModel::to_mta_model`. + class BaseModel { + public: + virtual ~BaseModel() = default; + + /// Get the capabilities of this model. + virtual ModelCapabilities capabilities() const = 0; + + /// Get metadata describing this model. + virtual ModelMetadata metadata() const = 0; + + /// List the pair lists (neighbor lists) this model needs as input. + virtual std::vector requested_pair_lists() const = 0; + + /// List the additional per-system inputs this model needs. + virtual std::vector requested_inputs() const = 0; + + /// Run the model and compute the requested outputs. + /// + /// This method should not be used directly. It is intended to be used + /// through `mta_execute_model`. + /// + /// @param systems systems to run the model on + /// @param selected_atoms optional selection of atoms to compute outputs + /// for, or `nullptr` to use all atoms + /// @param requested_outputs outputs the model should compute + /// @return the computed outputs, one tensor map per requested output + virtual std::vector execute_inner( + const std::vector& systems, + const metatensor::Labels* selected_atoms, + const std::vector& requested_outputs + ) = 0; + + /// Convert a C++ model to a `mta_model_t`. + /// + /// The returned `mta_model_t` takes ownership of the model and will + /// delete it when the `unload` callback is called. + /// + /// @param model model to convert + /// @return a `mta_model_t` model + static mta_model_t to_mta_model(std::unique_ptr model); + + /// Build a `mta_model_t` pointing at `model`, without taking + /// ownership of it. + /// + /// The `unload` callback of the returned `mta_model_t` is left as + /// `nullptr`, and none of the other callbacks free the model. This is + /// the counterpart of `BaseModel::to_mta_model` for cases where the + /// model must stay owned by the caller, such as `execute_model`. + /// + /// @warning The returned `mta_model_t` is a view of `model`: it stores + /// a plain pointer to it and does nothing to keep it alive. It is + /// the caller's responsibility to ensure `model` outlives every use + /// of the returned `mta_model_t`, and to never pass the result to + /// an API that takes ownership of the model (i.e. one that would + /// call `unload`). + /// `BaseModel::to_mta_model` whenever ownership can be transferred. + /// + /// @param model model to take a view of + /// @return a non-owning `mta_model_t` view of `model` + static mta_model_t mta_model_view(BaseModel& model); + }; + + /// RAII wrapper around an existing `mta_model_t`. + /// + /// This class wraps a model loaded from a plugin and exposes it through the + /// same `BaseModel` interface. It owns the underlying `mta_model_t` and + /// calls its `unload` callback on destruction. + class ExternalModel final: public BaseModel { + public: + /// Wrap an existing `mta_model_t`, taking ownership of it. + /// + /// @param model model to wrap + explicit ExternalModel(mta_model_t model): + model_(model) {} + + ~ExternalModel() override { + if (model_.unload != nullptr) { + model_.unload(model_.data); + } + } + + ExternalModel(const ExternalModel&) = delete; + ExternalModel& operator=(const ExternalModel&) = delete; + + ExternalModel(ExternalModel&& other) noexcept { + *this = std::move(other); + } + + ExternalModel& operator=(ExternalModel&& other) noexcept { + if (this != &other) { + if (model_.unload != nullptr) { + model_.unload(model_.data); + } + + model_ = other.model_; + + other.model_ = mta_model_t{}; + } + + return *this; + } + + /// Get the capabilities of this model. + ModelCapabilities capabilities() const override { + this->check_callback("capabilities", model_.capabilities); + + mta_string_t output = nullptr; + auto status = model_.capabilities(model_.data, &output); + details::check_status(status); + + auto json_str = details::string_from_mta(output); + return nlohmann::json::parse(json_str).get(); + } + + /// Get metadata describing this model. + ModelMetadata metadata() const override { + this->check_callback("metadata", model_.metadata); + + mta_string_t output = nullptr; + auto status = model_.metadata(model_.data, &output); + details::check_status(status); + + auto json_str = details::string_from_mta(output); + return nlohmann::json::parse(json_str).get(); + } + + /// List the pair lists (neighbor lists) this model needs as input. + std::vector requested_pair_lists() const override { + this->check_callback("requested_pair_lists", model_.requested_pair_lists); + + mta_string_t output = nullptr; + auto status = model_.requested_pair_lists(model_.data, &output); + details::check_status(status); + + auto json_str = details::string_from_mta(output); + return nlohmann::json::parse(json_str).get>(); + } + + /// List the additional per-system inputs this model needs. + std::vector requested_inputs() const override { + this->check_callback("requested_inputs", model_.requested_inputs); + + mta_string_t output = nullptr; + auto status = model_.requested_inputs(model_.data, &output); + details::check_status(status); + + auto json_str = details::string_from_mta(output); + return nlohmann::json::parse(json_str).get>(); + } + + /// Run the model and compute the requested outputs. + std::vector execute_inner( + const std::vector& systems, + const metatensor::Labels* selected_atoms, + const std::vector& requested_outputs + ) override { + this->check_callback("execute_inner", model_.execute_inner); + + std::vector systems_ptrs; + systems_ptrs.reserve(systems.size()); + for (const auto& system: systems) { + systems_ptrs.push_back(system.as_mta_system_t()); + } + + const mts_labels_t* selected_atoms_ptr = nullptr; + if (selected_atoms != nullptr) { + selected_atoms_ptr = selected_atoms->as_mts_labels_t(); + } + + nlohmann::json json = requested_outputs; + auto requested_outputs_str = json.dump(); + + std::vector outputs(requested_outputs.size(), nullptr); + + auto status = model_.execute_inner( + model_.data, + systems_ptrs.data(), + static_cast(systems_ptrs.size()), + selected_atoms_ptr, + requested_outputs_str.c_str(), + outputs.data(), + static_cast(outputs.size()) + ); + details::check_status(status); + + std::vector result; + result.reserve(outputs.size()); + for (auto* output: outputs) { + result.push_back(metatensor::TensorMap::unsafe_from_ptr(output)); + } + + return result; + } + + /// Get a pointer to the raw `mta_model_t` backing this wrapper. + /// + /// The `ExternalModel` keeps ownership of the underlying model. + mta_model_t* as_mta_model_t() & { + return &model_; + } + + /// Get a pointer to the raw `mta_model_t` backing this wrapper. + /// + /// The `ExternalModel` keeps ownership of the underlying model. + const mta_model_t* as_mta_model_t() const & { + return &model_; + } + + /// Getting the raw pointer from a temporary `ExternalModel` is forbidden, + /// as it would immediately dangle. + mta_model_t* as_mta_model_t() && = delete; + + /// Release ownership of the underlying `mta_model_t`. + /// + /// After this call, the `ExternalModel` is empty and will not call + /// the `unload` callback on destruction. The caller is responsible + /// for calling the `unload` callback. + mta_model_t release() { + auto model = model_; + model_ = mta_model_t{}; + return model; + } + + private: + template + void check_callback(const char* name, Callback callback) const { + if (callback == nullptr) { + throw Error( + "model is missing a '" + std::string(name) + "' callback" + ); + } + } + + mta_model_t model_ = mta_model_t{}; + }; + + inline mta_model_t BaseModel::mta_model_view(BaseModel& model) { + // Short-circuit if the model is already an ExternalModel, to avoid + // double wrapping. The `ExternalModel` keeps ownership of the + // underlying model, so we clear `unload`. + if (auto* ext = dynamic_cast(&model)) { + auto m = *ext->as_mta_model_t(); + m.unload = nullptr; + return m; + } + + mta_model_t m = mta_model_t{}; + + m.data = &model; + + m.capabilities = [](const void* model_data, mta_string_t* capabilities_json) -> mta_status_t { + return details::catch_exceptions([](const void* model_data, mta_string_t* capabilities_json) { + const auto* model = static_cast(model_data); + nlohmann::json json = model->capabilities(); + *capabilities_json = mta_string_create(json.dump().c_str()); + }, model_data, capabilities_json); + }; + + m.metadata = [](const void* model_data, mta_string_t* metadata_json) -> mta_status_t { + return details::catch_exceptions([](const void* model_data, mta_string_t* metadata_json) { + const auto* model = static_cast(model_data); + nlohmann::json json = model->metadata(); + *metadata_json = mta_string_create(json.dump().c_str()); + }, model_data, metadata_json); + }; + + m.requested_pair_lists = [](const void* model_data, mta_string_t* pair_options_json) -> mta_status_t { + return details::catch_exceptions([](const void* model_data, mta_string_t* pair_options_json) { + const auto* model = static_cast(model_data); + nlohmann::json json = model->requested_pair_lists(); + *pair_options_json = mta_string_create(json.dump().c_str()); + }, model_data, pair_options_json); + }; + + m.requested_inputs = [](const void* model_data, mta_string_t* inputs_json) -> mta_status_t { + return details::catch_exceptions([](const void* model_data, mta_string_t* inputs_json) { + const auto* model = static_cast(model_data); + nlohmann::json json = model->requested_inputs(); + *inputs_json = mta_string_create(json.dump().c_str()); + }, model_data, inputs_json); + }; + + m.execute_inner = []( + void* model_data, + const struct mta_system_t* const* systems, + uintptr_t systems_count, + const mts_labels_t* selected_atoms, + const char* requested_outputs_json, + mts_tensormap_t** outputs, + uintptr_t outputs_count + ) -> mta_status_t { + return details::catch_exceptions([]( + void* model_data, + const struct mta_system_t* const* systems, + uintptr_t systems_count, + const mts_labels_t* selected_atoms, + const char* requested_outputs_json, + mts_tensormap_t** outputs, + uintptr_t outputs_count + ) { + auto* model = static_cast(model_data); + + std::vector cpp_systems; + cpp_systems.reserve(systems_count); + for (uintptr_t i = 0; i < systems_count; ++i) { + cpp_systems.push_back(System::unsafe_view_from_ptr(systems[i])); + } + + std::optional selected_atoms_copy; + const metatensor::Labels* selected_atoms_cpp = nullptr; + if (selected_atoms != nullptr) { + selected_atoms_copy = metatensor::Labels::unsafe_from_ptr( + mts_labels_clone(selected_atoms) + ); + selected_atoms_cpp = &*selected_atoms_copy; + } + + nlohmann::json json = nlohmann::json::parse(requested_outputs_json); + auto requested_outputs = json.get>(); + + auto cpp_outputs = model->execute_inner( + cpp_systems, selected_atoms_cpp, requested_outputs + ); + + if (cpp_outputs.size() != outputs_count) { + throw Error( + "model returned " + std::to_string(cpp_outputs.size()) + + " outputs, but " + std::to_string(outputs_count) + + " were requested" + ); + } + + for (uintptr_t i = 0; i < outputs_count; ++i) { + outputs[i] = cpp_outputs[i].release(); + } + }, model_data, systems, systems_count, selected_atoms, requested_outputs_json, outputs, outputs_count); + }; + + return m; + } + + inline mta_model_t BaseModel::to_mta_model(std::unique_ptr model) { + // Short-circuit if the model is already an ExternalModel + // to avoid double wrapping + if (auto* ext = dynamic_cast(model.get())) { + return ext->release(); + } + + auto m = BaseModel::mta_model_view(*model); + + // mta_model_view returns a non-owning view of the model + // Here we add an `unload` callback to take ownership + m.unload = [](void* model_data) -> mta_status_t { + return details::catch_exceptions([](void* model_data) { + delete static_cast(model_data); + }, model_data); + }; + + model.release(); + + return m; + } + + /// Execute a model to compute the requested outputs for a set of systems. + /// + /// @param model the model to execute. A view of the model is created, + /// so the ownership of `model` remains with the caller. + /// @param systems systems to run the model on + /// @param selected_atoms optional selection of atoms to compute outputs + /// for, or `nullptr` to use all atoms + /// @param requested_outputs outputs the model should compute, one per + /// requested output + /// @param check_consistency if `true`, run additional checks on the inputs + /// and on the data produced by the model + /// @return the computed outputs, one tensor map per requested output + inline std::vector execute_model( + BaseModel& model, + const std::vector& systems, + const std::optional& selected_atoms, + const std::vector& requested_outputs, + bool check_consistency + ) { + // non-owning view of the model + // `model` is kept alive by the caller + auto raw_model = BaseModel::mta_model_view(model); + + std::vector systems_ptrs; + systems_ptrs.reserve(systems.size()); + for (const auto& system: systems) { + systems_ptrs.push_back(system.as_mta_system_t()); + } + + const mts_labels_t* selected_atoms_ptr = selected_atoms.has_value() ? selected_atoms->as_mts_labels_t() : nullptr; + + nlohmann::json json = requested_outputs; + auto requested_outputs_str = json.dump(); + + std::vector outputs(requested_outputs.size(), nullptr); + + auto status = mta_execute_model( + raw_model, + systems_ptrs.data(), + static_cast(systems_ptrs.size()), + selected_atoms_ptr, + requested_outputs_str.c_str(), + check_consistency, + outputs.data(), + static_cast(outputs.size()) + ); + details::check_status(status); + + std::vector result; + result.reserve(outputs.size()); + for (auto* output: outputs) { + result.push_back(metatensor::TensorMap::unsafe_from_ptr(output)); + } + + return result; + } +} // namespace metatomic diff --git a/metatomic-core/include/metatomic/plugin.hpp b/metatomic-core/include/metatomic/plugin.hpp new file mode 100644 index 000000000..5d4a33c34 --- /dev/null +++ b/metatomic-core/include/metatomic/plugin.hpp @@ -0,0 +1,156 @@ +#pragma once + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace metatomic { + /// Load the shared library at `path` and register the plugin contained + /// within. The library must export the symbols generated by the + /// `MTA_REGISTER_PLUGIN` macro. + /// + /// @param path path to the plugin shared library + inline void load_plugin(const std::string& path) { + details::check_status(mta_load_plugin(path.c_str())); + } + + /// Load a model from `load_from` with the given options. + /// + /// If `plugin_name` is `std::nullopt`, metatomic will try to determine the + /// correct plugin to use by checking the `load_from` parameter. If we can + /// not determine the correct plugin, we then try to load the model with + /// each registered plugin until one succeeds. + /// + /// If `plugin_name` is given, then we only try to load the model with the + /// specified plugin, and return an error if the plugin can not load the + /// model. + /// + /// @param load_from where to load the model from (e.g. a file path, a + /// model name, etc.) + /// @param plugin_name optional name of the plugin to use for loading the + /// model, or `std::nullopt` to let metatomic search + /// @param options_json optional JSON object containing string keys and + /// string values for loading the model + /// @return the loaded model + inline ExternalModel load_model( + const std::string& load_from, + std::optional options_json = std::nullopt, + std::optional plugin_name = std::nullopt + ) { + mta_model_t model; + const char* plugin_name_ptr = nullptr; + if (plugin_name.has_value()) { + plugin_name_ptr = plugin_name->c_str(); + } + + const char* options_json_ptr = nullptr; + if (options_json.has_value()) { + options_json_ptr = options_json->c_str(); + } + + details::check_status(mta_load_model( + load_from.c_str(), + options_json_ptr, + plugin_name_ptr, + &model + )); + + return ExternalModel(model); + } + + namespace details { + using load_model_t = std::unique_ptr (*)( + const std::string& load_from, + const std::map& options + ); + + /// `load_model` callback for a `mta_plugin_t` built by `MTA_REGISTER_CXX_PLUGIN` + inline mta_status_t cxx_plugin_load_model( + const char* load_from, + const char* options_json, + mta_model_t* model, + load_model_t load_model_fn + ) { + std::unique_ptr cpp_model; + auto status = metatomic::details::catch_exceptions([&]() { + auto options_obj = nlohmann::json::parse(options_json); + auto options = options_obj.get>(); + cpp_model = load_model_fn(std::string(load_from), options); + }); + + if (status != MTA_SUCCESS) { + // exception was caught while trying to load the model + // stop the plugin search and report the error + return status; + } + + if (cpp_model == nullptr) { + // the plugin could not load this model + return MTA_UNSUPPORTED_MODEL_ERROR; + } + + *model = BaseModel::to_mta_model(std::move(cpp_model)); + return MTA_SUCCESS; + } + } // namespace details +} // namespace metatomic + + +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunused-macros" +#endif + + +/// Plugin entry point for a C++ plugin. +/// +/// This should be used once in a C++ plugin shared library. `plugin_name` is +/// the name of the plugin, and `load_model_fn` must be a plain function with +/// the signature: +/// +/// ```cpp +/// std::unique_ptr load_model_fn( +/// const std::string& load_from, +/// const std::map& options +/// ); +/// ``` +/// +/// `load_model_fn` MUST return `nullptr` if it can not load the model +/// described by `load_from`, so metatomic can try another plugin. +/// +/// Only one `MTA_REGISTER_CXX_PLUGIN` can be used per shared library. +#define MTA_REGISTER_CXX_PLUGIN(plugin_name, load_model_fn) \ + MTA_REGISTER_PLUGIN(register_fn, { \ + static_assert( \ + std::is_convertible< \ + decltype(load_model_fn), metatomic::details::load_model_t \ + >::value, \ + "MTA_REGISTER_CXX_PLUGIN: load_model_fn must be callable as " \ + "std::unique_ptr(" \ + "std::string, std::map)" \ + ); \ + \ + static mta_plugin_t PLUGIN { \ + /*abi_version*/ MTA_ABI_VERSION, \ + /*name*/ plugin_name, \ + /*load_model*/ []( \ + const char* load_from, const char* options_json, mta_model_t* model \ + ) { \ + return metatomic::details::cxx_plugin_load_model( \ + load_from, options_json, model, load_model_fn \ + ); \ + } \ + }; \ + return register_fn(PLUGIN); \ + }) + +#ifdef __clang__ +#pragma clang diagnostic pop +#endif diff --git a/metatomic-core/include/metatomic/system.hpp b/metatomic-core/include/metatomic/system.hpp new file mode 100644 index 000000000..e6cc37d81 --- /dev/null +++ b/metatomic-core/include/metatomic/system.hpp @@ -0,0 +1,320 @@ +#pragma once + +#include +#include +#include + +#include +#include + +#include +#include +#include + +namespace metatomic { + /// A `System` contains all the information about an atomistic system, and is + /// used as the input of atomistic models. + /// + /// This is a RAII wrapper around the `mta_system_t` type from the C API. It + /// can either own the underlying system (in which case it is freed with the + /// `System`), or be a non-owning view into a system owned elsewhere (for + /// example a system passed to a model by the runtime). + class System final { + public: + /// Create a new `System` from DLPack tensors. + /// + /// Ownership of all four tensors is transferred to the new `System`. + /// + /// @param length_unit unit of length used by `positions` and `cell` + /// @param types tensor with shape `(n_atoms,)` of atomic types + /// @param positions tensor with shape `(n_atoms, 3)` of atomic positions + /// @param cell tensor with shape `(3, 3)` of the unit cell vectors + /// @param pbc tensor with shape `(3,)` of periodic boundary conditions + /// + /// The dtype and layout required for each tensor are validated by + /// `mta_system_create`; see the C API documentation for details. + System( + const std::string& length_unit, + DLPackTensor types, + DLPackTensor positions, + DLPackTensor cell, + DLPackTensor pbc + ) { + auto status = mta_system_create( + length_unit.c_str(), + types.release(), + positions.release(), + cell.release(), + pbc.release(), + &system_ + ); + details::check_status(status); + details::check_pointer(system_); + } + + ~System() { + if (!is_view_) { + // `mta_system_free` is a no-op on a null pointer + mta_system_free(system_); + } + } + + /// `System` is not copy-constructible + System(const System&) = delete; + /// `System` is not copy-assignable + System& operator=(const System&) = delete; + + /// `System` is move-constructible + System(System&& other) noexcept { + *this = std::move(other); + } + + /// `System` is move-assignable + System& operator=(System&& other) noexcept { + if (!is_view_) { + mta_system_free(system_); + } + + system_ = other.system_; + is_view_ = other.is_view_; + + other.system_ = nullptr; + other.is_view_ = true; + + return *this; + } + + /// Get the number of atoms in this system. + size_t size() const { + uintptr_t size = 0; + auto status = mta_system_size(system_, &size); + details::check_status(status); + return static_cast(size); + } + + /// Get the unit of length used by the positions and cell of this system. + std::string length_unit() const { + mta_string_t length_unit = nullptr; + auto status = mta_system_get_length_unit(system_, &length_unit); + details::check_status(status); + return details::string_from_mta(length_unit); + } + + /// Get the atomic types of all atoms in this system, as a tensor with + /// shape `(n_atoms,)`. + /// + /// @see `data` for the meaning of the returned tensor. + DLPackTensor types() const { + return this->data(MTA_SYSTEM_DATA_TYPES); + } + + /// Get the positions of all atoms in this system, as a tensor with shape + /// `(n_atoms, 3)`. + /// + /// @see `data` for the meaning of the returned tensor. + DLPackTensor positions() const { + return this->data(MTA_SYSTEM_DATA_POSITIONS); + } + + /// Get the unit cell of this system, as a tensor with shape `(3, 3)`. + /// + /// @see `data` for the meaning of the returned tensor. + DLPackTensor cell() const { + return this->data(MTA_SYSTEM_DATA_CELL); + } + + /// Get the periodic boundary conditions of this system, as a tensor with + /// shape `(3,)`. + /// + /// @see `data` for the meaning of the returned tensor. + DLPackTensor pbc() const { + return this->data(MTA_SYSTEM_DATA_PBC); + } + + /// Add a pair list (i.e. neighbor list) to this system. + /// + /// Ownership of `pairs` is transferred to this `System`. + /// + /// @param options options describing the pair list + /// @param pairs pairs data, stored as a metatensor block + void add_pairs(const PairListOptions& options, metatensor::TensorBlock pairs) { + nlohmann::json j = options; + this->add_pairs(j.dump(), std::move(pairs)); + } + + /// Add a pair list (i.e. neighbor list) to this system. + /// + /// Ownership of `pairs` is transferred to this `System`. + /// + /// @param options_json JSON-serialized `PairListOptions` describing the + /// pair list + /// @param pairs pairs data, stored as a metatensor block + void add_pairs(const std::string& options_json, metatensor::TensorBlock pairs) { + auto status = mta_system_add_pairs(system_, options_json.c_str(), pairs.release()); + details::check_status(status); + } + + /// Get a previously stored pair list matching the given `options_json`. + /// + /// The returned block is a non-owning view into data owned by this + /// `System`, and is only valid for as long as this `System` is alive. + /// + /// @param options options identifying the pair list to retrieve + metatensor::TensorBlock pairs(const PairListOptions& options) const { + nlohmann::json j = options; + return this->pairs(j.dump()); + } + + /// Get a previously stored pair list matching the given `options_json`. + /// + /// The returned block is a non-owning view into data owned by this + /// `System`, and is only valid for as long as this `System` is alive. + /// + /// @param options_json JSON-serialized `PairListOptions` identifying + /// the pair list to retrieve + metatensor::TensorBlock pairs(const std::string& options_json) const { + const mts_block_t* pairs = nullptr; + auto status = mta_system_get_pairs(system_, options_json.c_str(), &pairs); + details::check_status(status); + details::check_pointer(pairs); + return metatensor::TensorBlock::unsafe_view_from_ptr(const_cast(pairs)); + } + + /// Get the options of all pair lists registered with this `System` + std::vector known_pairs() const { + mta_string_t options = nullptr; + auto status = mta_system_known_pairs(system_, &options); + details::check_status(status); + nlohmann::json j = nlohmann::json::parse(mta_string_view(options)); + mta_string_free(options); + return j.get>(); + } + + /// Get the options of all pair lists registered with this `System`, as + /// a JSON-serialized array of `PairListOptions`. + std::vector known_pairs_json() const { + mta_string_t options = nullptr; + auto status = mta_system_known_pairs(system_, &options); + details::check_status(status); + nlohmann::json j = nlohmann::json::parse(mta_string_view(options)); + mta_string_free(options); + return j.get>(); + } + + /// Add custom data to this system, stored under the given `name`. + /// + /// Ownership of `data` is transferred to this `System`. + /// + /// @param name name used to identify the custom data + /// @param data custom data, stored as a metatensor tensor map + void add_custom_data(const std::string& name, metatensor::TensorMap data) { + auto status = mta_system_add_custom_data(system_, name.c_str(), data.release()); + details::check_status(status); + } + + /// Get the custom data previously stored under the given `name`. + /// + /// The returned tensor map is a non-owning view into data owned by this + /// `System`, and is only valid for as long as this `System` is alive. + /// + /// @param name name of the custom data to retrieve + metatensor::TensorMap custom_data(const std::string& name) const { + const mts_tensormap_t* data = nullptr; + auto status = mta_system_get_custom_data(system_, name.c_str(), &data); + details::check_status(status); + details::check_pointer(data); + return metatensor::TensorMap::unsafe_view_from_ptr(const_cast(data)); + } + + /// Get the names of all custom data registered with this `System` + std::vector known_custom_data() const { + mta_string_t names = nullptr; + auto status = mta_system_known_custom_data(system_, &names); + details::check_status(status); + nlohmann::json j = nlohmann::json::parse(mta_string_view(names)); + mta_string_free(names); + return j.get>(); + } + + /// Get the raw `mta_system_t` pointer backing this `System`. + /// + /// The `System` keeps ownership of the pointer, which is only valid for + /// as long as this `System` is alive. + mta_system_t* as_mta_system_t() & { + return system_; + } + + /// Get the raw `mta_system_t` pointer backing this `System`. + /// + /// The `System` keeps ownership of the pointer, which is only valid for + /// as long as this `System` is alive. + const mta_system_t* as_mta_system_t() const & { + return system_; + } + + /// Getting the raw pointer from a temporary `System` is forbidden, as it + /// would immediately dangle. + mta_system_t* as_mta_system_t() && = delete; + + /// Create an owning `System` from a raw `mta_system_t` pointer, taking + /// ownership of it. The system will be freed when the `System` is + /// destroyed. + /// + /// This is an advanced function, and the caller is responsible for + /// ensuring that `system` was allocated by the C API and is not used + /// anywhere else. + static System unsafe_from_ptr(mta_system_t* system) { + return System(system, /*is_view*/ false); + } + + /// Create a non-owning `System` view from a raw `mta_system_t` pointer. + /// The system will *not* be freed when the `System` is destroyed, and + /// must outlive it. + /// + /// This is an advanced function, mainly useful to wrap the systems given + /// to a model by the runtime. + static System unsafe_view_from_ptr(const mta_system_t* system) { + return System(const_cast(system), /*is_view*/ true); + } + + /// Release the raw `mta_system_t` pointer from this `System` without + /// freeing it, transferring ownership back to the caller. + mta_system_t* release() { + this->check_not_view("release"); + auto* system = system_; + system_ = nullptr; + is_view_ = true; + return system; + } + + private: + /// Wrap an existing `mta_system_t` pointer, see `unsafe_from_ptr` and + /// `unsafe_view_from_ptr`. + explicit System(mta_system_t* system, bool is_view): + system_(system), is_view_(is_view) {} + + void check_not_view(const char* method_name) const { + if (is_view_) { + throw Error( + "can not call System::" + std::string(method_name) + + " on this system since it is a view of a system owned elsewhere." + ); + } + } + + /// Get one of the always-present data tensors of this system. + /// + /// The returned `DLPackTensor` is a view sharing its data with the + /// system, which is kept alive for as long as the view exists. + DLPackTensor data(mta_system_data_kind request) const { + DLManagedTensorVersioned* data = nullptr; + auto status = mta_system_get_data(system_, request, &data); + details::check_status(status); + details::check_pointer(data); + return DLPackTensor(data); + } + + mta_system_t* system_ = nullptr; + bool is_view_ = false; + }; +} // namespace metatomic diff --git a/metatomic-core/include/metatomic/utils.hpp b/metatomic-core/include/metatomic/utils.hpp new file mode 100644 index 000000000..87cbd30de --- /dev/null +++ b/metatomic-core/include/metatomic/utils.hpp @@ -0,0 +1,122 @@ +#pragma once + +#include +#include + +#include +#include + +namespace metatomic { + /// RAII wrapper around a DLPack `DLManagedTensorVersioned*`. + /// + /// This owns the managed tensor and calls its deleter when the wrapper is + /// destroyed. It can be used to move ownership of DLPack tensors across the + /// metatomic C++ API. + class DLPackTensor final { + public: + /// Create an empty wrapper, not owning any tensor. + DLPackTensor() = default; + + /// Take ownership of an existing DLPack managed tensor. + explicit DLPackTensor(DLManagedTensorVersioned* tensor): tensor_(tensor) {} + + /// The managed tensor is freed through its own deleter on destruction. + ~DLPackTensor() = default; + + /// `DLPackTensor` is not copy-constructible + DLPackTensor(const DLPackTensor&) = delete; + /// `DLPackTensor` is not copy-assignable + DLPackTensor& operator=(const DLPackTensor&) = delete; + + /// `DLPackTensor` is move-constructible + DLPackTensor(DLPackTensor&&) noexcept = default; + /// `DLPackTensor` is move-assignable + DLPackTensor& operator=(DLPackTensor&&) noexcept = default; + + /// Check whether this wrapper currently owns a tensor. + explicit operator bool() const { + return static_cast(tensor_); + } + + /// Access the underlying `DLManagedTensorVersioned` without transferring + /// ownership. The pointer stays owned by this `DLPackTensor`. + DLManagedTensorVersioned* operator->() const { + return tensor_.get(); + } + + /// Get the underlying `DLManagedTensorVersioned` pointer. It stays owned + /// by this `DLPackTensor`, and is only valid for as long as it is alive. + DLManagedTensorVersioned* as_dlpack() const { + return tensor_.get(); + } + + /// Release the underlying `DLManagedTensorVersioned` without calling its + /// deleter, transferring ownership back to the caller. + DLManagedTensorVersioned* release() { + return tensor_.release(); + } + + private: + /// Deleter implementing the DLPack ownership protocol: invoke the managed + /// tensor's own `deleter` callback if it has one. + struct Deleter { + void operator()(DLManagedTensorVersioned* tensor) const noexcept { + if (tensor->deleter != nullptr) { + tensor->deleter(tensor); + } + } + }; + + std::unique_ptr tensor_; + }; + + namespace details { + /// Take ownership of an `mta_string_t` returned by the C API, copy its + /// contents into an owned `std::string`, and free the C string. + /// + /// The `unique_ptr` guard frees the C string on return, including if the + /// copy into the `std::string` throws. A null `mta_string_t` (as produced + /// by an empty output) yields an empty string. + inline std::string string_from_mta(mta_string_t string) { + struct Deleter { + void operator()(mta_string_t ptr) const noexcept { + mta_string_free(ptr); + } + }; + std::unique_ptr, Deleter> owned(string); + + if (string == nullptr) { + return std::string(); + } + return std::string(mta_string_view(string)); + } + } // namespace details + + /// Get the multiplicative conversion factor to use to convert from + /// `from_unit` to `to_unit`. Both units are parsed as expressions + /// (e.g. `kJ / mol / A^2`, `(eV * u)^(1/2)`) and their dimensions must + /// match. + /// + /// @verbatim embed:rst:leading-slashes + /// + /// .. seealso:: + /// + /// The general documentation for :ref:`units`, with the expression + /// syntax and list of supported base units. + /// + /// @endverbatim + /// + /// @param from_unit the unit to convert from + /// @param to_unit the unit to convert to + inline double unit_conversion_factor( + const std::string& from_unit, + const std::string& to_unit + ) { + double conversion = 0.0; + + auto status = mta_unit_conversion_factor(from_unit.c_str(), to_unit.c_str(), &conversion); + details::check_status(status); + + return conversion; + } +} // namespace metatomic diff --git a/metatomic-core/src/c_api/io.rs b/metatomic-core/src/c_api/io.rs new file mode 100644 index 000000000..a09918592 --- /dev/null +++ b/metatomic-core/src/c_api/io.rs @@ -0,0 +1,308 @@ +use std::ffi::{c_char, c_void, CStr}; +use std::fs::File; +use std::io::{BufReader, Cursor}; + +use metatensor::c_api::{mts_create_array_callback_t, mts_realloc_buffer_t}; + +use super::{catch_unwind, mta_status_t, mta_system_t}; +use crate::{Error, System}; + +/// Wrapper for an externally managed buffer, that can be grown to fit more data +struct ExternalBuffer { + data: *mut *mut u8, + writen: u64, + allocated: u64, + + realloc_user_data: *mut c_void, + realloc: unsafe extern "C" fn(*mut c_void, *mut u8, usize) -> *mut u8, + + current: u64, +} + +impl std::io::Write for ExternalBuffer { + #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let remaining_space = self.allocated.saturating_sub(self.current); + + if (remaining_space as usize) < buf.len() { + let required_size = self.current.saturating_add(buf.len() as u64); + let mut new_size = if self.allocated == 0 { 1024 } else { self.allocated }; + while new_size < required_size { + new_size = new_size.saturating_mul(2); + if new_size == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::OutOfMemory, + "requested allocation size overflow", + )); + } + } + + let new_ptr = unsafe { + (self.realloc)(self.realloc_user_data, *self.data, new_size as usize) + }; + + if new_ptr.is_null() { + return Err(std::io::Error::new( + std::io::ErrorKind::OutOfMemory, + "failed to allocate memory with the realloc callback" + )); + } + + unsafe { + *self.data = new_ptr; + } + + self.allocated = new_size; + } + + let mut output = unsafe { + let start = (*self.data).offset(self.current as isize); + // allocated >= current + buf.len() + std::slice::from_raw_parts_mut(start, buf.len()) + }; + + let count = output.write(buf).expect("failed to write to pre-allocated slice"); + assert_eq!(count, buf.len()); + self.current += count as u64; + + if self.current > self.writen { + self.writen = self.current; + } + return Ok(count); + } + + fn flush(&mut self) -> std::io::Result<()> { + return Ok(()); + } +} + + +#[allow(clippy::cast_sign_loss, clippy::cast_possible_wrap)] +impl std::io::Seek for ExternalBuffer { + fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result { + match pos { + std::io::SeekFrom::Start(offset) => { + if offset > self.writen { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, "tried to seek past the end of the buffer") + ); + } + + self.current = offset; + }, + + std::io::SeekFrom::End(offset) => { + if offset > 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, "tried to seek past the end of the buffer") + ); + } + + if -offset > self.writen as i64 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, "tried to seek past the beginning of the buffer") + ); + } + + self.current = (self.writen as i64 + offset) as u64; + }, + + std::io::SeekFrom::Current(offset) => { + let result = self.current as i64 + offset; + if result > self.writen as i64 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, "tried to seek past the end of the buffer") + ); + } + + if result < 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, "tried to seek past the beginning of the buffer") + ); + } + + self.current = result as u64; + }, + } + + return Ok(self.current); + } + + fn rewind(&mut self) -> std::io::Result<()> { + self.current = 0; + return Ok(()); + } + + fn stream_position(&mut self) -> std::io::Result { + return Ok(self.current); + } +} + + +/// Save a system to a file. +/// +/// The format consists of a zip archive containing NPY files for the system's +/// data (types, positions, cell, pbc), a `info.json` file for metadata, and +/// optional sub-directories for pair lists (`pairs//options.json` and +/// `pairs//data.mts`) and custom data (`data/.mts`). +/// +/// @param path A null-terminated C string containing the file path. Must not be +/// null. +/// @param system The system to save. Must not be null. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_save( + path: *const c_char, + system: *const mta_system_t +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(path, system); + + let path = unsafe { CStr::from_ptr(path) }.to_str() + .map_err(|_| Error::InvalidParameter("path is not valid UTF-8".into()))?; + + let file = File::create(path)?; + let system = unsafe { &*system.cast::() }; + crate::io::save(file, system)?; + + Ok(()) + }) +} + +/// Save a system to an in-memory buffer. +/// +/// The buffer is grown as needed using the provided `realloc` callback. On +/// success, `*buffer` points to the serialized data and `*buffer_count` +/// contains the number of bytes written. +/// +/// @param buffer Pointer to the buffer pointer. On input, `*buffer` may be NULL +/// (in which case `*buffer_count` must be 0). On output, `*buffer` is +/// updated to point to the serialized data. +/// @param buffer_count Pointer to the buffer size. On input, `*buffer_count` +/// must contain the current allocation size. On output, it is set to the +/// number of bytes written. +/// @param realloc_user_data User data passed as the first argument to +/// `realloc`. +/// @param realloc Callback to grow the buffer. Must not be NULL. +/// @param system The system to save. Must not be null. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation)] +pub unsafe extern "C" fn mta_save_buffer( + buffer: *mut *mut u8, + buffer_count: *mut usize, + realloc_user_data: *mut c_void, + realloc: mts_realloc_buffer_t, + system: *const mta_system_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(buffer, buffer_count, system); + + let realloc = if let Some(realloc) = realloc { + realloc + } else { + return Err(Error::InvalidParameter( + "realloc callback can not be NULL in mta_save_buffer".into() + )); + }; + + if unsafe { (*buffer).is_null() } { + // `ExternalBuffer.write` calls realloc with the current `*buffer` + // (which may be null) for the initial allocation. + unsafe { *buffer = std::ptr::null_mut(); } + } + + let system = unsafe { &*system.cast::() }; + let mut external_buffer = ExternalBuffer { + data: buffer, + allocated: unsafe { *buffer_count } as u64, + writen: 0, + realloc_user_data, + realloc, + current: 0, + }; + + crate::io::save(&mut external_buffer, system)?; + + unsafe { + *buffer_count = external_buffer.current as usize; + } + + Ok(()) + }) +} + +/// Load a system from a file. +/// +/// The file must have been written by `mta_save` and contain a valid metatomic +/// system. +/// +/// @param path A null-terminated C string containing the file path. Must not be +/// null. +/// @param create_array Callback to allocate arrays for the system's data. Must +/// not be NULL. +/// @param system Output parameter, set to the newly created system handle. +/// The caller takes ownership and must free it with `mta_system_free`. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_load( + path: *const c_char, + create_array: mts_create_array_callback_t, + system: *mut *mut mta_system_t, +) -> mta_status_t { + catch_unwind(move || { + check_pointers_non_null!(path); + + let path = unsafe { CStr::from_ptr(path) }.to_str() + .map_err(|_| Error::InvalidParameter("path is not valid UTF-8".into()))?; + + let file = BufReader::new(File::open(path)?); + let new_system = crate::io::load(file, create_array)?; + + unsafe { + *system = mta_system_t::into_raw(new_system); + } + + Ok(()) + }) +} + +/// Load a system from an in-memory buffer. +/// +/// The buffer must contain data serialized by `mta_save_buffer` (or the +/// equivalent Rust function). +/// +/// @param buffer Pointer to the serialized data. Must not be NULL. +/// @param buffer_size Number of bytes in `buffer`. +/// @param create_array Callback to allocate arrays for the system's data. Must +/// not be NULL. +/// @param system Output parameter, set to the newly created system handle. +/// The caller takes ownership and must free it with `mta_system_free`. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_load_buffer( + buffer: *const u8, + buffer_size: usize, + create_array: mts_create_array_callback_t, + system: *mut *mut mta_system_t, +) -> mta_status_t { + catch_unwind(move || { + check_pointers_non_null!(buffer); + + let slice = unsafe { + std::slice::from_raw_parts(buffer, buffer_size) + }; + let cursor = Cursor::new(slice); + let new_system = crate::io::load(cursor, create_array)?; + + unsafe { + *system = mta_system_t::into_raw(new_system); + } + + Ok(()) + }) +} diff --git a/metatomic-core/src/c_api/mod.rs b/metatomic-core/src/c_api/mod.rs new file mode 100644 index 000000000..dc705f81a --- /dev/null +++ b/metatomic-core/src/c_api/mod.rs @@ -0,0 +1,20 @@ +#![allow(clippy::doc_markdown)] + +#[macro_use] +mod status; +pub use self::status::{mta_status_t, catch_unwind}; + +mod utils; +pub use self::utils::mta_string_t; +pub use self::utils::{mta_string_create, mta_string_free, mta_string_view}; + +mod system; +pub use self::system::mta_system_t; + +mod model; +pub use self::model::mta_model_t; + +mod plugin; +pub use self::plugin::{mta_plugin_t, mta_register_plugin, mta_load_plugin, mta_load_model}; + +mod io; diff --git a/metatomic-core/src/c_api/model.rs b/metatomic-core/src/c_api/model.rs new file mode 100644 index 000000000..fcbd15f95 --- /dev/null +++ b/metatomic-core/src/c_api/model.rs @@ -0,0 +1,272 @@ +use std::ffi::{c_void, c_char, CStr}; +use std::sync::Arc; +use metatensor::c_api::{mts_labels_t, mts_tensormap_t}; + +use super::catch_unwind; +use crate::{Error, ModelMetadata, System}; +use crate::model::execute_model; + +use super::{mta_status_t, mta_string_t, mta_system_t}; + +/// A model that computes physical properties of atomistic systems. +/// +/// `mta_model_t` is a small virtual table: `data` holds the model's own state, +/// and the function pointers describe what the model can do. A model is usually +/// produced by a plugin's `load_model` callback (see `mta_load_model`) and then +/// executed with `mta_execute_model`. +/// +/// Every callback receives `data` as its first argument. metatomic treats +/// `data` as opaque and only hands it back to the callbacks. Callbacks should +/// report any error by saving it with `mta_set_last_error` and returning a +/// non-success `mta_status_t`. +#[repr(C)] +#[allow(non_camel_case_types)] +pub struct mta_model_t { + /// Opaque pointer to the model's internal state + /// + /// Its layout and meaning are private to the model implementation. It is + /// initialized by whoever creates the model (e.g. a plugin's `load_model`) + /// and released by `unload`. + pub data: *mut c_void, + + /// Release the resources owned by `model_data` + /// + /// Called exactly once when the model is no longer needed. May be `NULL` if + /// the model owns no resources. + /// + /// @param model_data the model's `data` pointer + /// @return `MTA_SUCCESS` on success, another status code on error + pub unload: Option mta_status_t>, + + /// Get the capabilities of the model as a JSON string. + /// + /// @verbatim embed:rst:leading-asterisk + /// The expected JSON structure is documented in :ref:`core-json-model-capabilities`. + /// @endverbatim + /// + /// @param model_data the model's `data` pointer + /// @param capabilities_json output string, set to a JSON-serialized + /// `ModelCapabilities` object. The caller takes ownership and must + /// free it with `mta_string_free`. + /// @return `MTA_SUCCESS` on success, another status code on error + pub capabilities: Option mta_status_t>, + + /// Get metadata describing the model (name, authors, references, ...) as a + /// JSON string. + /// + /// @verbatim embed:rst:leading-asterisk + /// The expected JSON structure is documented in :ref:`core-json-model-metadata`. + /// @endverbatim + /// + /// @param model_data the model's `data` pointer + /// @param metadata_json output string, set to a JSON-serialized + /// `ModelMetadata` object. The caller takes ownership and must + /// free it with `mta_string_free`. + /// @return `MTA_SUCCESS` on success, another status code on error + pub metadata: Option mta_status_t>, + + /// List the pair lists (neighbor lists) the model needs as input as a JSON + /// string. + /// + /// @verbatim embed:rst:leading-asterisk + /// + /// The engine is expected to compute these and attach them to every system + /// with :c:func:`mta_system_add_pairs` before calling + /// :c:func:`mta_execute_model`. + /// + /// The expected JSON structure for each pair list is documented in :ref:`core-json-pair-options`. + /// + /// @endverbatim + /// + /// @param model_data the model's `data` pointer + /// @param pair_options_json output string, set to a JSON array of + /// `PairListOptions` objects. The caller takes ownership and must + /// free it with `mta_string_free`. + /// @return `MTA_SUCCESS` on success, another status code on error + pub requested_pair_lists: Option mta_status_t>, + + /// List the additional per-system inputs the model needs as a JSON string. + /// + /// @verbatim embed:rst:leading-asterisk + /// + /// These correspond to custom data the engine should attach to every system + /// with :c:func:`mta_system_add_custom_data` before execution. + /// + /// The expected JSON structure for each input is documented in :ref:`core-json-quantity`. + /// + /// @endverbatim + /// + /// @param model_data the model's `data` pointer + /// @param inputs_json output string, set to a JSON array of `Quantity` + /// objects, one per requested input. The caller takes ownership and + /// must free it with `mta_string_free`. + /// @return `MTA_SUCCESS` on success, another status code on error + pub requested_inputs: Option mta_status_t>, + + /// Run the model and compute the requested outputs + /// + /// @verbatim embed:rst:leading-asterisk + /// + /// This performs the model's actual computation. This should not be called + /// directly, but rather through :c:func:`mta_execute_model`, which handles + /// unit conversion and can check inputs and output data for consistency. + /// + /// @endverbatim + /// + /// @param model_data the model's `data` pointer + /// @param systems array of `systems_count` systems to run the model on + /// @param systems_count number of entries in `systems` + /// @param selected_atoms optional labels selecting the subset of atoms to + /// compute outputs for, or `NULL` to use all atoms. When set, it has the + /// dimensions `"system"` and `"atom"` holding 0-based indices. + /// @param requested_outputs_json JSON string containing an array of + /// `Quantity`, one for each output the model should produce + /// @param outputs array of `outputs_count` tensor maps to fill, one per + /// requested output and in the same order + /// @param outputs_count number of entries in `outputs`, must equal + /// `requested_outputs_count` + /// @return `MTA_SUCCESS` on success, another status code on error + pub execute_inner: Option mta_status_t>, +} + +impl mta_model_t { + pub(crate) fn null() -> Self { + return mta_model_t { + data: std::ptr::null_mut(), + unload: None, + capabilities: None, + metadata: None, + requested_pair_lists: None, + requested_inputs: None, + execute_inner: None, + }; + } +} + +/// Execute a model to compute the requested outputs for a set of systems +/// +/// This is the main entry point to run a model loaded through the C API. It +/// validates the arguments and delegates the computation to the model's +/// `execute_inner` callback. +/// +/// @param model the model to execute +/// @param systems array of `systems_count` systems to run the model on +/// @param systems_count number of entries in `systems` +/// @param selected_atoms optional labels selecting the subset of atoms to +/// compute outputs for, or `NULL` to use all atoms +/// @param requested_outputs_json JSON string containing an array of +/// `Quantity`, one for each output the model should produce +/// @param check_consistency if `true`, run additional checks on the +/// inputs and on the data produced by the model +/// @param outputs array of `outputs_count` tensor maps to fill, one per +/// requested output and in the same order. The caller takes ownership of +/// the returned tensor maps. +/// @param outputs_count number of entries in `outputs`, must equal +/// `requested_outputs_count` +/// @return `MTA_SUCCESS` on success, another status code on error (the message +/// is available through `mta_last_error`) +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_execute_model( + model: mta_model_t, + systems: *const *const mta_system_t, + systems_count: usize, + selected_atoms: *const mts_labels_t, + requested_outputs_json: *const c_char, + check_consistency: bool, + outputs: *mut *mut mts_tensormap_t, + outputs_count: usize, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(systems, requested_outputs_json, outputs); + + let json = unsafe { CStr::from_ptr(requested_outputs_json) } + .to_str() + .map_err(|_| Error::InvalidParameter("requested_outputs_json is not valid UTF-8".into()))?; + + let model = crate::model::Model::from_ref(&model); + + // Recover Arc from each pointer, bumping the refcount. + // We use Arc::clone + forget so the caller's pointers remain valid. + let systems_vec: Vec> = { + let mut vec = Vec::with_capacity(systems_count); + for i in 0..systems_count { + let ptr = unsafe { *systems.add(i) }; + if ptr.is_null() { + return Err(Error::InvalidParameter( + format!("systems[{i}] is NULL") + )); + } + let system = unsafe { mta_system_t::from_raw(ptr) }; + vec.push(Arc::clone(&system)); + std::mem::forget(system); + } + vec + }; + + execute_model( + model, + &systems_vec, + selected_atoms, + json, + check_consistency, + outputs, + outputs_count, + )?; + + Ok(()) + }) +} + +/// Render model metadata as a human-readable string +/// +/// @param metadata a JSON-serialized `ModelMetadata` object as produced by a +/// model's `metadata` callback. Must not be null. +/// @param printed output string, set to a human-readable rendering of the +/// metadata. The caller takes ownership and must free it with +/// `mta_string_free`. +/// @return `MTA_SUCCESS` on success, another status code on error +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_format_metadata( + metadata: *const c_char, + printed: *mut mta_string_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(metadata, printed); + + let metadata = unsafe { std::ffi::CStr::from_ptr(metadata) }; + let metadata = metadata.to_str().map_err(|_| { + Error::InvalidParameter("metadata is not valid UTF-8".into()) + })?; + + let metadata = json::parse(metadata).map_err(|e| { + Error::Serialization(format!("invalid JSON for ModelMetadata: {e}")) + })?; + + let metadata = ModelMetadata::try_from(&metadata)?; + + unsafe { + *printed = mta_string_t::new(metadata.print()); + } + Ok(()) + }) +} diff --git a/metatomic-core/src/c_api/plugin.rs b/metatomic-core/src/c_api/plugin.rs new file mode 100644 index 000000000..f3df4ce3e --- /dev/null +++ b/metatomic-core/src/c_api/plugin.rs @@ -0,0 +1,176 @@ +use std::ffi::{CStr, c_char}; + +use super::catch_unwind; +use super::{mta_model_t, mta_status_t}; +use crate::Error; +use crate::Plugin; + +/// A metatomic plugin definition. +#[allow(non_camel_case_types)] +#[repr(C)] +pub struct mta_plugin_t { + /// ABI version this plugin was compiled against, this should be set to + /// `MTA_ABI_VERSION` when creating the plugin struct. + pub abi_version: i32, + + /// Name of the plugin, as a null-terminated UTF-8 string. This is the name + /// specified in `mta_load_model` when trying to load a model with a + /// specific plugin. The name must be unique among all registered plugins. + pub name: *const c_char, + + /// Callback function to load a model. This function should try to load a + /// model from `load_from` (which can be a file path, a model name, etc.) + /// and a set of key/values options passed as a JSON string. + /// + /// If the plugin can load the model, it should fill `model` with a pointer + /// to a valid `mta_model_t` struct and return `MTA_SUCCESS`. If the data in + /// `load_from` does not correspond to a model supported by the plugin, it + /// should return `MTA_UNSUPPORTED_MODEL_ERROR`. If an error occurs while + /// loading the model, it should return another status code and save an + /// error message with `mta_set_last_error`. + /// + /// @param load_from a null-terminated UTF-8 string describing where to load + /// the model from (e.g. a file path, a model name, etc.). The + /// interpretation of this string is up to the plugin. + /// @param options_json a null-terminated UTF-8 string containing a set of + /// string keys and string value options for loading the model. + /// @param model output pointer to the loaded model. The caller takes + /// ownership of the model and must unload it when the model is no + /// longer needed. + /// @return `MTA_SUCCESS` if the model was loaded successfully, + /// `MTA_UNSUPPORTED_MODEL_ERROR` if the plugin can not load the model, + /// or another status code if an error occurs. + pub load_model: Option mta_status_t>, +} + +unsafe impl Send for mta_plugin_t {} + +/// Register a plugin. This is passed as a callback to the `MTA_REGISTER_PLUGIN` +/// macro, and should not be called directly by C or C++ plugin implementations. +/// +/// @param plugin the plugin to register +/// @return `MTA_SUCCESS` if the plugin was registered successfully, or another +/// status code if an error occurs. You can get more details about the error +/// with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_register_plugin(plugin: mta_plugin_t) -> mta_status_t { + catch_unwind(move || { + let plugin = Plugin::new(plugin)?; + crate::plugin::register_plugin(plugin)?; + Ok(()) + }) +} + +/// Load the shared library at `path` and register the plugin contained within. +/// +/// The library must export the symbols generated by the `MTA_REGISTER_PLUGIN` +/// macro. +/// +/// @param path a null-terminated UTF-8 string containing the path to the plugin +/// shared library, or `NULL` to load the plugin from the current binary +/// @return `MTA_SUCCESS` if the plugin was loaded successfully, or another +/// status code if an error occurs. You can get more details about the +/// error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_load_plugin(path: *const c_char) -> mta_status_t { + catch_unwind(move || { + let path = if path.is_null() { + None + } else { + let path = unsafe { CStr::from_ptr(path) }; + let path = path.to_str().map_err(|_| { + Error::InvalidParameter("invalid UTF-8 in plugin path".into()) + })?; + Some(path) + }; + + crate::plugin::load_plugin(path) + }) +} + +/// Load a model from `load_from` with the given options. +/// +/// If `plugin_name` is a NULL pointer, metatomic will try to determine the +/// correct plugin to use by checking the `load_from` parameter. If we can not +/// determine the correct plugin, we then try to load the model with each +/// registered plugin until one succeeds. +/// +/// If `plugin_name` is given, then we only try to load the model with the +/// specified plugin, and return an error if the plugin can not load the model. +/// +/// @param plugin_name optional null-terminated UTF-8 string containing the name +/// of the plugin to use for loading the model, or `NULL` to let metatomic +/// search for a correct plugin +/// @param load_from a null-terminated UTF-8 string describing where to load the +/// model from (e.g. a file path, a model name, etc.). The interpretation +/// of this string is up to the plugin. +/// @param options_json a null-terminated UTF-8 string containing a set of string +/// keys and string value options for loading the model. The interpretation +/// of these options is up to the plugin. +/// @param model output pointer to the loaded model. The caller takes ownership of +/// the model and must unload it when the model is no longer needed. +/// @return `MTA_SUCCESS` if the model was loaded successfully, or another +/// status code if an error occurs. You can get more details about the +/// error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_load_model( + load_from: *const c_char, + options_json: *const c_char, + plugin_name: *const c_char, + model: *mut mta_model_t, +) -> mta_status_t { + let unwind_wrapper = std::panic::AssertUnwindSafe(model); + + catch_unwind(move || { + check_pointers_non_null!(load_from, model); + + let plugin_name = if plugin_name.is_null() { + None + } else { + let cstr = unsafe { CStr::from_ptr(plugin_name) }; + Some(cstr.to_str().map_err(|_| { + Error::InvalidParameter("invalid UTF-8 in plugin name".into()) + })?) + }; + + let options_json = if options_json.is_null() { + c"{}" + } else { + unsafe { CStr::from_ptr(options_json) } + }; + + let options_str = options_json.to_str().map_err(|_| { + Error::InvalidParameter("invalid UTF-8 in options JSON".into()) + })?; + + let options = json::parse(options_str).map_err( + |e| Error::Serialization(format!("JSON parsing error: {}", e)) + )?; + if !options.is_object() { + return Err(Error::Serialization("JSON options must be an object in `mta_load_model`".into())) + } + + // just some validation, we pass the raw JSON down to the plugins + for (key, value) in options.entries() { + if !value.is_string() { + return Err(Error::InvalidParameter(format!( + "JSON option '{}' has a non-string value in `mta_load_model`", + key + ))); + } + } + + let load_from = unsafe { CStr::from_ptr(load_from) }; + let loaded = crate::plugin::load_model(load_from, options_json, plugin_name)?; + + let _ = &unwind_wrapper; + unsafe { + *unwind_wrapper.0 = loaded.into_raw(); + } + Ok(()) + }) +} diff --git a/metatomic-core/src/c_api/status.rs b/metatomic-core/src/c_api/status.rs new file mode 100644 index 000000000..f589c3087 --- /dev/null +++ b/metatomic-core/src/c_api/status.rs @@ -0,0 +1,219 @@ +use std::cell::RefCell; +use std::ffi::{c_char, c_void, CStr, CString}; +use std::panic::UnwindSafe; + +use crate::Error; + +#[derive(Debug)] +struct LastError { + message: CString, + origin: CString, + custom_data: *mut c_void, + custom_data_deleter: Option, +} + +// Save the last error message in thread local storage. +thread_local! { + pub static LAST_ERROR: RefCell = RefCell::new(LastError { + message: CString::new("").expect("invalid C string"), + origin: CString::new("").expect("invalid C string"), + custom_data: std::ptr::null_mut(), + custom_data_deleter: None, + }); +} + +/// Status type returned by all functions in the C API. +/// +/// The value 0 (`MTA_SUCCESS`) indicates success, while any non-zero value indicates an error. +#[allow(non_camel_case_types)] +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum mta_status_t { + /// Status code indicating success + MTA_SUCCESS = 0, + /// Status code indicating invalid function parameters + MTA_INVALID_PARAMETER_ERROR = 1, + /// Status code indicating I/O errors + MTA_IO_ERROR = 2, + /// Status code indicating memory allocation errors + MTA_MEMORY_ERROR = 3, + /// Status code indicating serialization/deserialization errors + MTA_SERIALIZATION_ERROR = 4, + /// Status code indicating dlpack errors + MTA_DLPACK_ERROR = 5, + /// Status code indicating metatensor errors + MTA_METATENSOR_ERROR = 6, + /// Status code used by plugins when a model is not supported by the + /// current plugin + MTA_UNSUPPORTED_MODEL_ERROR = 7, + /// Status code used by model for any error that does not fit the cases above + MTA_MODEL_ERROR = 8, + /// Status code used when there is an internal error + MTA_INTERNAL_ERROR = 255, +} + +/// `std::panic::catch_unwind` that automatically transform +/// the error into `mta_status_t`. +pub fn catch_unwind(function: F) -> mta_status_t +where + F: FnOnce() -> Result<(), Error> + UnwindSafe, +{ + match std::panic::catch_unwind(function) { + Ok(Ok(())) => mta_status_t::MTA_SUCCESS, + Ok(Err(error)) => error.into(), + Err(error) => Error::from(error).into(), + } +} + +/// Check that pointers (used as C API function parameters) are not null. +#[macro_export] +#[doc(hidden)] +macro_rules! check_pointers_non_null { + ($pointer: ident) => { + if $pointer.is_null() { + return Err($crate::Error::InvalidParameter( + format!( + "got invalid NULL pointer for {} at {}:{}", + stringify!($pointer), file!(), line!() + ) + )); + } + }; + ($($pointer: ident),* $(,)?) => { + $(check_pointers_non_null!($pointer);)* + } +} + +impl From for mta_status_t { + fn from(error: Error) -> mta_status_t { + if let Error::CallbackError(status) = error { + // If the error is already a CallbackError, we can directly return the corresponding status code. + return status; + } + + LAST_ERROR.with(|last_error| { + let mut last_error = last_error.borrow_mut(); + + // If there is a custom data deleter, + // use it to free the custom data before overwriting it with the new error. + if let Some(deleter) = last_error.custom_data_deleter { + unsafe { + deleter(last_error.custom_data); + } + } + + *last_error = LastError { + message: CString::new(format!("{}", error)) + .expect("error message contains a null byte"), + origin: CString::new("metatomic-core").expect("invalid C string"), + custom_data: std::ptr::null_mut(), + custom_data_deleter: None, + }; + }); + + match error { + Error::InvalidParameter(_) => mta_status_t::MTA_INVALID_PARAMETER_ERROR, + Error::Io(_) => mta_status_t::MTA_IO_ERROR, + Error::Serialization(_) => mta_status_t::MTA_SERIALIZATION_ERROR, + Error::Dlpack(_) => mta_status_t::MTA_DLPACK_ERROR, + Error::Metatensor(_) => mta_status_t::MTA_METATENSOR_ERROR, + Error::CallbackError(_) => unreachable!("already handled above"), + Error::Internal(_) => mta_status_t::MTA_INTERNAL_ERROR, + + } + } +} + +/// Get last error message that was created on the current thread. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_last_error( + message: *mut *const c_char, + origin: *mut *const c_char, + data: *mut *mut c_void, +) -> mta_status_t { + let status = std::panic::catch_unwind(|| { + LAST_ERROR.with(|last_error| { + unsafe { + let last_error = last_error.borrow(); + if !message.is_null() { + *message = last_error.message.as_ptr(); + } + if !origin.is_null() { + *origin = last_error.origin.as_ptr(); + } + if !data.is_null() { + *data = last_error.custom_data; + } + } + }); + }); + + match status { + Ok(()) => mta_status_t::MTA_SUCCESS, + Err(error) => { + let last_error_debug = + LAST_ERROR.with(|last_error| format!("{:?}", last_error.borrow())); + if error.is::() { + eprintln!( + "panic in mta_last_error: {:?}, last_error: {:?}", + error.downcast_ref::(), + last_error_debug + ); + } else if error.is::<&str>() { + eprintln!( + "panic in mta_last_error: {:?}, last_error: {:?}", + error.downcast_ref::<&str>(), + last_error_debug + ); + } else { + eprintln!( + "panic in mta_last_error: unknown panic error type. last_error: {:?}", + last_error_debug + ); + } + mta_status_t::MTA_INTERNAL_ERROR + } + } +} + +/// Set last error message for the current thread. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_set_last_error( + message: *const c_char, + origin: *const c_char, + data: *mut c_void, + data_deleter: Option, +) -> mta_status_t { + catch_unwind(move || { + let message = if message.is_null() { + CString::new("").expect("invalid C string") + } else { + unsafe { CString::from(CStr::from_ptr(message)) } + }; + + let origin = if origin.is_null() { + CString::new("").expect("invalid C string") + } else { + unsafe { CString::from(CStr::from_ptr(origin)) } + }; + + LAST_ERROR.with(|last_error| { + let mut last_error = last_error.borrow_mut(); + + // Call custom data deleter before overwriting the custom data with the new one, to avoid memory leaks. + if let Some(deleter) = last_error.custom_data_deleter { + unsafe { + deleter(last_error.custom_data); + } + } + + *last_error = LastError { + message: message, + origin: origin, + custom_data: data, + custom_data_deleter: data_deleter, + }; + }); + Ok(()) + }) +} diff --git a/metatomic-core/src/c_api/system.rs b/metatomic-core/src/c_api/system.rs new file mode 100644 index 000000000..7659dc533 --- /dev/null +++ b/metatomic-core/src/c_api/system.rs @@ -0,0 +1,479 @@ +use std::ffi::{c_char, CStr}; +use std::sync::Arc; + +use dlpk::sys::DLManagedTensorVersioned; +use dlpk::{DLPackTensor, DLPackVersion}; +use metatensor::c_api::{mts_block_t, mts_tensormap_t}; +use metatensor::{TensorBlock, TensorMap}; + +use crate::{Error, PairListOptions, System}; +use super::{catch_unwind, mta_status_t, mta_string_t}; + +/// Opaque handle to an atomistic system. +/// +/// The system owns DLPack tensors for types, positions, cell, and PBC, as well +/// as metatensor blocks for pair lists and tensor maps for custom data. +#[repr(transparent)] +#[allow(non_camel_case_types)] +pub struct mta_system_t(pub(crate) System); + +impl mta_system_t { + /// Convert a `System` into a pointer inside an `Arc`, to be passed + /// through the C API as an `mta_system_t`. + pub(crate) fn into_raw(system: Arc) -> *mut mta_system_t { + Arc::into_raw(system).cast::().cast_mut() + } + + /// Recover the `Arc` from a pointer created with + /// [`mta_system_t::into_raw`] + pub(crate) unsafe fn from_raw(ptr: *const mta_system_t) -> Arc { + unsafe { Arc::from_raw(ptr.cast::()) } + } +} + +/// Create a new system from raw DLPack tensors. +/// +/// This function **takes ownership** of `types`, `positions`, `cell`, and +/// `pbc`. The caller must not use these tensors after calling this function. +/// +/// @param length_unit A null-terminated C string containing the length unit +/// (e.g. "Angstrom", "nanometer"). Must not be null. +/// @param types A DLPack managed tensor with shape `(n_atoms,)` and dtype +/// `int32`. Ownership is transferred. +/// @param positions A DLPack managed tensor with shape `(n_atoms, 3)` and +/// dtype `float32` or `float64`. Ownership is transferred. +/// @param cell A DLPack managed tensor with shape `(3, 3)` and the same dtype +/// as `positions`. Ownership is transferred. +/// @param pbc A DLPack managed tensor with shape `(3,)` and dtype `bool`. +/// Ownership is transferred. +/// @param system Output parameter, set to the newly created system handle. +/// The caller takes ownership and must free it with `mta_system_free`. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_create( + length_unit: *const c_char, + types: *mut DLManagedTensorVersioned, + positions: *mut DLManagedTensorVersioned, + cell: *mut DLManagedTensorVersioned, + pbc: *mut DLManagedTensorVersioned, + system: *mut *mut mta_system_t, +) -> mta_status_t { + let unwind_wrapper = std::panic::AssertUnwindSafe(system); + catch_unwind(move || { + check_pointers_non_null!(length_unit, types, positions, cell, pbc, system); + + unsafe { + let length_unit = CStr::from_ptr(length_unit) + .to_str() + .map_err(|_| Error::InvalidParameter("length_unit is not valid UTF-8".into()))? + .to_string(); + + + let types = DLPackTensor::from_ptr(types); + let positions = DLPackTensor::from_ptr(positions); + let cell = DLPackTensor::from_ptr(cell); + let pbc = DLPackTensor::from_ptr(pbc); + + let system = Arc::new(System::new(length_unit, types, positions, cell, pbc)?); + + let _ = &unwind_wrapper; + *unwind_wrapper.0 = mta_system_t::into_raw(system); + } + Ok(()) + }) +} + +/// Free a system previously created by `mta_system_create`. +/// +/// If there are outstanding borrowed views (from `mta_system_get_data`), the +/// system's data will remain alive until all views are released. +/// +/// @param system The system handle to free. Can be null, in which case this +/// function is a no-op. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_free(system: *mut mta_system_t) -> mta_status_t { + catch_unwind(|| { + if system.is_null() { + return Ok(()); + } + + let system = unsafe { mta_system_t::from_raw(system.cast_const()) }; + std::mem::drop(system); + Ok(()) + }) +} + +/// Get the number of atoms in a system. +/// +/// @param system The system handle. Must not be null. +/// @param size Output parameter, set to the number of atoms. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_size( + system: *const mta_system_t, + size: *mut usize, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, size); + + unsafe { + let system = &*system.cast::(); + *size = system.size(); + } + Ok(()) + }) +} + +/// Kind of data always stored in a system. +/// +/// Other kinds of data can be stored with `mta_system_add_custom_data` and +/// retrieved with `mta_system_get_custom_data`. +#[allow(non_camel_case_types)] +#[repr(C)] +#[non_exhaustive] +pub enum mta_system_data_kind { + MTA_SYSTEM_DATA_TYPES = 0, + MTA_SYSTEM_DATA_POSITIONS = 1, + MTA_SYSTEM_DATA_CELL = 2, + MTA_SYSTEM_DATA_PBC = 3, +} + +/// Custom deleter for borrowed DLPack tensors returned by `mta_system_get_data`. +/// +/// Releases the `Arc` reference stored in `manager_ctx` and +/// frees the heap-allocated `DLManagedTensorVersioned`. +unsafe extern "C" fn borrowed_tensor_deleter( + tensor: *mut DLManagedTensorVersioned, +) { + let system = unsafe { + mta_system_t::from_raw((*tensor).manager_ctx.cast()) + }; + std::mem::drop(system); + unsafe { + std::mem::drop(Box::from_raw(tensor)); + } +} + +/// Get a DLPack tensor from a system for the requested data. +/// +/// This function **returns a borrowed view** of the system's internal data. +/// The returned `DLManagedTensorVersioned` has a custom deleter that decrements +/// the system's reference count, keeping the system alive as long as the +/// borrowed view exists. +/// +/// The caller is responsible for calling the deleter on the returned tensor +/// when it is no longer needed. The tensor shares the data pointer with the +/// system; do **not** modify it. +/// +/// @param system The system handle. Must not be null. +/// @param request Which data to retrieve (types, positions, cell, or PBC). +/// @param data Output parameter, set to a pointer to a newly allocated +/// `DLManagedTensorVersioned` containing the requested data. The caller +/// takes ownership and must call the deleter when done. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_get_data( + system: *const mta_system_t, + request: mta_system_data_kind, + data: *mut *mut DLManagedTensorVersioned, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, data); + unsafe { + *data = std::ptr::null_mut(); + } + + // increase the reference count of the system so that it stays alive as + // long as the returned tensor is alive. We do this by creating a + // temporary Arc from the raw pointer, cloning it and storing the clone + // in the manager_ctx. + let system = unsafe { mta_system_t::from_raw(system) }; + let arc_clone = system.clone(); + + let tensor_ref = match request { + mta_system_data_kind::MTA_SYSTEM_DATA_TYPES => system.types(), + mta_system_data_kind::MTA_SYSTEM_DATA_POSITIONS => system.positions(), + mta_system_data_kind::MTA_SYSTEM_DATA_CELL => system.cell(), + mta_system_data_kind::MTA_SYSTEM_DATA_PBC => system.pbc(), + }; + + let packed = Box::new(DLManagedTensorVersioned { + version: DLPackVersion::current(), + manager_ctx: Arc::into_raw(arc_clone) as *mut std::ffi::c_void, + deleter: Some(borrowed_tensor_deleter), + flags: dlpk::sys::DLPACK_FLAG_BITMASK_READ_ONLY, + dl_tensor: tensor_ref.raw, + }); + + // do not drop the system, it is still owned by the caller. + std::mem::forget(system); + + unsafe { + *data = Box::into_raw(packed); + } + Ok(()) + }) +} + +/// Get the length unit of a system. +/// +/// This function returns a new `mta_string_t` that the caller must free with +/// `mta_string_free`. +/// +/// @param system The system handle. Must not be null. +/// @param length_unit Output parameter, set to the length unit string. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_get_length_unit( + system: *const mta_system_t, + length_unit: *mut mta_string_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, length_unit); + + unsafe { + let system = &*system.cast::(); + *length_unit = mta_string_t::new(system.length_unit()); + } + Ok(()) + }) +} + +/// Add a pair list (neighbor list) to a system. +/// +/// This function **takes ownership** of `pairs`. The caller must not use the +/// block after calling this function. +/// +/// @param system The system handle. Must not be null. +/// @param options A JSON-serialized `PairListOptions` object. Must not be null. +/// @param pairs A `mts_block_t` containing the pair data. Ownership is +/// transferred. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_add_pairs( + system: *mut mta_system_t, + options: *const c_char, + pairs: *mut mts_block_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, options, pairs); + + let options_str = unsafe { CStr::from_ptr(options) } + .to_str() + .map_err(|_| Error::InvalidParameter("options is not valid UTF-8".into()))?; + + let options_json = json::parse(options_str) + .map_err(|e| Error::Serialization(format!("invalid JSON for PairListOptions: {e}")))?; + + let options = PairListOptions::try_from(&options_json)?; + + let pairs = unsafe { TensorBlock::from_raw(pairs) }; + + let mut system = unsafe { mta_system_t::from_raw(system.cast_const()) }; + system.add_pairs(options, pairs)?; + + // do not drop the system, it is still owned by the caller. + std::mem::forget(system); + + Ok(()) + }) +} + +/// Get a pair list from a system. +/// +/// **Returns a borrowed view** of the pair list. The system must outlive the +/// returned pointer. Do **not** free the returned block. +/// +/// @param system The system handle. Must not be null. +/// @param options A JSON-serialized `PairListOptions` object identifying which +/// pair list to retrieve. Must not be null. +/// @param pairs Output parameter, set to a pointer to the pair list block, or +/// NULL if no pair list matches the options. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_get_pairs( + system: *const mta_system_t, + options: *const c_char, + pairs: *mut *const mts_block_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, options, pairs); + + let options_str = unsafe { CStr::from_ptr(options) } + .to_str() + .map_err(|_| Error::InvalidParameter("options is not valid UTF-8".into()))?; + + let options_json = json::parse(options_str) + .map_err(|e| Error::Serialization(format!("invalid JSON for PairListOptions: {e}")))?; + + let options = PairListOptions::try_from(&options_json)?; + + let system = unsafe { &*system.cast::() }; + match system.get_pairs(&options) { + Some(block) => { + unsafe { + *pairs = block.as_ptr(); + } + } + None => { + return Err(Error::InvalidParameter( + "no pair list found for the given options".into(), + )); + } + } + + Ok(()) + }) +} + +/// Get all pair list options known by a system. +/// +/// This function returns a new `mta_string_t` containing a JSON array of +/// `PairListOptions` objects. The caller must free it with `mta_string_free`. +/// +/// @param system The system handle. Must not be null. +/// @param pairs_options Output parameter, set to a JSON string containing an +/// array of `PairListOptions` objects. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_known_pairs( + system: *const mta_system_t, + pairs_options: *mut mta_string_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, pairs_options); + + let system = unsafe { &*system.cast::() }; + let known = system.known_pairs(); + let mut json_array = json::JsonValue::new_array(); + for options in known { + json_array.push(json::JsonValue::from(options.clone())).map_err(|_| { + Error::Internal("failed to build JSON array".into()) + })?; + } + + unsafe { + *pairs_options = mta_string_t::new(json::stringify(json_array)); + } + Ok(()) + }) +} + +/// Add custom data to a system. +/// +/// This function **takes ownership** of `data`. The caller must not use the +/// tensor map after calling this function. +/// +/// @param system The system handle. Must not be null. +/// @param name A null-terminated C string containing the name of the custom +/// data. Must not be null. +/// @param data A `mts_tensormap_t` containing the custom data. Ownership is +/// transferred. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_add_custom_data( + system: *mut mta_system_t, + name: *const c_char, + data: *mut mts_tensormap_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, name, data); + + let name = unsafe { CStr::from_ptr(name) } + .to_str() + .map_err(|_| Error::InvalidParameter("name is not valid UTF-8".into()))? + .to_string(); + + let data = unsafe { TensorMap::from_raw(data) }; + + let mut system = unsafe { mta_system_t::from_raw(system.cast_const()) }; + system.add_custom_data(name, data, false)?; + + // do not drop the system, it is still owned by the caller. + std::mem::forget(system); + + Ok(()) + }) +} + +/// Get custom data from a system by name. +/// +/// **Returns a borrowed view** of the custom data. The system must outlive the +/// returned pointer. Do **not** free the returned tensor map. +/// +/// @param system The system handle. Must not be null. +/// @param name A null-terminated C string containing the name of the custom +/// data. Must not be null. +/// @param data Output parameter, set to a pointer to the custom data tensor +/// map, or an error if no data with the given name exists. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_get_custom_data( + system: *const mta_system_t, + name: *const c_char, + data: *mut *const mts_tensormap_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, name, data); + + let name = unsafe { CStr::from_ptr(name) } + .to_str() + .map_err(|_| Error::InvalidParameter("name is not valid UTF-8".into()))?; + + let system = unsafe { &*system.cast::() }; + let result = system.get_custom_data(name)?; + + unsafe { + *data = result.as_ptr(); + } + + Ok(()) + }) +} + +/// Get all custom data names known by a system. +/// +/// **Returns a new** `mta_string_t` containing a JSON array of strings. The +/// caller must free it with `mta_string_free`. +/// +/// @param system The system handle. Must not be null. +/// @param names Output parameter, set to a JSON string containing an array of +/// custom data names. +/// @return `MTA_SUCCESS` on success, or another status code if an error occurs. +/// You can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_system_known_custom_data( + system: *const mta_system_t, + names: *mut mta_string_t, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(system, names); + + let system = unsafe { &*system.cast::() }; + let known = system.known_custom_data(); + let mut json_array = json::JsonValue::new_array(); + for name in known { + json_array.push(name).map_err(|_| { + Error::Internal("failed to build JSON array".into()) + })?; + } + + unsafe { + *names = mta_string_t::new(json::stringify(json_array)); + } + Ok(()) + }) +} + +// TODO: mta_system_to(device, dtype) diff --git a/metatomic-core/src/c_api/utils.rs b/metatomic-core/src/c_api/utils.rs new file mode 100644 index 000000000..22fabe38a --- /dev/null +++ b/metatomic-core/src/c_api/utils.rs @@ -0,0 +1,196 @@ +use std::ffi::{CString, c_char}; + +use std::sync::LazyLock; + +use super::{mta_status_t, catch_unwind}; +use crate::Error; + +static VERSION: LazyLock = LazyLock::new(|| { + CString::new(env!("METATOMIC_FULL_VERSION")).expect("version contains NULL byte") +}); + + +/// Get the runtime version of the metatomic library as a string. +/// +/// This version follows the `..[-]` format. +#[unsafe(no_mangle)] +pub extern "C" fn mta_version() -> *const c_char { + return VERSION.as_ptr(); +} + +/// Heap-allocated backing storage for `mta_string_t`, opaque to C users. +#[allow(non_camel_case_types)] +#[repr(transparent)] +pub struct mta_opaque_string_t(c_char); + +/// An heap-allocated UTF-8 string passed across the C API boundary. +/// +/// This is used whenever a C API function or callback needs to return a string. +/// +/// A null pointer represents an absent or empty string. Use `mta_string_create` +/// to allocate, `mta_string_free` to release, and `mta_string_view` to get a +/// pointer to the inner C string. +#[allow(non_camel_case_types)] +#[repr(transparent)] +pub struct mta_string_t(*mut mta_opaque_string_t); + +impl std::fmt::Debug for mta_string_t { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut builder = f.debug_tuple("mta_string_t"); + + if self.0.is_null() { + builder.field(&"NULL"); + } else { + builder.field(&self.as_str()); + } + builder.finish() + } +} + +impl mta_string_t { + /// Create a new `mta_string_t` from a Rust string. + pub fn new(value: impl Into) -> Self { + let cstring = CString::new(value.into()).expect("string contains NULL byte"); + let ptr = CString::into_raw(cstring); + return mta_string_t(ptr.cast()); + } + + /// Create a null `mta_string_t`, representing an absent string. + pub fn null() -> Self { + mta_string_t(std::ptr::null_mut()) + } + + /// View the string as a `&str`. Returns `""` for a null string. + pub fn as_str(&self) -> &str { + if self.0.is_null() { + return ""; + } + unsafe { + let cstr = std::ffi::CStr::from_ptr(self.0.cast()); + return cstr.to_str().expect("invalid UTF-8 in mta_string_t"); + } + } +} + +/// Allocate a new `mta_string_t` by copying the null-terminated C string +/// `string`. +/// +/// The returned string must be freed with `mta_string_free`. +/// +/// @param string A pointer to a null-terminated C string. Must not be null. +/// @return A new `mta_string_t` containing a copy of `string`, or null if an +/// error occurred. You can check the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_string_create( + string: *const c_char, +) -> mta_string_t { + let mut result = mta_string_t::null(); + let unwind_wrapper = std::panic::AssertUnwindSafe(&mut result); + + catch_unwind(move || { + check_pointers_non_null!(string); + + let cstr = unsafe { std::ffi::CStr::from_ptr(string) }; + let string = CString::from(cstr); + + let ptr = CString::into_raw(string); + + let _ = &unwind_wrapper; + *unwind_wrapper.0 = mta_string_t(ptr.cast()); + Ok(()) + }); + + return result; +} + +/// Free a `mta_string_t` previously created by `mta_string_create`. +/// +/// @param string A `mta_string_t` to free. Can be null, in which case this function is a no-op. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_string_free(string: mta_string_t) { + catch_unwind(|| { + if string.0.is_null() { + return Ok(()); + } + + let ptr = string.0.cast::(); + let cstring = unsafe { CString::from_raw(ptr) }; + std::mem::drop(cstring); + + Ok(()) + }); +} + +/// Return a pointer to the null-terminated string data inside `string`. +/// +/// The pointer is valid only for the lifetime of `string`. +/// +/// @param string A `mta_string_t` containing the string to view. Must not be null. +/// @return A pointer to the null-terminated C string inside `string` +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_string_view( + string: mta_string_t, +) -> *const c_char { + let mut result = std::ptr::null(); + let unwind_wrapper = std::panic::AssertUnwindSafe(&mut result); + + catch_unwind(move || { + let string = string.0; + check_pointers_non_null!(string); + + let _ = &unwind_wrapper; + *unwind_wrapper.0 = string.cast(); + + Ok(()) + }); + + return result; +} + +/// Get the multiplicative conversion factor to use to convert from `from_unit` +/// to `to_unit`. Both units are parsed as expressions (e.g. `kJ / mol / A^2`, +/// `(eV * u)^(1/2)`) and their dimensions must match. +/// +/// @verbatim embed:rst:leading-asterisk +/// +/// .. seealso:: +/// +/// The general documentation for :ref:`units`, with the expression +/// syntax and list of supported base units. +/// +/// @endverbatim +/// +/// @param from_unit A null-terminated C string containing the unit to convert from. +/// @param to_unit A null-terminated C string containing the unit to convert to. +/// @param conversion A pointer to a `double` where the conversion factor will be stored. +/// @return The status code of the operation. If this code is not `MTA_SUCCESS`, +/// you can get more details about the error with `mta_last_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn mta_unit_conversion_factor( + from_unit: *const c_char, + to_unit: *const c_char, + conversion: *mut f64, +) -> mta_status_t { + catch_unwind(|| { + check_pointers_non_null!(from_unit, to_unit, conversion); + + let from_cstr = unsafe { std::ffi::CStr::from_ptr(from_unit) }; + let to_cstr = unsafe { std::ffi::CStr::from_ptr(to_unit) }; + + let from_str = from_cstr.to_str().map_err(|_| { + Error::InvalidParameter("from_unit is not valid UTF-8".into()) + })?; + let to_str = to_cstr.to_str().map_err(|_| { + Error::InvalidParameter("to_unit is not valid UTF-8".into()) + })?; + + unsafe { + *conversion = crate::unit_conversion_factor(from_str, to_str)?; + } + + Ok(()) + }) +} + + +// TODO: logging & warnings? diff --git a/metatomic-core/src/io/mod.rs b/metatomic-core/src/io/mod.rs new file mode 100644 index 000000000..b0cfa2530 --- /dev/null +++ b/metatomic-core/src/io/mod.rs @@ -0,0 +1,34 @@ +use crate::Error; + +mod npy_header; + +mod tensor; +mod system; + +pub use system::{load, save}; + +pub trait ReadAndSeek: std::io::Read + std::io::Seek {} +impl ReadAndSeek for T {} + +pub enum PathOrBuffer<'a> { + Path(&'a str), + Buffer(&'a mut dyn ReadAndSeek), +} + +/// Byte order for multi-byte values in NPY files. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Endianness { + Little, + Big, + Native, +} + +// returns an error if the given reader contains any more data +fn check_for_extra_bytes(reader: &mut R) -> Result<(), Error> { + let extra = reader.read_to_end(&mut Vec::new())?; + if extra == 0 { + Ok(()) + } else { + Err(Error::Serialization(format!("found {} extra bytes after the expected end of data", extra))) + } +} diff --git a/metatomic-core/src/io/npy_header.rs b/metatomic-core/src/io/npy_header.rs new file mode 100644 index 000000000..3b789aafe --- /dev/null +++ b/metatomic-core/src/io/npy_header.rs @@ -0,0 +1,673 @@ +// This file was initially taken from https://github.com/jturner314/ndarray-npy, +// version 0.8.1. It is Copyright 2018–2021 Jim Turner and ndarray-npy +// developers, released under MIT and Apache Licenses. +use std::convert::TryFrom; +use std::sync::Arc; +use std::error::Error; +use std::fmt::Write as FmtWrite; +use std::io::Write as IoWrite; + +use byteorder::{ByteOrder, LittleEndian, ReadBytesExt}; + +/// Magic string to indicate npy format. +const MAGIC_STRING: &[u8] = b"\x93NUMPY"; + +/// The total header length (including magic string, version number, header +/// length value, array format description, padding, and final newline) must be +/// evenly divisible by this value. +// If this changes, update the docs of `ViewNpyExt` and `ViewMutNpyExt`. +const HEADER_DIVISOR: usize = 64; + +#[derive(Debug)] +pub enum ParseHeaderError { + MagicString, + Version { + major: u8, + minor: u8, + }, + /// Indicates that the `HEADER_LEN` doesn't fit in `usize`. + HeaderLengthOverflow(u32), + /// Indicates that the array format string contains non-ASCII characters. + /// This is an error for .npy format versions 1.0 and 2.0. + NonAscii, + /// Error parsing the array format string as UTF-8. This does not apply to + /// .npy format versions 1.0 and 2.0, which require the array format string + /// to be ASCII. + Utf8Parse(std::str::Utf8Error), + InvalidHeader(String), +} + +impl Error for ParseHeaderError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + ParseHeaderError::Utf8Parse(err) => Some(err), + ParseHeaderError::MagicString | + ParseHeaderError::Version { .. } | + ParseHeaderError::HeaderLengthOverflow(_) | + ParseHeaderError::NonAscii | + ParseHeaderError::InvalidHeader(_) => None, + } + } +} + +impl std::fmt::Display for ParseHeaderError { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + ParseHeaderError::MagicString => write!(f, "start does not match magic string"), + ParseHeaderError::Version { major, minor } => write!(f, "unknown version number: {}.{}", major, minor), + ParseHeaderError::HeaderLengthOverflow(header_len) => write!(f, "HEADER_LEN {} does not fit in `usize`", header_len), + ParseHeaderError::NonAscii => write!(f, "non-ascii in array format string; this is not supported in .npy format versions 1.0 and 2.0"), + ParseHeaderError::Utf8Parse(err) => write!(f, "error parsing array format string as UTF-8: {}", err), + ParseHeaderError::InvalidHeader(value) => write!(f, "invalid header in file: {}", value), + } + } +} + +impl From for ParseHeaderError { + fn from(err: std::str::Utf8Error) -> ParseHeaderError { + ParseHeaderError::Utf8Parse(err) + } +} + +impl From for ParseHeaderError { + fn from(e: std::num::ParseIntError) -> Self { + ParseHeaderError::InvalidHeader(format!("failed to parse an integer: {}", e)) + } +} + +#[derive(Debug)] +pub enum ReadHeaderError { + Io(std::io::Error), + Parse(ParseHeaderError), +} + +impl Error for ReadHeaderError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + ReadHeaderError::Io(err) => Some(err), + ReadHeaderError::Parse(err) => Some(err), + } + } +} + +impl std::fmt::Display for ReadHeaderError { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + ReadHeaderError::Io(err) => write!(f, "I/O error: {}", err), + ReadHeaderError::Parse(err) => write!(f, "error parsing header: {}", err), + } + } +} + +impl From for ReadHeaderError { + fn from(err: std::io::Error) -> ReadHeaderError { + ReadHeaderError::Io(err) + } +} + +impl From for ReadHeaderError { + fn from(err: ParseHeaderError) -> ReadHeaderError { + ReadHeaderError::Parse(err) + } +} + +#[derive(Clone, Copy)] +#[allow(non_camel_case_types)] +enum Version { + V1_0, + V2_0, + V3_0, +} + +impl Version { + /// Number of bytes taken up by version number (1 byte for major version, 1 + /// byte for minor version). + const VERSION_NUM_BYTES: usize = 2; + + fn from_bytes(bytes: &[u8]) -> Result { + debug_assert_eq!(bytes.len(), Self::VERSION_NUM_BYTES); + match (bytes[0], bytes[1]) { + (0x01, 0x00) => Ok(Version::V1_0), + (0x02, 0x00) => Ok(Version::V2_0), + (0x03, 0x00) => Ok(Version::V3_0), + (major, minor) => Err(ParseHeaderError::Version { major, minor }), + } + } + + /// Major version number. + fn major_version(self) -> u8 { + match self { + Version::V1_0 => 1, + Version::V2_0 => 2, + Version::V3_0 => 3, + } + } + + /// Major version number. + fn minor_version(self) -> u8 { + match self { + Version::V1_0 | Version::V2_0 | Version::V3_0 => 0, + } + } + + /// Number of bytes in representation of header length. + fn header_len_num_bytes(self) -> usize { + match self { + Version::V1_0 => 2, + Version::V2_0 | Version::V3_0 => 4, + } + } + + /// Read header length. + fn read_header_len(self, reader: &mut R) -> Result { + match self { + Version::V1_0 => Ok(usize::from(reader.read_u16::()?)), + Version::V2_0 | Version::V3_0 => { + let header_len: u32 = reader.read_u32::()?; + Ok(usize::try_from(header_len) + .map_err(|_| ParseHeaderError::HeaderLengthOverflow(header_len))?) + } + } + } + + /// Format header length as bytes for writing to file. + /// + /// Returns `None` if the value of `header_len` is too large for this .npy version. + fn format_header_len(self, header_len: usize) -> Option> { + match self { + Version::V1_0 => { + let header_len: u16 = u16::try_from(header_len).ok()?; + let mut out = vec![0; self.header_len_num_bytes()]; + LittleEndian::write_u16(&mut out, header_len); + Some(out) + } + Version::V2_0 | Version::V3_0 => { + let header_len: u32 = u32::try_from(header_len).ok()?; + let mut out = vec![0; self.header_len_num_bytes()]; + LittleEndian::write_u32(&mut out, header_len); + Some(out) + } + } + } + + /// Computes the total header length, formatted `HEADER_LEN` value, and + /// padding length for this .npy version. + /// + /// `unpadded_arr_format` is the Python literal describing the array + /// format, formatted as an ASCII string without any padding. + /// + /// Returns `None` if the total header length overflows `usize` or if the + /// value of `HEADER_LEN` is too large for this .npy version. + fn compute_lengths(self, unpadded_arr_format: &[u8]) -> Option { + /// Length of a '\n' char in bytes. + const NEWLINE_LEN: usize = 1; + + let prefix_len: usize = + MAGIC_STRING.len() + Version::VERSION_NUM_BYTES + self.header_len_num_bytes(); + let unpadded_total_len: usize = prefix_len + .checked_add(unpadded_arr_format.len())? + .checked_add(NEWLINE_LEN)?; + let padding_len: usize = HEADER_DIVISOR - unpadded_total_len % HEADER_DIVISOR; + let total_len: usize = unpadded_total_len.checked_add(padding_len)?; + let header_len: usize = total_len - prefix_len; + let formatted_header_len = self.format_header_len(header_len)?; + Some(HeaderLengthInfo { + total_len, + formatted_header_len, + }) + } +} + +struct HeaderLengthInfo { + /// Total header length (including magic string, version number, header + /// length value, array format description, padding, and final newline). + total_len: usize, + /// Formatted `HEADER_LEN` value. (This is the number of bytes in the array + /// format description, padding, and final newline.) + formatted_header_len: Vec, +} + +#[derive(Debug)] +pub enum WriteHeaderError { + Io(std::io::Error), + Format(String), +} + +impl Error for WriteHeaderError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + WriteHeaderError::Io(err) => Some(err), + WriteHeaderError::Format(_) => None, + } + } +} + +impl std::fmt::Display for WriteHeaderError { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + WriteHeaderError::Io(err) => write!(f, "I/O error: {}", err), + WriteHeaderError::Format(err) => write!(f, "error formatting header: {}", err), + } + } +} + +impl From for WriteHeaderError { + fn from(err: std::io::Error) -> WriteHeaderError { + WriteHeaderError::Io(err) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DataType { + Scalar(String), + Compound(Vec<(String, String)>), +} + +impl std::fmt::Display for DataType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DataType::Scalar(v) => write!(f, "'{}'", v), + DataType::Compound(list) => { + write!(f, "[")?; + for (k, v) in list { + write!(f, "('{}', '{}'), ", k, v)?; + } + write!(f, "]") + } + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Header { + pub type_descriptor: DataType, + pub fortran_order: bool, + pub shape: Vec, +} + +#[derive(Clone, Debug)] +struct HeaderParser { + data: Vec, + position: usize, +} + +impl HeaderParser { + fn done(&self) -> bool { + return self.position >= self.data.len(); + } + + fn current(&self) -> char { + return self.data[self.position]; + } + + fn advance(&mut self) -> char { + let value = self.current(); + self.position += 1; + return value; + } + + fn expects(&mut self, c: char) -> Result<(), ParseHeaderError> { + if self.current() == c { + self.advance(); + return Ok(()); + } else { + return Err(ParseHeaderError::InvalidHeader(format!( + "expected '{}', got '{}'", c, self.current() + ))); + } + } + + fn skip_whitespaces(&mut self) { + let mut c = self.current(); + while !self.done() && (c == ' ' || c == '\t' || c == '\x0C') { + self.advance(); + c = self.current(); + } + } + + fn parse_string(&mut self) -> Result { + let mut value = String::new(); + if self.current() == '\'' { + self.advance(); + while self.current() != '\'' { + value.push(self.advance()); + } + self.advance(); + + } else if self.current() == '"' { + self.advance(); + while self.current() != '"' { + value.push(self.advance()); + } + self.advance(); + } else { + return Err(ParseHeaderError::InvalidHeader(format!( + "expected a string, got '{}'", self.current() + ))); + } + + return Ok(value); + } + + fn parse_integer(&mut self) -> Result { + let mut value = String::new(); + loop { + if self.current().is_ascii_digit() { + value.push(self.advance()); + } else { + break; + } + } + + if value.is_empty() { + return Err(ParseHeaderError::InvalidHeader(format!( + "expected an integer, got '{}'", self.current() + ))); + } + + return Ok(value.parse()?); + } + + fn parse_data_type(&mut self) -> Result { + if self.current() == '\'' || self.current() == '"' { + let value = self.parse_string()?; + return Ok(DataType::Scalar(value)); + } else if self.current() == '[' { + self.advance(); + + let mut data_type = Vec::new(); + loop { + self.skip_whitespaces(); + self.expects('(')?; + self.skip_whitespaces(); + + let name = self.parse_string()?; + + self.skip_whitespaces(); + self.expects(',')?; + self.skip_whitespaces(); + + let value = self.parse_string()?; + + self.skip_whitespaces(); + self.expects(')')?; + self.skip_whitespaces(); + + data_type.push((name, value)); + + if self.current() == ',' { + self.advance(); + self.skip_whitespaces(); + } else { + self.expects(']')?; + break; + } + + if self.current() == ']' { + self.advance(); + break; + } + } + + return Ok(DataType::Compound(data_type)); + } else { + return Err(ParseHeaderError::InvalidHeader(format!( + "expected a string or a list, got '{}'", self.current() + ))); + } + } + + fn parse_bool(&mut self) -> Result { + if self.current() == 'T' { + self.advance(); + self.expects('r')?; + self.expects('u')?; + self.expects('e')?; + return Ok(true); + } else if self.current() == 'F' { + self.advance(); + self.expects('a')?; + self.expects('l')?; + self.expects('s')?; + self.expects('e')?; + return Ok(false); + } else { + return Err(ParseHeaderError::InvalidHeader(format!( + "expected a bool, got '{}'", self.current() + ))); + } + } + + fn parse_shape(&mut self) -> Result, ParseHeaderError> { + let mut shape = Vec::new(); + self.expects('(')?; + loop { + self.skip_whitespaces(); + shape.push(self.parse_integer()?); + self.skip_whitespaces(); + + + if self.current() == ',' { + self.advance(); + self.skip_whitespaces(); + } else { + self.expects(')')?; + break; + } + + if self.current() == ')' { + self.advance(); + break; + } + } + + return Ok(shape); + } + + fn parse(&mut self) -> Result { + let mut type_descriptor: Option = None; + let mut fortran_order: Option = None; + let mut shape: Option> = None; + + self.skip_whitespaces(); + self.expects('{')?; + self.skip_whitespaces(); + + loop { + let key = self.parse_string()?; + self.skip_whitespaces(); + self.expects(':')?; + self.skip_whitespaces(); + + if key == "descr" { + type_descriptor = Some(self.parse_data_type()?); + } else if key == "fortran_order" { + fortran_order = Some(self.parse_bool()?); + } else if key == "shape" { + shape = Some(self.parse_shape()?); + } else { + return Err(ParseHeaderError::InvalidHeader(format!( + "unknown key: '{}'", key + ))); + } + + self.skip_whitespaces(); + if self.current() == ',' { + self.advance(); + self.skip_whitespaces(); + } else { + self.expects('}')?; + break; + } + + if self.current() == '}' { + self.advance(); + break; + } + } + + match (type_descriptor, fortran_order, shape) { + (Some(type_descriptor), Some(fortran_order), Some(shape)) => Ok(Header { + type_descriptor, + fortran_order, + shape, + }), + (None, _, _) => Err(ParseHeaderError::InvalidHeader("missing 'descr' key".into())), + (_, None, _) => Err(ParseHeaderError::InvalidHeader("missing 'fortran_order' key".into())), + (_, _, None) => Err(ParseHeaderError::InvalidHeader("missing 'shape' key".into())), + } + } +} + +impl Header { + fn from_str(value: &str) -> Result { + let mut parser = HeaderParser { data: value.chars().collect(), position: 0 }; + return parser.parse(); + } + + pub fn from_reader(reader: &mut R) -> Result { + // Check for magic string. + let mut buf = vec![0; MAGIC_STRING.len()]; + reader.read_exact(&mut buf)?; + if buf != MAGIC_STRING { + return Err(ParseHeaderError::MagicString.into()); + } + + // Get version number. + let mut buf = [0; Version::VERSION_NUM_BYTES]; + reader.read_exact(&mut buf)?; + let version = Version::from_bytes(&buf)?; + + // Get `HEADER_LEN`. + let header_len = version.read_header_len(reader)?; + + // Parse the dictionary describing the array's format. + let mut buf = vec![0; header_len]; + reader.read_exact(&mut buf)?; + let without_newline = match buf.split_last() { + Some((&b'\n', rest)) => rest, + Some(_) | None => return Err(ParseHeaderError::InvalidHeader("missing new line".into()))?, + }; + let header_str = match version { + Version::V1_0 | Version::V2_0 => { + if without_newline.is_ascii() { + // ASCII strings are always valid UTF-8. + unsafe { std::str::from_utf8_unchecked(without_newline) } + } else { + return Err(ParseHeaderError::NonAscii.into()); + } + } + Version::V3_0 => { + std::str::from_utf8(without_newline).map_err(ParseHeaderError::from)? + } + }; + + Ok(Header::from_str(header_str)?) + } + + fn to_dict_literal(&self) -> String { + let mut result = String::new(); + write!(&mut result, "{{ 'descr': {}, ", self.type_descriptor).expect("failed to write"); + + let order = if self.fortran_order { + "True" + } else { + "False" + }; + write!(&mut result, "'fortran_order': {}, ", order).expect("failed to write"); + + write!(&mut result, "'shape': (").expect("failed to write"); + for s in &self.shape { + write!(&mut result, "{}, ", s).expect("failed to write"); + } + write!(&mut result, ") }}").expect("failed to write"); + return result; + } + + pub fn to_bytes(&self) -> Result, WriteHeaderError> { + // Metadata describing array's format as ASCII string. + let mut arr_format = Vec::new(); + + write!(&mut arr_format, "{}", self.to_dict_literal())?; + + // Determine appropriate version based on header length, and compute + // length information. + let (version, length_info) = [Version::V1_0, Version::V2_0] + .iter() + .find_map(|&version| Some((version, version.compute_lengths(&arr_format)?))) + .ok_or_else(|| WriteHeaderError::Format("header too long".into()))?; + + // Write the header. + let mut out = Vec::with_capacity(length_info.total_len); + out.extend_from_slice(MAGIC_STRING); + out.push(version.major_version()); + out.push(version.minor_version()); + out.extend_from_slice(&length_info.formatted_header_len); + out.extend_from_slice(&arr_format); + out.resize(length_info.total_len - 1, b' '); + out.push(b'\n'); + + // Verify the length of the header. + debug_assert_eq!(out.len(), length_info.total_len); + debug_assert_eq!(out.len() % HEADER_DIVISOR, 0); + + Ok(out) + } + + pub fn write(&self, mut writer: W) -> Result<(), WriteHeaderError> { + let bytes = self.to_bytes()?; + writer.write_all(&bytes)?; + Ok(()) + } +} + +/******************************************************************************/ + +impl From for crate::Error { + fn from(error: ReadHeaderError) -> Self { + match error { + ReadHeaderError::Io(e) => crate::Error::Io(Arc::new(e)), + ReadHeaderError::Parse(e) => crate::Error::Serialization(e.to_string()), + } + } +} + +impl From for crate::Error { + fn from(error: WriteHeaderError) -> Self { + match error { + WriteHeaderError::Io(e) => crate::Error::Io(Arc::new(e)), + WriteHeaderError::Format(e) => crate::Error::Serialization(e), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn npy_header_parsing() { + let header = " \t{'descr': [('a', '(reader: R, create_array: mts_create_array_callback_t) -> Result, Error> + where R: std::io::Read + std::io::Seek +{ + let mut archive = ZipArchive::new(reader).map_err(|e| ("", e))?; + + let mut length_unit = String::new(); + if let Some(index) = archive.index_for_name("info.json") { + let mut info_file = archive.by_index(index).map_err(|e| ("info.json", e))?; + let mut info_content = String::new(); + info_file.read_to_string(&mut info_content)?; + let info: JsonValue = json::parse(&info_content)?; + + + if info["format"].as_str() != Some("metatomic_system") { + return Err(Error::Serialization(format!( + "invalid format in info.json, expected 'metatomic_system', found {:?}", + info["format"] + ))); + } + + if info["version"].as_u8() != Some(1) { + return Err(Error::Serialization(format!( + "unsupported version in info.json, expected 1, found {:?}", + info["version"] + ))); + } + + if !info.has_key("length_unit") || !info["length_unit"].is_string() { + return Err(Error::Serialization( + "missing or invalid 'length_unit' field in info.json".into() + )); + } + length_unit = info["length_unit"].as_str().unwrap().to_string(); + } else { + // this is a legacy file from metatomic-torch + } + + let data_file = archive.by_name("types.npy").map_err(|e| ("types.npy", e))?; + let types = read_tensor(data_file, create_array)?; + + let data_file = archive.by_name("positions.npy").map_err(|e| ("positions.npy", e))?; + let position = read_tensor(data_file, create_array)?; + + let data_file = archive.by_name("cell.npy").map_err(|e| ("cell.npy", e))?; + let cell = read_tensor(data_file, create_array)?; + + let data_file = archive.by_name("pbc.npy").map_err(|e| ("pbc.npy", e))?; + let pbc = read_tensor(data_file, create_array)?; + + let mut system = Arc::new(System::new(length_unit, types, position, cell, pbc)?); + + let pairs_paths: Vec = archive.file_names() + .filter(|path| path.starts_with("pairs/") && path.ends_with("/options.json")) + .map(|path| path.to_string()) + .collect(); + + let mut buffer = Vec::new(); + for path in pairs_paths { + let options: PairListOptions = { + let mut options_file = archive.by_name(&path).map_err(|e| (&path, e))?; + let mut options_content = String::new(); + options_file.read_to_string(&mut options_content)?; + let options_json: &JsonValue = &json::parse(&options_content)?; + + options_json.try_into()? + }; + + let data_path = path.strip_suffix("/options.json").unwrap().to_string() + "/data.mts"; + let mut data_file = archive.by_name(&data_path).map_err(|e| (data_path, e))?; + + buffer.clear(); + data_file.read_to_end(&mut buffer)?; + + let pairs = metatensor::io::load_block_buffer_custom_array(&buffer, create_array)?; + + system.add_pairs(options, pairs)?; + } + + let data_paths: Vec = archive.file_names() + .filter(|path| path.starts_with("data/")) + .map(|path| path.to_string()) + .collect(); + + for path in data_paths { + let name = path.strip_prefix("data/").expect("data path should start with 'data/'") + .strip_suffix(".mts").expect("data path should end with '.mts'").to_string(); + + let mut data_file = archive.by_name(&path).map_err(|e| (&path, e))?; + + buffer.clear(); + data_file.read_to_end(&mut buffer)?; + + let data = metatensor::io::load_buffer_custom_array(&buffer, create_array)?; + + system.add_custom_data(name, data, /*override*/ true)?; + } + + return Ok(system); +} + +/// Save the given system to a file (or any other writer). +/// +/// The format consists of a zip archive containing NPY files for the system's +/// data (types, positions, cell, pbc), a `info.json` file for metadata, and +/// optional sub-directories for pair lists (`pairs//options.json` and +/// `pairs//data.mts`) and custom data (`data/.mts`). +/// +/// The recommended file extension is `.mta`. +pub fn save(writer: W, system: &System) -> Result<(), Error> { + let mut archive = ZipWriter::new(writer); + + let options = zip::write::FileOptions::<'_, ()>::default() + .with_alignment(16) + .compression_method(zip::CompressionMethod::Stored) + .large_file(true) + .last_modified_time(zip::DateTime::from_date_and_time(2000, 1, 1, 0, 0, 0).expect("invalid datetime")); + + archive.start_file("info.json", options).map_err(|e| ("info.json", e))?; + let info = json::object! { + "format": "metatomic_system", + "version": 1, + "length_unit": system.length_unit(), + }; + info.write(&mut archive)?; + + archive.start_file("types.npy", options).map_err(|e| ("types.npy", e))?; + write_tensor(&mut archive, system.types())?; + + archive.start_file("positions.npy", options).map_err(|e| ("positions.npy", e))?; + write_tensor(&mut archive, system.positions())?; + + archive.start_file("cell.npy", options).map_err(|e| ("cell.npy", e))?; + write_tensor(&mut archive, system.cell())?; + + archive.start_file("pbc.npy", options).map_err(|e| ("pbc.npy", e))?; + write_tensor(&mut archive, system.pbc())?; + + let mut buffer = Vec::new(); + for (i, &pairs_options) in system.known_pairs().iter().enumerate() { + let path = format!("pairs/{}/options.json", i); + archive.start_file(&path, options).map_err(|e| (path, e))?; + let json: JsonValue = pairs_options.clone().into(); + json.write(&mut archive)?; + + + let pairs_block = system.get_pairs(pairs_options).expect("pairs block should exist"); + buffer.clear(); + pairs_block.save_buffer(&mut buffer)?; + + let path = format!("pairs/{}/data.mts", i); + archive.start_file(&path, options).map_err(|e| (path, e))?; + archive.write_all(&buffer)?; + } + + for name in system.known_custom_data() { + let tensor = system.get_custom_data(name).expect("custom data should exist"); + buffer.clear(); + tensor.save_buffer(&mut buffer)?; + let path = format!("data/{}.mts", name); + archive.start_file(&path, options).map_err(|e| (path, e))?; + archive.write_all(&buffer)?; + } + + archive.finish().map_err(|e| ("", e))?; + + return Ok(()); +} + + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn load_legacy() { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/data/legacy.mta"); + + let file = std::fs::File::open(&path).unwrap(); + let system = load(file, Some(metatensor::io::create_ndarray)).unwrap(); + + assert_eq!(system.length_unit(), ""); + + let types: ndarray::ArrayView1 = system.types().try_into().unwrap(); + let positions: ndarray::ArrayView2 = system.positions().try_into().unwrap(); + let cell: ndarray::ArrayView2 = system.cell().try_into().unwrap(); + let pbc: ndarray::ArrayView1 = system.pbc().try_into().unwrap(); + + assert_eq!(types, ndarray::arr1(&[1, 6, 7, 8])); + assert_eq!( + positions, + ndarray::arr2(&[[0.0, 0.0, 0.0], [1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + ); + assert_eq!( + cell, + ndarray::arr2(&[[6.0, 0.0, 0.0], [0.0, 4.3, 0.0], [0.0, 0.0, 0.0]]) + ); + assert_eq!(pbc, ndarray::arr1(&[true, true, false])); + + let options = PairListOptions { cutoff: 5.5, full_list: true, strict: true, requestors: vec![] }; + let pairs = system.get_pairs(&options).unwrap(); + assert_eq!(pairs.samples().names(), ["first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"]); + assert_eq!(pairs.samples().count(), 28); + assert_eq!(pairs.values().shape().unwrap(), [28, 3, 1]); + + let options = system.known_pairs(); + assert_eq!(options.len(), 1); + // requestors are not used when looking up pairs, but are stored in the file + assert_eq!(options[0].requestors, ["some requestor", "another one with UTF8 Θµ"]); + + assert_eq!(system.known_custom_data(), vec!["custom::data"]); + let custom = system.get_custom_data("custom::data").unwrap(); + assert_eq!(custom.keys().count(), 2); + } + + #[test] + fn save_load_system() { + let system = crate::system::test_system("f32"); + + let path = std::env::temp_dir().join(format!("system-{}.mta", std::process::id())); + { + let file = std::fs::File::create(&path).unwrap(); + save(file, &system).unwrap(); + } + + { + let file = std::fs::File::open(&path).unwrap(); + let mut archive = zip::ZipArchive::new(file).unwrap(); + assert!(archive.by_name("types.npy").is_ok()); + assert!(archive.by_name("positions.npy").is_ok()); + assert!(archive.by_name("cell.npy").is_ok()); + assert!(archive.by_name("pbc.npy").is_ok()); + assert!(archive.by_name("pairs/0/data.mts").is_ok()); + assert!(archive.by_name("data/custom::data/name.mts").is_ok()); + + let options_file = archive.by_name("pairs/0/options.json").unwrap(); + let options_json = std::io::read_to_string(options_file).unwrap(); + let options_json: JsonValue = json::parse(&options_json).unwrap(); + + assert_eq!(options_json["type"].as_str(), Some("metatomic_pair_list_options")); + assert_eq!(options_json["cutoff"].as_str(), Some(&*format!("0x{:x}", 3.5_f64.to_bits()))); + assert_eq!(options_json["full_list"].as_bool(), Some(true)); + assert_eq!(options_json["strict"].as_bool(), Some(false)); + } + + let loaded = { + let file = std::fs::File::open(&path).unwrap(); + let loaded = load(file, Some(metatensor::io::create_ndarray)).unwrap(); + std::fs::remove_file(&path).unwrap(); + loaded + }; + + assert_eq!(loaded.length_unit(), "Angstrom"); + + let types: ndarray::ArrayView1 = loaded.types().try_into().unwrap(); + let positions: ndarray::ArrayView2 = loaded.positions().try_into().unwrap(); + let cell: ndarray::ArrayView2 = loaded.cell().try_into().unwrap(); + let pbc: ndarray::ArrayView1 = loaded.pbc().try_into().unwrap(); + + assert_eq!(types, ndarray::arr1(&[1, 6, 8])); + assert_eq!( + positions, + ndarray::arr2(&[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]) + ); + assert_eq!( + cell, + ndarray::arr2(&[[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) + ); + assert_eq!(pbc, ndarray::arr1(&[true, true, true])); + + let options = PairListOptions { + cutoff: 3.5, + full_list: true, + strict: false, + requestors: vec![], + }; + assert!(loaded.get_pairs(&options).is_some()); + assert!(loaded.get_custom_data("custom::data/name").is_ok()); + } +} diff --git a/metatomic-core/src/io/tensor.rs b/metatomic-core/src/io/tensor.rs new file mode 100644 index 000000000..81c8987cd --- /dev/null +++ b/metatomic-core/src/io/tensor.rs @@ -0,0 +1,333 @@ +use byteorder::{BigEndian, LittleEndian, NativeEndian, ReadBytesExt, WriteBytesExt}; + +use dlpk::{DLDataType, DLDataTypeCode, DLDevice, DLPackTensor, DLPackTensorRef, DLPackVersion}; +use metatensor::MtsArray; +use metatensor::c_api::{MTS_SUCCESS, mts_array_t, mts_create_array_callback_t}; + +use crate::Error; + +use super::{Endianness, check_for_extra_bytes}; +use super::npy_header::{Header, DataType}; + +/// Parse an NPY type descriptor string (e.g. `" Result<(DLDataTypeCode, u8, Endianness), Error> { + if descr.len() < 3 { + return Err(Error::Serialization(format!("invalid type descriptor: {}", descr))); + } + + let endian = match &descr[0..1] { + "<" => Endianness::Little, + "=" | "|" => Endianness::Native, + ">" => Endianness::Big, + // not applicable for single-byte types + _ => return Err(Error::Serialization(format!("unknown endianness in type descriptor: {}", descr))), + }; + + let type_char = &descr[1..2]; + let size_str = &descr[2..]; + let size: u8 = size_str.parse().map_err(|_| { + Error::Serialization(format!("invalid size in type descriptor: {}", descr)) + })?; + + let (code, bits) = match (type_char, size) { + ("f", 4) => (DLDataTypeCode::kDLFloat, 32), + ("f", 8) => (DLDataTypeCode::kDLFloat, 64), + ("i", 1) => (DLDataTypeCode::kDLInt, 8), + ("i", 2) => (DLDataTypeCode::kDLInt, 16), + ("i", 4) => (DLDataTypeCode::kDLInt, 32), + ("i", 8) => (DLDataTypeCode::kDLInt, 64), + ("u", 1) => (DLDataTypeCode::kDLUInt, 8), + ("u", 2) => (DLDataTypeCode::kDLUInt, 16), + ("u", 4) => (DLDataTypeCode::kDLUInt, 32), + ("u", 8) => (DLDataTypeCode::kDLUInt, 64), + ("b", 1) => (DLDataTypeCode::kDLBool, 8), + ("c", 8) => (DLDataTypeCode::kDLComplex, 64), + ("c", 16) => (DLDataTypeCode::kDLComplex, 128), + ("f", 2) => (DLDataTypeCode::kDLFloat, 16), + _ => return Err(Error::Serialization(format!("unsupported type descriptor: {}", descr))), + }; + + Ok((code, bits, endian)) +} + + +fn read_as(reader: &mut R, tensor: dlpk::DLPackTensorRefMut<'_>, cb: impl Fn(&mut R, &mut T) -> Result<(), std::io::Error>) -> Result<(), Error> +where R: std::io::Read, + T: dlpk::DLPackPointerCast + 'static +{ + let mut view: ndarray::ArrayViewMutD = tensor.try_into() + .map_err(|e| Error::Serialization(format!("failed to convert DLPack to ndarray mutable view: {}", e)))?; + + for value in &mut view { + cb(reader, value)?; + } + + Ok(()) +} + +// Read a data array from the given reader, using numpy's NPY format +#[allow(clippy::too_many_lines)] +pub fn read_tensor(mut reader: R, create_array: mts_create_array_callback_t) -> Result + where R: std::io::Read +{ + let create_array = create_array.ok_or_else(|| Error::InvalidParameter("create_array callback is NULL".into()))?; + let header = super::npy_header::Header::from_reader(&mut reader)?; + if header.fortran_order { + return Err(Error::Serialization("data can not be loaded from fortran-order arrays".into())); + } + + let descr = if let super::npy_header::DataType::Scalar(s) = &header.type_descriptor { + s.as_str() + } else { + return Err(Error::Serialization("structured arrays are not supported".into())); + }; + + let (file_code, file_bits, endian) = npy_descr_to_dtype(descr)?; + + let dl_dtype = DLDataType { code: file_code, bits: file_bits, lanes: 1 }; + + let shape = header.shape; + let mut array = mts_array_t::null(); + let status = unsafe { + create_array(shape.as_ptr(), shape.len(), dl_dtype, &mut array) + }; + + let array = if status == MTS_SUCCESS { + MtsArray::from_raw(array) + } else { + // TODO: how can we propagate the error from the callback? + return Err(Error::Serialization("failed to create array".into())); + }; + + let device = DLDevice::cpu(); + let version = DLPackVersion::current(); + let mut dl_tensor = array.as_dlpack(device, None, version)?; + + let num_elements: usize = shape.iter().product(); + if num_elements == 0 { + check_for_extra_bytes(&mut reader)?; + return Ok(dl_tensor); + } + + let tensor = dl_tensor.as_mut(); + + // Endianness is handled inside each arm to avoid tripling the number of + // match arms (which inflates uncovered-line counts for big/native paths + // that are not exercised in tests on little-endian CI). + match (file_code, file_bits) { + // Standard Floats + (DLDataTypeCode::kDLFloat, 32) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => r.read_f32::()?, + Endianness::Big => r.read_f32::()?, + Endianness::Native => r.read_f32::()?, + }; + Ok(()) + }), + (DLDataTypeCode::kDLFloat, 64) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => r.read_f64::()?, + Endianness::Big => r.read_f64::()?, + Endianness::Native => r.read_f64::()?, + }; + Ok(()) + }), + + // Standard Ints + (DLDataTypeCode::kDLInt, 8) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = r.read_i8()?; + Ok(()) + }), + (DLDataTypeCode::kDLInt, 16) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => r.read_i16::()?, + Endianness::Big => r.read_i16::()?, + Endianness::Native => r.read_i16::()?, + }; + Ok(()) + }), + (DLDataTypeCode::kDLInt, 32) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => r.read_i32::()?, + Endianness::Big => r.read_i32::()?, + Endianness::Native => r.read_i32::()?, + }; + Ok(()) + }), + (DLDataTypeCode::kDLInt, 64) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => r.read_i64::()?, + Endianness::Big => r.read_i64::()?, + Endianness::Native => r.read_i64::()?, + }; + Ok(()) + }), + + // Unsigned Ints + (DLDataTypeCode::kDLUInt, 8) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = r.read_u8()?; + Ok(()) + }), + (DLDataTypeCode::kDLUInt, 16) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => r.read_u16::()?, + Endianness::Big => r.read_u16::()?, + Endianness::Native => r.read_u16::()?, + }; + Ok(()) + }), + (DLDataTypeCode::kDLUInt, 32) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => r.read_u32::()?, + Endianness::Big => r.read_u32::()?, + Endianness::Native => r.read_u32::()?, + }; + Ok(()) + }), + (DLDataTypeCode::kDLUInt, 64) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => r.read_u64::()?, + Endianness::Big => r.read_u64::()?, + Endianness::Native => r.read_u64::()?, + }; + Ok(()) + }), + + // Boolean (Read as u8) + (DLDataTypeCode::kDLBool, 8) => read_as::(&mut reader, tensor, |r: &mut R, v| { + *v = r.read_u8()? != 0; + Ok(()) + }), + + // Complex Numbers (Read as array of 2 floats) + (DLDataTypeCode::kDLComplex, 64) => read_as::<[f32; 2], _>(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => [r.read_f32::()?, r.read_f32::()?], + Endianness::Big => [r.read_f32::()?, r.read_f32::()?], + Endianness::Native => [r.read_f32::()?, r.read_f32::()?], + }; + Ok(()) + }), + (DLDataTypeCode::kDLComplex, 128) => read_as::<[f64; 2], _>(&mut reader, tensor, |r: &mut R, v| { + *v = match endian { + Endianness::Little => [r.read_f64::()?, r.read_f64::()?], + Endianness::Big => [r.read_f64::()?, r.read_f64::()?], + Endianness::Native => [r.read_f64::()?, r.read_f64::()?], + }; + Ok(()) + }), + + _ => Err(Error::Serialization(format!( + "unsupported dtype for reading: {:?} {} bits", file_code, file_bits + ))), + }?; + + check_for_extra_bytes(&mut reader)?; + Ok(dl_tensor) +} + +fn dlpack_to_npy_descr(code: DLDataTypeCode, bits: u8) -> Result { + let endian = if cfg!(target_endian = "little") { "<" } else { ">" }; + + let (type_char, type_size) = match (code, bits) { + (DLDataTypeCode::kDLInt, 8) => ("i", 1), + (DLDataTypeCode::kDLInt, 16) => ("i", 2), + (DLDataTypeCode::kDLInt, 32) => ("i", 4), + (DLDataTypeCode::kDLInt, 64) => ("i", 8), + (DLDataTypeCode::kDLUInt, 8) => ("u", 1), + (DLDataTypeCode::kDLUInt, 16) => ("u", 2), + (DLDataTypeCode::kDLUInt, 32) => ("u", 4), + (DLDataTypeCode::kDLUInt, 64) => ("u", 8), + (DLDataTypeCode::kDLFloat, 32) => ("f", 4), + (DLDataTypeCode::kDLFloat, 64) => ("f", 8), + (DLDataTypeCode::kDLBool, 8) => ("b", 1), + (DLDataTypeCode::kDLComplex, 64) => ("c", 8), + (DLDataTypeCode::kDLComplex, 128) => ("c", 16), + (DLDataTypeCode::kDLFloat, 16) => ("f", 2), + _ => return Err(Error::Serialization( + format!("unsupported DLPack dtype: code {:?}, bits {:?}", code, bits) + ) + ), + }; + + Ok(format!("{}{}{}", endian, type_char, type_size)) +} + + +fn write_as(writer: &mut W, tensor: dlpk::DLPackTensorRef<'_>, cb: impl Fn(&mut W, T) -> Result<(), std::io::Error>) -> Result<(), Error> +where W: std::io::Write, + T: Copy + dlpk::DLPackPointerCast + 'static +{ + let view: ndarray::ArrayViewD = tensor.try_into() + .map_err(|e| Error::Serialization(format!("failed to convert DLPack to ndarray view: {}", e)))?; + + for &value in &view { + cb(writer, value)?; + } + + Ok(()) +} + +// Write an array to the given writer, using numpy's NPY format +#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] +pub fn write_tensor(writer: &mut W, tensor: DLPackTensorRef<'_>) -> Result<(), Error> { + let dtype = tensor.raw.dtype; + let (code, bits) = (dtype.code, dtype.bits); + + // Validate Lanes + if dtype.lanes != 1 { + return Err(Error::Serialization(format!( + "unsupported DLPack dtype: lanes != 1 ({})", dtype.lanes + ))); + } + + // Write Header + let tdesc = dlpack_to_npy_descr(code, bits)?; + let header = Header { + type_descriptor: DataType::Scalar(tdesc), + fortran_order: false, + shape: tensor.shape().iter().map(|&s| s as usize).collect(), + }; + + header.write(&mut *writer)?; + + // Get metadata for size and pointer for data + let num_elements: usize = header.shape.iter().product(); + if num_elements == 0 { + return Ok(()); + } + + match (code, bits) { + // Standard Floats + (DLDataTypeCode::kDLFloat, 32) => write_as::(writer, tensor, |w: &mut W, v| w.write_f32::(v)), + (DLDataTypeCode::kDLFloat, 64) => write_as::(writer, tensor, |w: &mut W, v| w.write_f64::(v)), + + // Standard Ints + (DLDataTypeCode::kDLInt, 8) => write_as::(writer, tensor, |w: &mut W, v| w.write_i8(v)), + (DLDataTypeCode::kDLInt, 16) => write_as::(writer, tensor, |w: &mut W, v| w.write_i16::(v)), + (DLDataTypeCode::kDLInt, 32) => write_as::(writer, tensor, |w: &mut W, v| w.write_i32::(v)), + (DLDataTypeCode::kDLInt, 64) => write_as::(writer, tensor, |w: &mut W, v| w.write_i64::(v)), + + // Unsigned Ints + (DLDataTypeCode::kDLUInt, 8) => write_as::(writer, tensor, |w: &mut W, v| w.write_u8(v)), + (DLDataTypeCode::kDLUInt, 16) => write_as::(writer, tensor, |w: &mut W, v| w.write_u16::(v)), + (DLDataTypeCode::kDLUInt, 32) => write_as::(writer, tensor, |w: &mut W, v| w.write_u32::(v)), + (DLDataTypeCode::kDLUInt, 64) => write_as::(writer, tensor, |w: &mut W, v| w.write_u64::(v)), + + // Boolean, stored as u8 + (DLDataTypeCode::kDLBool, 8) => write_as::(writer, tensor, |w: &mut W, v| w.write_u8(u8::from(v))), + + // Complex Numbers + (DLDataTypeCode::kDLComplex, 64) => write_as::<[f32; 2], _>(writer, tensor, |w: &mut W, v| { + w.write_f32::(v[0])?; + w.write_f32::(v[1]) + }), + (DLDataTypeCode::kDLComplex, 128) => write_as::<[f64; 2], _>(writer, tensor, |w: &mut W, v| { + w.write_f64::(v[0])?; + w.write_f64::(v[1]) + }), + + _ => Err(Error::Serialization(format!("unsupported dtype for writing: {:?} {} bits", code, bits))), + } +} diff --git a/metatomic-core/src/kernels/cpu.rs b/metatomic-core/src/kernels/cpu.rs new file mode 100644 index 000000000..39580f34d --- /dev/null +++ b/metatomic-core/src/kernels/cpu.rs @@ -0,0 +1,475 @@ +use std::collections::BTreeSet; + +use dlpk::{DLPackTensor, DLPackTensorRef, DLPackTensorRefMut}; +use ndarray::{ArrayView1, ArrayView2, ArrayViewD, ArrayViewMutD}; + +use crate::Error; +use super::{ReferenceValue, StridedNDIndex}; + +/// Check that the values of an i32 DLPack tensor match the expected reference. +/// +/// The tensor is converted to an ndarray view and compared element-wise and +/// shape-wise against `reference`. The `description` is used verbatim in the +/// error message on mismatch. +/// +/// # Parameters +/// - `tensor`: DLPack tensor with i32 data type +/// - `reference`: expected values with the same shape as the tensor +pub(crate) fn is_equal_i32( + tensor: DLPackTensorRef<'_>, + reference: &ReferenceValue, +) -> Result { + let values: ArrayViewD = tensor.try_into()?; + return Ok(values == reference.cpu.view()); +} + +macro_rules! validate_cell { + ($T: ty, $pbc: expr, $cell: expr) => { + let pbc_array: ArrayView1 = $pbc.try_into()?; + let cell_array: ArrayView2<$T> = $cell.try_into()?; + for i in 0..3 { + if !pbc_array[i] && !cell_array.row(i).iter().all(|&x| x == 0.0) { + return Err(Error::InvalidParameter(format!( + "invalid cell: for non-periodic dimensions, the corresponding \ + cell vector must be zero, but cell[{}] contains non-zero values", + i + ))); + } + } + }; +} + +/// Validate that cell vectors are zero for non-periodic dimensions on CPU. +/// +/// Converts the DLPack tensors to ndarray views and checks that for every +/// dimension where `pbc` is false, the corresponding row of `cell` contains +/// only zeros. +/// +/// # Parameters +/// - `pbc`: 1D boolean tensor of length 3 (periodic boundary condition flags) +/// - `cell`: 3x3 tensor (unit cell vectors as rows) +pub(crate) fn validate_cell_pbc( + pbc: DLPackTensorRef<'_>, + cell: DLPackTensorRef<'_>, +) -> Result<(), Error> { + let dtype = cell.dtype(); + if dtype.bits == 32 { + validate_cell!(f32, pbc, cell); + } else { + assert_eq!(dtype.bits, 64); + validate_cell!(f64, pbc, cell); + } + return Ok(()); +} + +/// Scale all elements of `tensor` in place by `factor` on CPU. +/// +/// Supports 32-bit and 64-bit floating point tensors. The tensor is converted +/// to a mutable ndarray view and scaled element-wise. +#[allow(clippy::cast_possible_truncation)] +pub(crate) fn scale_inplace( + tensor: DLPackTensorRefMut<'_>, + factor: f64, +) -> Result<(), Error> { + let dtype = tensor.dtype(); + if dtype.code == dlpk::sys::DLDataTypeCode::kDLFloat && dtype.bits == 32 { + let mut view: ArrayViewMutD = tensor.try_into()?; + view *= factor as f32; + } else if dtype.code == dlpk::sys::DLDataTypeCode::kDLFloat && dtype.bits == 64 { + let mut view: ArrayViewMutD = tensor.try_into()?; + view *= factor; + } else { + return Err(Error::InvalidParameter(format!( + "scale_inplace only supports 32-bit or 64-bit floats, got {}-bit {:?}", + dtype.bits, dtype.code + ))); + } + Ok(()) +} + +/// Check that all atomic types in `types` are present in `valid_types`. +/// +/// Returns `Ok(())` if all types are valid, or `Err` listing all invalid +/// types found. +#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] +pub(crate) fn check_atomic_types( + types: DLPackTensorRef<'_>, + valid_types: &ReferenceValue, +) -> Result<(), Error> { + assert!( + valid_types.cpu.is_standard_layout(), + "valid_types reference must be C-contiguous" + ); + assert_eq!( + types.n_dims(), 1, + "check_atomic_types expects a 1D types tensor" + ); + + let n_atoms = types.shape()[0] as usize; + if n_atoms == 0 { + return Ok(()); + } + + let types_idx = StridedNDIndex::from_dlpack(types); + let ptr = types.data_ptr::() + .map_err(|_| Error::Internal("failed to get types data pointer as i32".into()))?; + + // Compute the total number of elements in the buffer (including gaps from + // non-contiguous strides) so we can create a valid slice. + let n_elements = match types.strides() { + None => n_atoms, + Some(strides) => { + let max_offset: i64 = types.shape().iter() + .zip(strides.iter()) + .map(|(&s, &st)| (s - 1) * st) + .sum(); + + max_offset as usize + 1 + } + }; + let buffer = unsafe { + std::slice::from_raw_parts(ptr, n_elements) + }; + + check_atomic_types_buffer(buffer, &types_idx, n_atoms, valid_types) +} + +/// Check that all atomic types in a raw buffer are present in `valid_types`. +/// +/// This is used as a fallback by the CUDA and Metal backends when invalid types +/// are detected on-device. The `buffer` contains the raw i32 data (possibly +/// non-contiguous), and `types_idx` describes how to index into it. +#[allow(clippy::cast_possible_wrap, clippy::cast_possible_truncation, clippy::cast_sign_loss)] +pub(crate) fn check_atomic_types_buffer( + types: &[i32], + types_idx: &StridedNDIndex, + n_atoms: usize, + valid_types: &ReferenceValue, +) -> Result<(), Error> { + let valid = valid_types.cpu.as_slice().expect("reference should be contiguous"); + let mut invalid: BTreeSet = BTreeSet::new(); + for i in 0..n_atoms { + let offset = types_idx.offset(i as i64) as usize; + let atom_type = types[offset]; + if !valid.contains(&atom_type) { + invalid.insert(atom_type); + } + } + if !invalid.is_empty() { + let types: Vec = invalid.iter().map(|t| t.to_string()).collect(); + return Err(Error::InvalidParameter(format!( + "this model does not support the following atomic types which are present in the input systems: {}", + types.join(", ") + ))); + } + Ok(()) +} + +/// Clone a DLPack tensor on CPU, copying the underlying data. +/// +/// The returned `DLPackTensor` owns its own memory and is independent of the +/// original tensor. Supports all data types that can be viewed as an ndarray. +/// The clone is always C-contiguous, even when the original tensor is not. +#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] +pub(crate) fn clone_tensor(tensor: DLPackTensorRef<'_>) -> Result { + let dtype = tensor.dtype(); + + macro_rules! clone_as { + ($T: ty) => {{ + let view: ArrayViewD<$T> = tensor.try_into()?; + // `to_owned` would keep the memory order of the original tensor, we + // always want a C-contiguous clone + let cloned = view.as_standard_layout().into_owned(); + DLPackTensor::try_from(cloned) + .map_err(|e| Error::Internal(format!("failed to create DLPack tensor from ndarray: {e}"))) + }}; + } + + match (dtype.code, dtype.bits) { + (dlpk::sys::DLDataTypeCode::kDLInt, 32) => clone_as!(i32), + (dlpk::sys::DLDataTypeCode::kDLInt, 64) => clone_as!(i64), + (dlpk::sys::DLDataTypeCode::kDLUInt, 8) => clone_as!(u8), + (dlpk::sys::DLDataTypeCode::kDLUInt, 32) => clone_as!(u32), + (dlpk::sys::DLDataTypeCode::kDLUInt, 64) => clone_as!(u64), + (dlpk::sys::DLDataTypeCode::kDLFloat, 32) => clone_as!(f32), + (dlpk::sys::DLDataTypeCode::kDLFloat, 64) => clone_as!(f64), + (dlpk::sys::DLDataTypeCode::kDLBool, 8) => clone_as!(bool), + _ => Err(Error::InvalidParameter(format!( + "clone_tensor does not support {}-bit {:?} tensors", + dtype.bits, dtype.code + ))), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use super::super::ReferenceValue; + + use dlpk::DLPackTensor; + use ndarray::{Array1, Array2, ArrayD}; + + #[test] + fn test_is_equal_i32() { + let data = ArrayD::::from_shape_vec(vec![2, 3], vec![1, 2, 3, 4, 5, 6]).unwrap(); + let tensor: DLPackTensor = data.try_into().unwrap(); + let reference = ReferenceValue::new(ArrayD::::from_shape_vec( + vec![2, 3], + vec![1, 2, 3, 4, 5, 6], + ).unwrap()); + + assert!(is_equal_i32(tensor.as_ref(), &reference).unwrap()); + + let data = ArrayD::::from_shape_vec(vec![2, 3], vec![1, 2, 3, 42, 5, 6]).unwrap(); + let tensor: DLPackTensor = data.try_into().unwrap(); + assert!(!is_equal_i32(tensor.as_ref(), &reference).unwrap()); + + // shape mismatch + let data = ArrayD::::from_shape_vec(vec![2, 2], vec![1, 2, 3, 4]).unwrap(); + let tensor: DLPackTensor = data.try_into().unwrap(); + let reference = ReferenceValue::new(ArrayD::::from_shape_vec( + vec![4], + vec![1, 2, 3, 4], + ).unwrap()); + + assert!(!is_equal_i32(tensor.as_ref(), &reference).unwrap()); + + // empty arrays + let data = ArrayD::::from_shape_vec(vec![0], vec![]).unwrap(); + let tensor: DLPackTensor = data.try_into().unwrap(); + let reference = ReferenceValue::new(ArrayD::::from_shape_vec(vec![0], vec![]).unwrap()); + assert!(is_equal_i32(tensor.as_ref(), &reference).unwrap()); + } + + #[test] + fn test_validate_cell_pbc() { + // helper: pbc flags + 3x3 cell (row-major) => expected Ok or error substring + fn check(pbc: &[bool], cell: &[f64]) -> Result<(), Error> { + let pbc = Array1::::from_vec(pbc.to_vec()); + let cell = Array2::::from_shape_vec((3, 3), cell.to_vec()).unwrap(); + + let pbc: DLPackTensor = pbc.try_into().unwrap(); + let cell: DLPackTensor = cell.try_into().unwrap(); + + validate_cell_pbc(pbc.as_ref(), cell.as_ref()) + } + + // fully periodic — any cell is fine + check(&[true, true, true], &[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0]).unwrap(); + + // fully periodic — non-diagonal cell is fine too + check(&[true, true, true], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]).unwrap(); + + // non-periodic dim with zero cell vector — ok + check(&[true, false, true], &[10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 10.0]).unwrap(); + + // non-periodic dim with nonzero cell vector — error + let err = check( + &[true, false, true], + &[10.0, 0.0, 0.0, 5.0, 5.0, 5.0, 0.0, 0.0, 10.0], + ).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid cell: for non-periodic dimensions, \ + the corresponding cell vector must be zero, but cell[1] contains non-zero values" + ); + + // first dim non-periodic with nonzero cell + let err = check( + &[false, true, true], + &[1.0, 2.0, 3.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid cell: for non-periodic dimensions, \ + the corresponding cell vector must be zero, but cell[0] contains non-zero values" + ); + + // last dim non-periodic with nonzero cell + let err = check( + &[true, true, false], + &[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 7.0, 8.0, 9.0], + ).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid cell: for non-periodic dimensions, \ + the corresponding cell vector must be zero, but cell[2] contains non-zero values" + ); + + // all non-periodic with zero cell — ok + check(&[false, false, false], &[0.0; 9]).unwrap(); + + // f32 path + { + let pbc = Array1::::from_vec(vec![true, false, true]); + let cell = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap(); + + let pbc: DLPackTensor = pbc.try_into().unwrap(); + let cell: DLPackTensor = cell.try_into().unwrap(); + + let err = validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap_err(); + assert!(err.to_string().contains("cell[1] contains non-zero values")); + } + } + + #[test] + fn test_scale_inplace() { + // f32 2D + { + let data = ArrayD::::from_shape_vec(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); + let mut tensor: DLPackTensor = data.try_into().unwrap(); + + scale_inplace(tensor.as_mut(), 2.5).unwrap(); + + let view: ArrayViewD = tensor.as_ref().try_into().unwrap(); + assert_eq!(view, ndarray::arr2(&[[2.5_f32, 5.0, 7.5], [10.0, 12.5, 15.0]]).into_dyn()); + } + + // f64 2D + { + let data = ArrayD::::from_shape_vec(vec![2, 2], vec![1.0, -2.0, 3.5, 0.0]).unwrap(); + let mut tensor: DLPackTensor = data.try_into().unwrap(); + + scale_inplace(tensor.as_mut(), 0.5).unwrap(); + + let view: ArrayViewD = tensor.as_ref().try_into().unwrap(); + assert_eq!(view, ndarray::arr2(&[[0.5_f64, -1.0], [1.75, 0.0]]).into_dyn()); + } + + // zero factor + { + let data = ArrayD::::from_shape_vec(vec![3], vec![1.0, 2.0, 3.0]).unwrap(); + let mut tensor: DLPackTensor = data.try_into().unwrap(); + + scale_inplace(tensor.as_mut(), 0.0).unwrap(); + + let view: ArrayViewD = tensor.as_ref().try_into().unwrap(); + assert_eq!(view, ndarray::arr1(&[0.0_f32, 0.0, 0.0]).into_dyn()); + } + + // 1D f64 + { + let data = ArrayD::::from_shape_vec(vec![4], vec![2.0, 4.0, 8.0, 16.0]).unwrap(); + let mut tensor: DLPackTensor = data.try_into().unwrap(); + + scale_inplace(tensor.as_mut(), 0.25).unwrap(); + + let view: ArrayViewD = tensor.as_ref().try_into().unwrap(); + assert_eq!(view, ndarray::arr1(&[0.5_f64, 1.0, 2.0, 4.0]).into_dyn()); + } + } + + #[test] + fn test_check_atomic_types() { + let valid = ReferenceValue::new( + // this contains duplicated valid types, which is not expected but + // should be fine + ArrayD::::from_shape_vec(vec![4], vec![1, 6, 8, 1]).unwrap() + ); + + // all types valid + let types = ArrayD::::from_shape_vec(vec![5], vec![1, 6, 6, 6, 1]).unwrap(); + let tensor: DLPackTensor = types.try_into().unwrap(); + assert!(check_atomic_types(tensor.as_ref(), &valid).is_ok()); + + + // invalid types + let types = ArrayD::::from_shape_vec(vec![6], vec![1, 3, 8, 3, 4, 1]).unwrap(); + let tensor: DLPackTensor = types.try_into().unwrap(); + let err = check_atomic_types(tensor.as_ref(), &valid).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model does not support the following atomic \ + types which are present in the input systems: 3, 4" + ); + + // empty types — ok + let types = ArrayD::::from_shape_vec(vec![0], vec![]).unwrap(); + let tensor: DLPackTensor = types.try_into().unwrap(); + assert!(check_atomic_types(tensor.as_ref(), &valid).is_ok()); + } + + #[test] + fn test_clone_tensor() { + // f32 + { + let data = ArrayD::::from_shape_vec(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); + let tensor: DLPackTensor = data.try_into().unwrap(); + + let mut cloned = clone_tensor(tensor.as_ref()).unwrap(); + + // same values + let orig_view: ArrayViewD = tensor.as_ref().try_into().unwrap(); + let clone_view: ArrayViewD = cloned.as_ref().try_into().unwrap(); + assert_eq!(orig_view, clone_view); + + // modifying the clone does not affect the original + scale_inplace(cloned.as_mut(), 10.0).unwrap(); + let orig_after: ArrayViewD = tensor.as_ref().try_into().unwrap(); + assert_eq!(orig_after, orig_view); + } + + // i32 + { + let data = ArrayD::::from_shape_vec(vec![3], vec![10, 20, 30]).unwrap(); + let tensor: DLPackTensor = data.try_into().unwrap(); + + let cloned = clone_tensor(tensor.as_ref()).unwrap(); + + let orig_view: ArrayViewD = tensor.as_ref().try_into().unwrap(); + let clone_view: ArrayViewD = cloned.as_ref().try_into().unwrap(); + assert_eq!(orig_view, clone_view); + } + + // bool + { + let data = ArrayD::::from_shape_vec(vec![3], vec![true, false, true]).unwrap(); + let tensor: DLPackTensor = data.try_into().unwrap(); + + let cloned = clone_tensor(tensor.as_ref()).unwrap(); + + let orig_view: ArrayViewD = tensor.as_ref().try_into().unwrap(); + let clone_view: ArrayViewD = cloned.as_ref().try_into().unwrap(); + assert_eq!(orig_view, clone_view); + } + + // f64 2D + { + let data = ArrayD::::from_shape_vec(vec![2, 2], vec![1.5, -2.0, 3.0, 0.0]).unwrap(); + let expected = data.clone(); + let tensor: DLPackTensor = data.try_into().unwrap(); + + let cloned = clone_tensor(tensor.as_ref()).unwrap(); + + let clone_view: ArrayViewD = cloned.as_ref().try_into().unwrap(); + assert_eq!(clone_view, expected); + } + + // empty tensor + { + let data = ArrayD::::from_shape_vec(vec![0], vec![]).unwrap(); + let tensor: DLPackTensor = data.try_into().unwrap(); + + let cloned = clone_tensor(tensor.as_ref()).unwrap(); + assert_eq!(cloned.shape(), &[0]); + } + + // non-contiguous tensor: the clone is C-contiguous + { + let data = ArrayD::::from_shape_vec(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); + let transposed = data.view().reversed_axes(); + let tensor: DLPackTensorRef = (&transposed).try_into().unwrap(); + + let cloned = clone_tensor(tensor).unwrap(); + + assert_eq!(cloned.shape(), &[3, 2]); + assert_eq!(cloned.strides(), Some(&[2, 1][..])); + + let clone_view: ArrayViewD = cloned.as_ref().try_into().unwrap(); + assert_eq!(clone_view, transposed); + } + } +} diff --git a/metatomic-core/src/kernels/cuda.rs b/metatomic-core/src/kernels/cuda.rs new file mode 100644 index 000000000..45a6ffd05 --- /dev/null +++ b/metatomic-core/src/kernels/cuda.rs @@ -0,0 +1,935 @@ +use std::collections::hash_map::Entry; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, LazyLock}; + +use cudarc::driver::safe::DeviceRepr; +use cudarc::driver::safe::{ + CudaContext, CudaFunction, CudaModule, CudaStream, LaunchConfig, PushKernelArg, +}; +use cudarc::driver::sys; +use cudarc::nvrtc::compile_ptx; +use dlpk::{DLDevice, DLPackTensor, DLPackTensorRef, DLPackTensorRefMut}; + +use crate::Error; +use super::{ReferenceValue, StridedNDIndex}; + +// CUDA kernel source compiled at runtime via NVRTC for the exact GPU +const KERNEL_SRC: &str = include_str!("cuda_kernels.cu"); + +unsafe impl DeviceRepr for StridedNDIndex {} + +/// Create a [`LaunchConfig`] for `n_elements` with 64-bit element counts. +/// +/// This replaces `LaunchConfig::for_num_elems` which only accepts `u32`. +/// CUDA's gridDim.x supports up to 2^31 - 1 blocks; with a block size of +/// 1024 this covers up to ~2.2 × 10¹² elements. +#[allow(clippy::cast_possible_truncation)] +fn launch_config_for_elems(n_elements: u64) -> LaunchConfig { + const NUM_THREADS: u64 = 1024; + const MAX_GRID_X: u64 = (1u64 << 31) - 1; + let num_blocks = std::cmp::min(n_elements.div_ceil(NUM_THREADS), MAX_GRID_X); + LaunchConfig { + grid_dim: (num_blocks as u32, 1, 1), + block_dim: (NUM_THREADS as u32, 1, 1), + shared_mem_bytes: 0, + } +} + +/// Zero-cost wrapper to pass an existing device pointer as a CUDA kernel +/// argument. +/// +/// Does NOT own the memory — the caller (DLPack tensor) is responsible for +/// lifetime and must ensure the pointer remains valid for the duration of the +/// kernel launch. +/// +/// The `#[repr(transparent)]` wrapper over `cudarc::driver::sys::CUdeviceptr` +/// is passed to `PushKernelArg::arg()` which pushes the address of this struct +/// on the host stack. CUDA reads 8 bytes from that address as the kernel +/// parameter value, giving the kernel the correct device pointer. +#[repr(transparent)] +struct DLPackDevicePtr<'a> { + ptr: cudarc::driver::sys::CUdeviceptr, + _phantom: std::marker::PhantomData<&'a [u8]>, +} + +unsafe impl DeviceRepr for DLPackDevicePtr<'_> {} + +impl<'a> DLPackDevicePtr<'a> { + /// Wrap a CUDA-resident DLPack tensor's device pointer for use as a kernel + /// argument. + /// + /// The returned `DlpackDevicePtr` borrows the tensor's lifetime, ensuring the + /// backing memory stays alive as long as the argument is in use. + fn from_ref(tensor: DLPackTensorRef<'a>) -> Self { + Self { + ptr: unsafe { dlpack_to_device_ptr(tensor) }, + _phantom: std::marker::PhantomData, + } + } + + /// Wrap a CUDA-resident DLPack tensor's device pointer for use as a kernel + /// argument (mutable variant). + /// + /// The returned `DlpackDevicePtr` borrows the tensor's lifetime, ensuring the + /// backing memory stays alive as long as the argument is in use. + fn from_mut(tensor: DLPackTensorRefMut<'_>) -> Self { + Self { + ptr: unsafe { dlpack_to_device_ptr(tensor.as_ref()) }, + _phantom: std::marker::PhantomData, + } + } +} + +/// Per-device cached resources: context, module, and kernel function handles. +struct CudaKernelCache { + ctx: Arc, + module: Arc, + is_equal_i32: CudaFunction, + validate_cell_pbc_f32: CudaFunction, + validate_cell_pbc_f64: CudaFunction, + scale_f32: CudaFunction, + scale_f64: CudaFunction, + copy_to_contiguous_8bit: CudaFunction, + copy_to_contiguous_16bit: CudaFunction, + copy_to_contiguous_32bit: CudaFunction, + copy_to_contiguous_64bit: CudaFunction, + check_atomic_types: CudaFunction, +} + +impl CudaKernelCache { + fn new(device_id: usize) -> Result { + let ctx = CudaContext::new(device_id) + .map_err(|e| Error::Internal(format!("CudaContext::new({device_id}): {e}")))?; + + let ptx = compile_ptx(KERNEL_SRC) + .map_err(|e| Error::Internal(format!("NVRTC compile failed: {e}")))?; + + let module = ctx.load_module(ptx) + .map_err(|e| Error::Internal(format!("PTX load failed: {e}")))?; + + let is_equal_i32 = module.load_function("is_equal_i32") + .map_err(|e| Error::Internal(format!("load_function(is_equal_i32): {e}")))?; + + let validate_cell_pbc_f32 = module.load_function("validate_cell_pbc_f32") + .map_err(|e| Error::Internal(format!("load_function(validate_cell_pbc_f32): {e}")))?; + + let validate_cell_pbc_f64 = module.load_function("validate_cell_pbc_f64") + .map_err(|e| Error::Internal(format!("load_function(validate_cell_pbc_f64): {e}")))?; + + let scale_f32 = module.load_function("scale_f32") + .map_err(|e| Error::Internal(format!("load_function(scale_f32): {e}")))?; + + let scale_f64 = module.load_function("scale_f64") + .map_err(|e| Error::Internal(format!("load_function(scale_f64): {e}")))?; + + let copy_to_contiguous_8bit = module.load_function("copy_to_contiguous_8bit") + .map_err(|e| Error::Internal(format!("load_function(copy_to_contiguous_8bit): {e}")))?; + + let copy_to_contiguous_16bit = module.load_function("copy_to_contiguous_16bit") + .map_err(|e| Error::Internal(format!("load_function(copy_to_contiguous_16bit): {e}")))?; + + let copy_to_contiguous_32bit = module.load_function("copy_to_contiguous_32bit") + .map_err(|e| Error::Internal(format!("load_function(copy_to_contiguous_32bit): {e}")))?; + + let copy_to_contiguous_64bit = module.load_function("copy_to_contiguous_64bit") + .map_err(|e| Error::Internal(format!("load_function(copy_to_contiguous_64bit): {e}")))?; + let check_atomic_types = module.load_function("check_atomic_types") + .map_err(|e| Error::Internal(format!("load_function(check_atomic_types): {e}")))?; + + Ok(Self { + ctx, + module, + is_equal_i32, + validate_cell_pbc_f32, + validate_cell_pbc_f64, + scale_f32, + scale_f64, + copy_to_contiguous_8bit, + copy_to_contiguous_16bit, + copy_to_contiguous_32bit, + copy_to_contiguous_64bit, + check_atomic_types, + }) + } +} + +static CUDA_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new())); + +fn get_or_init(device_id: usize) -> Result, Error> { + let mut cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE"); + let entry = match cache.entry(device_id) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => entry.insert(CudaKernelCache::new(device_id)?), + }; + Ok(entry.ctx.default_stream()) +} + +fn check_valid_device(function: &str, device: DLDevice) { + assert_eq!( + device.device_type, dlpk::sys::DLDeviceType::kDLCUDA, + "{} called on non-CUDA tensor", function + ); + assert!(device.device_id >= 0, "{} called on invalid device_id", function); +} + +/// Extract a `CUdeviceptr` from a DLPack tensor's raw `data` + `byte_offset`. +/// +/// # Safety +/// +/// The returned `CUdeviceptr` is only valid as long as the DLPack tensor's +/// backing memory is alive. The caller must ensure the tensor is not dropped +/// before the kernel finishes execution. +unsafe fn dlpack_to_device_ptr(tensor: DLPackTensorRef<'_>) -> cudarc::driver::sys::CUdeviceptr { + debug_assert_eq!( + tensor.device().device_type, dlpk::sys::DLDeviceType::kDLCUDA, + "dlpack_to_device_ptr called on non-CUDA tensor" + ); + let raw_ptr = tensor.raw.data as u64; + (raw_ptr + tensor.raw.byte_offset) as cudarc::driver::sys::CUdeviceptr +} + +/// Check that the values of a CUDA-resident i32 DLPack tensor match an expected +/// reference array. +/// +/// The comparison is performed entirely on-device: the existing GPU pointer +/// from `tensor` is wrapped as a `DlpackDevicePtr`, the reference is uploaded to +/// the GPU (and cached for subsequent calls), and a single-element result flag +/// (`0` = ok, `1` = mismatch) is read back. +#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)] +pub(super) fn is_equal_i32(tensor: DLPackTensorRef<'_>, reference: &ReferenceValue) -> Result { + check_valid_device("is_equal_i32", tensor.device()); + + let device_id = tensor.device().device_id as usize; + let stream = get_or_init(device_id)?; + let cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE"); + let entry = &cache[&device_id]; + + let n_elements: i64 = tensor.shape().iter().product(); + + // Build strided index from the DLPack tensor (preserves actual strides) + let values_idx = StridedNDIndex::from_dlpack(tensor); + + // Wrap the existing GPU-allocated tensor pointer + let tensor_ptr = DLPackDevicePtr::from_ref(tensor); + + // Upload reference values to GPU (cached after first call, per device) + let (ref_dev, reference_idx) = reference.cuda_data(device_id, &stream)?; + + // Allocate result flag (initialized to 0 = no mismatch) + let mut result = stream.alloc_zeros::(1) + .map_err(|e| Error::Internal(format!("alloc_zeros: {e}")))?; + + unsafe { + stream.launch_builder(&entry.is_equal_i32) + .arg(&tensor_ptr) + .arg(&values_idx) + .arg(ref_dev) + .arg(reference_idx) + .arg(&n_elements) + .arg(&mut result) + .launch(launch_config_for_elems(n_elements as u64)) + .map_err(|e| Error::Internal(format!("kernel launch (is_equal_i32): {e}")))?; + } + + stream.synchronize() + .map_err(|e| Error::Internal(format!("device sync: {e}")))?; + + let host = stream.clone_dtoh(&result) + .map_err(|e| Error::Internal(format!("clone_dtoh result: {e}")))?; + + return Ok(host[0] == 0); +} + +/// Validate that cell vectors are zero for non-periodic dimensions, on CUDA device. +#[allow(clippy::cast_sign_loss)] +pub(super) fn validate_cell_pbc( + pbc: DLPackTensorRef<'_>, + cell: DLPackTensorRef<'_>, +) -> Result<(), Error> { + debug_assert_eq!(cell.device(), pbc.device(), "pbc and cell must be on the same device"); + check_valid_device("validate_cell_pbc", pbc.device()); + + let device_id = pbc.device().device_id as usize; + let stream = get_or_init(device_id)?; + let cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE"); + let entry = &cache[&device_id]; + + let pbc_ptr = DLPackDevicePtr::from_ref(pbc); + let cell_ptr = DLPackDevicePtr::from_ref(cell); + + let pbc_idx = StridedNDIndex::from_dlpack(pbc); + let cell_idx = StridedNDIndex::from_dlpack(cell); + + let mut result = stream.alloc_zeros::(1) + .map_err(|e| Error::Internal(format!("alloc_zeros: {e}")))?; + + if cell.dtype().bits == 32 { + unsafe { + stream.launch_builder(&entry.validate_cell_pbc_f32) + .arg(&pbc_ptr) + .arg(&pbc_idx) + .arg(&cell_ptr) + .arg(&cell_idx) + .arg(&mut result) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (3, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| Error::Internal(format!("kernel launch (f32): {e}")))?; + } + } else { + assert_eq!(cell.dtype().bits, 64, "validate_cell_pbc: unsupported cell dtype"); + unsafe { + stream.launch_builder(&entry.validate_cell_pbc_f64) + .arg(&pbc_ptr) + .arg(&pbc_idx) + .arg(&cell_ptr) + .arg(&cell_idx) + .arg(&mut result) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (3, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| Error::Internal(format!("kernel launch (f64): {e}")))?; + } + } + + stream.synchronize() + .map_err(|e| Error::Internal(format!("device sync: {e}")))?; + + let host = stream.clone_dtoh(&result) + .map_err(|e| Error::Internal(format!("clone_dtoh result: {e}")))?; + + if host[0] != 0 { + let dim = host[0] - 1; + return Err(Error::InvalidParameter(format!( + "invalid cell: for non-periodic dimensions, the corresponding \ + cell vector must be zero, but cell[{}] contains non-zero values", + dim + ))); + } + Ok(()) +} + + +/// Scale all elements of `tensor` in place by `factor`, on CUDA device. +/// +/// The tensor must be a 32-bit or 64-bit floating point tensor residing on a +/// CUDA device. The scaling is performed entirely on-device, in place. +#[allow(clippy::cast_sign_loss)] +pub(super) fn scale_inplace( + tensor: DLPackTensorRefMut<'_>, + factor: f64, +) -> Result<(), Error> { + check_valid_device("scale_inplace", tensor.device()); + + let device_id = tensor.device().device_id as usize; + let stream = get_or_init(device_id)?; + let cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE"); + let entry = &cache[&device_id]; + + let n_elements: i64 = tensor.shape().iter().product(); + if n_elements == 0 { + return Ok(()); + } + + let dtype = tensor.dtype(); + let tensor_idx = StridedNDIndex::from_dlpack(tensor.as_ref()); + let tensor_ptr = DLPackDevicePtr::from_mut(tensor); + + if dtype.code == dlpk::sys::DLDataTypeCode::kDLFloat && dtype.bits == 32 { + unsafe { + stream.launch_builder(&entry.scale_f32) + .arg(&tensor_ptr) + .arg(&tensor_idx) + .arg(&n_elements) + .arg(&factor) + .launch(launch_config_for_elems(n_elements as u64)) + .map_err(|e| Error::Internal(format!("kernel launch (scale_f32): {e}")))?; + } + } else if dtype.code == dlpk::sys::DLDataTypeCode::kDLFloat && dtype.bits == 64 { + unsafe { + stream.launch_builder(&entry.scale_f64) + .arg(&tensor_ptr) + .arg(&tensor_idx) + .arg(&n_elements) + .arg(&factor) + .launch(launch_config_for_elems(n_elements as u64)) + .map_err(|e| Error::Internal(format!("kernel launch (scale_f64): {e}")))?; + } + } else { + return Err(Error::InvalidParameter(format!( + "scale_inplace only supports 32-bit or 64-bit floats, got {}-bit {:?}", + dtype.bits, dtype.code + ))); + } + + stream.synchronize() + .map_err(|e| Error::Internal(format!("device sync: {e}")))?; + + Ok(()) +} + +/// Check that all atomic types in `types` are present in `valid_types`, on CUDA. +/// +/// The check runs on-device. If invalid types are found (count > 0), a CPU +/// fallback scan identifies the specific invalid type for the error message. +#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)] +pub(super) fn check_atomic_types( + types: DLPackTensorRef<'_>, + valid_types: &ReferenceValue, +) -> Result<(), Error> { + check_valid_device("check_atomic_types", types.device()); + assert!( + valid_types.cpu.is_standard_layout(), + "valid_types reference must be C-contiguous" + ); + assert_eq!( + types.n_dims(), 1, + "check_atomic_types expects a 1D types tensor" + ); + + let device_id = types.device().device_id as usize; + let stream = get_or_init(device_id)?; + let cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE"); + let entry = &cache[&device_id]; + + let n_atoms: i64 = types.shape().iter().product(); + if n_atoms == 0 { + return Ok(()); + } + + let types_idx = StridedNDIndex::from_dlpack(types); + let types_ptr = DLPackDevicePtr::from_ref(types); + + let (valid_types_device, _) = valid_types.cuda_data(device_id, &stream)?; + let n_valid_types = i64::try_from(valid_types.cpu.len()).expect("could not cast n_valid_types to i64"); + + // Allocate result counter (initialized to 0) + let mut result = stream.alloc_zeros::(1) + .map_err(|e| Error::Internal(format!("alloc_zeros: {e}")))?; + + unsafe { + stream.launch_builder(&entry.check_atomic_types) + .arg(&types_ptr) + .arg(&types_idx) + .arg(&n_atoms) + .arg(valid_types_device) + .arg(&n_valid_types) + .arg(&mut result) + .launch(launch_config_for_elems(n_atoms as u64)) + .map_err(|e| Error::Internal(format!("kernel launch (check_atomic_types): {e}")))?; + } + + stream.synchronize() + .map_err(|e| Error::Internal(format!("device sync: {e}")))?; + + let host = stream.clone_dtoh(&result) + .map_err(|e| Error::Internal(format!("clone_dtoh result: {e}")))?; + + if host[0] > 0 { + // Invalid types found — copy types to CPU and scan for the specific + // invalid types. The tensor may be non-contiguous, so we copy the full + // byte span and index using types_idx. + let n_atoms_usize = n_atoms as usize; + let elem_size = std::mem::size_of::(); + let n_bytes = match types.strides() { + None => n_atoms_usize * elem_size, + Some(strides) => { + let max_offset: i64 = types.shape().iter() + .zip(strides.iter()) + .map(|(&s, &st)| (s - 1) * st) + .sum(); + (max_offset as usize + 1) * elem_size + } + }; + let n_elements = n_bytes / elem_size; + let mut host_types = vec![0i32; n_elements]; + unsafe { + cudarc::driver::result::memcpy_dtoh_sync(host_types.as_mut_slice(), dlpack_to_device_ptr(types)) + }.map_err(|e| Error::Internal(format!("memcpy_dtoh_sync types: {e}")))?; + + super::cpu::check_atomic_types_buffer(&host_types, &types_idx, n_atoms_usize, valid_types)?; + } + + Ok(()) +} + +/// Context held by the deleter of a cloned CUDA `DLManagedTensorVersioned`. +/// +/// Stores the `CUdeviceptr` and the stream it was allocated on, so it can be +/// freed when the DLPack tensor is dropped. `ptr` is null for empty tensors, +/// for which no memory is allocated. +struct CudaCloneContext { + ptr: sys::CUdeviceptr, + stream: Arc, + shape: Vec, + strides: Vec, +} + +/// Deleter for a cloned CUDA DLPack tensor. +/// +/// Frees the device memory and the boxed `DLManagedTensorVersioned`. +unsafe extern "C" fn cuda_clone_deleter(tensor: *mut dlpk::sys::DLManagedTensorVersioned) { + unsafe { + let ctx = (*tensor).manager_ctx.cast::(); + let ctx = Box::from_raw(ctx); + + // free the device memory + if ctx.ptr != 0 { + let _ = sys::cuMemFreeAsync(ctx.ptr, ctx.stream.cu_stream()); + } + + // also drop the tensor itself + let _ = Box::from_raw(tensor); + } +} + +/// Clone a DLPack tensor on CUDA, copying the underlying device memory. +/// +/// The returned `DLPackTensor` owns its own CUDA memory allocation and is +/// independent of the original tensor. The clone is always C-contiguous, even +/// when the original tensor is not: the data is gathered with the +/// `copy_to_contiguous` kernel instead of a plain device to device copy. +#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss, clippy::cast_possible_wrap)] +pub(super) fn clone_tensor(tensor: &DLPackTensorRef<'_>) -> Result { + check_valid_device("clone_tensor", tensor.device()); + + let device_id = tensor.device().device_id as usize; + let stream = get_or_init(device_id)?; + + let cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE"); + let entry = &cache[&device_id]; + + let element_size_bits = super::element_size(tensor.dtype())?; + let n_elements: i64 = tensor.shape().iter().product(); + let num_bytes = n_elements as usize * element_size_bits / 8; + + let shape: Vec = tensor.shape().to_vec(); + // the clone stores the data contiguously, regardless of the strides used by + // the original tensor + let strides = super::contiguous_strides(&shape); + + // pick the kernel matching the element size before allocating anything + let kernel = match element_size_bits { + 8 => &entry.copy_to_contiguous_8bit, + 16 => &entry.copy_to_contiguous_16bit, + 32 => &entry.copy_to_contiguous_32bit, + 64 => &entry.copy_to_contiguous_64bit, + _ => { + return Err(Error::InvalidParameter(format!( + "clone_tensor does not support {} tensors on CUDA", + tensor.dtype() + ))); + } + }; + + // allocate device memory, only big enough for the contiguous data + let mut dst_ptr: sys::CUdeviceptr = 0; + stream.context().bind_to_thread() + .map_err(|e| Error::Internal(format!("bind_to_thread: {e}")))?; + + if num_bytes > 0 { + unsafe { + sys::cuMemAllocAsync(&mut dst_ptr, num_bytes, stream.cu_stream()) + .result() + .map_err(|e| Error::Internal(format!("cuMemAllocAsync: {e}")))?; + } + + // gather the (possibly strided) data from the original tensor into the + // contiguous allocation + let src_idx = StridedNDIndex::from_dlpack(*tensor); + let src_ptr = DLPackDevicePtr::from_ref(*tensor); + + unsafe { + stream.launch_builder(kernel) + .arg(&src_ptr) + .arg(&src_idx) + // `dst_ptr` is an `u64`, which is passed to the kernel by value + // and interpreted as a device pointer, as with `DLPackDevicePtr` + .arg(&dst_ptr) + .arg(&n_elements) + .launch(launch_config_for_elems(n_elements as u64)) + .map_err(|e| Error::Internal(format!("kernel launch (copy_to_contiguous): {e}")))?; + } + } + + stream.synchronize().map_err(|e| Error::Internal(format!("device sync: {e}")))?; + + // build the DLManagedTensorVersioned + let ctx = Box::new(CudaCloneContext { + ptr: dst_ptr, + stream: stream.clone(), + shape: shape, + strides: strides, + }); + + let ndim = ctx.shape.len() as i32; + let dl_tensor = dlpk::sys::DLTensor { + data: dst_ptr as *mut std::ffi::c_void, + device: tensor.device(), + ndim, + dtype: tensor.dtype(), + shape: ctx.shape.as_ptr().cast_mut(), + strides: ctx.strides.as_ptr().cast_mut(), + byte_offset: 0, + }; + + let managed = Box::new(dlpk::sys::DLManagedTensorVersioned { + version: dlpk::sys::DLPackVersion::current(), + manager_ctx: Box::into_raw(ctx).cast(), + deleter: Some(cuda_clone_deleter), + flags: dlpk::sys::DLPACK_FLAG_BITMASK_IS_COPIED, + dl_tensor, + }); + + let ptr = Box::into_raw(managed); + Ok(unsafe { DLPackTensor::from_ptr(ptr) }) +} + + +#[cfg(test)] +mod tests { + use super::*; + + use dlpk::{DLDevice, GetDLPackDataType}; + use ndarray::ArrayD; + + /// Check whether a CUDA device is available on this machine. The tests + /// below are skipped when there is none. + fn cuda_available() -> bool { + // this requires the `dynamic-loading` feature of cudarc, without which + // the tests would fail to link on machines without CUDA anyway + if !unsafe {cudarc::driver::sys::is_culib_present() } { + return false; + } + return CudaContext::device_count().unwrap_or(0) > 0; + } + + macro_rules! skip_without_cuda { + () => { + if !cuda_available() { + eprintln!("no CUDA device available, skipping this test"); + return; + } + }; + } + + /// A DLPack tensor with CUDA-resident data, used to test the kernels above. + struct CudaTensor { + ptr: cudarc::driver::sys::CUdeviceptr, + shape: Vec, + strides: Vec, + dtype: dlpk::sys::DLDataType, + } + + impl CudaTensor { + /// Create a new CUDA tensor with the given `shape` and `strides`, + /// containing a copy of `data`. + /// + /// `data` is the full memory span of the tensor, including any gap + /// between the elements actually part of the tensor. + fn new(data: &[T], shape: &[i64], strides: &[i64]) -> Self { + let stream = get_or_init(0).expect("failed to initialize CUDA device 0"); + stream.context().bind_to_thread().expect("bind_to_thread failed"); + + assert!(!data.is_empty()); + let ptr = unsafe { + cudarc::driver::result::malloc_sync(std::mem::size_of_val(data)) + }.expect("malloc_sync failed"); + + unsafe { + cudarc::driver::result::memcpy_htod_sync(ptr, data) + }.expect("memcpy_htod_sync failed"); + + CudaTensor { + ptr, + shape: shape.to_vec(), + strides: strides.to_vec(), + dtype: T::get_dlpack_data_type(), + } + } + + #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] + fn dl_tensor(&self) -> dlpk::sys::DLTensor { + dlpk::sys::DLTensor { + data: self.ptr as *mut std::ffi::c_void, + device: DLDevice { + device_type: dlpk::sys::DLDeviceType::kDLCUDA, + device_id: 0, + }, + ndim: self.shape.len() as i32, + dtype: self.dtype, + shape: self.shape.as_ptr().cast_mut(), + strides: self.strides.as_ptr().cast_mut(), + byte_offset: 0, + } + } + + fn as_ref(&self) -> DLPackTensorRef<'_> { + unsafe { DLPackTensorRef::from_raw(self.dl_tensor()) } + } + + fn as_mut(&mut self) -> DLPackTensorRefMut<'_> { + unsafe { DLPackTensorRefMut::from_raw(self.dl_tensor()) } + } + + /// Read the first `n` elements of this tensor's memory span back to the + /// CPU + fn data(&self, n: usize) -> Vec { + return read_cuda(self.as_ref(), n); + } + + /// Overwrite the data in this tensor's device allocation + fn overwrite(&self, data: &[T]) { + unsafe { + cudarc::driver::result::memcpy_htod_sync(self.ptr, data) + }.expect("memcpy_htod_sync failed"); + } + } + + /// Read the first `n` elements of the data of any CUDA-resident tensor back + /// to the CPU + fn read_cuda(tensor: DLPackTensorRef<'_>, n: usize) -> Vec { + let mut host = vec![T::default(); n]; + unsafe { + cudarc::driver::result::memcpy_dtoh_sync(host.as_mut_slice(), dlpack_to_device_ptr(tensor)) + }.expect("memcpy_dtoh_sync failed"); + + return host; + } + + impl Drop for CudaTensor { + fn drop(&mut self) { + unsafe { + let _ = cudarc::driver::result::free_sync(self.ptr); + } + } + } + + #[test] + fn is_equal_i32_kernel() { + skip_without_cuda!(); + + let reference = ReferenceValue::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![0, 1, 2]).unwrap() + ); + + // matching values + let tensor = CudaTensor::new(&[0_i32, 1, 2], &[3, 1], &[1, 1]); + assert!(is_equal_i32(tensor.as_ref(), &reference).unwrap()); + + // mismatching values + let tensor = CudaTensor::new(&[0_i32, 42, 2], &[3, 1], &[1, 1]); + assert!(!is_equal_i32(tensor.as_ref(), &reference).unwrap()); + + // matching values in a non-contiguous tensor: every other element of + // [0, -1, 1, -1, 2, -1] + let tensor = CudaTensor::new(&[0_i32, -1, 1, -1, 2, -1], &[3, 1], &[2, 1]); + assert!(is_equal_i32(tensor.as_ref(), &reference).unwrap()); + } + + #[test] + fn validate_cell_pbc_kernel() { + skip_without_cuda!(); + + // fully periodic: any cell is valid + let pbc = CudaTensor::new(&[true, true, true], &[3], &[1]); + let cell = CudaTensor::new( + &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], &[3, 3], &[3, 1] + ); + validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap(); + + // non-periodic dimension with a zero cell vector: valid + let pbc = CudaTensor::new(&[true, false, true], &[3], &[1]); + let cell = CudaTensor::new( + &[10.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 10.0], &[3, 3], &[3, 1] + ); + validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap(); + + // non-periodic dimension with a non-zero cell vector: invalid + let cell = CudaTensor::new( + &[10.0_f32, 0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 10.0], &[3, 3], &[3, 1] + ); + let err = validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap_err(); + assert!(err.to_string().contains("cell[1] contains non-zero values"), "{err}"); + + // the same checks with f64 data, using the last dimension as the + // non-periodic one + let pbc = CudaTensor::new(&[true, true, false], &[3], &[1]); + let cell = CudaTensor::new( + &[10.0_f64, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 0.0], &[3, 3], &[3, 1] + ); + validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap(); + + let cell = CudaTensor::new( + &[10.0_f64, 0.0, 0.0, 0.0, 10.0, 0.0, 3.0, 0.0, 10.0], &[3, 3], &[3, 1] + ); + let err = validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap_err(); + assert!(err.to_string().contains("cell[2] contains non-zero values"), "{err}"); + } + + #[test] + fn scale_inplace_kernel() { + skip_without_cuda!(); + + // f32 + let mut tensor = CudaTensor::new(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3], &[3, 1]); + scale_inplace(tensor.as_mut(), 2.5).unwrap(); + assert_eq!(tensor.data::(6), [2.5, 5.0, 7.5, 10.0, 12.5, 15.0]); + + // f64 + let mut tensor = CudaTensor::new(&[1.0_f64, -2.0, 3.5, 0.0], &[2, 2], &[2, 1]); + scale_inplace(tensor.as_mut(), 0.5).unwrap(); + assert_eq!(tensor.data::(4), [0.5, -1.0, 1.75, 0.0]); + + // non-contiguous tensor: only the 2x2 block in the top left corner of + // this 3x4 array is scaled, the rest of the data is left alone + let data: Vec = (0..12_i16).map(f32::from).collect(); + let mut tensor = CudaTensor::new(&data, &[2, 2], &[4, 1]); + scale_inplace(tensor.as_mut(), 10.0).unwrap(); + assert_eq!(tensor.data::(12), [ + 0.0, 10.0, 2.0, 3.0, + 40.0, 50.0, 6.0, 7.0, + 8.0, 9.0, 10.0, 11.0, + ]); + + // empty tensors are left alone (the allocation still has one element, + // since CUDA does not allow zero-sized allocations, and it should not + // be touched by the kernel) + let mut tensor = CudaTensor::new(&[3.0_f32], &[0], &[1]); + scale_inplace(tensor.as_mut(), 2.0).unwrap(); + assert_eq!(tensor.data::(1), [3.0]); + + // integers are not supported + let mut tensor = CudaTensor::new(&[1_i32, 2, 3], &[3], &[1]); + let err = scale_inplace(tensor.as_mut(), 2.0).unwrap_err(); + assert!(err.to_string().contains("only supports 32-bit or 64-bit floats"), "{err}"); + } + + #[test] + fn clone_contiguous() { + skip_without_cuda!(); + + let data = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + let tensor = CudaTensor::new(&data, &[2, 3], &[3, 1]); + + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + + assert_eq!(cloned.shape(), [2, 3]); + assert_eq!(cloned.strides(), Some(&[3, 1][..])); + assert_eq!(read_cuda::(cloned.as_ref(), 6), data); + + // the clone is independent from the original + tensor.overwrite(&[42.0_f32; 6]); + assert_eq!(read_cuda::(cloned.as_ref(), 6), data); + } + + #[test] + fn clone_non_contiguous() { + skip_without_cuda!(); + + // 2x2 block in the top left corner of a 3x4 array + let data: Vec = (0..12_i16).map(f32::from).collect(); + let tensor = CudaTensor::new(&data, &[2, 2], &[4, 1]); + + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + + assert_eq!(cloned.shape(), [2, 2]); + assert_eq!(cloned.strides(), Some(&[2, 1][..])); + assert_eq!(read_cuda::(cloned.as_ref(), 4), [0.0, 1.0, 4.0, 5.0]); + } + + #[test] + fn clone_transposed() { + skip_without_cuda!(); + + // 2x3 array in column-major order (i.e. the transpose of a 3x2 array) + let data: Vec = (0..6).collect(); + let tensor = CudaTensor::new(&data, &[2, 3], &[1, 2]); + + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + + assert_eq!(cloned.shape(), [2, 3]); + assert_eq!(cloned.strides(), Some(&[3, 1][..])); + assert_eq!(read_cuda::(cloned.as_ref(), 6), [0, 2, 4, 1, 3, 5]); + } + + #[test] + fn clone_element_sizes() { + skip_without_cuda!(); + + // 8-bit elements, every other one + let data: Vec = (0..6).collect(); + let tensor = CudaTensor::new(&data, &[3], &[2]); + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + assert_eq!(read_cuda::(cloned.as_ref(), 3), [0, 2, 4]); + + // 16-bit elements, every other one + let data: Vec = (0..6).collect(); + let tensor = CudaTensor::new(&data, &[3], &[2]); + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + assert_eq!(read_cuda::(cloned.as_ref(), 3), [0, 2, 4]); + + // bool elements + let data = vec![true, false, true, true]; + let tensor = CudaTensor::new(&data, &[2], &[2]); + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + assert_eq!(read_cuda::(cloned.as_ref(), 2), [true, true]); + } + + #[test] + fn clone_empty() { + skip_without_cuda!(); + + // the allocation still has one element, since CUDA does not allow + // zero-sized allocations, but the tensor itself is empty + let tensor = CudaTensor::new(&[3.0_f32], &[0], &[1]); + + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + + assert_eq!(cloned.shape(), [0]); + } + + #[test] + fn check_atomic_types_kernel() { + skip_without_cuda!(); + + let valid_types = ReferenceValue::new( + ArrayD::::from_shape_vec(vec![3], vec![1, 6, 8]).unwrap() + ); + + // all the types are valid + let types = CudaTensor::new(&[1_i32, 6, 6, 8, 1], &[5], &[1]); + check_atomic_types(types.as_ref(), &valid_types).unwrap(); + + // some of the types are invalid + let types = CudaTensor::new(&[1_i32, 3, 8, 3, 4, 1], &[6], &[1]); + let err = check_atomic_types(types.as_ref(), &valid_types).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model does not support the following atomic \ + types which are present in the input systems: 3, 4" + ); + + // non-contiguous types: only every other element is part of the tensor, + // so the invalid types in between should be ignored + let types = CudaTensor::new(&[1_i32, 12, 6, 12, 8, 12], &[3], &[2]); + check_atomic_types(types.as_ref(), &valid_types).unwrap(); + + // non-contiguous types, with an invalid type inside the tensor + let types = CudaTensor::new(&[1_i32, 12, 4, 12, 8, 12], &[3], &[2]); + let err = check_atomic_types(types.as_ref(), &valid_types).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model does not support the following atomic \ + types which are present in the input systems: 4" + ); + + // empty types are always valid + let types = CudaTensor::new(&[42_i32], &[0], &[1]); + check_atomic_types(types.as_ref(), &valid_types).unwrap(); + } +} diff --git a/metatomic-core/src/kernels/cuda_kernels.cu b/metatomic-core/src/kernels/cuda_kernels.cu new file mode 100644 index 000000000..1d74960ed --- /dev/null +++ b/metatomic-core/src/kernels/cuda_kernels.cu @@ -0,0 +1,215 @@ +typedef signed long long i64; +typedef unsigned char u8; +typedef unsigned short u16; +typedef unsigned int u32; +typedef unsigned long long u64; + + +#define MAX_NDIM 7 + +/// Multi-dimensional strided index (up to MAX_NDIM dimensions). +/// Decomposes a flat linear index into multi-dimensional coordinates from the +/// shape, then computes the strided memory offset using the stride array. +/// +/// WARNING: any change here needs to be reflected in the Rust and Metal sources. +struct StridedNDIndex { + i64 ndim; + i64 shape[MAX_NDIM]; + i64 strides[MAX_NDIM]; + + /// Get the offset from the start of the array for a given flat index + __device__ i64 offset(i64 flat_idx) const { + i64 off = 0; + for (int d = this->ndim - 1; d >= 0; d--) { + i64 coord = flat_idx % this->shape[d]; + flat_idx /= this->shape[d]; + off += coord * this->strides[d]; + } + return off; + } +}; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +extern "C" __global__ void is_equal_i32( + const int* values, + StridedNDIndex values_idx, + const int* reference, + StridedNDIndex reference_idx, + i64 n, + int* mismatch +) { + i64 i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + i64 value_offset = values_idx.offset(i); + i64 reference_offset = reference_idx.offset(i); + if (values[value_offset] != reference[reference_offset]) { + atomicMax(mismatch, 1); + } + } +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +__device__ void validate_cell_pbc_impl( + const bool* pbc, + StridedNDIndex pbc_idx, + const T* cell, + StridedNDIndex cell_idx, + int* mismatch_idx +) { + int i = threadIdx.x; + if (i < 3) { + if (!pbc[pbc_idx.offset(i)]) { + if ( + cell[cell_idx.offset(i * 3 + 0)] != T(0) || + cell[cell_idx.offset(i * 3 + 1)] != T(0) || + cell[cell_idx.offset(i * 3 + 2)] != T(0) + ) { + atomicMax(mismatch_idx, i + 1); + } + } + } +} + +extern "C" __global__ void validate_cell_pbc_f32( + const bool* pbc, + StridedNDIndex pbc_idx, + const float* cell, + StridedNDIndex cell_idx, + int* mismatch_idx +) { + validate_cell_pbc_impl(pbc, pbc_idx, cell, cell_idx, mismatch_idx); +} + +extern "C" __global__ void validate_cell_pbc_f64( + const bool* pbc, + StridedNDIndex pbc_idx, + const double* cell, + StridedNDIndex cell_idx, + int* mismatch_idx +) { + validate_cell_pbc_impl(pbc, pbc_idx, cell, cell_idx, mismatch_idx); +} + +//////////////////////////////////////////////////////////////////////////////// + +template +__device__ void scale_inplace_impl( + T* tensor, + StridedNDIndex tensor_idx, + i64 n, + double factor +) { + i64 i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + i64 offset = tensor_idx.offset(i); + tensor[offset] = static_cast(static_cast(tensor[offset]) * factor); + } +} + +extern "C" __global__ void scale_f32( + float* tensor, + StridedNDIndex tensor_idx, + i64 n, + double factor +) { + scale_inplace_impl(tensor, tensor_idx, n, factor); +} + +extern "C" __global__ void scale_f64( + double* tensor, + StridedNDIndex tensor_idx, + i64 n, + double factor +) { + scale_inplace_impl(tensor, tensor_idx, n, factor); +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Copy `n` elements from `src` (which can use arbitrary strides, described by +/// `src_idx`) to `dst`, which must be able to store `n` contiguous elements. +/// +/// The kernels below are instantiated for each element size instead of each +/// data type, since only the size of the elements matters when moving data +/// around. +template +__device__ void copy_to_contiguous_impl( + const T* src, + StridedNDIndex src_idx, + T* dst, + i64 n +) { + i64 i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + dst[i] = src[src_idx.offset(i)]; + } +} + +extern "C" __global__ void copy_to_contiguous_8bit( + const u8* src, + StridedNDIndex src_idx, + u8* dst, + i64 n +) { + copy_to_contiguous_impl(src, src_idx, dst, n); +} + +extern "C" __global__ void copy_to_contiguous_16bit( + const u16* src, + StridedNDIndex src_idx, + u16* dst, + i64 n +) { + copy_to_contiguous_impl(src, src_idx, dst, n); +} + +extern "C" __global__ void copy_to_contiguous_32bit( + const u32* src, + StridedNDIndex src_idx, + u32* dst, + i64 n +) { + copy_to_contiguous_impl(src, src_idx, dst, n); +} + +extern "C" __global__ void copy_to_contiguous_64bit( + const u64* src, + StridedNDIndex src_idx, + u64* dst, + i64 n +) { + copy_to_contiguous_impl(src, src_idx, dst, n); +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +extern "C" __global__ void check_atomic_types( + const int* types, + StridedNDIndex types_idx, + i64 n_atoms, + const int* valid_types, + i64 n_valid_types, + int* invalid_count +) { + i64 i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n_atoms) { + int atom_type = types[types_idx.offset(i)]; + bool found = false; + for (i64 j = 0; j < n_valid_types; j++) { + if (valid_types[j] == atom_type) { + found = true; + break; + } + } + if (!found) { + atomicAdd(invalid_count, 1); + } + } +} diff --git a/metatomic-core/src/kernels/metal.rs b/metatomic-core/src/kernels/metal.rs new file mode 100644 index 000000000..12809e4c8 --- /dev/null +++ b/metatomic-core/src/kernels/metal.rs @@ -0,0 +1,1014 @@ +use std::collections::{HashMap, hash_map::Entry}; +use std::ptr::NonNull; +use std::sync::Mutex; +use std::sync::LazyLock; + +use objc2::rc::Retained; +use objc2::runtime::ProtocolObject; +use objc2_foundation::ns_string; + +use objc2_metal::{ + MTLBuffer, MTLCommandBuffer, MTLCommandEncoder, MTLCommandQueue, + MTLComputeCommandEncoder, MTLComputePipelineState, + MTLCopyAllDevices, MTLCompileOptions, + MTLDevice, MTLLibrary, MTLResourceOptions, MTLSize, +}; + +use dlpk::{DLDevice, DLPackTensor, DLPackTensorRef, DLPackTensorRefMut}; + +use crate::Error; +use super::{ReferenceValue, StridedNDIndex}; + +// Small wrapper around MTLBuffer to implement Send and Sync, since the data is +// read-only after initialization. +pub(crate) struct MetalBuffer(pub(super) Retained>); + +unsafe impl Send for MetalBuffer {} +unsafe impl Sync for MetalBuffer {} + +impl std::ops::Deref for MetalBuffer { + type Target = ProtocolObject; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +/// Get the size of a memory page, which is the alignment required by +/// `newBufferWithBytesNoCopy`. +fn page_size() -> usize { + unsafe extern "C" { + fn getpagesize() -> std::ffi::c_int; + } + + let size = unsafe { getpagesize() }; + return usize::try_from(size).expect("got a negative page size"); +} + +/// A Metal buffer that borrows the lifetime of the data it points to. +/// +/// Created by wrapping an existing memory region (e.g. a DLPack tensor's data) +/// with `newBufferWithBytesNoCopy`, so the buffer does not own the memory and +/// must not outlive it. +/// +/// The buffer starts at the beginning of the memory page containing the data, +/// so [`MetalBufferRef::offset`] must be used when binding it to a kernel. +pub(crate) struct MetalBufferRef<'a> { + buffer: MetalBuffer, + offset: usize, + _phantom: std::marker::PhantomData<&'a [u8]>, +} + +impl std::ops::Deref for MetalBufferRef<'_> { + type Target = MetalBuffer; + fn deref(&self) -> &Self::Target { + &self.buffer + } +} + +impl<'a> MetalBufferRef<'a> { + /// Wrap a DLPack tensor's existing memory in a Metal buffer without copying. + /// + /// Uses `newBufferWithBytesNoCopy:length:options:deallocator:` with no + /// deallocator, since the DLPack tensor (or its owner) retains ownership of + /// the memory. The returned buffer borrows the tensor's lifetime and must + /// not outlive the tensor's backing memory. + /// + /// `newBufferWithBytesNoCopy` requires a page-aligned pointer, but the data + /// of a tensor can start anywhere: Metal itself sub-allocates small buffers + /// inside a single page, and tensors can be views inside a larger + /// allocation. We thus wrap the whole page-aligned memory range containing + /// the data, and the offset of the data inside this range is available with + /// [`MetalBufferRef::offset`]. + pub(crate) fn from_dlpack( + device: &ProtocolObject, + tensor: DLPackTensorRef<'a>, + ) -> Result { + let ptr = dlpack_data_ptr(tensor); + if ptr.is_null() { + return Err(Error::Internal("tensor data pointer is null".into())); + } + + let page_size = page_size(); + let offset = ptr as usize % page_size; + // the length must also be a multiple of the page size + let length = std::cmp::max( + (offset + dlpack_num_bytes(tensor)).next_multiple_of(page_size), + page_size, + ); + + let base = unsafe { ptr.cast::().sub(offset) }; + let nonnull = NonNull::new(base.cast_mut()) + .expect("the start of the page can not be null") + .cast(); + + let buffer = unsafe { + device.newBufferWithBytesNoCopy_length_options_deallocator( + nonnull, + length, + MTLResourceOptions::empty(), + None, + ) + }; + + let buffer = buffer.ok_or_else(|| Error::Internal( + "failed to create Metal buffer from DLPack tensor (newBufferWithBytesNoCopy returned nil)".into() + ))?; + + Ok(Self { + buffer: MetalBuffer(buffer), + offset, + _phantom: std::marker::PhantomData, + }) + } + + /// Offset in bytes of the tensor data inside this buffer + pub(crate) fn offset(&self) -> usize { + self.offset + } +} + +const KERNEL_SRC: &str = include_str!("metal_kernels.metal"); + +/// Cached metal ressources: device, command queue, and pipeline states for kernels. +struct MetalKernelCache { + device: Retained>, + queue: Retained>, + is_equal_i32: Retained>, + validate_cell_pbc_f32: Retained>, + scale_f32: Retained>, + copy_to_contiguous_8bit: Retained>, + copy_to_contiguous_16bit: Retained>, + copy_to_contiguous_32bit: Retained>, + copy_to_contiguous_64bit: Retained>, + check_atomic_types: Retained>, +} + +/// All Metal devices on this system, queried once on first access. +static METAL_DEVICES: LazyLock>>> = LazyLock::new(|| MTLCopyAllDevices().to_vec()); + +impl MetalKernelCache { + fn new(device_id: usize) -> Result { + let device = METAL_DEVICES + .get(device_id) + .ok_or_else(|| Error::Internal(format!("no Metal device with id {device_id}")))? + .clone(); + + let library = device + .newLibraryWithSource_options_error( + ns_string!(KERNEL_SRC), + Some(&MTLCompileOptions::new()), + ) + .map_err(|e| Error::Internal(format!("MSL compile failed: {e}")))?; + + let is_equal_i32 = make_pipeline(&device, &library, "is_equal_i32")?; + let validate_cell_pbc_f32 = make_pipeline(&device, &library, "validate_cell_pbc_f32")?; + let scale_f32 = make_pipeline(&device, &library, "scale_f32")?; + let copy_to_contiguous_8bit = make_pipeline(&device, &library, "copy_to_contiguous_8bit")?; + let copy_to_contiguous_16bit = make_pipeline(&device, &library, "copy_to_contiguous_16bit")?; + let copy_to_contiguous_32bit = make_pipeline(&device, &library, "copy_to_contiguous_32bit")?; + let copy_to_contiguous_64bit = make_pipeline(&device, &library, "copy_to_contiguous_64bit")?; + let check_atomic_types = make_pipeline(&device, &library, "check_atomic_types")?; + + let queue = device + .newCommandQueue() + .ok_or_else(|| Error::Internal("failed to create command queue".into()))?; + + Ok(Self { + device, + queue, + is_equal_i32, + validate_cell_pbc_f32, + scale_f32, + copy_to_contiguous_8bit, + copy_to_contiguous_16bit, + copy_to_contiguous_32bit, + copy_to_contiguous_64bit, + check_atomic_types, + }) + } +} + +fn make_pipeline( + device: &ProtocolObject, + library: &ProtocolObject, + name: &str, +) -> Result>, Error> { + use objc2_foundation::NSString; + + let ns_name = NSString::from_str(name); + let function = library + .newFunctionWithName(&ns_name) + .ok_or_else(|| Error::Internal(format!("get_function({name}): not found")))?; + + device + .newComputePipelineStateWithFunction_error(&function) + .map_err(|e| Error::Internal(format!("pipeline state ({name}): {e}"))) +} + +static METAL_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new())); + +fn get_or_init(cache: &mut HashMap, device_id: usize) -> Result<&MetalKernelCache, Error> { + let entry = match cache.entry(device_id) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => entry.insert(MetalKernelCache::new(device_id)?), + }; + Ok(entry) +} + +fn check_valid_device(function: &str, device: DLDevice) { + assert_eq!( + device.device_type, dlpk::sys::DLDeviceType::kDLMetal, + "{} called on non-metal tensor", function + ); + assert!(device.device_id >= 0, "{} called on invalid device_id", function); +} + +/// Compute the byte span of a DLPack tensor's data (including gaps from +/// strides). +#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] +fn dlpack_num_bytes(tensor: DLPackTensorRef<'_>) -> usize { + let elem_size = tensor.dtype().bits as usize / 8; + let shape = tensor.shape(); + match tensor.strides() { + None => shape.iter().map(|&s| s as usize).product::() * elem_size, + Some(strides) => { + let max_idx: i64 = shape.iter() + .zip(strides.iter()) + .map(|(&s, &st)| (s - 1) * st) + .sum(); + (max_idx as usize + 1) * elem_size + } + } +} + +/// Extract a raw pointer to the tensor's data, accounting for byte_offset. +/// +/// # Safety +/// +/// The returned pointer is only valid as long as the DLPack tensor's backing +/// memory is alive. +#[allow(clippy::cast_possible_truncation)] +fn dlpack_data_ptr(tensor: DLPackTensorRef<'_>) -> *const std::ffi::c_void { + unsafe { + tensor.raw.data.cast::().add(tensor.raw.byte_offset as usize).cast() + } +} + +/// Check that the values of a Metal-resident i32 DLPack tensor match an expected +/// reference array. +#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)] +pub(super) fn is_equal_i32(tensor: DLPackTensorRef<'_>, reference: &ReferenceValue) -> Result { + check_valid_device("is_equal_i32", tensor.device()); + + let device_id = tensor.device().device_id as usize; + let mut lock = METAL_CACHE.lock().expect("failed to lock METAL_CACHE"); + let cache = get_or_init(&mut lock, device_id)?; + + let n_elements: usize = tensor.shape().iter().map(|&s| s as usize).product(); + let ref_bytes = n_elements * std::mem::size_of::(); + + // Build strided index for the values + let values_idx = StridedNDIndex::from_dlpack(tensor); + + // Upload reference values to Metal (cached after first call, per device) + let (ref_buf, reference_idx) = reference.metal_data(device_id, &cache.device)?; + + let values_buf = MetalBufferRef::from_dlpack(&cache.device, tensor)?; + let result_buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::from(&0i32).cast(), + std::mem::size_of::(), + MTLResourceOptions::empty(), + ).expect("failed to create result buffer") + }; + + objc2::rc::autoreleasepool(|_| { + let cmd_buf = cache.queue.commandBuffer().expect("failed to create command buffer"); + let encoder = cmd_buf.computeCommandEncoder().expect("failed to create compute encoder"); + + encoder.setComputePipelineState(&cache.is_equal_i32); + unsafe { + encoder.setBuffer_offset_atIndex(Some(&*values_buf), values_buf.offset(), 0); + + encoder.setBytes_length_atIndex( + NonNull::::from(&values_idx).cast(), + std::mem::size_of::(), + 1, + ); + + encoder.setBuffer_offset_atIndex(Some(&*ref_buf), 0, 2); + + encoder.setBytes_length_atIndex( + NonNull::::from(reference_idx).cast(), + std::mem::size_of::(), + 3, + ); + + encoder.setBytes_length_atIndex( + NonNull::from(&(n_elements as u64)).cast(), + std::mem::size_of::(), + 4, + ); + + encoder.setBuffer_offset_atIndex(Some(&*result_buf), 0, 5); + } + + let tg_size = 32; + let tg_count = n_elements.div_ceil(tg_size); + encoder.dispatchThreadgroups_threadsPerThreadgroup( + MTLSize { width: tg_count, height: 1, depth: 1 }, + MTLSize { width: tg_size, height: 1, depth: 1 }, + ); + encoder.endEncoding(); + cmd_buf.commit(); + cmd_buf.waitUntilCompleted(); + }); + + let result = unsafe { + *result_buf.contents().as_ptr().cast::() + }; + return Ok(result == 0); +} + +/// Validate that cell vectors are zero for non-periodic dimensions on Metal. +#[allow(clippy::cast_sign_loss)] +pub(super) fn validate_cell_pbc( + pbc: DLPackTensorRef<'_>, + cell: DLPackTensorRef<'_>, +) -> Result<(), Error> { + debug_assert_eq!(cell.device(), pbc.device(), "pbc and cell must be on the same device"); + check_valid_device("validate_cell_pbc", cell.device()); + + let device_id = pbc.device().device_id as usize; + let mut lock = METAL_CACHE.lock().expect("failed to lock METAL_CACHE"); + let cache = get_or_init(&mut lock, device_id)?; + + let pbc_idx = StridedNDIndex::from_dlpack(pbc); + let cell_idx = StridedNDIndex::from_dlpack(cell); + + let pbc_buf = MetalBufferRef::from_dlpack(&cache.device, pbc)?; + let cell_buf = MetalBufferRef::from_dlpack(&cache.device, cell)?; + let result_buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::from(&0i32).cast(), + std::mem::size_of::(), + MTLResourceOptions::empty(), + ).expect("failed to create result buffer") + }; + + objc2::rc::autoreleasepool(|_| { + let cmd_buf = cache.queue.commandBuffer().expect("failed to create command buffer"); + let encoder = cmd_buf.computeCommandEncoder().expect("failed to create compute encoder"); + + assert!(cell.dtype().bits == 32, "only float32 is supported on Metal"); + + encoder.setComputePipelineState(&cache.validate_cell_pbc_f32); + unsafe { + encoder.setBuffer_offset_atIndex(Some(&*pbc_buf), pbc_buf.offset(), 0); + + encoder.setBytes_length_atIndex( + NonNull::::from(&pbc_idx).cast(), + std::mem::size_of::(), + 1, + ); + + encoder.setBuffer_offset_atIndex(Some(&*cell_buf), cell_buf.offset(), 2); + + encoder.setBytes_length_atIndex( + NonNull::::from(&cell_idx).cast(), + std::mem::size_of::(), + 3, + ); + + encoder.setBuffer_offset_atIndex(Some(&*result_buf), 0, 4); + } + + encoder.dispatchThreadgroups_threadsPerThreadgroup( + MTLSize { width: 1, height: 1, depth: 1 }, + MTLSize { width: 3, height: 1, depth: 1 }, + ); + encoder.endEncoding(); + cmd_buf.commit(); + cmd_buf.waitUntilCompleted(); + }); + + let result = unsafe { + *result_buf.contents().as_ptr().cast::() + }; + + if result != 0 { + let dim = result - 1; + return Err(Error::InvalidParameter(format!( + "invalid cell: for non-periodic dimensions, the corresponding \ + cell vector must be zero, but cell[{}] contains non-zero values", + dim + ))); + } + Ok(()) +} + +/// Scale all elements of `tensor` in place by `factor`, on Metal device. +/// +/// Only 32-bit floating point tensors are supported on Metal. +#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)] +pub(super) fn scale_inplace( + tensor: DLPackTensorRefMut<'_>, + factor: f64, +) -> Result<(), Error> { + check_valid_device("scale_inplace", tensor.device()); + + let device_id = tensor.device().device_id as usize; + let mut lock = METAL_CACHE.lock().expect("failed to lock METAL_CACHE"); + let cache = get_or_init(&mut lock, device_id)?; + + let dtype = tensor.dtype(); + if dtype.code != dlpk::sys::DLDataTypeCode::kDLFloat || dtype.bits != 32 { + return Err(Error::InvalidParameter(format!( + "scale_inplace on Metal only supports 32-bit floats, got {}-bit {:?}", + dtype.bits, dtype.code + ))); + } + + let n_elements: usize = tensor.shape().iter().map(|&s| s as usize).product(); + if n_elements == 0 { + return Ok(()); + } + + let tensor_idx = StridedNDIndex::from_dlpack(tensor.as_ref()); + let factor_f32 = factor as f32; + + let tensor_buf = MetalBufferRef::from_dlpack(&cache.device, tensor.as_ref())?; + + objc2::rc::autoreleasepool(|_| { + let cmd_buf = cache.queue.commandBuffer().expect("failed to create command buffer"); + let encoder = cmd_buf.computeCommandEncoder().expect("failed to create compute encoder"); + + encoder.setComputePipelineState(&cache.scale_f32); + unsafe { + encoder.setBuffer_offset_atIndex(Some(&*tensor_buf), tensor_buf.offset(), 0); + + encoder.setBytes_length_atIndex( + NonNull::::from(&tensor_idx).cast(), + std::mem::size_of::(), + 1, + ); + + encoder.setBytes_length_atIndex( + NonNull::from(&(n_elements as u64)).cast(), + std::mem::size_of::(), + 2, + ); + + encoder.setBytes_length_atIndex( + NonNull::from(&factor_f32).cast(), + std::mem::size_of::(), + 3, + ); + } + + let tg_size = 32; + let tg_count = n_elements.div_ceil(tg_size); + encoder.dispatchThreadgroups_threadsPerThreadgroup( + MTLSize { width: tg_count, height: 1, depth: 1 }, + MTLSize { width: tg_size, height: 1, depth: 1 }, + ); + encoder.endEncoding(); + cmd_buf.commit(); + cmd_buf.waitUntilCompleted(); + }); + + return Ok(()); +} + +/// Check that all atomic types in `types` are present in `valid_types`, on Metal. +/// +/// The check runs on-device. If invalid types are found (count > 0), a CPU +/// fallback scan identifies the specific invalid type for the error message. +#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)] +pub(super) fn check_atomic_types( + types: DLPackTensorRef<'_>, + valid_types: &ReferenceValue, +) -> Result<(), Error> { + check_valid_device("check_atomic_types", types.device()); + assert!( + valid_types.cpu.is_standard_layout(), + "valid_types reference must be C-contiguous" + ); + assert_eq!( + types.n_dims(), 1, + "check_atomic_types expects a 1D types tensor" + ); + + let device_id = types.device().device_id as usize; + let mut lock = METAL_CACHE.lock().expect("failed to lock METAL_CACHE"); + let cache = get_or_init(&mut lock, device_id)?; + + let n_atoms: usize = types.shape().iter().map(|&s| s as usize).product(); + if n_atoms == 0 { + return Ok(()); + } + + let types_idx = StridedNDIndex::from_dlpack(types); + + // Upload valid types to Metal (cached after first call, per device) + let (valid_types_buffer, _) = valid_types.metal_data(device_id, &cache.device)?; + let n_valid_types = valid_types.cpu.len() as u64; + + let types_buf = MetalBufferRef::from_dlpack(&cache.device, types)?; + let result_buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::from(&0i32).cast(), + std::mem::size_of::(), + MTLResourceOptions::empty(), + ).expect("failed to create result buffer") + }; + + objc2::rc::autoreleasepool(|_| { + let cmd_buf = cache.queue.commandBuffer().expect("failed to create command buffer"); + let encoder = cmd_buf.computeCommandEncoder().expect("failed to create compute encoder"); + + encoder.setComputePipelineState(&cache.check_atomic_types); + unsafe { + encoder.setBuffer_offset_atIndex(Some(&*types_buf), 0, 0); + + encoder.setBytes_length_atIndex( + NonNull::::from(&types_idx).cast(), + std::mem::size_of::(), + 1, + ); + + encoder.setBytes_length_atIndex( + NonNull::from(&(n_atoms as u64)).cast(), + std::mem::size_of::(), + 2, + ); + + encoder.setBuffer_offset_atIndex(Some(&*valid_types_buffer), 0, 3); + + encoder.setBytes_length_atIndex( + NonNull::from(&n_valid_types).cast(), + std::mem::size_of::(), + 4, + ); + + encoder.setBuffer_offset_atIndex(Some(&*result_buf), 0, 5); + } + + let tg_size = 32; + let tg_count = n_atoms.div_ceil(tg_size); + encoder.dispatchThreadgroups_threadsPerThreadgroup( + MTLSize { width: tg_count, height: 1, depth: 1 }, + MTLSize { width: tg_size, height: 1, depth: 1 }, + ); + encoder.endEncoding(); + cmd_buf.commit(); + cmd_buf.waitUntilCompleted(); + }); + + let result = unsafe { + *result_buf.contents().as_ptr().cast::() + }; + + if result > 0 { + // Invalid types found — read types from the Metal buffer and scan on CPU. + let n_bytes = dlpack_num_bytes(types); + let n_elements = n_bytes / std::mem::size_of::(); + let host_types: Vec = unsafe { + std::slice::from_raw_parts( + types_buf.contents().as_ptr().cast::(), + n_elements, + ).to_vec() + }; + super::cpu::check_atomic_types_buffer(&host_types, &types_idx, n_atoms, valid_types)?; + } + + Ok(()) +} + +/// Context held by the deleter of a cloned Metal `DLManagedTensorVersioned`. +struct MetalCloneContext { + buffer: Retained>, + shape: Vec, + strides: Vec, +} + +unsafe impl Send for MetalCloneContext {} +unsafe impl Sync for MetalCloneContext {} + +/// Deleter for a cloned Metal DLPack tensor. +/// +/// Drops the context (which releases the MTLBuffer) and the boxed +/// `DLManagedTensorVersioned`. +unsafe extern "C" fn metal_clone_deleter(tensor: *mut dlpk::sys::DLManagedTensorVersioned) { + unsafe { + let ctx = (*tensor).manager_ctx.cast::(); + let _ = Box::from_raw(ctx); + let _ = Box::from_raw(tensor); + } +} + +/// Clone a DLPack tensor on Metal, copying the underlying device memory. +/// +/// The returned `DLPackTensor` owns its own Metal buffer and is independent of +/// the original tensor. The clone is always C-contiguous, even when the +/// original tensor is not: the data is gathered with the `copy_to_contiguous` +/// kernel instead of a plain buffer copy. +#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub(super) fn clone_tensor(tensor: &DLPackTensorRef<'_>) -> Result { + check_valid_device("clone_tensor", tensor.device()); + + let device_id = tensor.device().device_id as usize; + let mut lock = METAL_CACHE.lock().expect("failed to lock METAL_CACHE"); + let cache = get_or_init(&mut lock, device_id)?; + + let element_size_bits = super::element_size(tensor.dtype())?; + let n_elements: usize = tensor.shape().iter().map(|&s| s as usize).product(); + let num_bytes = n_elements * element_size_bits / 8; + + let shape: Vec = tensor.shape().to_vec(); + // the clone stores the data contiguously, regardless of the strides used by + // the original tensor + let strides = super::contiguous_strides(&shape); + + // pick the kernel matching the element size before allocating anything + let pipeline = match element_size_bits { + 8 => &cache.copy_to_contiguous_8bit, + 16 => &cache.copy_to_contiguous_16bit, + 32 => &cache.copy_to_contiguous_32bit, + 64 => &cache.copy_to_contiguous_64bit, + _ => { + return Err(Error::InvalidParameter(format!( + "clone_tensor does not support {} tensors on Metal", + tensor.dtype() + ))); + } + }; + + // allocate a new buffer, only big enough for the contiguous data. Metal + // does not allow zero-sized buffers, so we always allocate at least one + // byte (which is never read, since the tensor is then empty). + let buffer = cache.device + .newBufferWithLength_options(std::cmp::max(num_bytes, 1), MTLResourceOptions::empty()) + .ok_or_else(|| Error::Internal("failed to allocate Metal buffer for the clone".into()))?; + + if n_elements > 0 { + // gather the (possibly strided) data from the original tensor into the + // contiguous allocation + let src_idx = StridedNDIndex::from_dlpack(*tensor); + let src_buf = MetalBufferRef::from_dlpack(&cache.device, *tensor)?; + + objc2::rc::autoreleasepool(|_| { + let cmd_buf = cache.queue.commandBuffer().expect("failed to create command buffer"); + let encoder = cmd_buf.computeCommandEncoder().expect("failed to create compute encoder"); + + encoder.setComputePipelineState(pipeline); + unsafe { + encoder.setBuffer_offset_atIndex(Some(&*src_buf), src_buf.offset(), 0); + + encoder.setBytes_length_atIndex( + NonNull::::from(&src_idx).cast(), + std::mem::size_of::(), + 1, + ); + + encoder.setBuffer_offset_atIndex(Some(&*buffer), 0, 2); + + encoder.setBytes_length_atIndex( + NonNull::from(&(n_elements as u64)).cast(), + std::mem::size_of::(), + 3, + ); + } + + let tg_size = 32; + let tg_count = n_elements.div_ceil(tg_size); + encoder.dispatchThreadgroups_threadsPerThreadgroup( + MTLSize { width: tg_count, height: 1, depth: 1 }, + MTLSize { width: tg_size, height: 1, depth: 1 }, + ); + encoder.endEncoding(); + cmd_buf.commit(); + cmd_buf.waitUntilCompleted(); + }); + } + + let ctx = Box::new(MetalCloneContext { + buffer, + shape: shape, + strides: strides, + }); + + let ndim = ctx.shape.len() as i32; + let data_ptr = ctx.buffer.contents().as_ptr(); + + let dl_tensor = dlpk::sys::DLTensor { + data: data_ptr.cast::(), + device: tensor.device(), + ndim, + dtype: tensor.dtype(), + shape: ctx.shape.as_ptr().cast_mut(), + strides: ctx.strides.as_ptr().cast_mut(), + byte_offset: 0, + }; + + let managed = Box::new(dlpk::sys::DLManagedTensorVersioned { + version: dlpk::sys::DLPackVersion::current(), + manager_ctx: Box::into_raw(ctx).cast(), + deleter: Some(metal_clone_deleter), + flags: dlpk::sys::DLPACK_FLAG_BITMASK_IS_COPIED, + dl_tensor, + }); + + let ptr = Box::into_raw(managed); + Ok(unsafe { DLPackTensor::from_ptr(ptr) }) +} + + +#[cfg(test)] +mod tests { + use super::*; + + use dlpk::{DLDevice, GetDLPackDataType}; + use ndarray::ArrayD; + + /// A DLPack tensor with Metal-resident data, used to test the kernels above. + struct MetalTensor { + buffer: Retained>, + shape: Vec, + strides: Vec, + dtype: dlpk::sys::DLDataType, + } + + impl MetalTensor { + /// Create a new Metal tensor with the given `shape` and `strides`, + /// containing a copy of `data`. + /// + /// `data` is the full memory span of the tensor, including any gap + /// between the elements actually part of the tensor. + fn new(data: &[T], shape: &[i64], strides: &[i64]) -> Self { + let device = METAL_DEVICES.first().expect("no Metal device available"); + + // Metal does not allow zero-sized buffers + assert!(!data.is_empty()); + let buffer = device + .newBufferWithLength_options(std::mem::size_of_val(data), MTLResourceOptions::empty()) + .expect("failed to allocate Metal buffer"); + + unsafe { + std::ptr::copy_nonoverlapping( + data.as_ptr(), + buffer.contents().as_ptr().cast::(), + data.len(), + ); + } + + MetalTensor { + buffer, + shape: shape.to_vec(), + strides: strides.to_vec(), + dtype: T::get_dlpack_data_type(), + } + } + + #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] + fn dl_tensor(&self) -> dlpk::sys::DLTensor { + dlpk::sys::DLTensor { + data: self.buffer.contents().as_ptr(), + device: DLDevice { + device_type: dlpk::sys::DLDeviceType::kDLMetal, + device_id: 0, + }, + ndim: self.shape.len() as i32, + dtype: self.dtype, + shape: self.shape.as_ptr().cast_mut(), + strides: self.strides.as_ptr().cast_mut(), + byte_offset: 0, + } + } + + fn as_ref(&self) -> DLPackTensorRef<'_> { + unsafe { DLPackTensorRef::from_raw(self.dl_tensor()) } + } + + fn as_mut(&mut self) -> DLPackTensorRefMut<'_> { + unsafe { DLPackTensorRefMut::from_raw(self.dl_tensor()) } + } + + /// Read the first `n` elements of this tensor's memory span + fn data(&self, n: usize) -> Vec { + return read_metal(self.as_ref(), n); + } + + /// Overwrite the data in this tensor's Metal buffer + fn overwrite(&self, data: &[T]) { + unsafe { + std::ptr::copy_nonoverlapping( + data.as_ptr(), + self.buffer.contents().as_ptr().cast::(), + data.len(), + ); + } + } + } + + /// Read the first `n` elements of the data of any Metal-resident tensor + fn read_metal(tensor: DLPackTensorRef<'_>, n: usize) -> Vec { + unsafe { + std::slice::from_raw_parts(dlpack_data_ptr(tensor).cast::(), n).to_vec() + } + } + + #[test] + fn is_equal_i32_kernel() { + let reference = ReferenceValue::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![0, 1, 2]).unwrap() + ); + + // matching values + let tensor = MetalTensor::new(&[0_i32, 1, 2], &[3, 1], &[1, 1]); + assert!(is_equal_i32(tensor.as_ref(), &reference).unwrap()); + + // mismatching values + let tensor = MetalTensor::new(&[0_i32, 42, 2], &[3, 1], &[1, 1]); + assert!(!is_equal_i32(tensor.as_ref(), &reference).unwrap()); + + // matching values in a non-contiguous tensor: every other element of + // [0, -1, 1, -1, 2, -1] + let tensor = MetalTensor::new(&[0_i32, -1, 1, -1, 2, -1], &[3, 1], &[2, 1]); + assert!(is_equal_i32(tensor.as_ref(), &reference).unwrap()); + } + + #[test] + fn validate_cell_pbc_kernel() { + // fully periodic: any cell is valid + let pbc = MetalTensor::new(&[true, true, true], &[3], &[1]); + let cell = MetalTensor::new( + &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], &[3, 3], &[3, 1] + ); + validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap(); + + // non-periodic dimension with a zero cell vector: valid + let pbc = MetalTensor::new(&[true, false, true], &[3], &[1]); + let cell = MetalTensor::new( + &[10.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 10.0], &[3, 3], &[3, 1] + ); + validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap(); + + // non-periodic dimension with a non-zero cell vector: invalid + let cell = MetalTensor::new( + &[10.0_f32, 0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 10.0], &[3, 3], &[3, 1] + ); + let err = validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap_err(); + assert!(err.to_string().contains("cell[1] contains non-zero values"), "{err}"); + + // the last dimension being non-periodic is reported correctly too + let pbc = MetalTensor::new(&[true, true, false], &[3], &[1]); + let cell = MetalTensor::new( + &[10.0_f32, 0.0, 0.0, 0.0, 10.0, 0.0, 3.0, 0.0, 10.0], &[3, 3], &[3, 1] + ); + let err = validate_cell_pbc(pbc.as_ref(), cell.as_ref()).unwrap_err(); + assert!(err.to_string().contains("cell[2] contains non-zero values"), "{err}"); + } + + #[test] + fn scale_inplace_kernel() { + let mut tensor = MetalTensor::new(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3], &[3, 1]); + scale_inplace(tensor.as_mut(), 2.5).unwrap(); + assert_eq!(tensor.data::(6), [2.5, 5.0, 7.5, 10.0, 12.5, 15.0]); + + // non-contiguous tensor: only the 2x2 block in the top left corner of + // this 3x4 array is scaled, the rest of the data is left alone + let data: Vec = (0..12_i16).map(f32::from).collect(); + let mut tensor = MetalTensor::new(&data, &[2, 2], &[4, 1]); + scale_inplace(tensor.as_mut(), 10.0).unwrap(); + assert_eq!(tensor.data::(12), [ + 0.0, 10.0, 2.0, 3.0, + 40.0, 50.0, 6.0, 7.0, + 8.0, 9.0, 10.0, 11.0, + ]); + + // empty tensors are left alone (the buffer still has one element, + // since Metal does not allow zero-sized buffers, and it should not be + // touched by the kernel) + let mut tensor = MetalTensor::new(&[3.0_f32], &[0], &[1]); + scale_inplace(tensor.as_mut(), 2.0).unwrap(); + assert_eq!(tensor.data::(1), [3.0]); + + // only 32-bit floats are supported on Metal + let mut tensor = MetalTensor::new(&[1.0_f64, 2.0], &[2], &[1]); + let err = scale_inplace(tensor.as_mut(), 2.0).unwrap_err(); + assert!(err.to_string().contains("only supports 32-bit floats"), "{err}"); + } + + #[test] + fn clone_contiguous() { + let data = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + let tensor = MetalTensor::new(&data, &[2, 3], &[3, 1]); + + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + + assert_eq!(cloned.shape(), [2, 3]); + assert_eq!(cloned.strides(), Some(&[3, 1][..])); + assert_eq!(read_metal::(cloned.as_ref(), 6), data); + + // the clone is independent from the original + tensor.overwrite(&[42.0_f32; 6]); + assert_eq!(read_metal::(cloned.as_ref(), 6), data); + } + + #[test] + fn clone_non_contiguous() { + // 2x2 block in the top left corner of a 3x4 array + let data: Vec = (0..12_i16).map(f32::from).collect(); + let tensor = MetalTensor::new(&data, &[2, 2], &[4, 1]); + + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + + assert_eq!(cloned.shape(), [2, 2]); + assert_eq!(cloned.strides(), Some(&[2, 1][..])); + assert_eq!(read_metal::(cloned.as_ref(), 4), [0.0, 1.0, 4.0, 5.0]); + } + + #[test] + fn clone_transposed() { + // 2x3 array in column-major order (i.e. the transpose of a 3x2 array) + let data: Vec = (0..6).collect(); + let tensor = MetalTensor::new(&data, &[2, 3], &[1, 2]); + + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + + assert_eq!(cloned.shape(), [2, 3]); + assert_eq!(cloned.strides(), Some(&[3, 1][..])); + assert_eq!(read_metal::(cloned.as_ref(), 6), [0, 2, 4, 1, 3, 5]); + } + + #[test] + fn clone_element_sizes() { + // 8-bit elements, every other one + let data: Vec = (0..6).collect(); + let tensor = MetalTensor::new(&data, &[3], &[2]); + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + assert_eq!(read_metal::(cloned.as_ref(), 3), [0, 2, 4]); + + // 16-bit elements, every other one + let data: Vec = (0..6).collect(); + let tensor = MetalTensor::new(&data, &[3], &[2]); + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + assert_eq!(read_metal::(cloned.as_ref(), 3), [0, 2, 4]); + + // bool elements + let data = vec![true, false, true, true]; + let tensor = MetalTensor::new(&data, &[2], &[2]); + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + assert_eq!(read_metal::(cloned.as_ref(), 2), [true, true]); + } + + #[test] + fn clone_empty() { + // the buffer still has one element, since Metal does not allow + // zero-sized buffers, but the tensor itself is empty + let tensor = MetalTensor::new(&[3.0_f32], &[0], &[1]); + + let cloned = clone_tensor(&tensor.as_ref()).unwrap(); + + assert_eq!(cloned.shape(), [0]); + } + + #[test] + fn check_atomic_types_kernel() { + let valid_types = ReferenceValue::new( + ArrayD::::from_shape_vec(vec![3], vec![1, 6, 8]).unwrap() + ); + + // all the types are valid + let types = MetalTensor::new(&[1_i32, 6, 6, 8, 1], &[5], &[1]); + check_atomic_types(types.as_ref(), &valid_types).unwrap(); + + // some of the types are invalid + let types = MetalTensor::new(&[1_i32, 3, 8, 3, 4, 1], &[6], &[1]); + let err = check_atomic_types(types.as_ref(), &valid_types).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model does not support the following atomic \ + types which are present in the input systems: 3, 4" + ); + + // non-contiguous types: only every other element is part of the tensor, + // so the invalid types in between should be ignored + let types = MetalTensor::new(&[1_i32, 12, 6, 12, 8, 12], &[3], &[2]); + check_atomic_types(types.as_ref(), &valid_types).unwrap(); + + // non-contiguous types, with an invalid type inside the tensor + let types = MetalTensor::new(&[1_i32, 12, 4, 12, 8, 12], &[3], &[2]); + let err = check_atomic_types(types.as_ref(), &valid_types).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model does not support the following atomic \ + types which are present in the input systems: 4" + ); + + // empty types are always valid + let types = MetalTensor::new(&[42_i32], &[0], &[1]); + check_atomic_types(types.as_ref(), &valid_types).unwrap(); + } +} diff --git a/metatomic-core/src/kernels/metal_kernels.metal b/metatomic-core/src/kernels/metal_kernels.metal new file mode 100644 index 000000000..44e75d177 --- /dev/null +++ b/metatomic-core/src/kernels/metal_kernels.metal @@ -0,0 +1,191 @@ +#include +using namespace metal; + +// --------------------------------------------------------------------------- +// Multi-dimensional strided index helper (up to MAX_NDIM dimensions). +// +// Decomposes a flat linear index into multi-dimensional coordinates based on +// the shape and then computes the strided memory offset using the stride +// array. +// +// WARNING: the layout of this struct must match both the CUDA +// (cuda_kernels.cu) and Rust (kernels/mod.rs) definitions. +// --------------------------------------------------------------------------- +constant long MAX_NDIM [[maybe_unused]] = 7; + +struct StridedNDIndex { + long ndim; + long shape[MAX_NDIM]; + long strides[MAX_NDIM]; +}; + +/// Get the offset from the start of the array for a given flat index. +/// +/// This is a free function instead of a member function of `StridedNDIndex`, +/// since MSL does not allow calling member functions on objects living in the +/// `constant` address space. +static long strided_offset(constant StridedNDIndex& index, long flat_idx) { + long off = 0; + for (int d = index.ndim - 1; d >= 0; d--) { + long coord = flat_idx % index.shape[d]; + flat_idx /= index.shape[d]; + off += coord * index.strides[d]; + } + return off; +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +kernel void is_equal_i32( + [[buffer(0)]] device const int* values, + [[buffer(1)]] constant StridedNDIndex& values_idx, + [[buffer(2)]] device const int* reference, + [[buffer(3)]] constant StridedNDIndex& reference_idx, + [[buffer(4)]] constant uint64_t& n, + [[buffer(5)]] device atomic_int* mismatch, + [[thread_position_in_grid]] uint gid +) { + if (gid < n) { + long v_off = strided_offset(values_idx, gid); + long r_off = strided_offset(reference_idx, gid); + if (values[v_off] != reference[r_off]) { + atomic_fetch_max_explicit(mismatch, 1, memory_order_relaxed); + } + } +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Validate cell vectors against PBC flags (f32 only on Metal). +kernel void validate_cell_pbc_f32( + [[buffer(0)]] device const bool* pbc, + [[buffer(1)]] constant StridedNDIndex& pbc_idx, + [[buffer(2)]] device const float* cell, + [[buffer(3)]] constant StridedNDIndex& cell_idx, + [[buffer(4)]] device atomic_int* mismatch_idx, + [[thread_position_in_threadgroup]] uint tid +) { + if (tid < 3) { + if (!pbc[strided_offset(pbc_idx, tid)]) { + if ( + cell[strided_offset(cell_idx, tid * 3 + 0)] != 0.0f || + cell[strided_offset(cell_idx, tid * 3 + 1)] != 0.0f || + cell[strided_offset(cell_idx, tid * 3 + 2)] != 0.0f + ) { + atomic_fetch_max_explicit(mismatch_idx, int(tid + 1), memory_order_relaxed); + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Scale all elements of a tensor in place by `factor` (f32 only on Metal). +kernel void scale_f32( + [[buffer(0)]] device float* tensor, + [[buffer(1)]] constant StridedNDIndex& tensor_idx, + [[buffer(2)]] constant uint64_t& n, + [[buffer(3)]] constant float& factor, + [[thread_position_in_grid]] uint gid +) { + if (gid < n) { + long offset = strided_offset(tensor_idx, gid); + tensor[offset] = tensor[offset] * factor; + } +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Copy `n` elements from `src` (which can use arbitrary strides, described by +/// `src_idx`) to `dst`, which must be able to store `n` contiguous elements. +/// +/// The kernels below are instantiated for each element size instead of each +/// data type, since only the size of the elements matters when moving data +/// around. +template +static void copy_to_contiguous_impl( + device const T* src, + constant StridedNDIndex& src_idx, + device T* dst, + uint64_t n, + uint gid +) { + if (gid < n) { + dst[gid] = src[strided_offset(src_idx, gid)]; + } +} + +kernel void copy_to_contiguous_8bit( + [[buffer(0)]] device const uint8_t* src, + [[buffer(1)]] constant StridedNDIndex& src_idx, + [[buffer(2)]] device uint8_t* dst, + [[buffer(3)]] constant uint64_t& n, + [[thread_position_in_grid]] uint gid +) { + copy_to_contiguous_impl(src, src_idx, dst, n, gid); +} + +kernel void copy_to_contiguous_16bit( + [[buffer(0)]] device const uint16_t* src, + [[buffer(1)]] constant StridedNDIndex& src_idx, + [[buffer(2)]] device uint16_t* dst, + [[buffer(3)]] constant uint64_t& n, + [[thread_position_in_grid]] uint gid +) { + copy_to_contiguous_impl(src, src_idx, dst, n, gid); +} + +kernel void copy_to_contiguous_32bit( + [[buffer(0)]] device const uint32_t* src, + [[buffer(1)]] constant StridedNDIndex& src_idx, + [[buffer(2)]] device uint32_t* dst, + [[buffer(3)]] constant uint64_t& n, + [[thread_position_in_grid]] uint gid +) { + copy_to_contiguous_impl(src, src_idx, dst, n, gid); +} + +kernel void copy_to_contiguous_64bit( + [[buffer(0)]] device const uint64_t* src, + [[buffer(1)]] constant StridedNDIndex& src_idx, + [[buffer(2)]] device uint64_t* dst, + [[buffer(3)]] constant uint64_t& n, + [[thread_position_in_grid]] uint gid +) { + copy_to_contiguous_impl(src, src_idx, dst, n, gid); +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + + +/// Check that all atomic types in `types` are present in `valid_types`. +/// Increments `invalid_count` for each invalid type found. +/// `invalid_count` is initialized to 0 by the caller. +kernel void check_atomic_types( + [[buffer(0)]] device const int* types, + [[buffer(1)]] constant StridedNDIndex& types_idx, + [[buffer(2)]] constant uint64_t& n_atoms, + [[buffer(3)]] device const int* valid_types, + [[buffer(4)]] constant uint64_t& n_valid_types, + [[buffer(5)]] device atomic_int* invalid_count, + [[thread_position_in_grid]] uint gid +) { + if (gid < n_atoms) { + int atom_type = types[strided_offset(types_idx, gid)]; + bool found = false; + for (uint64_t j = 0; j < n_valid_types; j++) { + if (valid_types[j] == atom_type) { + found = true; + break; + } + } + if (!found) { + atomic_fetch_add_explicit(invalid_count, 1, memory_order_relaxed); + } + } +} diff --git a/metatomic-core/src/kernels/mod.rs b/metatomic-core/src/kernels/mod.rs new file mode 100644 index 000000000..1e7565540 --- /dev/null +++ b/metatomic-core/src/kernels/mod.rs @@ -0,0 +1,420 @@ +use std::sync::{Arc, OnceLock}; + +use cudarc::driver::safe::{CudaContext, CudaStream, DeviceRepr}; +use cudarc::driver::CudaSlice; +use dlpk::sys::DLDeviceType; +use dlpk::{DLPackTensor, DLPackTensorRef, DLPackTensorRefMut}; +use ndarray::{ArrayD, ArrayViewD}; + +use crate::Error; + +mod cpu; +mod cuda; + +#[cfg(target_os = "macos")] +mod metal; + +const MAX_NDIM: usize = 7; + +/// Multi-dimensional strided index (up to MAX_NDIM dimensions). +/// +/// Decomposes a flat linear index into multi-dimensional coordinates from the +/// shape, then computes the strided memory offset using the stride array. +/// +/// WARNING: any change here needs to be reflected in the CUDA and Metal sources. +#[repr(C)] +pub(crate) struct StridedNDIndex { + pub(crate) ndim: i64, + pub(crate) shape: [i64; MAX_NDIM], + pub(crate) strides: [i64; MAX_NDIM], +} + +#[allow(clippy::cast_possible_wrap)] +impl StridedNDIndex { + /// Create a `StridedNDIndex` from a DLPack tensor's shape and strides. + pub(crate) fn from_dlpack(tensor: DLPackTensorRef<'_>) -> Self { + Self::from_shape_strides(tensor.shape(), tensor.strides()) + } + + /// Create a `StridedNDIndex` from an ndarray view's shape and strides. + pub(crate) fn from_ndarray(array: &ArrayViewD<'_, T>) -> Self { + let shape: Vec = array.shape().iter().map(|&s| s as i64).collect(); + let strides: Vec = array.strides().iter().map(|&s| s as i64).collect(); + Self::from_shape_strides(&shape, Some(&strides)) + } + + /// Create a `StridedNDIndex` from shape and optional strides. + /// + /// If strides is `None`, the strides are computed as if the array were + /// contiguous (row-major / C-contiguous). + pub(crate) fn from_shape_strides(shape: &[i64], strides: Option<&[i64]>) -> Self { + let ndim = shape.len(); + assert!( + ndim <= MAX_NDIM, + "StridedNDIndex only supports up to {MAX_NDIM} dimensions, got {ndim}" + ); + let mut shape_arr = [0i64; MAX_NDIM]; + let mut strides_arr = [0i64; MAX_NDIM]; + + // Contiguous fallback strides (row-major / C-contiguous) + let mut acc: i64 = 1; + for i in (0..ndim).rev() { + shape_arr[i] = shape[i]; + strides_arr[i] = acc; + acc *= shape[i]; + } + + if let Some(strides) = strides { + strides_arr[..ndim].copy_from_slice(&strides[..ndim]); + } + StridedNDIndex { ndim: ndim as i64, shape: shape_arr, strides: strides_arr } + } + + /// Compute the strided memory offset for a given flat index. + pub(crate) fn offset(&self, flat_idx: i64) -> i64 { + let mut off = 0i64; + let mut idx = flat_idx; + let ndim = usize::try_from(self.ndim).expect("ndim must be >=0"); + for d in (0..ndim).rev() { + let coord = idx % self.shape[d]; + idx /= self.shape[d]; + off += coord * self.strides[d]; + } + off + } +} + +/// Compute C-contiguous (row-major) strides for the given `shape`. +/// +/// The returned strides are in number of elements, as expected by DLPack. +pub(crate) fn contiguous_strides(shape: &[i64]) -> Vec { + let mut strides = vec![0i64; shape.len()]; + let mut acc: i64 = 1; + for i in (0..shape.len()).rev() { + strides[i] = acc; + acc *= shape[i]; + } + return strides; +} + +/// Get the size in bits of a single element of a tensor with this `dtype`. +pub(crate) fn element_size(dtype: dlpk::sys::DLDataType) -> Result { + if dtype.lanes != 1 { + return Err(Error::InvalidParameter(format!( + "vector data types are not supported, got {dtype} with {} lanes", + dtype.lanes + ))); + } + + if dtype.bits == 0 || !dtype.bits.is_multiple_of(8) { + return Err(Error::InvalidParameter(format!( + "only data types with a whole number of bytes are supported, got {dtype}" + ))); + } + + return Ok(dtype.bits as usize); +} + +type CudaArray = (CudaSlice, StridedNDIndex); +#[cfg(target_os = "macos")] +type MetalArray = (metal::MetalBuffer, StridedNDIndex); + +/// Store and cache reference values for different backends (CPU, CUDA, Metal). +/// +/// The CPU copy is always present. Device-resident copies are lazily uploaded +/// on first use per device: the outer `OnceLock` initializes a `Vec` with one +/// entry per device (sized from the device count), and each inner `OnceLock` +/// is independently initialized on first access for that specific device. +pub struct ReferenceValue { + /// The reference values stored on the CPU, always there + pub(crate) cpu: ArrayD, + /// Reference values stored on CUDA, one `OnceLock` per device (lazily sized) + pub(crate) cuda: OnceLock>>>, + #[cfg(target_os = "macos")] + /// Reference values stored on Metal, one `OnceLock` per device (lazily sized) + pub(crate) metal: OnceLock>>, +} + +impl std::fmt::Debug for ReferenceValue { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ReferenceValue") + .field("cpu", &self.cpu) + .finish() + } +} + +impl ReferenceValue { + pub(crate) fn new(cpu: ArrayD) -> Self { + Self { + cpu, + cuda: OnceLock::new(), + #[cfg(target_os = "macos")] + metal: OnceLock::new(), + } + } +} + +impl ReferenceValue { + /// Get the CUDA-resident copy of the reference values for `device_id`, + /// uploading from CPU on first use for this device. + /// + /// The returned reference is tied to `&self` and valid for the lifetime of + /// this `ReferenceValue`. + pub(crate) fn cuda_data( + &self, + device_id: usize, + stream: &Arc, + ) -> Result<&(CudaSlice, StridedNDIndex), Error> { + let entries = self.cuda.get_or_init(|| { + let count = usize::try_from(CudaContext::device_count().unwrap_or(0)).expect("got negative device count"); + (0..count).map(|_| OnceLock::new()).collect() + }); + + if device_id >= entries.len() { + return Err(Error::Internal(format!( + "CUDA device {device_id} does not exist (only {} devices available)", + entries.len() + ))); + } + + Ok(entries[device_id].get_or_init(|| { + let slice = stream + .clone_htod(self.cpu.as_slice().expect("reference should be contiguous")) + .expect("clone_htod reference failed"); + let idx = StridedNDIndex::from_ndarray(&self.cpu.view()); + (slice, idx) + })) + } +} + +#[cfg(target_os = "macos")] +impl ReferenceValue { + /// Get the Metal-resident copy of the reference values for `device_id`, + /// uploading from CPU on first use for this device. + /// + /// The returned reference is tied to `&self` and valid for the lifetime of + /// this `ReferenceValue`. + pub(crate) fn metal_data( + &self, + device_id: usize, + device: &objc2::runtime::ProtocolObject, + ) -> Result<&(metal::MetalBuffer, StridedNDIndex), Error> { + use objc2_metal::{MTLCopyAllDevices, MTLDevice}; + + let entries = self.metal.get_or_init(|| { + let count = MTLCopyAllDevices().count(); + (0..count).map(|_| OnceLock::new()).collect() + }); + + if device_id >= entries.len() { + return Err(Error::Internal(format!( + "Metal device {device_id} does not exist (only {} devices available)", + entries.len() + ))); + } + + Ok(entries[device_id].get_or_init(|| { + let ref_bytes = self.cpu.len() * std::mem::size_of::(); + let ref_ptr: *const std::ffi::c_void = self.cpu + .as_slice() + .expect("reference should be contiguous") + .as_ptr() + .cast(); + let buf = unsafe { + use std::ptr::NonNull; + use objc2_metal::MTLResourceOptions; + device.newBufferWithBytes_length_options( + NonNull::new(ref_ptr.cast_mut()).expect("reference pointer must not be null"), + ref_bytes, + MTLResourceOptions::empty(), + ).expect("failed to create reference buffer") + }; + let idx = StridedNDIndex::from_ndarray(&self.cpu.view()); + (metal::MetalBuffer(buf), idx) + })) + } +} + +/// Check that the values of an i32 DLPack tensor match the expected reference. +/// +/// This dispatches to the appropriate backend based on the device of `tensor`. +/// +/// # Parameters +/// - `tensor`: DLPack tensor with i32 data type +/// - `reference`: expected values with the same shape as the tensor +pub(crate) fn is_equal_i32(tensor: DLPackTensorRef<'_>, reference: &ReferenceValue) -> Result { + match tensor.device().device_type { + DLDeviceType::kDLCPU | DLDeviceType::kDLCUDAHost | DLDeviceType::kDLROCMHost => { + cpu::is_equal_i32(tensor, reference) + } + DLDeviceType::kDLCUDA | DLDeviceType::kDLCUDAManaged => { + cuda::is_equal_i32(tensor, reference) + } + DLDeviceType::kDLMetal => { + #[cfg(target_os = "macos")] { + metal::is_equal_i32(tensor, reference) + } + #[cfg(not(target_os = "macos"))] { + Err(Error::Internal( + "Metal backend is only available on macOS".into(), + )) + } + } + _ => { + eprintln!( + "is_equal_i32 for device {:?} is not implemented", + tensor.device() + ); + Ok(true) + } + } +} + +/// Validate that cell vectors are zero for non-periodic dimensions. +/// +/// This dispatches to the appropriate backend based on the device of `pbc`. +/// +/// # Parameters +/// - `pbc`: 1D boolean tensor of length 3 (periodic boundary condition flags) +/// - `cell`: 3x3 tensor (unit cell vectors as rows) +pub(crate) fn validate_cell_pbc(pbc: DLPackTensorRef<'_>, cell: DLPackTensorRef<'_>) -> Result<(), Error> { + debug_assert!( + pbc.device() == cell.device(), + "pbc and cell must be on the same device" + ); + + match pbc.device().device_type { + DLDeviceType::kDLCPU | DLDeviceType::kDLCUDAHost | DLDeviceType::kDLROCMHost => { + cpu::validate_cell_pbc(pbc, cell) + } + DLDeviceType::kDLCUDA | DLDeviceType::kDLCUDAManaged => { + cuda::validate_cell_pbc(pbc, cell) + } + DLDeviceType::kDLMetal => { + #[cfg(target_os = "macos")] { + metal::validate_cell_pbc(pbc, cell) + } + #[cfg(not(target_os = "macos"))] { + Err(Error::Internal( + "Metal backend is only available on macOS".into(), + )) + } + } + _ => { + eprintln!( + "Cell/PBC validation for device {:?} is not implemented", + pbc.device() + ); + Ok(()) + } + } +} + +/// Scale all elements of `tensor` in place by `factor`. +/// +/// This dispatches to the appropriate backend based on the device of `tensor`. +/// Only 32-bit and 64-bit floating point tensors are supported. +/// +/// # Parameters +/// - `tensor`: a mutable DLPack tensor with f32 or f64 data type +/// - `factor`: the multiplicative factor to apply to every element +pub(crate) fn scale_inplace(tensor: DLPackTensorRefMut<'_>, factor: f64) -> Result<(), Error> { + match tensor.device().device_type { + DLDeviceType::kDLCPU | DLDeviceType::kDLCUDAHost | DLDeviceType::kDLROCMHost => { + cpu::scale_inplace(tensor, factor) + } + DLDeviceType::kDLCUDA | DLDeviceType::kDLCUDAManaged => { + cuda::scale_inplace(tensor, factor) + } + DLDeviceType::kDLMetal => { + #[cfg(target_os = "macos")] { + metal::scale_inplace(tensor, factor) + } + #[cfg(not(target_os = "macos"))] { + Err(Error::Internal( + "Metal backend is only available on macOS".into(), + )) + } + } + _ => { + Err(Error::Internal(format!( + "scale_inplace is not implemented for device {:?}", + tensor.device() + ))) + } + } +} + +/// Clone a DLPack tensor, copying the underlying data to a new allocation. +/// +/// The returned `DLPackTensor` owns its own memory and is independent of the +/// original tensor. The clone is on the same device as the original. +/// +/// # Parameters +/// - `tensor`: the DLPack tensor to clone +pub(crate) fn clone_tensor(tensor: &DLPackTensorRef<'_>) -> Result { + match tensor.device().device_type { + DLDeviceType::kDLCPU | DLDeviceType::kDLCUDAHost | DLDeviceType::kDLROCMHost => { + cpu::clone_tensor(*tensor) + } + DLDeviceType::kDLCUDA | DLDeviceType::kDLCUDAManaged => { + cuda::clone_tensor(tensor) + } + DLDeviceType::kDLMetal => { + #[cfg(target_os = "macos")] { + metal::clone_tensor(tensor) + } + #[cfg(not(target_os = "macos"))] { + Err(Error::Internal( + "Metal backend is only available on macOS".into(), + )) + } + } + _ => { + Err(Error::Internal(format!( + "clone_tensor is not implemented for device {:?}", + tensor.device() + ))) + } + } +} + +/// Check that all atomic types in `types` are present in `valid_types`. +/// +/// This dispatches to the appropriate backend based on the device of `types`. +/// On CPU, the check is done directly. On CUDA/Metal, the check runs on-device +/// and a result count is read back; if invalid types are found, a CPU fallback +/// scan identifies the specific invalid type for the error message. +/// +/// # Parameters +/// - `types`: 1D i32 DLPack tensor of atomic types +/// - `valid_types`: device-resident reference of valid atomic types +pub(crate) fn check_atomic_types( + types: DLPackTensorRef<'_>, + valid_types: &ReferenceValue, +) -> Result<(), Error> { + match types.device().device_type { + DLDeviceType::kDLCPU | DLDeviceType::kDLCUDAHost | DLDeviceType::kDLROCMHost => { + cpu::check_atomic_types(types, valid_types) + } + DLDeviceType::kDLCUDA | DLDeviceType::kDLCUDAManaged => { + cuda::check_atomic_types(types, valid_types) + } + DLDeviceType::kDLMetal => { + #[cfg(target_os = "macos")] { + metal::check_atomic_types(types, valid_types) + } + #[cfg(not(target_os = "macos"))] { + Err(Error::Internal( + "Metal backend is only available on macOS".into(), + )) + } + } + _ => { + Err(Error::Internal(format!( + "check_atomic_types is not implemented for device {:?}", + types.device() + ))) + } + } +} diff --git a/metatomic-core/src/lib.rs b/metatomic-core/src/lib.rs new file mode 100644 index 000000000..537722152 --- /dev/null +++ b/metatomic-core/src/lib.rs @@ -0,0 +1,145 @@ +#![warn(clippy::all, clippy::pedantic)] + +// disable some style lints +#![allow(clippy::needless_return, clippy::must_use_candidate, clippy::comparison_chain)] +#![allow(clippy::redundant_field_names, clippy::redundant_closure_for_method_calls, clippy::redundant_else)] +#![allow(clippy::unreadable_literal, clippy::option_if_let_else, clippy::module_name_repetitions)] +#![allow(clippy::missing_errors_doc, clippy::missing_panics_doc, clippy::missing_safety_doc)] +#![allow(clippy::similar_names, clippy::borrow_as_ptr, clippy::uninlined_format_args)] +#![allow(clippy::doc_markdown, clippy::needless_continue)] +#![allow(clippy::let_underscore_untyped, clippy::manual_let_else, clippy::empty_line_after_doc_comments)] + +// To be removed later +#![allow(unused_variables, dead_code, clippy::needless_pass_by_value)] + +use std::sync::Arc; + +#[doc(hidden)] +pub mod c_api; + +mod utils; + +mod metadata; +use crate::c_api::mta_status_t; + +pub use self::metadata::{Device, DType, ModelCapabilities, ModelMetadata, PairListOptions}; + +mod quantity; +pub use self::quantity::{QuantityName, Quantity, SampleKind, Gradients}; + +mod kernels; + +mod system; +pub use self::system::System; + +mod io; + +mod model; +pub use self::model::Model; + +mod plugin; +pub use self::plugin::Plugin; + +mod units; +pub use self::units::unit_conversion_factor; + +/// The possible sources of error in metatomic +#[derive(Debug, Clone)] +pub enum Error { + /// Error while serializing data to or deserializing data + Serialization(String), + /// Invalid parameters passed to a function + InvalidParameter(String), + /// I/O error + Io(Arc), + /// Error related to dlpack tensors, such as invalid tensor shapes or types + Dlpack(Arc), + /// Error coming from metatensor + Metatensor(metatensor::Error), + /// Error coming from an external function used as a callback + CallbackError(mta_status_t), + /// Any other internal error, usually these are internal bugs. + Internal(String), +} + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Error::Serialization(e) => write!(f, "serialization error: {}", e), + Error::InvalidParameter(e) => write!(f, "invalid parameter: {}", e), + Error::Io(e) => write!(f, "io error: {}", e), + Error::Dlpack(e) => write!(f, "dlpack error: {}", e), + Error::Metatensor(e) => write!(f, "metatensor error: {}", e), + Error::CallbackError(e) => write!(f, "callback error, status code: {:?}", e), + Error::Internal(e) => write!(f, + "internal metatomic error (this is likely a bug, please report it): {}", e + ), + } + } +} + +impl std::error::Error for Error { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Error::InvalidParameter(_) + | Error::Serialization(_) + | Error::Internal(_) + | Error::CallbackError(_) => None, + Error::Io(e) => Some(e), + Error::Dlpack(e) => Some(e), + Error::Metatensor(e) => Some(e), + } + } + + fn cause(&self) -> Option<&dyn std::error::Error> { + self.source() + } +} + +// Box is the error type in std::panic::catch_unwind +impl From> for Error { + fn from(error: Box) -> Error { + if error.is::() { + Error::Internal(*error.downcast::().expect("should be a String")) + } else if error.is::<&str>() { + Error::Internal((*error.downcast::<&str>().expect("should be an &str")).to_owned()) + } else if error.is::() { + return *error.downcast::().expect("it should be an Error"); + } else { + panic!("panic message is not a string, something is very wrong") + } + } +} + +impl From for Error { + fn from(error: std::io::Error) -> Self { + Error::Io(Arc::new(error)) + } +} + +impl From for Error { + fn from(error: dlpk::ndarray::DLPackNDarrayError) -> Self { + Error::Dlpack(Arc::new(error)) + } +} + +impl From for Error { + fn from(error: metatensor::Error) -> Self { + Error::Metatensor(error) + } +} + +impl From for Error { + fn from(error: json::Error) -> Self { + Error::Serialization(format!("json error: {}", error)) + } +} + +impl> From<(T, zip::result::ZipError)> for Error { + fn from((path, error): (T, zip::result::ZipError)) -> Self { + match error { + zip::result::ZipError::Io(e) => Error::Io(Arc::new(e)), + error => Error::Serialization(format!("{}: at '{}'", error, path.as_ref())), + } + } +} diff --git a/metatomic-core/src/metadata.rs b/metatomic-core/src/metadata.rs new file mode 100644 index 000000000..64a7ff5be --- /dev/null +++ b/metatomic-core/src/metadata.rs @@ -0,0 +1,1083 @@ +use std::collections::BTreeMap; +use std::fmt::Write; + +use json::JsonValue; + +use crate::metadata::DType::Float32; +use crate::{Error, Quantity}; +use crate::units::validate_unit; + +/// Options for the calculation of a pair list (neighbor list) +#[derive(Debug, Clone)] +pub struct PairListOptions { + /// Cutoff radius for this pair list in the length unit of the model + pub cutoff: f64, + /// Whether the list is a full list (contains both the pair `i -> j` and `j -> i`) + /// or a half list (contains only `i -> j`) + pub full_list: bool, + /// Whether the list guarantees that only atoms within the cutoff are + /// included (strict) or may also include pairs slightly beyond the cutoff + /// (non-strict) + pub strict: bool, + /// List of strings describing who requested this pair list + pub requestors: Vec, +} + +impl std::cmp::PartialEq for PairListOptions { + fn eq(&self, other: &Self) -> bool { + self.cutoff == other.cutoff + && self.full_list == other.full_list + && self.strict == other.strict + } +} + +impl std::cmp::Eq for PairListOptions {} + +impl std::cmp::PartialOrd for PairListOptions { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl std::cmp::Ord for PairListOptions { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.cutoff.partial_cmp(&other.cutoff).expect("cutoff is NaN") + .then_with(|| self.full_list.cmp(&other.full_list)) + .then_with(|| self.strict.cmp(&other.strict)) + } +} + +impl From for JsonValue { + fn from(value: PairListOptions) -> Self { + let mut result = JsonValue::new_object(); + result["type"] = "metatomic_pair_list_options".into(); + // store the bit pattern so the float round-trips exactly + result["cutoff"] = format!("{:#x}", value.cutoff.to_bits()).into(); + result["full_list"] = value.full_list.into(); + result["strict"] = value.strict.into(); + result["requestors"] = value.requestors.into(); + return result; + } +} + +impl<'a> TryFrom<&'a JsonValue> for PairListOptions { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + if !value.is_object() { + return Err(Error::Serialization( + "invalid JSON data for PairListOptions, expected an object".into() + )); + } + + let cutoff = if value.has_key("class") { + // this is the legacy format from metatomic-torch, which can be used + // to load serialized PairListOptions + if value["class"].as_str() != Some("NeighborListOptions") { + return Err(Error::Serialization( + "'class' in legacy JSON for PairListOptions must be 'NeighborListOptions'".into() + )); + } + + let cutoff_bits = value["cutoff"].as_u64().ok_or_else(|| Error::Serialization( + "'cutoff' in legacy JSON for PairListOptions must be an integer".into() + ))?; + + f64::from_bits(cutoff_bits) + } else { + if value["type"].as_str() != Some("metatomic_pair_list_options") { + return Err(Error::Serialization( + "'type' in JSON for PairListOptions must be 'metatomic_pair_list_options'".into() + )); + } + + let cutoff_str = value["cutoff"].as_str().ok_or_else(|| Error::Serialization( + "'cutoff' in JSON for PairListOptions must be a hex-encoded string".into() + ))?; + let cutoff_bits = u64::from_str_radix(cutoff_str.strip_prefix("0x").unwrap_or(cutoff_str), 16) + .map_err(|_| Error::Serialization( + "'cutoff' in JSON for PairListOptions must be a hex-encoded string".into() + ))?; + + f64::from_bits(cutoff_bits) + }; + + if !cutoff.is_finite() || cutoff <= 0.0 { + return Err(Error::Serialization( + "'cutoff' in JSON for PairListOptions must be a finite positive number".into() + )); + } + + let full_list = value["full_list"].as_bool().ok_or_else(|| Error::Serialization( + "'full_list' in JSON for PairListOptions must be a boolean".into() + ))?; + + let strict = value["strict"].as_bool().ok_or_else(|| Error::Serialization( + "'strict' in JSON for PairListOptions must be a boolean".into() + ))?; + + let mut requestors = Vec::new(); + if value.has_key("requestors") { + if !value["requestors"].is_array() { + return Err(Error::Serialization( + "'requestors' in JSON for PairListOptions must be an array".into() + )); + } + + for requestor in value["requestors"].members() { + let requestor = requestor.as_str().ok_or_else(|| Error::Serialization( + "'requestors' in JSON for PairListOptions must be an array of strings".into() + ))?; + // ignore empty strings and duplicates, keeping first-seen order + if !requestor.is_empty() && !requestors.iter().any(|r| r == requestor) { + requestors.push(requestor.to_string()); + } + } + } + + return Ok(PairListOptions { cutoff, full_list, strict, requestors }); + } +} + +// ========================================================================== // +// ========================================================================== // +// ========================================================================== // + +/// References for a model, divided into three categories: references about the +/// model as a whole, references about the architecture of the model, and +/// references about the implementation of the model. Each category is a list of +/// strings, which can be DOIs, URLs, or any other format the model author finds +/// useful. +#[derive(Debug, Clone)] +pub struct References { + /// The references about the model as a whole, e.g. a paper describing the + /// model or a website presenting it. + model: Vec, + /// The references about the architecture of the model, e.g. papers + /// describing the mathematical form of the model. + architecture: Vec, + /// The references about the implementation of the model, e.g. a link to + /// the source code repository or a paper describing the software. + implementation: Vec, +} + +impl From for JsonValue { + fn from(value: References) -> Self { + let mut result = JsonValue::new_object(); + result["model"] = value.model.into(); + result["architecture"] = value.architecture.into(); + result["implementation"] = value.implementation.into(); + return result; + } +} + + +fn read_references(object: &JsonValue, key: &str) -> Result, Error> { + let mut references = Vec::new(); + if !object[key].is_array() { + return Err(Error::Serialization( + format!("'{}' in references of ModelMetadata must be an array", key) + )); + } + for reference in object[key].members() { + let reference = reference.as_str().ok_or_else(|| Error::Serialization( + format!("'{}' in references of ModelMetadata must be an array of strings", key) + ))?; + references.push(reference.to_string()); + } + Ok(references) +} + +impl<'a> TryFrom<&'a JsonValue> for References { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + if !value.is_object() { + return Err(Error::Serialization( + "invalid JSON data for references in ModelMetadata, expected an object".into() + )); + } + + let model = read_references(value, "model")?; + let architecture = read_references(value, "architecture")?; + let implementation = read_references(value, "implementation")?; + + Ok(References { model, architecture, implementation }) + } +} + + +fn normalize_whitespace(data: &str) -> String { + let mut normalized_string = String::new(); + for c in data.chars() { + if c == '\n' || c == '\r' || c == '\t' { + normalized_string.push(' '); + } else { + normalized_string.push(c); + } + } + normalized_string +} + + +fn wrap_80_chars(output: &mut String, data: &str, indent: usize) { + let string = normalize_whitespace(data); + assert!(indent < 30); + let line_length = 80 - indent; + assert!(line_length > 50); + let mut first_line = true; + let mut start = 0; + + loop { + let remaining = &string[start..]; + + if remaining.len() <= line_length { + if !first_line { + output.push_str(&" ".repeat(indent)); + } + output.push_str(remaining); + break; + } + + // byte offset of the character just past the first `line_length` chars + let end = remaining.char_indices().nth(line_length).map_or(remaining.len(), |(i, _)| i); + + if let Some(space_pos) = remaining[..end].rfind(' ') { + if !first_line { + output.push_str(&" ".repeat(indent)); + } + output.push_str(&remaining[..space_pos]); + output.push('\n'); + start += space_pos + 1; + first_line = false; + } else { + let word_end = remaining.find(' ').unwrap_or(remaining.len()); + if !first_line { + output.push_str(&" ".repeat(indent)); + } + output.push_str(&remaining[..word_end]); + output.push('\n'); + first_line = false; + if word_end < remaining.len() { + start += word_end + 1; + } else { + break; + } + } + } +} + +/// Metadata about a model +#[derive(Debug, Clone)] +pub struct ModelMetadata { + /// The name of the model, e.g. `"MyCoolModel v1.2"` + pub name: String, + /// The authors of the model, e.g. `["Alice Smith", "Bob Johnson + /// "]` + pub authors: Vec, + /// A description of the model + pub description: String, + /// References for the model that should be cited when using it + pub references: References, + /// Any other key-value pairs the model author wants to include in the + /// metadata. This can be used for any purpose. + pub extra: BTreeMap, +} + +impl From for JsonValue { + fn from(value: ModelMetadata) -> Self { + let mut result = JsonValue::new_object(); + result["type"] = "metatomic_model_metadata".into(); + result["name"] = value.name.into(); + result["authors"] = value.authors.into(); + result["description"] = value.description.into(); + result["references"] = value.references.into(); + result["extra"] = value.extra.into(); + return result; + } +} + +impl<'a> TryFrom<&'a JsonValue> for ModelMetadata { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + if !value.is_object() { + return Err(Error::Serialization( + "invalid JSON data for ModelMetadata, expected an object".into() + )); + } + + if value["type"].as_str() != Some("metatomic_model_metadata") { + return Err(Error::Serialization( + "'type' in JSON for ModelMetadata must be 'metatomic_model_metadata'".into() + )); + } + + let name = value["name"].as_str().ok_or_else(|| Error::Serialization( + "'name' in JSON for ModelMetadata must be a string".into() + ))?; + + if !value["authors"].is_array() { + return Err(Error::Serialization( + "'authors' in JSON for ModelMetadata must be an array".into() + )); + } + + let authors = value["authors"].members().map(|author| { + author.as_str().ok_or_else(|| Error::Serialization( + "'authors' in JSON for ModelMetadata must be an array of strings".into() + )).map(|s| s.to_string()) + }).collect::, Error>>()?; + + let description = value["description"].as_str().ok_or_else(|| Error::Serialization( + "'description' in JSON for ModelMetadata must be a string".into() + ))?.to_string(); + + let references = References::try_from(&value["references"])?; + + if !value["extra"].is_object() { + return Err(Error::Serialization( + "'extra' in JSON for ModelMetadata must be an object".into() + )); + } + + let mut extra = BTreeMap::new(); + for (key, value) in value["extra"].entries() { + let value = value.as_str().ok_or_else(|| Error::Serialization( + "'extra' in JSON for ModelMetadata must be an object with string values".into() + ))?; + extra.insert(key.to_string(), value.to_string()); + } + + // Validate the contents of `authors` and `references` + for author in &authors { + if author.is_empty() { + return Err(Error::InvalidParameter("author can not be empty string in ModelMetadata".into())); + } + } + + for model_ref in &references.model { + if model_ref.is_empty() { + return Err(Error::InvalidParameter("reference can not be empty string (in 'model' section)".into())); + } + } + + for architecture_ref in &references.architecture { + if architecture_ref.is_empty() { + return Err(Error::InvalidParameter("reference can not be empty string (in 'architecture' section)".into())); + } + } + + for implementation_ref in &references.implementation { + if implementation_ref.is_empty() { + return Err(Error::InvalidParameter("reference can not be empty string (in 'implementation' section)".into())); + } + } + + let metadata = ModelMetadata { + name: name.to_string(), + authors: authors, + description: description, + references: references, + extra: extra, + }; + Ok(metadata) + } +} + +impl ModelMetadata{ + pub fn print(&self) -> String { + let mut output = String::new(); + if self.name.is_empty() { + let _ = writeln!(output, "This is an unnamed model"); + let _ = writeln!(output, "========================"); + } else { + let _ = writeln!(output, "This is the {} model", self.name); + let _ = writeln!(output, "============{}======", "=".repeat(self.name.len())); + } + + if !self.description.is_empty() { + let _ = writeln!(output); + wrap_80_chars(&mut output, &(self.description), 0); + let _ = writeln!(output); + } + + if !self.authors.is_empty() { + let _ = writeln!(output, "\nModel authors\n-------------\n"); + for author in &self.authors { + let _ = write!(output, "- "); + wrap_80_chars(&mut output, author, 2); + output.push('\n'); + } + } + + let mut references_output = String::new(); + if !self.references.model.is_empty() { + references_output.push_str("- about this specific model:\n"); + for reference in &self.references.model { + references_output.push_str(" * "); + wrap_80_chars(&mut references_output, reference, 4); + references_output.push('\n'); + } + } + + if !self.references.architecture.is_empty() { + references_output.push_str("- about the architecture of this model:\n"); + for reference in &self.references.architecture { + references_output.push_str(" * "); + wrap_80_chars(&mut references_output, reference, 4); + references_output.push('\n'); + } + } + + if !self.references.implementation.is_empty() { + references_output.push_str("- about the implementation of this model:\n"); + for reference in &self.references.implementation { + references_output.push_str(" * "); + wrap_80_chars(&mut references_output, reference, 4); + references_output.push('\n'); + } + } + + if !references_output.is_empty() { + output.push_str("\nModel references\n----------------\n\n"); + output.push_str("Please cite the following references when using this model:\n"); + output.push_str(&references_output); + } + + return output; + } +} + +/// The data type of a model, used for all inputs and outputs. The model can +/// still internally use a different data type for its calculations, but it will +/// get inputs in this type and must produce outputs in this type. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DType { + /// 32-bit floating point, following the IEEE 754 standard + Float32, + /// 64-bit floating point, following the IEEE 754 standard + Float64, +} + +impl std::fmt::Display for DType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if *self == Float32 { + write!(f, "float32") + } else { + write!(f, "float64") + } + } +} + +impl From for JsonValue { + fn from(value: DType) -> Self { + match value { + DType::Float32 => "float32".into(), + DType::Float64 => "float64".into(), + } + } +} + +impl<'a> TryFrom<&'a JsonValue> for DType { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + if let Some(s) = value.as_str() { + match s { + "float32" => Ok(DType::Float32), + "float64" => Ok(DType::Float64), + _ => Err(Error::Serialization( + "invalid string for dtype in JSON for ModelCapabilities, expected 'float32' or 'float64'".into() + )), + } + } else { + Err(Error::Serialization( + "dtype in JSON for ModelCapabilities must be a string".into() + )) + } + } +} + +/// A device on which a model can run. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Device(dlpk::DLDeviceType); + +impl Device { + /// Create a `Device` representing the CPU. + pub fn cpu() -> Self { + Device(dlpk::DLDeviceType::kDLCPU) + } + + /// Check wether a dlpack device is valid to use with this model device. + pub fn is_valid(&self, device: dlpk::DLDevice) -> bool { + // this should probably handle the "Host" and "Managed" variants, but + // should be enough for now. + return device.device_type == self.0; + } +} + +impl From for JsonValue { + fn from(value: Device) -> Self { + match value.0 { + dlpk::DLDeviceType::kDLCPU => "cpu".into(), + dlpk::DLDeviceType::kDLCUDA => "cuda".into(), + dlpk::DLDeviceType::kDLROCM => "rocm".into(), + dlpk::DLDeviceType::kDLMetal => "metal".into(), + dlpk::DLDeviceType::kDLCUDAHost | dlpk::DLDeviceType::kDLCUDAManaged => { + // These refer to memory devices more than execution devices + panic!("Do not use kDLCUDAHost or kDLCUDAManaged, use kDLCUDA instead."); + } + dlpk::DLDeviceType::kDLROCMHost => { + // This refers to a memory device more than an execution device + panic!("Do not use kDLROCMHost, use kDLROCM instead."); + } + _ => { + // We don't want to expose other device types until we have a + // use case for them, and we don't want to accidentally leak + // them if they're added in the future + panic!("unsupported device type: {:?}", value.0); + } + } + } +} + +impl<'a> TryFrom<&'a JsonValue> for Device { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + if let Some(s) = value.as_str() { + match s { + "cpu" => Ok(Device(dlpk::DLDeviceType::kDLCPU)), + "cuda" => Ok(Device(dlpk::DLDeviceType::kDLCUDA)), + "rocm" => Ok(Device(dlpk::DLDeviceType::kDLROCM)), + "metal" => Ok(Device(dlpk::DLDeviceType::kDLMetal)), + _ => Err(Error::Serialization( + "invalid string for device in JSON for ModelCapabilities, expected 'cpu', 'cuda', 'rocm', or 'metal'".into() + )), + } + } else { + Err(Error::Serialization( + "device in JSON for ModelCapabilities must be a string".into() + )) + } + } +} + +/// Capabilities about a model: which outputs it provides, which atoms it +/// supports, etc. +#[derive(Debug, Clone)] +pub struct ModelCapabilities { + /// The outputs this model can provide + pub outputs: Vec, + /// The atomic types this model supports. The meaning of the integers in + /// this list is up to the model, and is not required to be the atomic + /// numbers. + pub atomic_types: Vec, + /// The interaction range of the model (in the length unit of the model), + /// i.e. the maximum distance between two atoms for which the model's output + /// can depend on their relative position. + pub interaction_range: f64, + /// The length unit of the model, e.g. "angstrom" or "nanometer". This is + /// used to interpret the `interaction_range` and convert the inputs. + pub length_unit: String, + /// The devices on which the model can run, e.g. `["cpu", "cuda"]`. + pub supported_devices: Vec, + /// The data type of the model, used for all inputs and outputs. + pub dtype: DType, +} + +impl ModelCapabilities { + /// Find the declared output quantity matching the given `request` name and + /// sample kind, if any. + pub fn find_output(&self, request: &Quantity) -> Option<&Quantity> { + self.outputs.iter().find(|q| + q.name == request.name && q.sample_kind == request.sample_kind + ) + } +} + +impl From for JsonValue { + fn from(value: ModelCapabilities) -> Self { + let mut result = JsonValue::new_object(); + result["type"] = "metatomic_model_capabilities".into(); + result["outputs"] = value.outputs.into(); + result["atomic_types"] = value.atomic_types.into(); + result["interaction_range"] = value.interaction_range.into(); + result["length_unit"] = value.length_unit.into(); + result["supported_devices"] = value.supported_devices.into(); + result["dtype"] = value.dtype.into(); + return result; + } +} + +impl<'a> TryFrom<&'a JsonValue> for ModelCapabilities { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + if !value.is_object() { + return Err(Error::Serialization( + "invalid JSON data for ModelCapabilities, expected an object".into() + )); + } + + if value["type"].as_str() != Some("metatomic_model_capabilities") { + return Err(Error::Serialization( + "'type' in JSON for ModelCapabilities must be 'metatomic_model_capabilities'".into() + )); + } + + let mut outputs = Vec::new(); + if !value["outputs"].is_array() { + return Err(Error::Serialization( + "'outputs' in JSON for ModelCapabilities must be an array".into() + )); + } + for output in value["outputs"].members() { + outputs.push(Quantity::try_from(output)?); + } + + + let mut atomic_types = Vec::new(); + if !value["atomic_types"].is_array() { + return Err(Error::Serialization( + "'atomic_types' in JSON for ModelCapabilities must be an array".into() + )); + } + + for atomic_type in value["atomic_types"].members() { + let atomic_type = atomic_type.as_i32().ok_or_else(|| Error::Serialization( + "'atomic_types' in JSON for ModelCapabilities must be an array of integers".into() + ))?; + atomic_types.push(atomic_type); + } + + let interaction_range = value["interaction_range"].as_f64().ok_or_else(|| Error::Serialization( + "'interaction_range' in JSON for ModelCapabilities must be a number".into() + ))?; + if interaction_range < 0.0 { + return Err(Error::Serialization( + "'interaction_range' in JSON for ModelCapabilities must be non-negative".into() + )); + } + + let length_unit = value["length_unit"].as_str().ok_or_else(|| Error::Serialization( + "'length_unit' in JSON for ModelCapabilities must be a string".into() + ))?.to_string(); + validate_unit(&length_unit, "m", Some("'length_unit' in JSON for ModelCapabilities"))?; + + let mut supported_devices = Vec::new(); + if !value["supported_devices"].is_array() { + return Err(Error::Serialization( + "'supported_devices' in JSON for ModelCapabilities must be an array".into() + )); + } + for device in value["supported_devices"].members() { + supported_devices.push(Device::try_from(device)?); + } + + let dtype = DType::try_from(&value["dtype"])?; + + Ok(ModelCapabilities { + outputs, + atomic_types, + interaction_range, + length_unit, + supported_devices, + dtype, + }) + } +} + + +#[cfg(test)] +mod tests { + mod pair_list_options { + use super::super::*; + + fn example() -> PairListOptions { + PairListOptions { + cutoff: 3.5, + full_list: true, + strict: false, + requestors: vec!["nl-1".to_string(), "nl-2".to_string()], + } + } + + #[test] + fn roundtrip() { + let options = example(); + let json: JsonValue = options.clone().into(); + + assert_eq!(json["type"].as_str(), Some("metatomic_pair_list_options")); + assert_eq!(json["cutoff"].as_str(), Some(format!("{:#x}", 3.5_f64.to_bits()).as_str())); + assert_eq!(json["full_list"].as_bool(), Some(true)); + assert_eq!(json["strict"].as_bool(), Some(false)); + + let parsed = PairListOptions::try_from(&json).unwrap(); + assert_eq!(parsed.cutoff.to_bits(), options.cutoff.to_bits()); + assert_eq!(parsed.full_list, options.full_list); + assert_eq!(parsed.strict, options.strict); + assert_eq!(parsed.requestors, options.requestors); + } + + #[test] + fn cutoff_keeps_full_precision() { + let mut options = example(); + options.cutoff = 1.0 / 3.0; + let parsed = PairListOptions::try_from(&JsonValue::from(options.clone())).unwrap(); + assert_eq!(parsed.cutoff.to_bits(), options.cutoff.to_bits()); + } + + #[test] + fn requestors_are_optional() { + let mut json: JsonValue = example().into(); + json.remove("requestors"); + let parsed = PairListOptions::try_from(&json).unwrap(); + assert!(parsed.requestors.is_empty()); + } + + #[test] + fn rejects_invalid_json() { + // each case corrupts exactly one field of an otherwise valid object + let with_cutoff = |value: f64| { + let mut json = JsonValue::from(example()); + json["cutoff"] = format!("{:#x}", value.to_bits()).into(); + json + }; + + let mut wrong_type = JsonValue::from(example()); + wrong_type["type"] = "something-else".into(); + + let mut missing_cutoff = JsonValue::from(example()); + missing_cutoff.remove("cutoff"); + + let mut non_hex_cutoff = JsonValue::from(example()); + non_hex_cutoff["cutoff"] = "not-hex".into(); + + let mut non_boolean_flag = JsonValue::from(example()); + non_boolean_flag["full_list"] = "yes".into(); + + let mut non_array_requestors = JsonValue::from(example()); + non_array_requestors["requestors"] = "nl-1".into(); + + let mut non_string_requestor = JsonValue::from(example()); + non_string_requestor["requestors"] = json::array![ "nl-1", 42 ]; + + let cases = [ + (JsonValue::from("not an object"), + "serialization error: invalid JSON data for PairListOptions, expected an object"), + (wrong_type, + "serialization error: 'type' in JSON for PairListOptions must be 'metatomic_pair_list_options'"), + (missing_cutoff, + "serialization error: 'cutoff' in JSON for PairListOptions must be a hex-encoded string"), + (non_hex_cutoff, + "serialization error: 'cutoff' in JSON for PairListOptions must be a hex-encoded string"), + (with_cutoff(f64::NAN), + "serialization error: 'cutoff' in JSON for PairListOptions must be a finite positive number"), + (with_cutoff(f64::INFINITY), + "serialization error: 'cutoff' in JSON for PairListOptions must be a finite positive number"), + (with_cutoff(-1.0), + "serialization error: 'cutoff' in JSON for PairListOptions must be a finite positive number"), + (with_cutoff(0.0), + "serialization error: 'cutoff' in JSON for PairListOptions must be a finite positive number"), + (non_boolean_flag, + "serialization error: 'full_list' in JSON for PairListOptions must be a boolean"), + (non_array_requestors, + "serialization error: 'requestors' in JSON for PairListOptions must be an array"), + (non_string_requestor, + "serialization error: 'requestors' in JSON for PairListOptions must be an array of strings"), + ]; + + for (json, expected) in cases { + let error = PairListOptions::try_from(&json).expect_err("expected an error"); + assert_eq!(error.to_string(), expected); + } + } + + #[test] + fn requestors_skip_empty_and_duplicates() { + let mut json: JsonValue = example().into(); + json["requestors"] = json::array![ "a", "", "b", "a" ]; + + let parsed = PairListOptions::try_from(&json).unwrap(); + assert_eq!(parsed.requestors, vec!["a".to_string(), "b".to_string()]); + } + } + + mod model_metadata { + +use super::super::*; + + fn example() -> ModelMetadata { + ModelMetadata { + name: "test-model".into(), + authors: vec!["Alice".into(), "Bob ".into()], + description: "A test model".into(), + references: References { + model: vec!["doi:10.1234/test".into()], + architecture: vec!["doi:10.1234/arch".into()], + implementation: vec!["https://github.com/test".into()], + }, + extra: BTreeMap::from([ + ("key1".into(), "value1".into()), + ("key2".into(), "value2".into()), + ]), + } + } + + #[test] + fn roundtrip() { + let metadata = example(); + let json: JsonValue = metadata.clone().into(); + + assert_eq!(json["type"].as_str(), Some("metatomic_model_metadata")); + assert_eq!(json["name"].as_str(), Some("test-model")); + assert_eq!(json["authors"][0].as_str(), Some("Alice")); + assert_eq!(json["authors"][1].as_str(), Some("Bob ")); + assert_eq!(json["description"].as_str(), Some("A test model")); + assert_eq!(json["references"]["model"][0].as_str(), Some("doi:10.1234/test")); + assert_eq!(json["references"]["architecture"][0].as_str(), Some("doi:10.1234/arch")); + assert_eq!(json["references"]["implementation"][0].as_str(), Some("https://github.com/test")); + assert_eq!(json["extra"]["key1"].as_str(), Some("value1")); + assert_eq!(json["extra"]["key2"].as_str(), Some("value2")); + + let parsed = ModelMetadata::try_from(&json).unwrap(); + assert_eq!(parsed.name, metadata.name); + assert_eq!(parsed.authors, metadata.authors); + assert_eq!(parsed.description, metadata.description); + assert_eq!(parsed.references.model, metadata.references.model); + assert_eq!(parsed.references.architecture, metadata.references.architecture); + assert_eq!(parsed.references.implementation, metadata.references.implementation); + assert_eq!(parsed.extra, metadata.extra); + } + + #[test] + fn rejects_invalid_json() { + let mut wrong_type = JsonValue::from(example()); + wrong_type["type"] = "something-else".into(); + + let mut missing_name = JsonValue::from(example()); + missing_name.remove("name"); + + let mut non_string_name = JsonValue::from(example()); + non_string_name["name"] = 42.into(); + + let mut non_array_authors = JsonValue::from(example()); + non_array_authors["authors"] = "Alice".into(); + + let mut non_string_author = JsonValue::from(example()); + non_string_author["authors"] = json::array!["Alice", 42]; + + let mut missing_description = JsonValue::from(example()); + missing_description.remove("description"); + + let mut non_object_extra = JsonValue::from(example()); + non_object_extra["extra"] = "not-an-object".into(); + + let mut non_string_extra_value = JsonValue::from(example()); + non_string_extra_value["extra"] = json::object!{ "key" => 42 }; + + let mut non_object_references = JsonValue::from(example()); + non_object_references["references"] = "not-an-object".into(); + + let cases = [ + (JsonValue::from("not an object"), + "serialization error: invalid JSON data for ModelMetadata, expected an object"), + (wrong_type, + "serialization error: 'type' in JSON for ModelMetadata must be 'metatomic_model_metadata'"), + (missing_name, + "serialization error: 'name' in JSON for ModelMetadata must be a string"), + (non_string_name, + "serialization error: 'name' in JSON for ModelMetadata must be a string"), + (non_array_authors, + "serialization error: 'authors' in JSON for ModelMetadata must be an array"), + (non_string_author, + "serialization error: 'authors' in JSON for ModelMetadata must be an array of strings"), + (missing_description, + "serialization error: 'description' in JSON for ModelMetadata must be a string"), + (non_object_extra, + "serialization error: 'extra' in JSON for ModelMetadata must be an object"), + (non_string_extra_value, + "serialization error: 'extra' in JSON for ModelMetadata must be an object with string values"), + (non_object_references, + "serialization error: invalid JSON data for references in ModelMetadata, expected an object"), + ]; + + for (json, expected) in cases { + let error = ModelMetadata::try_from(&json).expect_err("expected an error"); + assert_eq!(error.to_string(), expected); + } + } + + #[test] + fn printing() { + let metadata = example(); + let output = metadata.print(); + let expected = String::from( + "This is the test-model model +============================ + +A test model + +Model authors +------------- + +- Alice +- Bob + +Model references +---------------- + +Please cite the following references when using this model: +- about this specific model: + * doi:10.1234/test +- about the architecture of this model: + * doi:10.1234/arch +- about the implementation of this model: + * https://github.com/test +" +); + + assert_eq!(output, expected); + } + } + + mod model_capabilities { + use crate::QuantityName; + use super::super::*; + + fn example() -> ModelCapabilities { + ModelCapabilities { + outputs: vec![ + Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: Some("total energy".into()), + gradients: vec![crate::Gradients::Positions], + sample_kind: crate::SampleKind::System, + }, + Quantity { + name: QuantityName::new("custom::charge/with_variant".into()).unwrap(), + unit: "e".into(), + description: None, + gradients: vec![], + sample_kind: crate::SampleKind::Atom, + }, + ], + atomic_types: vec![1, 6, 8], + interaction_range: 5.0, + length_unit: "Angstrom".into(), + supported_devices: vec![Device(dlpk::DLDeviceType::kDLCPU), Device(dlpk::DLDeviceType::kDLCUDA)], + dtype: DType::Float32, + } + } + + #[test] + fn roundtrip() { + let capabilities = example(); + let json: JsonValue = capabilities.clone().into(); + + assert_eq!(json["type"].as_str(), Some("metatomic_model_capabilities")); + assert_eq!(json["outputs"][0]["name"].as_str(), Some("energy")); + assert_eq!(json["outputs"][1]["name"].as_str(), Some("custom::charge/with_variant")); + assert_eq!(json["atomic_types"][0].as_i32(), Some(1)); + assert_eq!(json["atomic_types"][1].as_i32(), Some(6)); + assert_eq!(json["atomic_types"][2].as_i32(), Some(8)); + assert_eq!(json["interaction_range"].as_f64(), Some(5.0)); + assert_eq!(json["length_unit"].as_str(), Some("Angstrom")); + assert_eq!(json["supported_devices"][0].as_str(), Some("cpu")); + assert_eq!(json["supported_devices"][1].as_str(), Some("cuda")); + assert_eq!(json["dtype"].as_str(), Some("float32")); + + let parsed = ModelCapabilities::try_from(&json).unwrap(); + assert_eq!(parsed.outputs.len(), 2); + assert_eq!(parsed.outputs[0].name.namespace(), None); + assert_eq!(parsed.outputs[0].name.base(), "energy"); + assert_eq!(parsed.outputs[0].name.variant(), None); + + assert_eq!(parsed.outputs[1].name.namespace(), Some("custom")); + assert_eq!(parsed.outputs[1].name.base(), "charge"); + assert_eq!(parsed.outputs[1].name.variant(), Some("with_variant")); + + assert_eq!(parsed.atomic_types, vec![1, 6, 8]); + assert_eq!(parsed.interaction_range.to_bits(), 5.0_f64.to_bits()); + assert_eq!(parsed.length_unit, "Angstrom"); + assert_eq!(parsed.supported_devices.len(), 2); + assert_eq!(parsed.dtype, DType::Float32); + } + + #[test] + fn rejects_invalid_json() { + let mut wrong_type = JsonValue::from(example()); + wrong_type["type"] = "something-else".into(); + + let mut non_array_outputs = JsonValue::from(example()); + non_array_outputs["outputs"] = "energy".into(); + + let mut non_array_atomic_types = JsonValue::from(example()); + non_array_atomic_types["atomic_types"] = "1".into(); + + let mut non_integer_atomic_type = JsonValue::from(example()); + non_integer_atomic_type["atomic_types"] = json::array![1, "x"]; + + let mut missing_interaction_range = JsonValue::from(example()); + missing_interaction_range.remove("interaction_range"); + + let mut negative_interaction_range = JsonValue::from(example()); + negative_interaction_range["interaction_range"] = (-1.0).into(); + + let mut missing_length_unit = JsonValue::from(example()); + missing_length_unit.remove("length_unit"); + + let mut wrong_dimension_length_unit = JsonValue::from(example()); + wrong_dimension_length_unit["length_unit"] = "eV".into(); + + let mut non_array_supported_devices = JsonValue::from(example()); + non_array_supported_devices["supported_devices"] = "cpu".into(); + + let mut invalid_device = JsonValue::from(example()); + invalid_device["supported_devices"] = json::array!["cpu", "wat"]; + + let mut missing_dtype = JsonValue::from(example()); + missing_dtype.remove("dtype"); + + let mut invalid_dtype = JsonValue::from(example()); + invalid_dtype["dtype"] = "float16".into(); + + let cases: Vec<(JsonValue, &str)> = vec![ + (JsonValue::from("not an object"), + "serialization error: invalid JSON data for ModelCapabilities, expected an object"), + (wrong_type, + "serialization error: 'type' in JSON for ModelCapabilities must be 'metatomic_model_capabilities'"), + (non_array_outputs, + "serialization error: 'outputs' in JSON for ModelCapabilities must be an array"), + (non_array_atomic_types, + "serialization error: 'atomic_types' in JSON for ModelCapabilities must be an array"), + (non_integer_atomic_type, + "serialization error: 'atomic_types' in JSON for ModelCapabilities must be an array of integers"), + (missing_interaction_range, + "serialization error: 'interaction_range' in JSON for ModelCapabilities must be a number"), + (negative_interaction_range, + "serialization error: 'interaction_range' in JSON for ModelCapabilities must be non-negative"), + (missing_length_unit, + "serialization error: 'length_unit' in JSON for ModelCapabilities must be a string"), + (wrong_dimension_length_unit, + "invalid parameter: dimension mismatch in 'length_unit' in JSON for ModelCapabilities: 'eV' has dimension [L^2 T^-2 M] but expected dimension [L]"), + (non_array_supported_devices, + "serialization error: 'supported_devices' in JSON for ModelCapabilities must be an array"), + (invalid_device, + "serialization error: invalid string for device in JSON for ModelCapabilities, expected 'cpu', 'cuda', 'rocm', or 'metal'"), + (missing_dtype, + "serialization error: dtype in JSON for ModelCapabilities must be a string"), + (invalid_dtype, + "serialization error: invalid string for dtype in JSON for ModelCapabilities, expected 'float32' or 'float64'"), + ]; + + for (json, expected) in cases { + let error = ModelCapabilities::try_from(&json).expect_err("expected an error"); + assert_eq!(error.to_string(), expected); + } + } + } +} diff --git a/metatomic-core/src/model/execute.rs b/metatomic-core/src/model/execute.rs new file mode 100644 index 000000000..018344f8b --- /dev/null +++ b/metatomic-core/src/model/execute.rs @@ -0,0 +1,251 @@ +use std::ffi::CString; +use std::sync::{Arc, LazyLock, Mutex}; + +use lru::LruCache; +use metatensor::TensorMap; +use metatensor::c_api::{mts_labels_t, mts_tensormap_t}; + +use crate::{Error, Quantity, System}; +use crate::c_api::{mta_system_t, mta_status_t}; +use crate::kernels::{check_atomic_types, ReferenceValue}; +use crate::quantity::check_quantity; +use crate::unit_conversion_factor; +use crate::utils::scale_tensormap; + +use super::Model; +use super::inputs::{check_inputs, check_requested_outputs}; + +#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)] +struct AtomicTypes(Vec); + +impl AtomicTypes { + fn new(mut types: Vec) -> Self { + types.sort_unstable(); + return AtomicTypes(types); + } +} + +impl std::ops::Deref for AtomicTypes { + type Target = [i32]; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +/// LRU cache of `ReferenceValue` for atomic types, keyed by the sorted +/// vector of valid types. This avoids re-uploading the same set of valid types +/// to the device on every `execute_model` call. +static ATOMIC_TYPES_CACHE: LazyLock>>> = LazyLock::new( + || Mutex::new(LruCache::new(std::num::NonZero::new(8).unwrap())) +); + +/// Run a model on a set of systems, computing the requested outputs. +/// +/// This is the main entry point for executing a model. It validates the +/// arguments (optionally), converts units, delegates the computation to the +/// model's `execute_inner` callback, and converts the outputs to the requested +/// units. +#[allow(clippy::too_many_lines)] +pub fn execute_model( + model: &Model, + systems: &[Arc], + selected_atoms: *const mts_labels_t, + requested_outputs_json: &str, + check_consistency: bool, + outputs: *mut *mut mts_tensormap_t, + outputs_count: usize, +) -> Result<(), Error> { + // TODO: measure the overhead of getting both of these information at every + // model call. If it is too high, we can add a caching layer for them. + let capabilities = model.capabilities()?; + let requested_inputs = model.requested_inputs()?; + + // parse requested outputs for validation and unit conversion + let requested_outputs: Vec = { + let json = json::parse(requested_outputs_json).map_err(|e| { + Error::Serialization(format!("invalid JSON for requested_outputs: {e}")) + })?; + if !json.is_array() { + return Err(Error::InvalidParameter( + "requested_outputs_json must contain a JSON array".into() + )); + } + let mut result = Vec::new(); + for item in json.members() { + result.push(Quantity::try_from(item)?); + } + result + }; + + if requested_outputs.len() != outputs_count { + return Err(Error::InvalidParameter(format!( + "the number of requested outputs ({}) does not match the outputs buffer length ({})", + requested_outputs.len(), outputs_count + ))); + } + + let mut selected_atoms_labels = None; + if check_consistency { + selected_atoms_labels = if selected_atoms.is_null() { + None + } else { + // increase the refcount on Labels, and only keep one for ourself + // (we do not own the data passed by the caller) + let labels = unsafe { metatensor::Labels::from_raw(selected_atoms) }; + let clone = labels.clone(); + std::mem::forget(labels); + Some(clone) + }; + + check_requested_outputs(&capabilities, &requested_outputs)?; + + let requested_pair_lists = model.requested_pair_lists()?; + check_inputs( + &capabilities, + &requested_pair_lists, + &requested_inputs, + systems, + selected_atoms_labels.as_ref(), + )?; + } + + // always check atomic types (even when check_consistency is false) + { + let mut cache = ATOMIC_TYPES_CACHE.lock().expect("ATOMIC_TYPES_CACHE lock poisoned"); + let atomic_types = AtomicTypes::new(capabilities.atomic_types.clone()); + let valid_types = cache.get_or_insert(atomic_types, || { + ReferenceValue::new( + ndarray::ArrayD::from_shape_vec( + ndarray::IxDyn(&[capabilities.atomic_types.len()]), + capabilities.atomic_types.clone(), + ).expect("atomic_types should be contiguous") + ) + }); + for system in systems { + check_atomic_types(system.types(), valid_types)?; + } + } + + // Convert systems from engine to model units. This returns a new + // Arc for each system — either a refcount bump (no conversion + // needed) or a deep copy that has been scaled. + let model_length_unit = &capabilities.length_unit; + let converted_systems: Vec> = systems.iter() + .map(|s| s.clone().convert_units(model_length_unit, &requested_inputs)) + .collect::, _>>()?; + + // Marshal systems as *const mta_system_t for execute_inner. + // We borrow from the Arcs (Arc::as_ptr) instead of consuming them, so the + // converted_systems Vec keeps the Arcs alive for the duration of the call. + let system_ptrs: Vec<*const mta_system_t> = converted_systems.iter() + .map(|s| Arc::as_ptr(s).cast::()) + .collect(); + + // call execute_inner + let execute_inner = model.0.execute_inner.ok_or_else(|| { + Error::Internal("model is missing an 'execute_inner' callback".into()) + })?; + + let requested_outputs_cstr = CString::new(requested_outputs_json).map_err(|e| { + Error::InvalidParameter(format!("requested_outputs_json contains a null byte: {e}")) + })?; + + let status = unsafe { + execute_inner( + model.0.data, + system_ptrs.as_ptr(), + system_ptrs.len(), + selected_atoms, + requested_outputs_cstr.as_ptr(), + outputs, + outputs_count, + ) + }; + + // drop any outputs that where allocated so far + let drop_partial_outputs = |outputs: *mut *mut mts_tensormap_t, outputs_count: usize| { + for i in 0..outputs_count { + unsafe { + // Set the pointer to null so that the caller doesn't double-free + let ptr = std::mem::replace(&mut *outputs.add(i), std::ptr::null_mut()); + + if !ptr.is_null() { + std::mem::drop(TensorMap::from_raw(ptr)); + } + } + } + }; + + if status != mta_status_t::MTA_SUCCESS { + // the model reported an error; free any outputs it may have partially + // filled to avoid leaking them, then propagate the error + drop_partial_outputs(outputs, outputs_count); + return Err(Error::CallbackError(status)); + } + + // verify all outputs were filled + for i in 0..outputs_count { + if unsafe { *outputs.add(i) }.is_null() { + // free any outputs that were filled so far + drop_partial_outputs(outputs, outputs_count); + return Err(Error::InvalidParameter( + "model's execute_inner did not fill all requested outputs".into() + )); + } + } + + // determine whether we need to take ownership of the outputs at all: + // only if check_consistency is true or any unit conversion factor != 1.0 + let mut needs_ownership = check_consistency; + for requested in &requested_outputs { + #[allow(clippy::float_cmp)] + if let Some(declared) = capabilities.find_output(requested) { + let factor = unit_conversion_factor(&declared.unit, &requested.unit)?; + needs_ownership |= factor != 1.0; + } + } + + if needs_ownership { + // take ownership of returned tensor maps for checking and unit conversion + let mut result: Vec = Vec::with_capacity(outputs_count); + for i in 0..outputs_count { + unsafe { + let ptr = std::mem::replace(&mut *outputs.add(i), std::ptr::null_mut()); + result.push(TensorMap::from_raw(ptr)); + } + } + + // check outputs + if check_consistency { + for (output, requested) in result.iter().zip(&requested_outputs) { + check_quantity( + requested, + output, + systems, + selected_atoms_labels.as_ref(), + )?; + } + } + + // convert output units: declared.unit → requested.unit + for (i, tensor) in result.into_iter().enumerate() { + let requested = &requested_outputs[i]; + + // find the declared quantity in capabilities + let tensor = if let Some(declared) = capabilities.find_output(requested) { + let factor = unit_conversion_factor(&declared.unit, &requested.unit)?; + scale_tensormap(tensor, factor)? + } else { + tensor + }; + + // write back into the outputs buffer + unsafe { + *outputs.add(i) = TensorMap::into_raw(tensor); + } + } + } + + Ok(()) +} diff --git a/metatomic-core/src/model/inputs.rs b/metatomic-core/src/model/inputs.rs new file mode 100644 index 000000000..48e36aa77 --- /dev/null +++ b/metatomic-core/src/model/inputs.rs @@ -0,0 +1,498 @@ +use std::sync::Arc; + +use dlpk::DLDataTypeCode::kDLFloat; +use metatensor::Labels; + +use crate::{Error, ModelCapabilities, PairListOptions, Quantity, System}; +use crate::metadata::DType; + +/// Validate that the requested outputs match the model's capabilities. +/// +/// This checks that each requested output is in the model's capabilities, with +/// compatible `sample_kind` and `explicit_gradients`. +pub(crate) fn check_requested_outputs( + capabilities: &ModelCapabilities, + requested_outputs: &[Quantity], +) -> Result<(), Error> { + for requested in requested_outputs { + // find all capability entries for this output name (a model may declare + // the same output with different sample_kinds, e.g. "energy" as both + // system-level and per-atom) + let possible = capabilities.outputs.iter() + .filter(|q| q.name == requested.name) + .find(|q| q.sample_kind == requested.sample_kind); + + let possible = if let Some(quantity) = possible { + quantity + } else { + // check if the output name exists at all (with any sample_kind) + let name_exists = capabilities.outputs.iter().any(|q| q.name == requested.name); + if name_exists { + let available_kinds: Vec<_> = capabilities.outputs.iter() + .filter(|q| q.name == requested.name) + .map(|q| q.sample_kind.to_string()) + .collect(); + return Err(Error::InvalidParameter(format!( + "this model can not compute '{}' with sample kind '{}', only with sample kind{} [{}]", + requested.name, + requested.sample_kind, + if available_kinds.len() > 1 { "s" } else { "" }, + available_kinds.join(", ") + ))); + } else { + return Err(Error::InvalidParameter(format!( + "this model can not compute '{}', the implemented outputs are [{}]", + requested.name, + capabilities.outputs.iter().map(|q| q.name.full()).collect::>().join(", ") + ))); + } + }; + + // check explicit gradients + for gradient in &requested.gradients { + if !possible.gradients.contains(gradient) { + return Err(Error::InvalidParameter(format!( + "this model can not compute explicit gradients of '{}' with respect to '{}'", + requested.name, gradient + ))); + } + } + } + + Ok(()) +} + +/// Validate that the inputs to a model are consistent with its capabilities. +/// +/// This checks that: +/// - all systems are on the same device and have the same dtype +/// - the systems device and dtype match what the model supports +/// - `selected_atoms` (if provided) has the right names, device, and only +/// contains entries that correspond to actual atoms in the systems +/// - all requested neighbor lists are present on every system +/// - all requested inputs are present on every system +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +#[allow(clippy::too_many_lines)] +pub(crate) fn check_inputs( + capabilities: &ModelCapabilities, + requested_neighbor_lists: &[PairListOptions], + requested_inputs: &[Quantity], + systems: &[Arc], + selected_atoms: Option<&Labels>, +) -> Result<(), Error> { + if systems.is_empty() { + return Ok(()); + } + + let global_device = systems[0].device(); + let global_dtype = systems[0].dtype(); + + // check dtype matches the model's expected dtype + let expected_dlpack_dtype = match capabilities.dtype { + DType::Float32 => ::get_dlpack_data_type(), + DType::Float64 => ::get_dlpack_data_type(), + }; + + if global_dtype != expected_dlpack_dtype { + let actual_dtype = if global_dtype.code == kDLFloat && global_dtype.bits == 32 { + "float32" + } else if global_dtype.code == kDLFloat && global_dtype.bits == 64 { + "float64" + } else { + // unknown float type; display the raw DLDataType + &global_dtype.to_string() + }; + return Err(Error::InvalidParameter(format!( + "wrong dtype for the systems: the model wants {}, we got {}", + capabilities.dtype, actual_dtype + ))); + } + + let mut is_supported_device = false; + for supported_device in &capabilities.supported_devices { + if supported_device.is_valid(global_device) { + is_supported_device = true; + break; + } + } + + if !is_supported_device { + return Err(Error::InvalidParameter(format!( + "the systems device ({}) is not supported by this model, only {:?} are", + global_device, capabilities.supported_devices + ))); + } + + // check selected_atoms + if let Some(selected) = selected_atoms { + if selected.device() != global_device { + return Err(Error::InvalidParameter(format!( + "expected selected_atoms to be on the same device as the systems, got {} and {}", + selected.device(), global_device + ))); + } + + if selected.names() != ["system", "atom"] { + return Err(Error::InvalidParameter(format!( + "invalid names for selected_atoms: expected ['system', 'atom'], got {:?}", + selected.names() + ))); + } + + // build the set of all possible (system, atom) pairs + let total_atoms: usize = systems.iter().map(|s| s.size()).sum(); + let mut possible_values = ndarray::Array2::from_elem((total_atoms, 2), 0i32); + let mut index = 0; + for (system_i, system) in systems.iter().enumerate() { + for atom_i in 0..system.size() { + possible_values[[index, 0]] = system_i as i32; + possible_values[[index, 1]] = atom_i as i32; + index += 1; + } + } + let possible_atoms = metatensor::Labels::new_assume_unique(["system", "atom"], possible_values); + + let intersection = selected.intersection(&possible_atoms, None, None)?; + if intersection.count() != selected.count() { + return Err(Error::InvalidParameter( + "invalid selected_atoms: there are entries that are not possible for the current systems".into() + )); + } + } + + // check each system + for system in systems { + if system.device() != global_device { + return Err(Error::InvalidParameter(format!( + "expected all systems to be on the same device, got {} and {}", + global_device, system.device() + ))); + } + + if system.dtype() != global_dtype { + return Err(Error::InvalidParameter(format!( + "expected all systems to have the same dtype, got {} and {}", + global_dtype, system.dtype() + ))); + } + + // check neighbor lists + for request in requested_neighbor_lists { + if let Some(pairs) = system.get_pairs(request) { + if !pairs.as_ref().gradient_list().is_empty() { + return Err(Error::InvalidParameter(format!( + "neighbors list for {:?} contain gradients, which are \ + not supported", + request + ))); + } + } else { + return Err(Error::InvalidParameter(format!( + "missing neighbors list in the system: the model requested \ + a list for {:?}, but it was not provided in the system", + request + ))); + } + + + } + + // check additional inputs + let known_data = system.known_custom_data(); + for request in requested_inputs { + let found = known_data.iter().any(|known| *known == request.name.full()); + if !found { + return Err(Error::InvalidParameter(format!( + "missing additional input in the system: the model requested \ + '{}' as an extra input, but it was not provided in the system", + request.name + ))); + } + } + } + + Ok(()) +} + + +#[cfg(test)] +mod tests { + use super::*; + + use crate::quantity::{Gradients, Quantity, QuantityName, SampleKind}; + use crate::system::test_system; + + #[test] + #[allow(clippy::too_many_lines)] + fn test_check_requested_outputs() { + let capabilities = ModelCapabilities { + outputs: vec![Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![Gradients::Positions], + sample_kind: SampleKind::System, + }], + atomic_types: vec![1, 6, 8], + interaction_range: 5.0, + length_unit: "nm".into(), + supported_devices: vec![crate::Device::cpu()], + dtype: DType::Float32, + }; + + // happy path + let requested_outputs = vec![Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::System, + }]; + check_requested_outputs(&capabilities, &requested_outputs).unwrap(); + + // requested output not in capabilities + let bad_outputs = vec![Quantity { + name: QuantityName::new("custom::forces".into()).unwrap(), + unit: "eV/A".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + }]; + let err = check_requested_outputs(&capabilities, &bad_outputs).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model can not compute 'custom::forces', \ + the implemented outputs are [energy]" + ); + + // requested gradient not in capabilities + let bad_grad_outputs = vec![Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![Gradients::Strain], + sample_kind: SampleKind::System, + }]; + let err = check_requested_outputs(&capabilities, &bad_grad_outputs).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model can not compute explicit gradients \ + of 'energy' with respect to 'strain'" + ); + + // sample kind mismatch: requesting atom when model only offers system + let bad_sample_outputs = vec![Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + }]; + let err = check_requested_outputs(&capabilities, &bad_sample_outputs).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model can not compute 'energy' with sample \ + kind 'atom', only with sample kind [system]" + ); + + // model with multiple sample_kinds for the same output — requesting + // one that exists should pass + let multi_capabilities = ModelCapabilities { + outputs: vec![ + Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::System, + }, + Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + }, + ], + atomic_types: vec![1, 6, 8], + interaction_range: 5.0, + length_unit: "Angstrom".into(), + supported_devices: vec![crate::Device::cpu()], + dtype: DType::Float32, + }; + let atom_outputs = vec![Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + }]; + check_requested_outputs(&multi_capabilities, &atom_outputs).unwrap(); + + // requesting a sample kind that doesn't match any of the available ones + let pair_outputs = vec![Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::AtomPair, + }]; + let err = check_requested_outputs(&multi_capabilities, &pair_outputs).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: this model can not compute 'energy' with \ + sample kind 'atom_pair', only with sample kinds [system, atom]" + ); + } + + #[test] + #[allow(clippy::too_many_lines)] + fn test_check_inputs() { + // build capabilities matching the test_system() helper + let capabilities = ModelCapabilities { + outputs: vec![Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![Gradients::Positions], + sample_kind: SampleKind::System, + }], + atomic_types: vec![1, 6, 8], + interaction_range: 5.0, + length_unit: "nm".into(), + supported_devices: vec![crate::Device::cpu()], + dtype: DType::Float32, + }; + + let requested_neighbor_lists = [PairListOptions { + cutoff: 3.5, + full_list: true, + strict: false, + requestors: vec![], + }]; + + let requested_inputs = vec![Quantity { + name: QuantityName::new("custom::data/name".into()).unwrap(), + unit: String::new(), + description: None, + gradients: vec![], + sample_kind: SampleKind::System, + }]; + + let systems = vec![test_system("f32")]; + + // happy path — everything matches + check_inputs( + &capabilities, + &requested_neighbor_lists, + &requested_inputs, + &systems, + None, + ).unwrap(); + + // wrong dtype + let f64_capabilities = ModelCapabilities { + outputs: capabilities.outputs.clone(), + atomic_types: capabilities.atomic_types.clone(), + interaction_range: capabilities.interaction_range, + length_unit: capabilities.length_unit.clone(), + supported_devices: capabilities.supported_devices.clone(), + dtype: DType::Float64, + }; + let err = check_inputs( + &f64_capabilities, + &requested_neighbor_lists, + &requested_inputs, + &systems, + None, + ).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: wrong dtype for the systems: the model wants float64, we got float32" + ); + + // missing neighbor list + let bad_nl = vec![PairListOptions { + cutoff: 5.0, + full_list: false, + strict: true, + requestors: vec![], + }]; + let err = check_inputs( + &capabilities, + &bad_nl, + &requested_inputs, + &systems, + None, + ).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: missing neighbors list in the system: \ + the model requested a list for PairListOptions { cutoff: 5.0, \ + full_list: false, strict: true, requestors: [] }, but it was not \ + provided in the system" + ); + + // missing requested input + let bad_inputs = vec![Quantity { + name: QuantityName::new("custom::missing".into()).unwrap(), + unit: String::new(), + description: None, + gradients: vec![], + sample_kind: SampleKind::System, + }]; + let err = check_inputs( + &capabilities, + &requested_neighbor_lists, + &bad_inputs, + &systems, + None, + ).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: missing additional input in the system: \ + the model requested 'custom::missing' as an extra input, but it was \ + not provided in the system" + ); + + // invalid selected_atoms names + let bad_selected = metatensor::Labels::new(["foo", "bar"], [[0i32, 0]]); + let err = check_inputs( + &capabilities, + &requested_neighbor_lists, + &requested_inputs, + &systems, + Some(&bad_selected), + ).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid names for selected_atoms: expected \ + ['system', 'atom'], got [\"foo\", \"bar\"]" + ); + + // selected_atoms with out-of-range atom + let bad_selected = metatensor::Labels::new(["system", "atom"], [[0, 99]]); + let err = check_inputs( + &capabilities, + &requested_neighbor_lists, + &requested_inputs, + &systems, + Some(&bad_selected), + ).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid selected_atoms: there are entries that \ + are not possible for the current systems" + ); + + // valid selected_atoms + let good_selected = metatensor::Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + check_inputs( + &capabilities, + &requested_neighbor_lists, + &requested_inputs, + &systems, + Some(&good_selected), + ).unwrap(); + } +} diff --git a/metatomic-core/src/model/mod.rs b/metatomic-core/src/model/mod.rs new file mode 100644 index 000000000..e920f8de9 --- /dev/null +++ b/metatomic-core/src/model/mod.rs @@ -0,0 +1,8 @@ +mod inputs; + +#[allow(clippy::module_inception)] +mod model; +pub use self::model::Model; + +mod execute; +pub use self::execute::execute_model; diff --git a/metatomic-core/src/model/model.rs b/metatomic-core/src/model/model.rs new file mode 100644 index 000000000..4144d7daf --- /dev/null +++ b/metatomic-core/src/model/model.rs @@ -0,0 +1,283 @@ +use std::ffi::c_void; + +use crate::{Error, ModelCapabilities, ModelMetadata, PairListOptions, Quantity}; +use crate::c_api::{mta_model_t, mta_status_t, mta_string_t, mta_string_free}; + +/// A loaded atomistic model, ready to be executed on a set of systems. +/// +/// `Model` wraps a [`mta_model_t`] vtable provided by a plugin. It gives +/// access to the model's metadata and capabilities, and can be run with +/// [`execute_model`]. +#[repr(transparent)] +pub struct Model(pub(crate) mta_model_t); + +impl Drop for Model { + fn drop(&mut self) { + if let Some(unload) = self.0.unload { + unsafe { unload(self.0.data) }; + } + } +} + +fn call_string_callback( + callback: unsafe extern "C" fn(*const c_void, *mut mta_string_t) -> mta_status_t, + data: *const c_void, +) -> Result { + let mut output = mta_string_t::null(); + let status = unsafe { callback(data, &mut output) }; + if status != mta_status_t::MTA_SUCCESS { + unsafe { mta_string_free(output) }; + return Err(Error::CallbackError(status)); + } + let json_str = output.as_str().to_owned(); + unsafe { mta_string_free(output) }; + return Ok(json_str); +} + +impl Model { + /// Create a new `Model` from the corresponding C API struct. + /// + /// The `Model` takes ownership of `model` and will call its `unload` + /// callback when dropped. + pub fn new(model: mta_model_t) -> Self { + return Model(model); + } + + /// Create a `&Model` from a `&mta_model_t` without taking ownership. + /// + /// This is used by the C API to call [`execute_model`] on a model that is + /// owned by the caller (e.g. passed by value to `mta_execute_model`). The + /// returned reference does not own the model and will not call `unload` when + /// dropped. + pub fn from_ref(model: &mta_model_t) -> &Self { + // SAFETY: `Model` is repr(transparent) over mta_model_t + unsafe { &*std::ptr::from_ref(model).cast::() } + } + + /// Extract the underlying C API struct, transferring ownership to the caller. + /// + /// The caller is responsible for eventually calling the `unload` callback + /// on the returned [`mta_model_t`] to free its resources. The `Model`'s + /// own `Drop` implementation is skipped. + pub fn into_raw(self) -> mta_model_t { + let model = std::mem::ManuallyDrop::new(self); + return unsafe { std::ptr::read(&model.0) }; + } + + /// Get the metadata describing this model (name, authors, description, + /// references, ...). + pub fn metadata(&self) -> Result { + let callback = self.0.metadata.ok_or_else(|| { + Error::Internal("model is missing a 'metadata' callback".into()) + })?; + let json_str = call_string_callback(callback, self.0.data)?; + let json = json::parse(&json_str).map_err(|e| { + Error::Serialization(format!("model returned invalid JSON for metadata: {}", e)) + })?; + return ModelMetadata::try_from(&json); + } + + /// Get the capabilities of this model: which outputs it can compute, which + /// atomic types it supports, its interaction range, length unit, supported + /// devices, and data type. + pub fn capabilities(&self) -> Result { + let callback = self.0.capabilities.ok_or_else(|| { + Error::Internal("model is missing a 'capabilities' callback".into()) + })?; + let json_str = call_string_callback(callback, self.0.data)?; + let json = json::parse(&json_str).map_err(|e| { + Error::Serialization(format!("model returned invalid JSON for capabilities: {}", e)) + })?; + return ModelCapabilities::try_from(&json); + } + + /// Get the pair lists (neighbor lists) this model needs as input. + /// + /// The engine must compute these and attach them to every system with + /// `mta_system_add_pairs` before calling [`execute_model`]. + pub fn requested_pair_lists(&self) -> Result, Error> { + let callback = self.0.requested_pair_lists.ok_or_else(|| { + Error::Internal("model is missing a 'requested_pair_lists' callback".into()) + })?; + let json_str = call_string_callback(callback, self.0.data)?; + let json = json::parse(&json_str).map_err(|e| { + Error::Serialization(format!("model returned invalid JSON for requested_pair_lists: {}", e)) + })?; + if !json.is_array() { + return Err(Error::Serialization( + "model returned invalid JSON for requested_pair_lists, expected an array".into() + )); + } + let mut result = Vec::new(); + for item in json.members() { + result.push(PairListOptions::try_from(item)?); + } + return Ok(result); + } + + /// Get the additional per-system inputs this model needs. + /// + /// The engine must attach these to every system with + /// `mta_system_add_custom_data` before calling [`execute_model`]. + pub fn requested_inputs(&self) -> Result, Error> { + let callback = self.0.requested_inputs.ok_or_else(|| { + Error::Internal("model is missing a 'requested_inputs' callback".into()) + })?; + let json_str = call_string_callback(callback, self.0.data)?; + let json = json::parse(&json_str).map_err(|e| { + Error::Serialization(format!("model returned invalid JSON for requested_inputs: {}", e)) + })?; + if !json.is_array() { + return Err(Error::Serialization( + "model returned invalid JSON for requested_inputs, expected an array".into() + )); + } + let mut result = Vec::new(); + for item in json.members() { + result.push(Quantity::try_from(item)?); + } + return Ok(result); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::c_api::{mta_model_t, mta_status_t, mta_string_t}; + + // Each function below is a stand-in for what a real plugin would implement. + // They simply write a hard-coded JSON string into the output mta_string_t + // and return MTA_SUCCESS. + unsafe extern "C" fn metadata_impl( + _data: *const c_void, + out: *mut mta_string_t, + ) -> mta_status_t { + unsafe { + *out = mta_string_t::new(r#"{ + "type": "metatomic_model_metadata", + "name": "test-model", + "authors": ["Alice"], + "description": "A test model", + "references": {"model": [], "architecture": [], "implementation": []}, + "extra": {} + }"#); + } + return mta_status_t::MTA_SUCCESS; + } + + unsafe extern "C" fn + capabilities_impl( + _data: *const c_void, + out: *mut mta_string_t, + ) -> mta_status_t { + unsafe { + *out = mta_string_t::new(r#"{ + "type": "metatomic_model_capabilities", + "outputs": [{ + "type": "metatomic_quantity", + "name": "energy", + "unit": "eV", + "gradients": ["positions"], + "sample_kind": "system" + }, + { + "type": "metatomic_quantity", + "name": "custom::output", + "unit": "", + "gradients": [], + "sample_kind": "atom_pair" + }], + "atomic_types": [1, 6], + "interaction_range": 5.0, + "length_unit": "Angstrom", + "supported_devices": ["cpu"], + "dtype": "float32" + }"#); + } + return mta_status_t::MTA_SUCCESS; + } + + unsafe extern "C" fn requested_pair_lists_impl( + _data: *const c_void, + out: *mut mta_string_t, + ) -> mta_status_t { + unsafe { + *out = mta_string_t::new(format!(r#"[{{ + "type": "metatomic_pair_list_options", + "cutoff": "{:#x}", + "full_list": true, + "strict": true + }}]"#, 3.5_f64.to_bits())); + } + return mta_status_t::MTA_SUCCESS; + } + + unsafe extern "C" fn requested_inputs_impl( + _data: *const c_void, + out: *mut mta_string_t, + ) -> mta_status_t { + unsafe { + *out = mta_string_t::new(r#"[{ + "type": "metatomic_quantity", + "name": "charge", + "unit": "e", + "gradients": [], + "sample_kind": "atom" + }]"#); + } + return mta_status_t::MTA_SUCCESS; + } + + + fn test_model() -> Model { + Model(mta_model_t { + metadata: Some(metadata_impl), + capabilities: Some(capabilities_impl), + requested_pair_lists: Some(requested_pair_lists_impl), + requested_inputs: Some(requested_inputs_impl), + ..mta_model_t::null() + }) + } + + #[test] + fn metadata() { + let metadata = test_model().metadata().unwrap(); + assert_eq!(metadata.name, "test-model"); + assert_eq!(metadata.authors, vec!["Alice"]); + assert_eq!(metadata.description, "A test model"); + } + + + #[test] + fn capabilities() { + let capabilities = test_model().capabilities().unwrap(); + + assert_eq!(capabilities.outputs.len(), 2); + assert_eq!(capabilities.outputs[0].name.full(), "energy"); + assert_eq!(capabilities.outputs[0].unit, "eV"); + + assert_eq!(capabilities.outputs[1].name.full(), "custom::output"); + assert_eq!(capabilities.outputs[1].unit, ""); + + assert_eq!(capabilities.atomic_types, vec![1, 6]); + assert_eq!(capabilities.interaction_range.to_bits(), 5.0_f64.to_bits()); + assert_eq!(capabilities.length_unit, "Angstrom"); + } + + #[test] + fn requested_pair_lists() { + let options = test_model().requested_pair_lists().unwrap(); + assert_eq!(options.len(), 1); + assert_eq!(options[0].cutoff.to_bits(), 3.5_f64.to_bits()); + assert!(options[0].full_list); + assert!(options[0].strict); + } + + #[test] + fn requested_inputs() { + let inputs = test_model().requested_inputs().unwrap(); + assert_eq!(inputs.len(), 1); + assert_eq!(inputs[0].name.full(), "charge"); + assert_eq!(inputs[0].unit, "e"); + } +} diff --git a/metatomic-core/src/plugin.rs b/metatomic-core/src/plugin.rs new file mode 100644 index 000000000..4a7d485e8 --- /dev/null +++ b/metatomic-core/src/plugin.rs @@ -0,0 +1,221 @@ +use std::ffi::CStr; +use std::sync::{Mutex, LazyLock}; + +use libloading::Library; + +use crate::c_api::{mta_model_t, mta_plugin_t, mta_register_plugin, mta_status_t}; +use crate::{Error, Model}; + +/// ABI version of the metatomic plugin interface. +/// +/// This increases anytime the plugin or model C API changes in a non backward +/// compatible way. Plugins compiled with an incompatible ABI version will be +/// rejected at registration time. +pub const MTA_ABI_VERSION: i32 = 1; + +/// The list of registered plugins in the current process. +static PLUGINS: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); +/// Keep the loaded libraries alive for the entire process lifetime, to ensure +/// that the plugin code is not unloaded while it's still in use. +static LIBRARIES: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); + +pub struct Plugin(mta_plugin_t); + +impl Plugin { + /// Create a new plugin from the C struct + pub fn new(plugin: mta_plugin_t) -> Result { + if plugin.name.is_null() { + return Err(Error::InvalidParameter( + "can not register plugin: plugin `name` is NULL".into(), + )); + } + + let c_str_name = unsafe { CStr::from_ptr(plugin.name) }; + if c_str_name.to_str().is_err() { + return Err(Error::InvalidParameter(format!( + "can not register plugin: plugin `name` is not valid UTF-8: {}", + c_str_name.to_string_lossy() + ))); + } + + if plugin.load_model.is_none() { + return Err(Error::InvalidParameter( + "can not register plugin: plugin `load_model` callback is NULL".into(), + )); + } + + if plugin.abi_version != MTA_ABI_VERSION { + let name = unsafe { + CStr::from_ptr(plugin.name).to_string_lossy() + }; + + return Err(Error::InvalidParameter(format!( + "can not register plugin '{}': plugin ABI version is {}, but metatomic expects {}", + name, + plugin.abi_version, + MTA_ABI_VERSION, + ))); + } + + Ok(Plugin(plugin)) + } + + /// Get the plugin name. + pub fn name(&self) -> &str { + unsafe { + return CStr::from_ptr(self.0.name) + .to_str() + .expect("invalid UTF-8 in plugin name"); + } + } + + /// Try to load a model with this plugin. + pub fn load_model( + &self, + load_from: &CStr, + options_json: &CStr, + ) -> Result { + let load_model = self.0.load_model.expect("`load_model` is NULL"); + + let mut model = mta_model_t::null(); + let status = unsafe { + load_model(load_from.as_ptr(), options_json.as_ptr(), &mut model) + }; + + if status != mta_status_t::MTA_SUCCESS { + return Err(Error::CallbackError(status)); + } + + return Ok(Model::new(model)); + } +} + +/// Register a new plugin in the current process. +pub fn register_plugin(plugin: Plugin) -> Result<(), Error> { + let mut plugins = PLUGINS.lock().expect("plugin registry mutex was poisoned"); + if plugins.iter().any(|existing| existing.name() == plugin.name()) { + return Err(Error::InvalidParameter(format!( + "a plugin named '{}' is already registered", + plugin.name() + ))); + } + + plugins.push(plugin); + return Ok(()); +} + +/// Load a plugin from a shared library. +/// +/// The shared library must export the symbols generated by the +/// `MTA_REGISTER_PLUGIN` C macro. +pub fn load_plugin(path: Option<&str>) -> Result<(), Error> { + // this needs to be kept in sync with the definition in `MTA_REGISTER_PLUGIN` in build.rs + type PluginInitFn = unsafe extern "C" fn(abi: i32, data: *mut std::ffi::c_void) -> mta_status_t; + + let library = if let Some(path) = path { + let library = unsafe { Library::new(path) }; + library.map_err(|error| { + std::io::Error::other( + format!("failed to load plugin '{}': {}", path, error), + ) + })? + } else { + #[cfg(unix)] + { + libloading::os::unix::Library::this().into() + } + + #[cfg(windows)] + { + let library = libloading::os::windows::Library::this(); + library.map_err(|error| { + std::io::Error::other( + format!("failed to load as a plugin: {}", error), + ) + })?.into() + } + + #[cfg(not(any(unix, windows)))] + { + panic!("loading the current process library is not supported on this platform"); + } + }; + + let status = unsafe { + let init_plugin = library.get::(b"mta_plugin_init\0") + .map_err(|error| Error::InvalidParameter(format!( + "failed to load plugin registration symbol from '{}': {}", + path.unwrap_or(""), error + )))?; + init_plugin(MTA_ABI_VERSION, mta_register_plugin as *mut std::ffi::c_void) + }; + + if status != mta_status_t::MTA_SUCCESS { + return Err(Error::CallbackError(status)); + } + + LIBRARIES.lock().expect("loaded plugin registry mutex was poisoned").push(library); + + return Ok(()); +} + +/// Load a model from `load_from`, using the given options. +pub fn load_model( + load_from: &CStr, + options_json: &CStr, + plugin_name: Option<&str>, +) -> Result { + let plugins = PLUGINS.lock().expect("plugin registry mutex was poisoned"); + + if let Some(plugin_name) = plugin_name { + for plugin in plugins.iter() { + if plugin.name() == plugin_name { + return plugin.load_model(load_from, options_json).map_err(|e| { + if let Error::CallbackError(mta_status_t::MTA_UNSUPPORTED_MODEL_ERROR) = e { + Error::InvalidParameter(format!( + "failed to load model from '{}': plugin '{}' could not load the model", + load_from.to_string_lossy(), + plugin_name + )) + } else { + e + } + }); + } + } + + return Err(Error::InvalidParameter(format!( + "no plugin named '{}' is registered", + plugin_name + ))); + } + + for plugin in plugins.iter() { + match plugin.load_model(load_from, options_json) { + Ok(model) => return Ok(model), + Err(e) => { + if let Error::CallbackError(mta_status_t::MTA_UNSUPPORTED_MODEL_ERROR) = e { + // try the next plugin + continue; + } else { + return Err(e); + } + } + } + } + + let message = if plugins.is_empty() { + "no plugin is registered".into() + } else { + format!( + "tried the following plugins, but none could load the model: {}", + plugins.iter().map(|p| p.name()).collect::>().join(", ") + ) + }; + + return Err(Error::InvalidParameter(format!( + "failed to load model from '{}': {}", + load_from.to_string_lossy(), + message + ))); +} diff --git a/metatomic-core/src/quantity/charge.rs b/metatomic-core/src/quantity/charge.rs new file mode 100644 index 000000000..d133015e0 --- /dev/null +++ b/metatomic-core/src/quantity/charge.rs @@ -0,0 +1,392 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, SINGLE_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + + +/// Check the layout of the "charge" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "charge"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::System, SampleKind::Atom])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[])?; + + let expected_properties = ExpectedLabels { + names: &["charge"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("charge".into()).unwrap(), + unit: "e".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + } + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2]], + ); + let properties = Labels::new(["charge"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(); + TensorBlock::new(values, &samples, &[], &properties).unwrap() + } + + fn valid_charge() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_charge(), &[system(3)], None).unwrap(); + + // Also check SampleKind::System + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.5]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["charge"], [[0]]) + ).unwrap(); + + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &charge, &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + // Empty systems slice, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &charge, &[], None).unwrap(); + + // System with 0 atoms, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &charge, &[system(0)], None).unwrap(); + + // Empty systems slice, per-system output + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system"], + Array2::::from_shape_vec((0, 1), vec![]).unwrap(), + ), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &charge, &[], None).unwrap(); + } + + #[test] + fn selected_atoms() { + // Per-atom output with selected_atoms across multiple systems + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &charge, &systems, Some(&selected_atoms)).unwrap(); + + // Per-system values with selected_atoms across multiple systems + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 1], vec![5.0, 6.0]).unwrap(), + &Labels::new(["system"], [[0], [1]]), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &charge, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![4, 1], vec![1.0; 4]).unwrap(), + &Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2], [1, 0]], + ), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &charge, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::AtomPair; + let err = check(&request, &valid_charge(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'charge': expected one of [system, atom], got 'atom_pair'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let charge = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'charge': expected a single block, but found 0 blocks" + ); + + let charge = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'charge': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let charge = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'charge': expected a single block with key '_', but found key names [foo]" + ); + + let charge = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'charge': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'charge': expected names [charge], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["charge"], [[1]]), + ).unwrap(); + + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'charge': expected [[0]]" + ); + } + + #[test] + fn has_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: components for 'charge' should be empty" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["charge"], [[0]]) + ).unwrap(); + + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'charge': expected [system, atom], got [system]" + ); + + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0]]), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&request, &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'charge': expected [system], got [system, atom]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let charge = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &charge, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'charge': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + // samples that are not in the selected_atoms + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &charge, &[system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'charge', they do not match the `systems` and `selected_atoms`" + ); + + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 1], vec![3.0, 4.0]).unwrap(), + // systems that are not in the selected_atoms + &Labels::new(["system"], [[0], [1]]), + &[], + &Labels::new(["charge"], [[0]]), + ).unwrap(); + let charge = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&request, &charge, &[system(3), system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'charge', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/checks.rs b/metatomic-core/src/quantity/checks.rs new file mode 100644 index 000000000..95ba71070 --- /dev/null +++ b/metatomic-core/src/quantity/checks.rs @@ -0,0 +1,401 @@ +use std::sync::Arc; +use std::collections::BTreeSet; +use std::sync::LazyLock; + +use metatensor::{Labels, TensorBlockRef, TensorMap}; + +use super::Quantity; + +use crate::{Error, SampleKind, System}; +use crate::kernels::{is_equal_i32, ReferenceValue}; + + +pub(super) static XYZ_LABELS_REFERENCE: LazyLock> = LazyLock::new(|| { + ReferenceValue::new( + ndarray::ArrayD::from_shape_vec( + ndarray::IxDyn(&[3usize, 1]), + vec![0i32, 1, 2], + ).unwrap() + ) +}); + +pub(super) static SINGLE_LABELS_REFERENCE: LazyLock> = LazyLock::new(|| { + ReferenceValue::new( + ndarray::ArrayD::from_shape_vec( + ndarray::IxDyn(&[1usize, 1]), + vec![0i32], + ).unwrap() + ) +}); + +/// Check that the `sample_kind` is one of the valid kinds for the given quantity. +pub(super) fn it_should_have_valid_sample_kind( + context: &str, + sample_kind: SampleKind, + valid_kinds: &[SampleKind] +) -> Result<(), Error> { + if !valid_kinds.contains(&sample_kind) { + return Err(Error::InvalidParameter(format!( + "invalid sample_kind for {}: expected one of [{}], got '{}'", + context, + valid_kinds.iter().map(|k| k.to_string()).collect::>().join(", "), + sample_kind + ))); + } + + return Ok(()); +} + +/// Ensure the TensorMap has a single block with the expected key +pub(super) fn it_should_have_a_single_block(context: &str, value: &TensorMap) -> Result<(), Error> { + let keys = value.keys(); + if keys.count() != 1 { + return Err(Error::InvalidParameter(format!( + "invalid {}: expected a single block, but found {} blocks", + context, + keys.count() + ))); + } + + if keys.names() != ["_"] { + return Err(Error::InvalidParameter(format!( + "invalid {}: expected a single block with key '_', but found key names [{}]", + context, + keys.names().join(", ") + ))); + } + + let values = keys.values().as_dlpack(dlpk::DLDevice::cpu(), None, dlpk::DLPackVersion::current())?; + if !is_equal_i32(values.as_ref(), &SINGLE_LABELS_REFERENCE)? { + return Err(Error::InvalidParameter(format!( + "invalid {}: expected a single block with key value 0", + context, + ))); + } + + Ok(()) +} + +/// Validate the values for "system" samples +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +fn validate_system_samples( + context: &str, + samples: &Labels, + systems: &[Arc], + selected_atoms: Option<&Labels>, +) -> Result<(), Error> { + let values = if let Some(selected) = selected_atoms { + // only include the systems that are present in the selected_atoms + let mut values = BTreeSet::new(); + for [system_i, _] in selected.iter_fixed_size::<2>() { + values.insert(system_i.i32()); + } + ndarray::Array2::from_shape_vec( + (values.len(), 1), + values.into_iter().collect() + ).expect("created invalid array for system samples") + } else { + ndarray::Array2::from_shape_vec( + (systems.len(), 1), + (0..systems.len()).map(|s| s as i32).collect() + ).expect("created invalid array for system samples") + }; + + let expected = Labels::new_assume_unique(["system"], values); + + if expected.union(samples, None, None)?.count() != expected.count() { + return Err(Error::InvalidParameter(format!( + "invalid samples for {}, they do not match the \ + `systems` and `selected_atoms`", + context, + // TODO: add Labels::print to metatensor and use it here + ))); + } + + return Ok(()); +} + +/// Validate the values for "atom" samples +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +fn validate_atom_samples( + context: &str, + samples: &Labels, + systems: &[Arc], + selected_atoms: Option<&Labels>, +) -> Result<(), Error> { + let total_atoms: usize = systems.iter().map(|s| s.size()).sum(); + let mut values = ndarray::Array2::from_elem((total_atoms, 2), 0); + + let mut index = 0; + for (system_i, system) in systems.iter().enumerate() { + for atom_i in 0..system.size() { + values[[index, 0]] = system_i as i32; + values[[index, 1]] = atom_i as i32; + index += 1; + } + } + let mut expected = Labels::new_assume_unique(["system", "atom"], values); + if let Some(selected) = selected_atoms { + expected = expected.intersection(selected, None, None)?; + } + + if expected.union(samples, None, None)?.count() != expected.count() { + return Err(Error::InvalidParameter(format!( + "invalid samples for {}, they do not match the \ + `systems` and `selected_atoms`", + context, + // TODO: add Labels::print to metatensor and use it here + ))); + } + + return Ok(()); +} + +/// Validate the values for "atom_pair" samples +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap, clippy::cast_sign_loss)] +fn validate_atom_pair_samples( + context: &str, + samples: &Labels, + systems: &[Arc], + selected_atoms: Option<&Labels>, +) -> Result<(), Error> { + for [system, first_atom, second_atom, _, _, _] in samples.iter_fixed_size::<6>() { + let system = system.i32(); + let first_atom = first_atom.i32(); + let second_atom = second_atom.i32(); + + if system < 0 || system >= systems.len() as i32 { + return Err(Error::InvalidParameter(format!( + "invalid system index in samples for {}: {} is out of bounds", + context, + system + ))); + } + + let n_atoms = systems[system as usize].size() as i32; + if first_atom < 0 || first_atom >= n_atoms { + return Err(Error::InvalidParameter(format!( + "invalid first_atom index in samples for {}: {} is out of bounds for system {}", + context, + first_atom, + system + ))); + } + if second_atom < 0 || second_atom >= n_atoms { + return Err(Error::InvalidParameter(format!( + "invalid second_atom index in samples for {}: {} is out of bounds for system {}", + context, + second_atom, + system + ))); + } + } + + return Ok(()); +} + +/// Validates that the sample labels match the expected structure based on the +/// sample_kind and the systems/selected_atoms provided. +pub(super) fn it_should_have_valid_samples( + context: &str, + sample_kind: SampleKind, + block: TensorBlockRef<'_>, + systems: &[Arc], + selected_atoms: Option<&Labels>, +) -> Result<(), Error> { + let expected_samples_names: &[&str] = match sample_kind { + SampleKind::System => &["system"], + SampleKind::Atom => &["system", "atom"], + SampleKind::AtomPair => &[ + "system", + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + }; + + let samples = block.samples(); + if samples.names() != expected_samples_names { + return Err(Error::InvalidParameter(format!( + "invalid sample names for {}: expected [{}], got [{}]", + context, + expected_samples_names.join(", "), + samples.names().join(", ") + ))); + } + + // Check if the samples entries match the systems and selected_atoms + match sample_kind { + SampleKind::System => validate_system_samples(context, &samples, systems, selected_atoms), + SampleKind::Atom => validate_atom_samples(context, &samples, systems, selected_atoms), + SampleKind::AtomPair => validate_atom_pair_samples(context, &samples, systems, selected_atoms), + } +} + +#[derive(Debug, Clone, Copy)] +pub(super) struct ExpectedLabels<'a> { + /// Expected names of the labels + pub names: &'a [&'a str], + /// Expected values of the labels + pub values: &'a ReferenceValue, + /// Message to display if the values do not match, showing the expected values + pub values_message: &'a str, +} + +pub(super) fn it_should_have_expected_labels( + context: &str, + labels_kind: &str, + labels: &Labels, + expected: ExpectedLabels<'_>, +) -> Result<(), Error> { + + if labels.names() != expected.names { + return Err(Error::InvalidParameter(format!( + "invalid {} for {}: expected names [{}], got [{}]", + labels_kind, + context, + expected.names.join(", "), + labels.names().join(", ") + ))); + } + + let values = labels.values().as_dlpack(dlpk::DLDevice::cpu(), None, dlpk::DLPackVersion::current())?; + if !is_equal_i32(values.as_ref(), expected.values)? { + return Err(Error::InvalidParameter(format!( + "invalid {} values for {}: expected {}", + labels_kind, + context, + expected.values_message + ))); + } + Ok(()) +} + +pub(super) fn it_should_have_expected_components( + context: &str, + block: TensorBlockRef<'_>, + expected: &[ExpectedLabels<'_>], +) -> Result<(), Error> { + let components = block.components(); + if components.len() != expected.len() { + if expected.is_empty() { + return Err(Error::InvalidParameter(format!( + "components for {} should be empty", + context + ))); + } else { + return Err(Error::InvalidParameter(format!( + "invalid components for {}: expected {} component(s), got {}", + context, + expected.len(), + components.len() + ))); + } + } + + for (component, &expected) in components.iter().zip(expected) { + it_should_have_expected_labels(context, "components", component, expected)?; + } + + return Ok(()); +} + +pub(super) fn it_should_have_expected_gradients( + context: &str, + request: &Quantity, + block: TensorBlockRef<'_>, + potential_gradients: &[&str], +) -> Result<(), Error> { + if potential_gradients.is_empty() && block.gradients().len() > 0 { + return Err(Error::InvalidParameter(format!( + "invalid gradients for {}: expected no gradients, but found \ + gradients with respect to [{}]", + context, + block.gradient_list().join(", ") + ))); + } + + for (parameter, gradient) in block.gradients() { + if !potential_gradients.contains(¶meter) { + return Err(Error::InvalidParameter(format!( + "invalid gradient '{}' for {}: expected one of [{}]", + parameter, + context, + potential_gradients.join(", ") + ))); + } + + match parameter { + "strain" => { + if !request.gradients.contains(&super::Gradients::Strain) { + return Err(Error::InvalidParameter(format!( + "invalid gradient 'strain' for {}: these gradients were not requested", + context + ))); + } + + let context = format!("strain gradient of {}", context); + if gradient.samples().names() != ["sample"] { + return Err(Error::InvalidParameter(format!( + "invalid samples for {}: expected samples names ['sample'], got [{}]", + context, + gradient.samples().names().join(", ") + ))); + } + + it_should_have_expected_components( + &context, + gradient, + &[ + ExpectedLabels { + names: &["xyz_1"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ExpectedLabels { + names: &["xyz_2"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ] + )?; + }, + "positions" => { + if !request.gradients.contains(&super::Gradients::Positions) { + return Err(Error::InvalidParameter(format!( + "invalid gradient 'positions' for {}: these gradients were not requested", + context + ))); + } + + let context = format!("positions gradient of {}", context); + if gradient.samples().names() != ["sample", "system", "atom"] { + return Err(Error::InvalidParameter(format!( + "invalid samples for {}: expected samples names ['sample', 'system', 'atom'], got [{}]", + context, + gradient.samples().names().join(", ") + ))); + } + + it_should_have_expected_components( + &context, + gradient, + &[ + ExpectedLabels { + names: &["xyz"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ] + )?; + }, + _ => { + unreachable!("got unknown gradient parameter for {}: {}", context, parameter); + } + } + } + + return Ok(()); +} diff --git a/metatomic-core/src/quantity/energy.rs b/metatomic-core/src/quantity/energy.rs new file mode 100644 index 000000000..9ac0876e4 --- /dev/null +++ b/metatomic-core/src/quantity/energy.rs @@ -0,0 +1,559 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, SINGLE_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; +use crate::kernels::ReferenceValue; + + +/// Check the layout of one of the energy-related quantities ("energy", +/// "energy_ensemble", "energy_uncertainty"). +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + let name = &request.name; + assert!(!name.is_custom()); + assert!(name.base() == "energy" || name.base() == "energy_ensemble" || name.base() == "energy_uncertainty"); + + let context = format!("'{}'", name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::System, SampleKind::Atom])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[])?; + + if name.base() == "energy" || name.base() == "energy_uncertainty" { + checks::it_should_have_expected_labels( + &context, + "properties", + &block.properties(), + ExpectedLabels { + names: &["energy"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]", + } + )?; + } else { + let n_ensemble_members = *block.values().shape()?.last().expect("energy block has an empty shape"); + let reference = ReferenceValue::new(ndarray::ArrayD::from_shape_vec( + vec![n_ensemble_members, 1], (0..n_ensemble_members as i32).collect() + ).expect("created invalid array for energy_ensemble properties")); + checks::it_should_have_expected_labels( + &context, + "properties", + &block.properties(), + ExpectedLabels { + names: &["energy"], + values: &reference, + values_message: "[[0, ..., n]]", + } + )?; + } + + checks::it_should_have_expected_gradients(&context, request, block, &["strain", "positions"])?; + return Ok(()); +} + +#[cfg(test)] +mod tests { + // use a macro to generate the test code for all three energy-related quantities + macro_rules! energy_tests { + ($base_name: ident) => { + mod $base_name { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Gradients, Quantity, QuantityName, SampleKind, System}; + + use super::super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new(String::from(stringify!($base_name))).unwrap(), + unit: "eV".into(), + description: None, + gradients: vec![Gradients::Positions, Gradients::Strain], + sample_kind: SampleKind::Atom, + } + } + + fn n_properties() -> usize { + if stringify!($base_name) == "energy_ensemble" { 2 } else { 1 } + } + + fn property_labels() -> Labels { + if stringify!($base_name) == "energy_ensemble" { + Labels::new(["energy"], [[0], [1]]) + } else { + Labels::new(["energy"], [[0]]) + } + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2]], + ); + let n_props = n_properties(); + let values = ArrayD::::from_shape_vec( + vec![3, n_props], + vec![1.0; 3 * n_props], + ).unwrap(); + TensorBlock::new(values, &samples, &[], &property_labels()).unwrap() + } + + fn with_gradients(block: &mut TensorBlock) { + let n_props = n_properties(); + let props = property_labels(); + + let pos_gradient = TensorBlock::new( + ArrayD::::from_shape_vec( + vec![1, 3, n_props], + vec![0.1; 3 * n_props], + ).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &props, + ).unwrap(); + block.add_gradient("positions", pos_gradient).unwrap(); + + let strain_gradient = TensorBlock::new( + ArrayD::::from_shape_vec( + vec![1, 3, 3, n_props], + vec![0.1; 9 * n_props], + ).unwrap(), + &Labels::new(["sample"], [[0]]), + &[ + Labels::new(["xyz_1"], [[0], [1], [2]]), + Labels::new(["xyz_2"], [[0], [1], [2]]), + ], + &props, + ).unwrap(); + block.add_gradient("strain", strain_gradient).unwrap(); + } + + fn valid_energy() -> TensorMap { + let mut block = valid_block(); + with_gradients(&mut block); + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![block]).unwrap() + } + + fn system_energy() -> TensorMap { + let samples = Labels::new(["system"], [[0]]); + let n_props = n_properties(); + let values = ArrayD::::from_shape_vec( + vec![1, n_props], + vec![2.0; n_props], + ).unwrap(); + let mut block = TensorBlock::new(values, &samples, &[], &property_labels()).unwrap(); + with_gradients(&mut block); + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![block]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_energy(), &[system(3)], None).unwrap(); + + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + check(&request, &system_energy(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + // Empty systems slice, per-atom output + let n_props = n_properties(); + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, n_props], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[], + &property_labels(), + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &energy, &[], None).unwrap(); + + // System with 0 atoms, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, n_props], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[], + &property_labels(), + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &energy, &[system(0)], None).unwrap(); + + // Empty systems slice, per-system output + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, n_props], vec![]).unwrap(), + &Labels::new( + ["system"], + Array2::::from_shape_vec((0, 1), vec![]).unwrap(), + ), + &[], + &property_labels(), + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &energy, &[], None).unwrap(); + } + + #[test] + fn selected_atoms() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let n_props = n_properties(); + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, n_props], vec![1.0; 3 * n_props]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]), + &[], + &property_labels(), + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &energy, &systems, Some(&selected_atoms)).unwrap(); + + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, n_props], vec![2.0; 2 * n_props]).unwrap(), + &Labels::new(["system"], [[0], [1]]), + &[], + &property_labels(), + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &energy, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2], [1, 0]], + ); + let n_props = n_properties(); + let values = ArrayD::::from_shape_vec( + vec![4, n_props], + vec![1.0; 4 * n_props], + ).unwrap(); + let mut block = TensorBlock::new(values, &samples, &[], &property_labels()).unwrap(); + with_gradients(&mut block); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &energy, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::AtomPair; + let err = check(&request, &valid_energy(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid sample_kind for '{}': expected one of [system, atom], got 'atom_pair'", + stringify!($base_name) + ) + ); + } + + #[test] + fn wrong_number_of_blocks() { + let energy = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid '{}': expected a single block, but found 0 blocks", + stringify!($base_name) + ) + ); + + let energy = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid '{}': expected a single block, but found 2 blocks", + stringify!($base_name) + ) + ); + } + + #[test] + fn wrong_key() { + let energy = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid '{}': expected a single block with key '_', but found key names [foo]", + stringify!($base_name) + ) + ); + + let energy = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid '{}': expected a single block with key value 0", + stringify!($base_name) + ) + ); + } + + #[test] + fn wrong_property() { + let samples = Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]); + let n_props = n_properties(); + let values = ArrayD::::from_shape_vec( + vec![3, n_props], + vec![1.0; 3 * n_props], + ).unwrap(); + + let props_wrong = if stringify!($base_name) == "energy_ensemble" { + Labels::new(["wrong"], [[0], [1]]) + } else { + Labels::new(["wrong"], [[0]]) + }; + + let block = TensorBlock::new(values, &samples, &[], &props_wrong).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid properties for '{}': expected names [energy], got [wrong]", + stringify!($base_name) + ) + ); + + let props_wrong = if stringify!($base_name) == "energy_ensemble" { + Labels::new(["energy"], [[1], [0]]) + } else { + Labels::new(["energy"], [[1]]) + }; + let values = ArrayD::::from_shape_vec( + vec![3, n_props], + vec![1.0; 3 * n_props], + ).unwrap(); + + let block = TensorBlock::new(values, &samples, &[], &props_wrong).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + + let expected_msg = if stringify!($base_name) == "energy_ensemble" { + format!("invalid parameter: invalid properties values for '{}': expected [[0, ..., n]]", stringify!($base_name)) + } else { + format!("invalid parameter: invalid properties values for '{}': expected [[0]]", stringify!($base_name)) + }; + assert_eq!(err.to_string(), expected_msg); + } + + #[test] + fn has_components() { + let n_props = n_properties(); + let props = property_labels(); + let values = ArrayD::::from_shape_vec( + vec![3, 3, n_props], + vec![1.0; 9 * n_props], + ).unwrap(); + let block = TensorBlock::new( + values, + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &props, + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: components for '{}' should be empty", + stringify!($base_name) + ) + ); + } + + #[test] + fn wrong_sample_names() { + let n_props = n_properties(); + let props = property_labels(); + let values = ArrayD::::from_shape_vec( + vec![1, n_props], + vec![1.0; n_props], + ).unwrap(); + let block = TensorBlock::new( + values, + &Labels::new(["system"], [[0]]), + &[], + &props, + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid sample names for '{}': expected [system, atom], got [system]", + stringify!($base_name) + ) + ); + } + + #[test] + fn gradients_dummy() { + let n_props = n_properties(); + let props = property_labels(); + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec( + vec![3, n_props], + vec![1.0; 3 * n_props], + ).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &props, + ).unwrap(); + + let dummy_gradient = TensorBlock::new( + ArrayD::::from_shape_vec( + vec![1, n_props], + vec![0.1; n_props], + ).unwrap(), + &Labels::new(["sample"], [[0]]), + &[], + &props, + ).unwrap(); + block.add_gradient("dummy", dummy_gradient).unwrap(); + + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid gradient 'dummy' for '{}': expected one of [strain, positions]", + stringify!($base_name) + ) + ); + } + + #[test] + fn gradients_position_not_requested() { + let mut request = valid_request(); + request.gradients = vec![]; + + let n_props = n_properties(); + let props = property_labels(); + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec( + vec![3, n_props], + vec![1.0; 3 * n_props], + ).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &props, + ).unwrap(); + + let pos_gradient = TensorBlock::new( + ArrayD::::from_shape_vec( + vec![1, 3, n_props], + vec![0.1; 3 * n_props], + ).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &props, + ).unwrap(); + block.add_gradient("positions", pos_gradient).unwrap(); + + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&request, &energy, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid gradient 'positions' for '{}': these gradients were not requested", + stringify!($base_name) + ) + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let n_props = n_properties(); + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, n_props], vec![1.0; 3 * n_props]).unwrap(), + // samples that are not in the selected_atoms + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &property_labels(), + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &energy, &[system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid samples for '{}', they do not match the `systems` and `selected_atoms`", + stringify!($base_name) + ) + ); + + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, n_props], vec![2.0; 2 * n_props]).unwrap(), + // systems that are not in the selected_atoms + &Labels::new(["system"], [[0], [1]]), + &[], + &property_labels(), + ).unwrap(); + let energy = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&request, &energy, &[system(3), system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "invalid parameter: invalid samples for '{}', they do not match the `systems` and `selected_atoms`", + stringify!($base_name) + ) + ); + } + } + }; + } + + energy_tests!(energy); + energy_tests!(energy_ensemble); + energy_tests!(energy_uncertainty); +} diff --git a/metatomic-core/src/quantity/feature.rs b/metatomic-core/src/quantity/feature.rs new file mode 100644 index 000000000..80f91deee --- /dev/null +++ b/metatomic-core/src/quantity/feature.rs @@ -0,0 +1,307 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks; + +use crate::{Error, SampleKind, System}; + + +/// Check the layout of the "feature" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "feature"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::System, SampleKind::Atom])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[])?; + // no check on properties, as they can be anything for "feature" + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("feature".into()).unwrap(), + unit: String::new(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + } + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2]], + ); + let properties = Labels::new(["anything_goes_here"], [[-42], [5]]); + let values = ArrayD::::from_shape_vec(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); + TensorBlock::new(values, &samples, &[], &properties).unwrap() + } + + fn valid_feature() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_feature(), &[system(3)], None).unwrap(); + + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.5]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["something_else"], [[0]]) + ).unwrap(); + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &feature, &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + // Empty systems slice, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &feature, &[], None).unwrap(); + + // System with 0 atoms, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &feature, &[system(0)], None).unwrap(); + + // Empty systems slice, per-system output + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system"], + Array2::::from_shape_vec((0, 1), vec![]).unwrap(), + ), + &[], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &feature, &[], None).unwrap(); + } + + #[test] + fn selected_atoms() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]), + &[], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &feature, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![4, 1], vec![1.0; 4]).unwrap(), + &Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2], [1, 0]], + ), + &[], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &feature, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::AtomPair; + let err = check(&request, &valid_feature(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'feature': expected one of [system, atom], got 'atom_pair'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let feature = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &feature, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'feature': expected a single block, but found 0 blocks" + ); + + let feature = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &feature, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'feature': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let feature = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &feature, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'feature': expected a single block with key '_', but found key names [foo]" + ); + + let feature = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &feature, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'feature': expected a single block with key value 0" + ); + } + + #[test] + fn has_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &feature, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: components for 'feature' should be empty" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["feature"], [[0]]) + ).unwrap(); + + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &feature, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'feature': expected [system, atom], got [system]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let feature = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &feature, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'feature': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + // samples that are not in the selected_atoms + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["feature"], [[0]]), + ).unwrap(); + let feature = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &feature, &[system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'feature', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/heat_flux.rs b/metatomic-core/src/quantity/heat_flux.rs new file mode 100644 index 000000000..c2297e03c --- /dev/null +++ b/metatomic-core/src/quantity/heat_flux.rs @@ -0,0 +1,363 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, SINGLE_LABELS_REFERENCE, XYZ_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + + +/// Check the layout of the "heat_flux" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "heat_flux"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::System])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[ + ExpectedLabels { + names: &["xyz"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]" + } + ])?; + + let expected_properties = ExpectedLabels { + names: &["heat_flux"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("heat_flux".into()).unwrap(), + unit: "eV/ps".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::System, + } + } + + fn valid_xyz_component() -> Labels { + Labels::new(["xyz"], [[0], [1], [2]]) + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new(["system"], [[0]]); + let properties = Labels::new(["heat_flux"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(); + TensorBlock::new(values, &samples, &[valid_xyz_component()], &properties).unwrap() + } + + fn valid_heat_flux() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_heat_flux(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + // Empty systems slice, per-system output + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system"], + Array2::::from_shape_vec((0, 1), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &heat_flux, &[], None).unwrap(); + } + + #[test] + fn selected_atoms() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 3, 1], vec![1.0; 6]).unwrap(), + &Labels::new(["system"], [[0], [1]]), + &[valid_xyz_component()], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &heat_flux, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let samples = Labels::new( + ["system"], + [[0], [1]], + ); + let properties = Labels::new(["heat_flux"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![2, 3, 1], vec![1.0; 6]).unwrap(); + let block = TensorBlock::new(values, &samples, &[valid_xyz_component()], &properties).unwrap(); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &heat_flux, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::Atom; + let err = check(&request, &valid_heat_flux(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'heat_flux': expected one of [system], got 'atom'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let heat_flux = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'heat_flux': expected a single block, but found 0 blocks" + ); + + let heat_flux = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'heat_flux': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let heat_flux = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'heat_flux': expected a single block with key '_', but found key names [foo]" + ); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'heat_flux': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[valid_xyz_component()], + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'heat_flux': expected names [heat_flux], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[valid_xyz_component()], + &Labels::new(["heat_flux"], [[1]]), + ).unwrap(); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'heat_flux': expected [[0]]" + ); + } + + #[test] + fn missing_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'heat_flux': expected 1 component(s), got 0" + ); + } + + #[test] + fn wrong_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[Labels::new(["abc"], [[0], [1], [2]])], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'heat_flux': expected names [xyz], got [abc]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[Labels::new(["xyz"], [[1], [2], [3]])], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components values for 'heat_flux': expected [[0], [1], [2]]" + ); + } + + #[test] + fn extra_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system"], [[0]]), + &[ + valid_xyz_component(), + Labels::new(["abc"], [[0], [1], [2]]), + ], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'heat_flux': expected 1 component(s), got 2" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0]]), + &[valid_xyz_component()], + &Labels::new(["heat_flux"], [[0]]) + ).unwrap(); + + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'heat_flux': expected [system], got [system, atom]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[valid_xyz_component()], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample"], [[0]]), + &[valid_xyz_component()], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let heat_flux = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &heat_flux, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'heat_flux': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 3, 1], vec![1.0; 6]).unwrap(), + // systems that are not in the selected_atoms + &Labels::new(["system"], [[0], [1]]), + &[valid_xyz_component()], + &Labels::new(["heat_flux"], [[0]]), + ).unwrap(); + let heat_flux = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &heat_flux, &[system(3), system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'heat_flux', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/mass.rs b/metatomic-core/src/quantity/mass.rs new file mode 100644 index 000000000..444667557 --- /dev/null +++ b/metatomic-core/src/quantity/mass.rs @@ -0,0 +1,318 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, SINGLE_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + + +/// Check the layout of the "mass" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "mass"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::Atom])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[])?; + + let expected_properties = ExpectedLabels { + names: &["mass"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("mass".into()).unwrap(), + unit: "dalton".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + } + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2]], + ); + let properties = Labels::new(["mass"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(); + TensorBlock::new(values, &samples, &[], &properties).unwrap() + } + + fn valid_mass() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_mass(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + // Empty systems slice, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[], + &Labels::new(["mass"], [[0]]), + ).unwrap(); + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &mass, &[], None).unwrap(); + + // System with 0 atoms, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[], + &Labels::new(["mass"], [[0]]), + ).unwrap(); + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &mass, &[system(0)], None).unwrap(); + } + + #[test] + fn selected_atoms() { + // Per-atom output with selected_atoms across multiple systems + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]), + &[], + &Labels::new(["mass"], [[0]]), + ).unwrap(); + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &mass, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![4, 1], vec![1.0; 4]).unwrap(), + &Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2], [1, 0]], + ), + &[], + &Labels::new(["mass"], [[0]]), + ).unwrap(); + + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &mass, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let err = check(&request, &valid_mass(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'mass': expected one of [atom], got 'system'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let mass = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'mass': expected a single block, but found 0 blocks" + ); + + let mass = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'mass': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let mass = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'mass': expected a single block with key '_', but found key names [foo]" + ); + + let mass = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'mass': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'mass': expected names [mass], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["mass"], [[1]]), + ).unwrap(); + + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'mass': expected [[0]]" + ); + } + + #[test] + fn has_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &Labels::new(["mass"], [[0]]), + ).unwrap(); + + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: components for 'mass' should be empty" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["mass"], [[0]]) + ).unwrap(); + + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'mass': expected [system, atom], got [system]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["mass"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &Labels::new(["mass"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let mass = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &mass, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'mass': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + // samples that are not in the selected_atoms + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["mass"], [[0]]), + ).unwrap(); + let mass = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &mass, &[system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'mass', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/mod.rs b/metatomic-core/src/quantity/mod.rs new file mode 100644 index 000000000..4ccca6d3c --- /dev/null +++ b/metatomic-core/src/quantity/mod.rs @@ -0,0 +1,96 @@ +use std::sync::Arc; + +use metatensor::{Labels, TensorMap}; + +use crate::{Error, System}; + + +mod quantities; +pub use quantities::{QuantityName, Quantity, SampleKind, Gradients}; + +mod checks; + +mod energy; +mod feature; +mod non_conservative_force; +mod non_conservative_stress; +mod position; +mod momentum; +mod velocity; +mod mass; +mod charge; +mod heat_flux; +mod spin_multiplicity; + + +/// Check that the provided `TensorMap` matches the expected layout for the +/// given `Quantity`. +/// +/// Only standard quantities are checked, custom quantities are only validated +/// for device/dtype compatibility. +/// +/// `selected_atoms` can change the expected samples, and should be provided if +/// the `TensorMap` was computed for a subset of atoms. +pub fn check_quantity( + quantity: &Quantity, + values: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels>, +) -> Result<(), Error> { + let (device, dtype) = if systems.is_empty() { + (None, None) + } else { + debug_assert!(systems.iter().all(|s| s.dtype() == systems[0].dtype()), "all systems must have the same dtype"); + debug_assert!(systems.iter().all(|s| s.device() == systems[0].device()), "all systems must have the same device"); + + (Some(systems[0].device()), Some(systems[0].dtype())) + }; + + + if !values.keys().is_empty() { + if let Some(device) = device && values.device()? != device { + return Err(Error::InvalidParameter(format!( + "invalid device for quantity '{}': expected {}, got {}", + quantity.name, + device, + values.device()? + ))); + } + + if let Some(dtype) = dtype && values.dtype()? != dtype { + return Err(Error::InvalidParameter(format!( + "invalid dtype for quantity '{}': expected {}, got {}", + quantity.name, + dtype, + values.dtype()? + ))); + } + } + + if quantity.name.is_custom() { + // nothing to check + return Ok(()); + } + + match quantity.name.base() { + "energy" | "energy_ensemble" | "energy_uncertainty" => energy::check(quantity, values, systems, selected_atoms)?, + "feature" => feature::check(quantity, values, systems, selected_atoms)?, + "non_conservative_force" => non_conservative_force::check(quantity, values, systems, selected_atoms)?, + "non_conservative_stress" => non_conservative_stress::check(quantity, values, systems, selected_atoms)?, + "position" => position::check(quantity, values, systems, selected_atoms)?, + "momentum" => momentum::check(quantity, values, systems, selected_atoms)?, + "mass" => mass::check(quantity, values, systems, selected_atoms)?, + "velocity" => velocity::check(quantity, values, systems, selected_atoms)?, + "charge" => charge::check(quantity, values, systems, selected_atoms)?, + "heat_flux" => heat_flux::check(quantity, values, systems, selected_atoms)?, + "spin_multiplicity" => spin_multiplicity::check(quantity, values, systems, selected_atoms)?, + _ => { + return Err(Error::Internal(format!( + "invalid quantity name '{}': unknown standard quantity", + quantity.name + ))); + } + } + + Ok(()) +} diff --git a/metatomic-core/src/quantity/momentum.rs b/metatomic-core/src/quantity/momentum.rs new file mode 100644 index 000000000..ab5775ec8 --- /dev/null +++ b/metatomic-core/src/quantity/momentum.rs @@ -0,0 +1,377 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, SINGLE_LABELS_REFERENCE, XYZ_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + + +/// Check the layout of the "momentum" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "momentum"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::Atom])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[ + ExpectedLabels { + names: &["xyz"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ])?; + + let expected_properties = ExpectedLabels { + names: &["momentum"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("momentum".into()).unwrap(), + unit: "Angstrom*amu/ps".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + } + } + + fn valid_xyz_component() -> Labels { + Labels::new(["xyz"], [[0], [1], [2]]) + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2]], + ); + let properties = Labels::new(["momentum"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(); + TensorBlock::new(values, &samples, &[valid_xyz_component()], &properties).unwrap() + } + + fn valid_momentum() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_momentum(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &momentum, &[], None).unwrap(); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &momentum, &[system(0)], None).unwrap(); + } + + #[test] + fn selected_atoms() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &momentum, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![4, 3, 1], vec![1.0; 12]).unwrap(), + &Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2], [1, 0]], + ), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &momentum, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let err = check(&request, &valid_momentum(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'momentum': expected one of [atom], got 'system'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let momentum = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'momentum': expected a single block, but found 0 blocks" + ); + + let momentum = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'momentum': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let momentum = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'momentum': expected a single block with key '_', but found key names [foo]" + ); + + let momentum = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'momentum': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'momentum': expected names [momentum], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[1]]), + ).unwrap(); + + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'momentum': expected [[0]]" + ); + } + + #[test] + fn missing_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0; 3]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'momentum': expected 1 component(s), got 0" + ); + } + + #[test] + fn wrong_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["abc"], [[0], [1], [2]])], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'momentum': expected names [xyz], got [abc]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["xyz"], [[1], [2], [3]])], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components values for 'momentum': expected [[0], [1], [2]]" + ); + } + + #[test] + fn extra_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 3, 1], vec![1.0; 27]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[ + valid_xyz_component(), + Labels::new(["abc"], [[0], [1], [2]]), + ], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'momentum': expected 1 component(s), got 2" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[0]]) + ).unwrap(); + + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'momentum': expected [system, atom], got [system]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let momentum = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &momentum, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'momentum': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + // samples that are not in the selected_atoms + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["momentum"], [[0]]), + ).unwrap(); + let momentum = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &momentum, &[system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'momentum', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/non_conservative_force.rs b/metatomic-core/src/quantity/non_conservative_force.rs new file mode 100644 index 000000000..3ab6bbfa1 --- /dev/null +++ b/metatomic-core/src/quantity/non_conservative_force.rs @@ -0,0 +1,378 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, XYZ_LABELS_REFERENCE, SINGLE_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + + + +/// Check the layout of the "non_conservative_force" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "non_conservative_force"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::Atom])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[ + ExpectedLabels { + names: &["xyz"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ])?; + + let expected_properties = ExpectedLabels { + names: &["non_conservative_force"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("non_conservative_force".into()).unwrap(), + unit: "eV/Angstrom".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + } + } + + fn valid_xyz_component() -> Labels { + Labels::new(["xyz"], [[0], [1], [2]]) + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2]], + ); + let properties = Labels::new(["non_conservative_force"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(); + TensorBlock::new(values, &samples, &[valid_xyz_component()], &properties).unwrap() + } + + fn valid_non_conservative_force() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_non_conservative_force(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &non_conservative_force, &[], None).unwrap(); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap(); + } + + #[test] + fn selected_atoms() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &non_conservative_force, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![4, 3, 1], vec![1.0; 12]).unwrap(), + &Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2], [1, 0]], + ), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &non_conservative_force, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let err = check(&request, &valid_non_conservative_force(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'non_conservative_force': expected one of [atom], got 'system'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let non_conservative_force = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'non_conservative_force': expected a single block, but found 0 blocks" + ); + + let non_conservative_force = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'non_conservative_force': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let non_conservative_force = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'non_conservative_force': expected a single block with key '_', but found key names [foo]" + ); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'non_conservative_force': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'non_conservative_force': expected names [non_conservative_force], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[1]]), + ).unwrap(); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'non_conservative_force': expected [[0]]" + ); + } + + #[test] + fn missing_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0; 3]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'non_conservative_force': expected 1 component(s), got 0" + ); + } + + #[test] + fn wrong_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["abc"], [[0], [1], [2]])], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'non_conservative_force': expected names [xyz], got [abc]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["xyz"], [[1], [2], [3]])], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components values for 'non_conservative_force': expected [[0], [1], [2]]" + ); + } + + #[test] + fn extra_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 3, 1], vec![1.0; 27]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[ + valid_xyz_component(), + Labels::new(["abc"], [[0], [1], [2]]), + ], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'non_conservative_force': expected 1 component(s), got 2" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[0]]) + ).unwrap(); + + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'non_conservative_force': expected [system, atom], got [system]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let non_conservative_force = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &non_conservative_force, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'non_conservative_force': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + // samples that are not in the selected_atoms + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["non_conservative_force"], [[0]]), + ).unwrap(); + let non_conservative_force = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_force, &[system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'non_conservative_force', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/non_conservative_stress.rs b/metatomic-core/src/quantity/non_conservative_stress.rs new file mode 100644 index 000000000..fbda12aa1 --- /dev/null +++ b/metatomic-core/src/quantity/non_conservative_stress.rs @@ -0,0 +1,371 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, XYZ_LABELS_REFERENCE, SINGLE_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + +/// Check the layout of the "non_conservative_stress" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "non_conservative_stress"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::System])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[ + ExpectedLabels { + names: &["xyz_1"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ExpectedLabels { + names: &["xyz_2"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ])?; + + let expected_properties = ExpectedLabels { + names: &["non_conservative_stress"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("non_conservative_stress".into()).unwrap(), + unit: "eV/Angstrom^3".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::System, + } + } + + fn valid_xyz_components() -> Vec { + vec![ + Labels::new(["xyz_1"], [[0], [1], [2]]), + Labels::new(["xyz_2"], [[0], [1], [2]]), + ] + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new(["system"], [[0]]); + let properties = Labels::new(["non_conservative_stress"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(); + TensorBlock::new(values, &samples, &valid_xyz_components(), &properties).unwrap() + } + + fn valid_non_conservative_stress() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_non_conservative_stress(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + // Empty systems slice, per-system output + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system"], + Array2::::from_shape_vec((0, 1), vec![]).unwrap(), + ), + &valid_xyz_components(), + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &non_conservative_stress, &[], None).unwrap(); + } + + #[test] + fn selected_atoms() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 3, 3, 1], vec![1.0; 18]).unwrap(), + &Labels::new(["system"], [[0], [1]]), + &valid_xyz_components(), + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &non_conservative_stress, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let samples = Labels::new( + ["system"], + [[0], [1]], + ); + let properties = Labels::new(["non_conservative_stress"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![2, 3, 3, 1], vec![1.0; 18]).unwrap(); + let block = TensorBlock::new(values, &samples, &valid_xyz_components(), &properties).unwrap(); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &non_conservative_stress, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::Atom; + let err = check(&request, &valid_non_conservative_stress(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'non_conservative_stress': expected one of [system], got 'atom'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let non_conservative_stress = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'non_conservative_stress': expected a single block, but found 0 blocks" + ); + + let non_conservative_stress = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'non_conservative_stress': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let non_conservative_stress = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'non_conservative_stress': expected a single block with key '_', but found key names [foo]" + ); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'non_conservative_stress': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system"], [[0]]), + &valid_xyz_components(), + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'non_conservative_stress': expected names [non_conservative_stress], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system"], [[0]]), + &valid_xyz_components(), + &Labels::new(["non_conservative_stress"], [[1]]), + ).unwrap(); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'non_conservative_stress': expected [[0]]" + ); + } + + #[test] + fn missing_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'non_conservative_stress': expected 2 component(s), got 0" + ); + } + + #[test] + fn wrong_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system"], [[0]]), + &[Labels::new(["abc"], [[0], [1], [2]]), Labels::new(["xyz_2"], [[0], [1], [2]])], + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'non_conservative_stress': expected names [xyz_1], got [abc]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system"], [[0]]), + &[Labels::new(["xyz_1"], [[1], [2], [3]]), Labels::new(["xyz_2"], [[0], [1], [2]])], + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components values for 'non_conservative_stress': expected [[0], [1], [2]]" + ); + } + + #[test] + fn extra_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 3, 1], vec![1.0; 27]).unwrap(), + &Labels::new(["system"], [[0]]), + &[ + Labels::new(["xyz_1"], [[0], [1], [2]]), + Labels::new(["xyz_2"], [[0], [1], [2]]), + Labels::new(["abc"], [[0], [1], [2]]), + ], + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'non_conservative_stress': expected 2 component(s), got 3" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0]]), + &valid_xyz_components(), + &Labels::new(["non_conservative_stress"], [[0]]) + ).unwrap(); + + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'non_conservative_stress': expected [system], got [system, atom]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system"], [[0]]), + &valid_xyz_components(), + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["sample"], [[0]]), + &valid_xyz_components(), + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let non_conservative_stress = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &non_conservative_stress, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'non_conservative_stress': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 3, 3, 1], vec![1.0; 18]).unwrap(), + // systems that are not in the selected_atoms + &Labels::new(["system"], [[0], [1]]), + &valid_xyz_components(), + &Labels::new(["non_conservative_stress"], [[0]]), + ).unwrap(); + let non_conservative_stress = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &non_conservative_stress, &[system(3), system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'non_conservative_stress', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/position.rs b/metatomic-core/src/quantity/position.rs new file mode 100644 index 000000000..0014391be --- /dev/null +++ b/metatomic-core/src/quantity/position.rs @@ -0,0 +1,379 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, SINGLE_LABELS_REFERENCE, XYZ_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + +/// Check the layout of the "position" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "position"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::Atom])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[ + ExpectedLabels { + names: &["xyz"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ])?; + + let expected_properties = ExpectedLabels { + names: &["position"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("position".into()).unwrap(), + unit: "Angstrom".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + } + } + + fn valid_xyz_component() -> Labels { + Labels::new(["xyz"], [[0], [1], [2]]) + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2]], + ); + let properties = Labels::new(["position"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(); + TensorBlock::new(values, &samples, &[valid_xyz_component()], &properties).unwrap() + } + + fn valid_position() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_position(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + // Empty systems slice, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["position"], [[0]]), + ).unwrap(); + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &position, &[], None).unwrap(); + + // System with 0 atoms, per-atom output + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["position"], [[0]]), + ).unwrap(); + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &position, &[system(0)], None).unwrap(); + } + + #[test] + fn selected_atoms() { + // Per-atom output with selected_atoms across multiple systems + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]), + &[valid_xyz_component()], + &Labels::new(["position"], [[0]]), + ).unwrap(); + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &position, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![4, 3, 1], vec![1.0; 12]).unwrap(), + &Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2], [1, 0]], + ), + &[valid_xyz_component()], + &Labels::new(["position"], [[0]]), + ).unwrap(); + + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &position, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let err = check(&request, &valid_position(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'position': expected one of [atom], got 'system'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let position = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'position': expected a single block, but found 0 blocks" + ); + + let position = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'position': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let position = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'position': expected a single block with key '_', but found key names [foo]" + ); + + let position = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'position': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'position': expected names [position], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["position"], [[1]]), + ).unwrap(); + + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'position': expected [[0]]" + ); + } + + #[test] + fn missing_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0; 3]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["position"], [[0]]), + ).unwrap(); + + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'position': expected 1 component(s), got 0" + ); + } + + #[test] + fn wrong_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["abc"], [[0], [1], [2]])], + &Labels::new(["position"], [[0]]), + ).unwrap(); + + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'position': expected names [xyz], got [abc]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["xyz"], [[1], [2], [3]])], + &Labels::new(["position"], [[0]]), + ).unwrap(); + + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components values for 'position': expected [[0], [1], [2]]" + ); + } + + #[test] + fn extra_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 3, 1], vec![1.0; 27]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[ + valid_xyz_component(), + Labels::new(["abc"], [[0], [1], [2]]), + ], + &Labels::new(["position"], [[0]]), + ).unwrap(); + + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'position': expected 1 component(s), got 2" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[valid_xyz_component()], + &Labels::new(["position"], [[0]]) + ).unwrap(); + + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'position': expected [system, atom], got [system]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["position"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[valid_xyz_component()], + &Labels::new(["position"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let position = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &position, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'position': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + // samples that are not in the selected_atoms + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["position"], [[0]]), + ).unwrap(); + let position = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &position, &[system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'position', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/quantities.rs b/metatomic-core/src/quantity/quantities.rs new file mode 100644 index 000000000..6a97a64bf --- /dev/null +++ b/metatomic-core/src/quantity/quantities.rs @@ -0,0 +1,521 @@ +use json::JsonValue; + +use crate::Error; + +static STANDARD_QUANTITIES: &[&str] = &[ + "charge", + "energy_ensemble", + "energy_uncertainty", + "energy", + "feature", + "heat_flux", + "mass", + "momentum", + "non_conservative_force", + "non_conservative_stress", + "position", + "spin_multiplicity", + "velocity", +]; + +fn is_valid_identifier(s: &str) -> bool { + if s.is_empty() { + return false; + } + let first = s.chars().next().unwrap(); + if !(first.is_ascii_alphabetic() || first == '_') { + return false; + } + s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') +} + +/// The name of a quantity, which can be either a standard name or a custom name +/// with an optional variant. +/// +/// This struct enforces that the name is either a known standard name or a +/// custom name. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct QuantityName { + /// The full name of the quantity, including namespace and variant if present + full: String, + /// Optional namespace for custom quantity names. Standard quantity names do + /// not have a namespace. + namespace: Option, + /// The base name of the quantity + base: String, + /// Optional variant of the quantity, i.e. `pbe0` in `energy/pbe0` + variant: Option, +} + +impl QuantityName { + /// Parse and validate a quantity name. + /// + /// The name can be either a standard name or a custom name with the form + /// `::`, where the namespace can itself contain `::` to + /// define sub-namespaces. + /// + /// Both standard and custom names can also define a variant with the form + /// `/` or `::/`. + /// + /// All components (namespace, name, variant) must be non-empty if they are + /// present, and must be valid identifiers (alphanumeric + underscore, not + /// starting with a digit). + pub fn new(name: String) -> Result { + let (main_part, variant) = if let Some(pos) = name.find('/') { + (&name[..pos], Some(name[pos + 1..].to_string())) + } else { + (&*name, None) + }; + + let (namespace, base) = match main_part.rsplit_once("::") { + Some((ns, base)) => (Some(ns.to_string()), base.to_string()), + None => (None, main_part.to_string()), + }; + + if let Some(ref ns) = namespace { + for component in ns.split("::") { + if !is_valid_identifier(component) { + return Err(Error::InvalidParameter(format!( + "invalid namespace '{}' in '{}': must be a valid \ + identifier (alphanumeric or underscore, not starting with a digit)", + ns, name + ))); + } + } + } + + if base.is_empty() { + return Err(Error::InvalidParameter(format!( + "quantity name cannot be empty in '{}'", name + ))); + } + + if !is_valid_identifier(&base) { + return Err(Error::InvalidParameter(format!( + "invalid quantity name '{}' in '{}': \ + must be a valid identifier (alphanumeric or underscore, not starting with a digit)", + base, name + ))); + } + + if let Some(ref variant) = variant && !is_valid_identifier(variant) { + return Err(Error::InvalidParameter(format!( + "invalid quantity variant '{}' in '{}': \ + must be a valid identifier (alphanumeric or underscore, not starting with a digit)", + variant, name + ))); + } + + if namespace.is_none() && !STANDARD_QUANTITIES.contains(&&*base) { + return Err(Error::InvalidParameter(format!( + "'{}' is not a standard quantity name; custom quantity names must use '::'", + name + ))); + } + + return Ok(QuantityName { + full: name, + namespace, + base, + variant, + }) + } + + /// Is this a custom quantity name? + pub fn is_custom(&self) -> bool { + self.namespace.is_some() + } + + /// Get the base name of this quantity + pub fn base(&self) -> &str { + &self.base + } + + /// Get the namespace of this quantity, if any + pub fn namespace(&self) -> Option<&str> { + self.namespace.as_deref() + } + + /// Get the variant of this quantity, if any + pub fn variant(&self) -> Option<&str> { + self.variant.as_deref() + } + + /// Get the full name of this quantity, including namespace and variant if + /// present + pub fn full(&self) -> &str { + &self.full + } +} + +impl std::fmt::Display for QuantityName { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.full()) + } +} + +/// Different kind of samples a quantity can be associated with +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum SampleKind { + /// The quantity is defined for each atom (e.g. atomic energy, charge, ...) + Atom, + /// The quantity is defined for the whole system (e.g. total energy, ...) + System, + /// The quantity is defined for each pair of atoms (e.g. hamiltonian elements, ...) + AtomPair, +} + +impl From for JsonValue { + fn from(value: SampleKind) -> Self { + let s = match value { + SampleKind::Atom => "atom", + SampleKind::System => "system", + SampleKind::AtomPair => "atom_pair", + }; + JsonValue::from(s) + } +} + +impl<'a> TryFrom<&'a JsonValue> for SampleKind { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + let s = value.as_str().ok_or_else(|| Error::Serialization( + "'sample_kind' in JSON for Quantity must be a string".into() + ))?; + match s { + "atom" => Ok(SampleKind::Atom), + "system" => Ok(SampleKind::System), + "atom_pair" => Ok(SampleKind::AtomPair), + _ => Err(Error::Serialization(format!( + "'sample_kind' in JSON for Quantity must be 'atom', 'system' or 'atom_pair', got '{}'", s + ))), + } + } +} + +impl std::fmt::Display for SampleKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + SampleKind::Atom => write!(f, "atom"), + SampleKind::AtomPair => write!(f, "atom_pair"), + SampleKind::System => write!(f, "system"), + } + } +} + +/// Different gradients that a quantity can have +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Gradients { + /// Gradients with respect to atomic positions + Positions, + /// Gradients with respect to the strain (typically used for stress) + Strain, +} + +impl std::fmt::Display for Gradients { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Gradients::Positions => write!(f, "positions"), + Gradients::Strain => write!(f, "strain"), + } + } +} + +impl From for JsonValue { + fn from(value: Gradients) -> Self { + let s = match value { + Gradients::Positions => "positions", + Gradients::Strain => "strain", + }; + JsonValue::from(s) + } +} + +impl<'a> TryFrom<&'a JsonValue> for Gradients { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + let s = value.as_str().ok_or_else(|| Error::Serialization( + "'gradients' in JSON for Quantity must be a string".into() + ))?; + match s { + "positions" => Ok(Gradients::Positions), + "strain" => Ok(Gradients::Strain), + _ => Err(Error::Serialization(format!( + "'gradients' in JSON for Quantity must be 'positions' or 'strain', got '{}'", s + ))), + } + } +} + +/// A quantity that a model can use as input or output +#[derive(Debug, Clone)] +pub struct Quantity { + /// Name of the quantity, this can be a standard name from + /// , or + /// a custom name of the form `::[/]` + pub name: QuantityName, + /// Unit of the quantity + pub unit: String, + /// Description of the quantity, used to provide more details about the + /// quantity, especially when a model defines multiple variants of the same + /// quantity. + pub description: Option, + /// List of explicit gradients for this quantity, stored in the + /// corresponding `TensorMap` + pub gradients: Vec, + /// The kind of samples this quantity is associated with (e.g. per-atom, + /// per-system, ...) + pub sample_kind: SampleKind, +} + +impl From for JsonValue { + fn from(value: Quantity) -> Self { + let mut result = JsonValue::new_object(); + result["type"] = "metatomic_quantity".into(); + result["name"] = value.name.full().into(); + result["unit"] = value.unit.into(); + if let Some(description) = value.description { + result["description"] = description.into(); + } + result["gradients"] = value.gradients.into(); + result["sample_kind"] = value.sample_kind.into(); + return result; + } +} + + +impl<'a> TryFrom<&'a JsonValue> for Quantity { + type Error = Error; + + fn try_from(value: &'a JsonValue) -> Result { + if !value.is_object() { + return Err(Error::Serialization( + "invalid JSON data for Quantity, expected an object".into() + )); + } + + if value["type"].as_str() != Some("metatomic_quantity") { + return Err(Error::Serialization( + "'type' in JSON for Quantity must be 'metatomic_quantity'".into() + )); + } + + let name = value["name"].as_str().ok_or_else(|| Error::Serialization( + "'name' in JSON for Quantity must be a string".into() + ))?; + let name = QuantityName::new(name.to_string())?; + + let unit = value["unit"].as_str().ok_or_else(|| Error::Serialization( + "'unit' in JSON for Quantity must be a string".into() + ))?; + + let mut description = value["description"].as_str().map(|s| s.to_string()); + if description == Some(String::new()) { + // Treat empty description as None + description = None; + } + + let gradients = &value["gradients"]; + if !gradients.is_array() { + return Err(Error::Serialization( + "'gradients' in JSON for Quantity must be an array".into() + )); + } + let gradients = gradients.members() + .map(Gradients::try_from) + .collect::, _>>()?; + + let sample_kind = SampleKind::try_from(&value["sample_kind"])?; + + Ok(Quantity { + name: name, + unit: unit.to_string(), + description, + gradients, + sample_kind, + }) + } +} + + +#[cfg(test)] +mod tests { + use super::*; + + fn example() -> Quantity { + Quantity { + name: QuantityName::new("energy".into()).unwrap(), + unit: "eV".into(), + description: Some("total energy of the system".into()), + gradients: vec![Gradients::Positions], + sample_kind: SampleKind::Atom, + } + } + + #[test] + fn roundtrip() { + let quantity = example(); + let json: JsonValue = quantity.into(); + + assert_eq!(json["type"].as_str(), Some("metatomic_quantity")); + assert_eq!(json["name"].as_str(), Some("energy")); + assert_eq!(json["unit"].as_str(), Some("eV")); + assert_eq!(json["gradients"][0].as_str(), Some("positions")); + assert_eq!(json["sample_kind"].as_str(), Some("atom")); + + let parsed = Quantity::try_from(&json).unwrap(); + assert_eq!(parsed.name.base, "energy"); + assert_eq!(parsed.unit, "eV"); + assert_eq!(parsed.gradients, vec![Gradients::Positions]); + assert!(matches!(parsed.sample_kind, SampleKind::Atom)); + } + + #[test] + fn roundtrip_all_variants() { + for sample_kind in [SampleKind::Atom, SampleKind::System, SampleKind::AtomPair] { + for gradients in [ + vec![], + vec![Gradients::Positions], + vec![Gradients::Strain], + vec![Gradients::Positions, Gradients::Strain], + ] { + let quantity = Quantity { + name: QuantityName::new("test_ns::test".into()).unwrap(), + unit: "unit".into(), + description: Some("Hello".to_string()), + gradients: gradients.clone(), + sample_kind: sample_kind, + }; + let parsed = Quantity::try_from(&JsonValue::from(quantity.clone())).unwrap(); + assert_eq!(parsed.name, quantity.name); + assert_eq!(parsed.unit, quantity.unit); + assert_eq!(parsed.gradients, gradients); + assert_eq!(parsed.sample_kind, sample_kind); + } + } + } + + #[test] + fn rejects_invalid_json() { + let mut wrong_type = JsonValue::from(example()); + wrong_type["type"] = "something-else".into(); + + let mut missing_name = JsonValue::from(example()); + missing_name.remove("name"); + + let mut missing_unit = JsonValue::from(example()); + missing_unit.remove("unit"); + + let mut missing_gradients = JsonValue::from(example()); + missing_gradients.remove("gradients"); + + let mut non_array_gradients = JsonValue::from(example()); + non_array_gradients["gradients"] = "positions".into(); + + let mut invalid_gradient = JsonValue::from(example()); + invalid_gradient["gradients"] = json::array!["positions", "foo"]; + + let mut missing_sample_kind = JsonValue::from(example()); + missing_sample_kind.remove("sample_kind"); + + let mut invalid_sample_kind = JsonValue::from(example()); + invalid_sample_kind["sample_kind"] = "foo".into(); + + let cases: Vec<(JsonValue, &str)> = vec![ + (JsonValue::from("not an object"), + "serialization error: invalid JSON data for Quantity, expected an object"), + (wrong_type, + "serialization error: 'type' in JSON for Quantity must be 'metatomic_quantity'"), + (missing_name, + "serialization error: 'name' in JSON for Quantity must be a string"), + (missing_unit, + "serialization error: 'unit' in JSON for Quantity must be a string"), + (missing_gradients, + "serialization error: 'gradients' in JSON for Quantity must be an array"), + (non_array_gradients, + "serialization error: 'gradients' in JSON for Quantity must be an array"), + (invalid_gradient, + "serialization error: 'gradients' in JSON for Quantity must be 'positions' or 'strain', got 'foo'"), + (missing_sample_kind, + "serialization error: 'sample_kind' in JSON for Quantity must be a string"), + (invalid_sample_kind, + "serialization error: 'sample_kind' in JSON for Quantity must be 'atom', 'system' or 'atom_pair', got 'foo'"), + ]; + + for (json, expected) in cases { + let error = Quantity::try_from(&json).expect_err("expected an error"); + assert_eq!(error.to_string(), expected); + } + } + + #[test] + fn validate_names() { + for name in STANDARD_QUANTITIES { + QuantityName::new(name.to_string()).unwrap(); + } + + let custom = [ + "my_model::energy", + "org::my_model::custom_qty", + "ns1::ns2::ns3::energy", + "some_ns::name_with_underscores", + "_ns::_name", + ]; + for name in custom { + QuantityName::new(name.to_string()).unwrap(); + } + + let variants = [ + "energy/ensemble", + "my_ns::energy/raw", + "ns1::ns2::energy/some_variant", + ]; + for name in variants { + QuantityName::new(name.to_string()).unwrap(); + } + + let error = QuantityName::new(String::new()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: quantity name cannot be empty in ''"); + + let error = QuantityName::new("not_a_standard_name".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: 'not_a_standard_name' is not a standard quantity name; custom quantity names must use '::'"); + + let error = QuantityName::new("/variant".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: quantity name cannot be empty in '/variant'"); + + let error = QuantityName::new("name/".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid quantity variant '' in 'name/': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + + let error = QuantityName::new("::energy".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid namespace '' in '::energy': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + + let error = QuantityName::new("ns::".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: quantity name cannot be empty in 'ns::'"); + + let error = QuantityName::new("ns::/variant".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: quantity name cannot be empty in 'ns::/variant'"); + + let error = QuantityName::new("::".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid namespace '' in '::': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + + let error = QuantityName::new("123name".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid quantity name '123name' in '123name': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + + let error = QuantityName::new("my_ns::123name".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid quantity name '123name' in 'my_ns::123name': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + + let error = QuantityName::new("my_ns::name/123variant".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid quantity variant '123variant' in 'my_ns::name/123variant': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + + let error = QuantityName::new("has spaces".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid quantity name 'has spaces' in 'has spaces': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + + let error = QuantityName::new("my_ns::name/has spaces".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid quantity variant 'has spaces' in 'my_ns::name/has spaces': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + + let error = QuantityName::new("has-dash".into()).expect_err("expected an error"); + assert_eq!(error.to_string(), "invalid parameter: invalid quantity name 'has-dash' in 'has-dash': must be a valid identifier (alphanumeric or underscore, not starting with a digit)"); + } +} diff --git a/metatomic-core/src/quantity/spin_multiplicity.rs b/metatomic-core/src/quantity/spin_multiplicity.rs new file mode 100644 index 000000000..fdb9dc564 --- /dev/null +++ b/metatomic-core/src/quantity/spin_multiplicity.rs @@ -0,0 +1,304 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, SINGLE_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + + +/// Check the layout of the "spin_multiplicity" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "spin_multiplicity"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::System])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[])?; + + let expected_properties = ExpectedLabels { + names: &["spin_multiplicity"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("spin_multiplicity".into()).unwrap(), + unit: "dimensionless".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::System, + } + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new(["system"], [[0]]); + let properties = Labels::new(["spin_multiplicity"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(); + TensorBlock::new(values, &samples, &[], &properties).unwrap() + } + + fn valid_spin_multiplicity() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_spin_multiplicity(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + // Empty systems slice, per-system output + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 1], vec![]).unwrap(), + &Labels::new( + ["system"], + Array2::::from_shape_vec((0, 1), vec![]).unwrap(), + ), + &[], + &Labels::new(["spin_multiplicity"], [[0]]), + ).unwrap(); + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&request, &spin_multiplicity, &[], None).unwrap(); + } + + #[test] + fn selected_atoms() { + // Per-system values with selected_atoms across multiple systems + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 1], vec![5.0, 6.0]).unwrap(), + &Labels::new(["system"], [[0], [1]]), + &[], + &Labels::new(["spin_multiplicity"], [[0]]), + ).unwrap(); + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &spin_multiplicity, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 1], vec![1.0; 2]).unwrap(), + &Labels::new( + ["system"], + [[0], [1]], + ), + &[], + &Labels::new(["spin_multiplicity"], [[0]]), + ).unwrap(); + + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &spin_multiplicity, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::Atom; + let err = check(&request, &valid_spin_multiplicity(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'spin_multiplicity': expected one of [system], got 'atom'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let spin_multiplicity = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'spin_multiplicity': expected a single block, but found 0 blocks" + ); + + let spin_multiplicity = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'spin_multiplicity': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let spin_multiplicity = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'spin_multiplicity': expected a single block with key '_', but found key names [foo]" + ); + + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'spin_multiplicity': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'spin_multiplicity': expected names [spin_multiplicity], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["spin_multiplicity"], [[1]]), + ).unwrap(); + + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'spin_multiplicity': expected [[0]]" + ); + } + + #[test] + fn has_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0; 3]).unwrap(), + &Labels::new(["system"], [[0]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &Labels::new(["spin_multiplicity"], [[0]]), + ).unwrap(); + + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: components for 'spin_multiplicity' should be empty" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0]]), + &[], + &Labels::new(["spin_multiplicity"], [[0]]) + ).unwrap(); + + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'spin_multiplicity': expected [system], got [system, atom]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[], + &Labels::new(["spin_multiplicity"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample"], [[0]]), + &[Labels::new(["xyz"], [[0], [1], [2]])], + &Labels::new(["spin_multiplicity"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let spin_multiplicity = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &spin_multiplicity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'spin_multiplicity': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![2, 1], vec![3.0, 4.0]).unwrap(), + // systems that are not in the selected_atoms + &Labels::new(["system"], [[0], [1]]), + &[], + &Labels::new(["spin_multiplicity"], [[0]]), + ).unwrap(); + let spin_multiplicity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &spin_multiplicity, &[system(3), system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'spin_multiplicity', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/quantity/velocity.rs b/metatomic-core/src/quantity/velocity.rs new file mode 100644 index 000000000..b8aa607ea --- /dev/null +++ b/metatomic-core/src/quantity/velocity.rs @@ -0,0 +1,378 @@ +use std::sync::Arc; +use metatensor::{Labels, TensorMap}; + +use super::Quantity; +use super::checks::{self, ExpectedLabels, SINGLE_LABELS_REFERENCE, XYZ_LABELS_REFERENCE}; + +use crate::{Error, SampleKind, System}; + + + +/// Check the layout of the "velocity" quantity. +pub(super) fn check( + request: &Quantity, + value: &TensorMap, + systems: &[Arc], + selected_atoms: Option<&Labels> +) -> Result<(), Error> { + assert!(!request.name.is_custom() && request.name.base() == "velocity"); + + let context = format!("'{}'", request.name.full()); + checks::it_should_have_valid_sample_kind(&context, request.sample_kind, &[SampleKind::Atom])?; + + checks::it_should_have_a_single_block(&context, value)?; + let block = value.block_by_id(0); + + checks::it_should_have_valid_samples(&context, request.sample_kind, block, systems, selected_atoms)?; + checks::it_should_have_expected_components(&context, block, &[ + ExpectedLabels { + names: &["xyz"], + values: &XYZ_LABELS_REFERENCE, + values_message: "[[0], [1], [2]]", + }, + ])?; + + let expected_properties = ExpectedLabels { + names: &["velocity"], + values: &SINGLE_LABELS_REFERENCE, + values_message: "[[0]]" + }; + checks::it_should_have_expected_labels(&context, "properties", &block.properties(), expected_properties)?; + checks::it_should_have_expected_gradients(&context, request, block, &[])?; + + return Ok(()); +} + +#[cfg(test)] +mod tests { + use metatensor::{Labels, TensorBlock, TensorMap}; + use ndarray::{Array1, Array2, ArrayD}; + use dlpk::DLPackTensor; + use std::sync::Arc; + + use crate::{Quantity, QuantityName, SampleKind, System}; + + use super::check; + + fn system(n_atoms: usize) -> Arc { + let types: DLPackTensor = Array1::::from_vec(vec![1; n_atoms]).try_into().unwrap(); + let positions: DLPackTensor = Array2::::from_shape_vec((n_atoms, 3), vec![0.0; n_atoms * 3]).unwrap().try_into().unwrap(); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + + let pbc: DLPackTensor = Array1::::from_vec(vec![true, true, true]).try_into().unwrap(); + Arc::new(System::new("Angstrom".into(), types, positions, cell, pbc).unwrap()) + } + + fn valid_request() -> Quantity { + Quantity { + name: QuantityName::new("velocity".into()).unwrap(), + unit: "Angstrom/ps".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::Atom, + } + } + + fn valid_xyz_component() -> Labels { + Labels::new(["xyz"], [[0], [1], [2]]) + } + + fn valid_block() -> TensorBlock { + let samples = Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2]], + ); + let properties = Labels::new(["velocity"], [[0]]); + let values = ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(); + TensorBlock::new(values, &samples, &[valid_xyz_component()], &properties).unwrap() + } + + fn valid_velocity() -> TensorMap { + let keys = Labels::new(["_"], [[0]]); + TensorMap::new(keys, vec![valid_block()]).unwrap() + } + + #[test] + fn ok() { + check(&valid_request(), &valid_velocity(), &[system(3)], None).unwrap(); + } + + #[test] + fn empty_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &velocity, &[], None).unwrap(); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![0, 3, 1], vec![]).unwrap(), + &Labels::new( + ["system", "atom"], + Array2::::from_shape_vec((0, 2), vec![]).unwrap(), + ), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &velocity, &[system(0)], None).unwrap(); + } + + #[test] + fn selected_atoms() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]); + let systems = [system(3), system(1)]; + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [1, 0]]), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &velocity, &systems, Some(&selected_atoms)).unwrap(); + } + + #[test] + fn multiple_systems() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![4, 3, 1], vec![1.0; 12]).unwrap(), + &Labels::new( + ["system", "atom"], + [[0, 0], [0, 1], [0, 2], [1, 0]], + ), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + check(&valid_request(), &velocity, &[system(3), system(1)], None).unwrap(); + } + + #[test] + fn invalid_sample_kind() { + let mut request = valid_request(); + request.sample_kind = SampleKind::System; + let err = check(&request, &valid_velocity(), &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample_kind for 'velocity': expected one of [atom], got 'system'" + ); + } + + #[test] + fn wrong_number_of_blocks() { + let velocity = TensorMap::new(Labels::empty(vec!["_"]), vec![]).unwrap(); + + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'velocity': expected a single block, but found 0 blocks" + ); + + let velocity = TensorMap::new( + Labels::new(["_"], [[0], [1]]), + vec![valid_block(), valid_block()] + ).unwrap(); + + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'velocity': expected a single block, but found 2 blocks" + ); + } + + #[test] + fn wrong_key() { + let velocity = TensorMap::new(Labels::new(["foo"], [[0]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'velocity': expected a single block with key '_', but found key names [foo]" + ); + + let velocity = TensorMap::new(Labels::new(["_"], [[1]]), vec![valid_block()]).unwrap(); + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid 'velocity': expected a single block with key value 0" + ); + } + + #[test] + fn wrong_property() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["wrong"], [[0]]), + ).unwrap(); + + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties for 'velocity': expected names [velocity], got [wrong]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[1]]), + ).unwrap(); + + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid properties values for 'velocity': expected [[0]]" + ); + } + + #[test] + fn missing_components() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 1], vec![1.0; 3]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'velocity': expected 1 component(s), got 0" + ); + } + + #[test] + fn wrong_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["abc"], [[0], [1], [2]])], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'velocity': expected names [xyz], got [abc]" + ); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[Labels::new(["xyz"], [[1], [2], [3]])], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components values for 'velocity': expected [[0], [1], [2]]" + ); + } + + #[test] + fn extra_component() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 3, 1], vec![1.0; 27]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[ + valid_xyz_component(), + Labels::new(["abc"], [[0], [1], [2]]), + ], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid components for 'velocity': expected 1 component(s), got 2" + ); + } + + #[test] + fn wrong_sample_names() { + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.0, 2.0, 3.0]).unwrap(), + &Labels::new(["system"], [[0]]), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[0]]) + ).unwrap(); + + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid sample names for 'velocity': expected [system, atom], got [system]" + ); + } + + #[test] + fn gradients() { + let mut block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + + let gradient = TensorBlock::new( + ArrayD::::from_shape_vec(vec![1, 3, 1], vec![0.1, 0.2, 0.3]).unwrap(), + &Labels::new(["sample", "system", "atom"], [[0, 0, 0]]), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + + block.add_gradient("positions", gradient).unwrap(); + + let velocity = TensorMap::new( + Labels::new(["_"], [[0]]), + vec![block] + ).unwrap(); + + let err = check(&valid_request(), &velocity, &[system(3)], None).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid gradients for 'velocity': expected no gradients, but found gradients with respect to [positions]" + ); + } + + #[test] + fn selected_atoms_error() { + let selected_atoms = Labels::new(["system", "atom"], [[0, 0], [0, 1]]); + + let block = TensorBlock::new( + ArrayD::::from_shape_vec(vec![3, 3, 1], vec![1.0; 9]).unwrap(), + // samples that are not in the selected_atoms + &Labels::new(["system", "atom"], [[0, 0], [0, 1], [0, 2]]), + &[valid_xyz_component()], + &Labels::new(["velocity"], [[0]]), + ).unwrap(); + let velocity = TensorMap::new(Labels::new(["_"], [[0]]), vec![block]).unwrap(); + let err = check(&valid_request(), &velocity, &[system(3)], Some(&selected_atoms)).unwrap_err(); + assert_eq!( + err.to_string(), + "invalid parameter: invalid samples for 'velocity', they do not match the `systems` and `selected_atoms`" + ); + } +} diff --git a/metatomic-core/src/system.rs b/metatomic-core/src/system.rs new file mode 100644 index 000000000..c85f67657 --- /dev/null +++ b/metatomic-core/src/system.rs @@ -0,0 +1,1069 @@ +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::{Arc, LazyLock}; + +use dlpk::sys::{DLDataType, DLDevice}; +use dlpk::{DLPackTensor, DLPackTensorRef}; +use metatensor::{TensorBlock, TensorMap}; + +use crate::kernels::{self, ReferenceValue}; +use crate::quantity::check_quantity; +use crate::{Error, Gradients, PairListOptions, Quantity, QuantityName, SampleKind}; + +/// Names that can never be used as custom data in a system +static INVALID_DATA_NAMES: LazyLock> = LazyLock::new(|| { + HashSet::from(["types", "type", "positions", "position", "cell", "neighbors", "neighbor", "pair", "pairs"]) +}); + +static XYZ_REFERENCE: LazyLock> = LazyLock::new(|| { + ReferenceValue::new( + ndarray::ArrayD::from_shape_vec( + ndarray::IxDyn(&[3usize, 1]), + vec![0i32, 1, 2], + ).unwrap() + ) +}); + +static DISTANCE_REFERENCE: LazyLock> = LazyLock::new(|| { + ReferenceValue::new( + ndarray::ArrayD::from_shape_vec( + ndarray::IxDyn(&[1usize, 1]), + vec![0i32], + ).unwrap() + ) +}); + +/// Storage for an atomistic system. +/// +/// This owns the raw DLPack tensors and metatensor objects used at FFI +/// boundaries. +pub struct System { + length_unit: String, + types: DLPackTensor, + positions: DLPackTensor, + cell: DLPackTensor, + pbc: DLPackTensor, + + pairs: BTreeMap, + custom_data: HashMap, +} + +unsafe impl Send for System {} +unsafe impl Sync for System {} + +impl System { + /// Clone this system, deep-copying all underlying data (types, positions, + /// cell, pbc, pair lists, and custom data) to new device allocations. + /// + /// The cloned system is fully independent of the original: modifying one + /// does not affect the other. + pub fn try_clone(&self) -> Result { + let types = kernels::clone_tensor(&self.types.as_ref())?; + let positions = kernels::clone_tensor(&self.positions.as_ref())?; + let cell = kernels::clone_tensor(&self.cell.as_ref())?; + let pbc = kernels::clone_tensor(&self.pbc.as_ref())?; + + let mut pairs = BTreeMap::new(); + for (options, block) in &self.pairs { + pairs.insert(options.clone(), block.as_ref().try_clone()?); + } + + let mut custom_data = HashMap::new(); + for (name, data) in &self.custom_data { + custom_data.insert(name.clone(), data.try_clone()?); + } + + Ok(System { + length_unit: self.length_unit.clone(), + types, + positions, + cell, + pbc, + pairs, + custom_data, + }) + } + + /// Convert this system's data from the system's length unit to + /// `model_length_unit`, and with custom data matching `requested_inputs` + /// converted from their stored units to the requested units. + /// + /// Returns a new `Arc`. If no conversion is needed, this is just + /// a refcount bump (the data is shared). Otherwise, the system is + /// deep-copied and the copy is scaled in place, so the original data is + /// preserved. + /// + /// # Parameters + /// - `model_length_unit`: the length unit the model expects (e.g. + /// "Angstrom"). If empty, no length conversion is performed. + /// - `requested_inputs`: the list of quantities the model requested as + /// extra inputs. Custom data matching one of these will have its values + /// converted from the data's stored unit to the requested unit. + pub fn convert_units( + self: Arc, + model_length_unit: &str, + requested_inputs: &[Quantity], + ) -> Result, Error> { + // length conversion factor + let length_factor = crate::unit_conversion_factor(&self.length_unit, model_length_unit)?; + + // Clone the system before modifying any data, so the engine's original + // data is preserved. + #[allow(clippy::float_cmp)] + let mut system = if length_factor == 1.0 { + // Check if any custom data needs conversion before returning early + let mut needs_custom_conversion = false; + for requested in requested_inputs { + if let Ok(data) = self.get_custom_data(requested.name.full()) { + let data_unit = data.get_info("unit").map(|s| s.to_string()).unwrap_or_default(); + let factor = crate::unit_conversion_factor(&data_unit, &requested.unit)?; + if factor != 1.0 { + needs_custom_conversion = true; + break; + } + } + } + if !needs_custom_conversion { + return Ok(self); + } + self.try_clone()? + } else { + let mut system = self.try_clone()?; + + // scale positions and cell + let positions = system.positions.as_mut(); + kernels::scale_inplace(positions, length_factor)?; + + let cell = system.cell.as_mut(); + kernels::scale_inplace(cell, length_factor)?; + + // scale pair list values (distances). If the underlying DLPack + // tensor is read-only (e.g. backed by non-writable memory), we copy + // it, scale the copy, and rebuild the block instead of scaling in + // place. + for pairs in system.pairs.values_mut() { + let mut pairs_ref_mut = pairs.as_ref_mut(); + let values = pairs_ref_mut.values_mut(); + let device = values.device()?; + let mut dlpack = values.as_dlpack(device, None, dlpk::sys::DLPackVersion::current())?; + + if dlpack.is_read_only() { + let copy = pairs.values().copy(device)?; + + let samples = pairs.samples(); + let components = pairs.components(); + let properties = pairs.properties(); + let new_block = TensorBlock::new( + copy, + &samples, + &components, + &properties, + )?; + *pairs = new_block; + + let mut pairs_ref_mut = pairs.as_ref_mut(); + let values = pairs_ref_mut.values_mut(); + dlpack = values.as_dlpack(device, None, dlpk::sys::DLPackVersion::current())?; + } + + kernels::scale_inplace(dlpack.as_mut(), length_factor)?; + } + + system + }; + + // convert custom data units for requested inputs + let dlpack_version = dlpk::sys::DLPackVersion::current(); + let data_names: Vec = system.custom_data.keys().cloned().collect(); + for name in data_names { + // find if this data is a requested input + let requested = requested_inputs.iter().find(|q| q.name.full() == name); + + if let Some(requested) = requested { + let mut tensor = system.custom_data.remove(&name).unwrap(); + let data_unit = tensor.get_info("unit").map(|s| s.to_string()).unwrap_or_default(); + + let factor = crate::unit_conversion_factor(&data_unit, &requested.unit)?; + + tensor = crate::utils::scale_tensormap(tensor, factor)?; + + // update the unit info to the requested unit + tensor.set_info("unit", &requested.unit); + system.custom_data.insert(name, tensor); + } + } + + // update the system's length unit + system.length_unit = model_length_unit.to_string(); + + Ok(Arc::new(system)) + } +} + +impl std::fmt::Debug for System { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("System") + .field("length_unit", &self.length_unit) + .field("types", &self.types) + .field("positions", &self.positions) + .field("cell", &self.cell) + .field("pbc", &self.pbc) + .field("pairs", &self.pairs.keys().collect::>()) + .field("custom_data", &self.custom_data.keys().collect::>()) + .finish() + } +} + +impl System { + /// Create a `System` from raw DLPack tensors + pub fn new( + length_unit: String, + types: DLPackTensor, + positions: DLPackTensor, + cell: DLPackTensor, + pbc: DLPackTensor, + ) -> Result { + validate_system_tensors(&types, &positions, &cell, &pbc)?; + + let system = System { + length_unit, + types, + positions, + cell, + pbc, + pairs: BTreeMap::new(), + custom_data: HashMap::new(), + }; + + crate::kernels::validate_cell_pbc(system.pbc(), system.cell())?; + + return Ok(system); + } + + /// Get the length unit used by this system + pub fn length_unit(&self) -> &str { + &self.length_unit + } + + /// Get the number of atoms/particles in this system + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + pub fn size(&self) -> usize { + let size = self.types.shape()[0]; + debug_assert!(usize::try_from(size).is_ok()); + return size as usize; + } + + /// Get the particle types + pub fn types(&self) -> DLPackTensorRef<'_> { + self.types.as_ref() + } + + /// Get the particle positions + pub fn positions(&self) -> DLPackTensorRef<'_> { + self.positions.as_ref() + } + + /// Get the unit cell + pub fn cell(&self) -> DLPackTensorRef<'_> { + self.cell.as_ref() + } + + /// Get the periodic boundary condition flags + pub fn pbc(&self) -> DLPackTensorRef<'_> { + self.pbc.as_ref() + } + + /// Add a pair list to this system + pub fn add_pairs( + self: &mut Arc, + options: PairListOptions, + pairs: TensorBlock, + ) -> Result<(), Error> { + let system_mut = Arc::get_mut(self).ok_or_else(|| { + Error::InvalidParameter( + "cannot modify system while there are outstanding borrowed views".into(), + ) + })?; + + if system_mut.pairs.contains_key(&options) { + return Err(Error::InvalidParameter( + "the pair list for these options already exists in this system".into(), + )); + } + + let samples = pairs.samples(); + let samples_names = samples.names(); + if samples_names != ["first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"] { + return Err(Error::InvalidParameter( + "invalid samples for `pairs`: the samples names must be \ + 'first_atom', 'second_atom', 'cell_shift_a', 'cell_shift_b', \ + 'cell_shift_c'".into(), + )); + } + + let components = pairs.components(); + if components.len() != 1 || components[0].names() != ["xyz"] || components[0].count() != 3 { + return Err(Error::InvalidParameter( + "invalid components for `pairs`: there should be a \ + single 'xyz'=[0, 1, 2] component".into() + )); + } + + { + let mts_array = components[0].values(); + let dl_tensor = mts_array.as_dlpack( + components[0].device(), + None, + dlpk::sys::DLPackVersion::current(), + )?; + + if !crate::kernels::is_equal_i32(dl_tensor.as_ref(), &XYZ_REFERENCE)? { + return Err(Error::InvalidParameter( + "invalid components for `pairs`: the 'xyz' component should \ + contain [[0], [1], [2]]".into() + )); + } + } + + let properties = pairs.properties(); + if properties.names() != ["distance"] || properties.count() != 1 { + return Err(Error::InvalidParameter( + "invalid properties for `pairs`: there should be a single \ + 'distance'=0 property".into() + )); + } + + { + let mts_array = properties.values(); + let dl_tensor = mts_array.as_dlpack( + properties.device(), + None, + dlpk::sys::DLPackVersion::current(), + )?; + + if !crate::kernels::is_equal_i32(dl_tensor.as_ref(), &DISTANCE_REFERENCE)? { + return Err(Error::InvalidParameter( + "invalid properties for `pairs`: the 'distance' property \ + should contain [0]".into() + )); + } + } + + if !pairs.as_ref().gradient_list().is_empty() { + return Err(Error::InvalidParameter( + "`pairs` should not have any gradients".into() + )); + } + + if pairs.device()? != system_mut.device() { + return Err(Error::InvalidParameter(format!( + "`pairs` device ({}) does not match this system's device ({})", + pairs.device()?, system_mut.device(), + ))); + } + + if pairs.dtype()? != system_mut.dtype() { + return Err(Error::InvalidParameter(format!( + "`pairs` dtype ({}) does not match this system's dtype ({})", + pairs.dtype()?, system_mut.dtype(), + ))); + } + + system_mut.pairs.insert(options, pairs); + return Ok(()); + } + + /// Get a pair list from this system + pub fn get_pairs(&self, options: &PairListOptions) -> Option<&TensorBlock> { + return self.pairs.get(options); + } + + /// Get all pair list options known by this system + pub fn known_pairs(&self) -> Vec<&PairListOptions> { + return self.pairs.keys().collect(); + } + + /// Add custom data to this system + /// + /// If `override_` is `true`, existing data with the same name will be + /// replaced. + pub fn add_custom_data(self: &mut Arc, name: impl Into, data: TensorMap, override_: bool) -> Result<(), Error> { + let name = name.into(); + if INVALID_DATA_NAMES.contains(name.to_lowercase().as_str()) { + return Err(Error::InvalidParameter(format!( + "custom data can not be named '{}'", name + ))); + } + + if data.keys().is_empty() { + return Err(Error::InvalidParameter(format!( + "custom data '{}' has no blocks", name + ))); + } + + // validate the quantity + let name = QuantityName::new(name)?; + let quantity = quantity_for_data(name, &data)?; + check_quantity(&quantity, &data, std::slice::from_ref(self), None)?; + + let system_mut = Arc::get_mut(self).ok_or_else(|| { + Error::InvalidParameter( + "cannot modify system while there are outstanding borrowed views".into(), + ) + })?; + + if !override_ && system_mut.custom_data.contains_key(quantity.name.full()) { + return Err(Error::InvalidParameter(format!( + "custom data '{}' is already present in this system", + quantity.name + ))); + } + + system_mut.custom_data.insert(quantity.name.full().to_string(), data); + return Ok(()); + } + + /// Get custom data from this system. + pub fn get_custom_data(&self, name: &str) -> Result<&TensorMap, Error> { + let lower = name.to_lowercase(); + if INVALID_DATA_NAMES.contains(lower.as_str()) { + return Err(Error::InvalidParameter(format!( + "custom data can not be named '{}'", name + ))); + } + + return self.custom_data.get(name).ok_or_else(|| Error::InvalidParameter(format!( + "no custom data for '{}' found in this system", name + ))); + } + + /// Get all custom data names known by this system. + pub fn known_custom_data(&self) -> Vec<&str> { + return self.custom_data.keys().map(String::as_str).collect(); + } + + /// The device used for all tensors in this system + pub fn device(&self) -> DLDevice { + self.types.device() + } + + /// The data type used for the `positions` and `cell` tensors in this + /// system, as well as any pair lists and custom data added to this system. + pub fn dtype(&self) -> DLDataType { + self.positions.dtype() + } +} + +/// Guess the `SampleKind` corresponding to the provided `TensorMap`. +/// +/// If `allow_unknown` is `true`, this will return `SampleKind::System` when +/// unable to determine the sample kind. Otherwise, it will return an error. +fn sample_kind_from_sample_names(data: &TensorMap, allow_unknown: bool) -> Result { + assert!(!data.keys().is_empty()); + + let first_block = data.block_by_id(0); + let samples = first_block.samples(); + let sample_names = samples.names(); + + if sample_names == ["system"] { + Ok(SampleKind::System) + } else if sample_names == ["system", "atom"] { + Ok(SampleKind::Atom) + } else if sample_names == ["system", "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"] { + Ok(SampleKind::AtomPair) + } else if allow_unknown { + Ok(SampleKind::System) + } else { + Err(Error::InvalidParameter(format!( + "data has unknown sample names: [{}]", + sample_names.join(", ") + ))) + } +} + +/// Guess the `Quantity` corresponding to the provided custom data name and +/// `TensorMap`. +fn quantity_for_data(name: QuantityName, data: &TensorMap) -> Result { + assert!(!data.keys().is_empty()); + + if name.is_custom() { + return Ok(Quantity { + name: name, + unit: String::new(), + description: None, + gradients: vec![], + sample_kind: sample_kind_from_sample_names(data, true)?, + }); + } + + let mut gradients = Vec::new(); + let first_block = data.block_by_id(0); + for parameter in first_block.gradient_list() { + if parameter == "positions" { + gradients.push(Gradients::Positions); + } else if parameter == "cell" { + gradients.push(Gradients::Strain); + } else { + return Err(Error::InvalidParameter(format!( + "data '{}' has an unknown gradient '{}'", + name, parameter + ))); + } + } + + return Ok(Quantity { + name: name, + unit: data.get_info("unit").unwrap_or("").into(), + description: None, + gradients: gradients, + sample_kind: sample_kind_from_sample_names(data, false)?, + }); +} + +fn validate_system_tensors( + types: &DLPackTensor, + positions: &DLPackTensor, + cell: &DLPackTensor, + pbc: &DLPackTensor, +) -> Result<(), Error> { + let device = types.device(); + if positions.device() != device || cell.device() != device || pbc.device() != device { + return Err(Error::InvalidParameter( + "`types`, `positions`, `cell`, and `pbc` must be on the same device".into() + )); + } + + let dtype_i32 = ::get_dlpack_data_type(); + let dtype_f32 = ::get_dlpack_data_type(); + let dtype_f64 = ::get_dlpack_data_type(); + let dtype_bool = ::get_dlpack_data_type(); + + if types.dtype() != dtype_i32 { + return Err(Error::InvalidParameter( + "`types` must be a tensor of 32-bit integers".into() + )); + } + + let types_shape = types.shape(); + if types_shape.len() != 1 || types_shape[0] < 0 { + return Err(Error::InvalidParameter(format!( + "`types` must be a (n_atoms,) tensor, got a tensor with shape [{}]", + types_shape.iter().map(|dim| dim.to_string()).collect::>().join(", ") + ))); + } + + let n_atoms = types_shape[0]; + + let positions_shape = positions.shape(); + if positions_shape.len() != 2 || positions_shape[0] != n_atoms || positions_shape[1] != 3 { + return Err(Error::InvalidParameter(format!( + "`positions` must be a (n_atoms x 3) tensor, got a tensor with shape [{}]", + positions_shape.iter().map(|dim| dim.to_string()).collect::>().join(", ") + ))); + } + + if positions.dtype() != dtype_f32 && positions.dtype() != dtype_f64 { + return Err(Error::InvalidParameter( + "`positions` must be a tensor of 32 or 64-bit floating point data".into() + )); + } + + let cell_shape = cell.shape(); + if cell_shape.len() != 2 || cell_shape[0] != 3 || cell_shape[1] != 3 { + return Err(Error::InvalidParameter(format!( + "`cell` must be a (3 x 3) tensor, got a tensor with shape [{}]", + cell_shape.iter().map(|dim| dim.to_string()).collect::>().join(", ") + ))); + } + + if cell.dtype() != positions.dtype() { + return Err(Error::InvalidParameter(format!( + "`cell` must have the same dtype as `positions`, got {} and {}", + cell.dtype(), + positions.dtype() + ))); + } + + let pbc_shape = pbc.shape(); + if pbc_shape.len() != 1 || pbc_shape[0] != 3 { + return Err(Error::InvalidParameter(format!( + "`pbc` must contain 3 entries, got a tensor with shape [{}]", + pbc_shape.iter().map(|dim| dim.to_string()).collect::>().join(", ") + ))); + } + + if pbc.dtype() != dtype_bool { + return Err(Error::InvalidParameter( + "`pbc` must be a tensor of booleans".into() + )); + } + + return Ok(()); +} + +#[cfg(test)] +pub(crate) use tests::test_system; + +#[cfg(test)] +mod tests { + use super::*; + use metatensor::Labels; + use ndarray::{Array1, Array2, ArrayViewD}; + use approx::assert_relative_eq; + + // ----------------------------------------------------------------------- + // helpers to create DLPack tensors + // ----------------------------------------------------------------------- + fn type_tensor(data: &[i32]) -> DLPackTensor { + Array1::from_vec(data.to_vec()).try_into().unwrap() + } + + #[allow(clippy::cast_precision_loss)] + fn positions_tensor(n_atoms: usize, dtype: &str) -> DLPackTensor { + match dtype { + "f32" => { + let mut data = Vec::with_capacity(3 * n_atoms); + for i in 0..n_atoms { + data.extend_from_slice(&[i as f32, 0.0, 0.0]); + } + Array2::from_shape_vec((n_atoms, 3), data).unwrap().try_into().unwrap() + } + "f64" => { + let mut data = Vec::with_capacity(3 * n_atoms); + for i in 0..n_atoms { + data.extend_from_slice(&[i as f64, 0.0, 0.0]); + } + Array2::from_shape_vec((n_atoms, 3), data).unwrap().try_into().unwrap() + } + _ => panic!("unsupported dtype '{}'", dtype), + } + } + + #[allow(clippy::cast_possible_truncation)] + fn cell_tensor(size: f64, dtype: &str) -> DLPackTensor { + match dtype { + "f32" => { + Array2::::from_shape_vec( + (3, 3), + vec![ + size as f32, 0.0, 0.0, + 0.0, size as f32, 0.0, + 0.0, 0.0, size as f32, + ], + ).unwrap().try_into().unwrap() + } + "f64" => Array2::::from_shape_vec( + (3, 3), + vec![ + size, 0.0, 0.0, + 0.0, size, 0.0, + 0.0, 0.0, size, + ], + ).unwrap().try_into().unwrap(), + _ => panic!("unsupported dtype '{}'", dtype), + } + } + + fn pbc_tensor(data: &[bool]) -> DLPackTensor { + Array1::from_vec(data.to_vec()).try_into().unwrap() + } + + fn valid_pair_block(dtype: &str) -> TensorBlock { + let samples = Labels::new( + ["first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"], + [[0i32, 1, 0, 0, 0]], + ); + let components = vec![Labels::new(["xyz"], [[0i32], [1], [2]])]; + let properties = Labels::new(["distance"], [[0i32]]); + + match dtype { + "f32" => { + let values = ndarray::ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.5, 2.5, 3.5]).unwrap(); + TensorBlock::new(values, &samples, &components, &properties).unwrap() + } + "f64" => { + let values = ndarray::ArrayD::::from_shape_vec(vec![1, 3, 1], vec![1.5, 2.5, 3.5]).unwrap(); + TensorBlock::new(values, &samples, &components, &properties).unwrap() + } + _ => panic!("unsupported dtype '{}'", dtype), + } + } + + fn valid_custom_data(dtype: &str) -> TensorMap { + let keys = Labels::new(["key"], [[0i32]]); + let samples = Labels::new(["sample"], [[0i32]]); + let properties = Labels::new(["property"], [[0i32]]); + + let block = match dtype { + "f32" => { + let values = ndarray::ArrayD::::from_shape_vec(vec![1, 1], vec![42.0]).unwrap(); + TensorBlock::new(values, &samples, &[], &properties).unwrap() + } + "f64" => { + let values = ndarray::ArrayD::::from_shape_vec(vec![1, 1], vec![42.0]).unwrap(); + TensorBlock::new(values, &samples, &[], &properties).unwrap() + } + _ => panic!("unsupported dtype '{}'", dtype), + }; + + let mut tensor = TensorMap::new(keys, vec![block]).unwrap(); + tensor.set_info("unit", "eV"); + + return tensor; + } + + pub(crate) fn test_system(dtype: &str) -> Arc { + let mut system = Arc::new(System::new( + "Angstrom".into(), + tests::type_tensor(&[1, 6, 8]), + tests::positions_tensor(3, dtype), + tests::cell_tensor(10.0, dtype), + tests::pbc_tensor(&[true, true, true]), + ).unwrap()); + + system.add_custom_data("custom::data/name", valid_custom_data(dtype), true).unwrap(); + + let options = PairListOptions { + cutoff: 3.5, + full_list: true, + strict: false, + requestors: vec![], + }; + + system.add_pairs(options, valid_pair_block(dtype)).unwrap(); + + return system; + } + + #[test] + fn system() { + let system = System::new( + "Angstrom".into(), + type_tensor(&[1, 6, 8]), + positions_tensor(3, "f32"), + cell_tensor(10.0, "f32"), + pbc_tensor(&[true, true, true]), + ).unwrap(); + + assert_eq!(system.length_unit(), "Angstrom"); + assert_eq!(system.size(), 3); + assert_eq!(system.device(), DLDevice::cpu()); + assert_eq!(system.dtype().bits, 32); + + let system = System::new( + "Angstrom".into(), + type_tensor(&[1, 6, 8]), + positions_tensor(3, "f64"), + cell_tensor(10.0, "f64"), + pbc_tensor(&[true, true, true]), + ).unwrap(); + assert_eq!(system.length_unit(), "Angstrom"); + assert_eq!(system.size(), 3); + assert_eq!(system.device(), DLDevice::cpu()); + assert_eq!(system.dtype().bits, 64); + } + + #[test] + fn system_invalid_tensors() { + let length_unit = "Angstrom".to_string(); + + let bad_types: DLPackTensor = Array1::::from_vec(vec![1.0, 2.0]).try_into().unwrap(); + let positions = positions_tensor(2, "f32"); + let cell = cell_tensor(0.0, "f32"); + let pbc = pbc_tensor(&[true, true, true]); + + let err = System::new(length_unit.clone(), bad_types, positions, cell, pbc).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: `types` must be a tensor of 32-bit integers"); + + let bad_types: DLPackTensor = Array2::::from_shape_vec((2, 2), vec![1, 2, 3, 4]).unwrap().try_into().unwrap(); + let positions = positions_tensor(2, "f32"); + let cell = cell_tensor(0.0, "f32"); + let pbc = pbc_tensor(&[true, true, true]); + let err = System::new(length_unit.clone(), bad_types, positions, cell, pbc).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: `types` must be a (n_atoms,) tensor, got a tensor with shape [2, 2]"); + + let types = type_tensor(&[1]); + let bad_positions: DLPackTensor = Array2::::from_shape_vec((1, 3), vec![1, 2, 3]).unwrap().try_into().unwrap(); + let cell = cell_tensor(0.0, "f32"); + let pbc = pbc_tensor(&[true, true, true]); + let err = System::new(length_unit.clone(), types, bad_positions, cell, pbc).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: `positions` must be a tensor of 32 or 64-bit floating point data"); + + let types = type_tensor(&[1, 6]); + let bad_positions = Array2::::from_shape_vec((2, 2), vec![0.0; 4]).unwrap().try_into().unwrap(); + let cell = cell_tensor(0.0, "f32"); + let pbc = pbc_tensor(&[true, true, true]); + let err = System::new("Angstrom".into(), types, bad_positions, cell, pbc).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: `positions` must be a (n_atoms x 3) tensor, got a tensor with shape [2, 2]"); + + let types = type_tensor(&[1, 6]); + let positions = positions_tensor(2, "f32"); + let bad_cell = Array2::::from_shape_vec((2, 3), vec![0.0; 6]).unwrap().try_into().unwrap(); + let pbc = pbc_tensor(&[true, true, true]); + let err = System::new(length_unit.clone(), types, positions, bad_cell, pbc).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: `cell` must be a (3 x 3) tensor, got a tensor with shape [2, 3]"); + + let types = type_tensor(&[1, 6]); + let positions = positions_tensor(2, "f32"); + let cell = cell_tensor(0.0, "f64"); + let pbc = pbc_tensor(&[true, true, true]); + let err = System::new(length_unit.clone(), types, positions, cell, pbc).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: `cell` must have the same dtype as `positions`, got f64 and f32"); + + let bad_pbc_dtype: DLPackTensor = Array1::::from_vec(vec![1, 0, 1]).try_into().unwrap(); + let types = type_tensor(&[1, 6]); + let positions = positions_tensor(2, "f32"); + let cell = cell_tensor(0.0, "f32"); + let err = System::new(length_unit.clone(), types, positions, cell, bad_pbc_dtype).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: `pbc` must be a tensor of booleans"); + + let types = type_tensor(&[1, 6]); + let positions = positions_tensor(2, "f32"); + let cell = cell_tensor(0.0, "f32"); + let bad_pbc = pbc_tensor(&[true, true]); + let err = System::new(length_unit, types, positions, cell, bad_pbc).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: `pbc` must contain 3 entries, got a tensor with shape [2]"); + } + + #[test] + fn system_periodic() { + let length_unit = "Angstrom".to_string(); + + // valid periodicity combinations: (1) fully periodic + let types = type_tensor(&[1]); + let positions = positions_tensor(1, "f32"); + let cell = cell_tensor(10.0, "f32"); + let pbc = pbc_tensor(&[true, true, true]); + System::new(length_unit.clone(), types, positions, cell, pbc).unwrap(); + + // (2) fully non-periodic with zero cell + let types = type_tensor(&[1]); + let positions = positions_tensor(1, "f32"); + let cell = cell_tensor(0.0, "f32"); + let pbc = pbc_tensor(&[false, false, false]); + System::new(length_unit.clone(), types, positions, cell, pbc).unwrap(); + + // (3) mixed periodic/non-periodic + let types = type_tensor(&[1]); + let positions = positions_tensor(1, "f32"); + let cell: DLPackTensor = Array2::::from_shape_vec( + (3, 3), + vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 10.0], + ).unwrap().try_into().unwrap(); + let pbc = pbc_tensor(&[true, false, true]); + System::new(length_unit.clone(), types, positions, cell, pbc).unwrap(); + + // invalid periodicity/cell + let types = type_tensor(&[1]); + let positions = positions_tensor(1, "f32"); + let cell = cell_tensor(10.0, "f32"); + let pbc = pbc_tensor(&[true, false, true]); + let err = System::new(length_unit.clone(), types, positions, cell, pbc).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: invalid cell: for non-periodic dimensions, the corresponding cell vector must be zero, but cell[1] contains non-zero values"); + } + + #[test] + fn add_pairs() { + let mut system = Arc::new(System::new( + "Angstrom".into(), + type_tensor(&[1, 6, 8]), + positions_tensor(3, "f32"), + cell_tensor(10.0, "f32"), + pbc_tensor(&[true, true, true]), + ).unwrap()); + + let options = PairListOptions { cutoff: 3.5, full_list: true, strict: false, requestors: vec![] }; + let pairs = valid_pair_block("f32"); + let pairs_ptr = pairs.as_ptr(); + system.add_pairs(options.clone(), pairs).unwrap(); + assert_eq!(system.known_pairs().len(), 1); + assert_eq!(system.get_pairs(&options).unwrap().properties().names(), ["distance"]); + + let options_with_requestor = PairListOptions { + cutoff: 3.5, + full_list: true, + strict: false, + requestors: vec!["test-requestor".into()], + }; + + let pairs_from_system = system.get_pairs(&options_with_requestor).unwrap(); + assert_eq!(pairs_from_system.as_ptr(), pairs_ptr); + + system.add_pairs( + PairListOptions { cutoff: 5.0, full_list: false, strict: true, requestors: vec![] }, + valid_pair_block("f32"), + ).unwrap(); + assert_eq!(system.known_pairs().len(), 2); + } + + + #[test] + fn custom_data() { + let mut system = Arc::new(System::new( + "Angstrom".into(), + type_tensor(&[1, 6, 8]), + positions_tensor(3, "f32"), + cell_tensor(10.0, "f32"), + pbc_tensor(&[true, true, true]), + ).unwrap()); + + let data = valid_custom_data("f32"); + system.add_custom_data("test::my_data", data, false).unwrap(); + assert_eq!(system.known_custom_data(), vec!["test::my_data"]); + assert_eq!(system.get_custom_data("test::my_data").unwrap().keys().names(), ["key"]); + + let err = system.add_custom_data("test::my_data", valid_custom_data("f32"), false).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: custom data 'test::my_data' is already present in this system"); + + let replacement = valid_custom_data("f32"); + system.add_custom_data("test::my_data", replacement, true).unwrap(); + assert_eq!(system.known_custom_data(), vec!["test::my_data"]); + + let mut system = Arc::new(System::new( + "Angstrom".into(), + type_tensor(&[1, 6, 8]), + positions_tensor(3, "f32"), + cell_tensor(10.0, "f32"), + pbc_tensor(&[true, true, true]), + ).unwrap()); + + let test_data_a = valid_custom_data("f32"); + let test_data_a_ptr = test_data_a.as_ptr(); + system.add_custom_data("test::a", test_data_a, false).unwrap(); + + let test_data_b = valid_custom_data("f32"); + let test_data_b_ptr = test_data_b.as_ptr(); + system.add_custom_data("test::b", test_data_b, false).unwrap(); + + let mut names = system.known_custom_data(); + names.sort_unstable(); + assert_eq!(names, vec!["test::a", "test::b"]); + + let data_a = system.get_custom_data("test::a").unwrap(); + assert_eq!(data_a.as_ptr(), test_data_a_ptr); + + let data_b = system.get_custom_data("test::b").unwrap(); + assert_eq!(data_b.as_ptr(), test_data_b_ptr); + + let err = system.get_custom_data("no_such_data").unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: no custom data for 'no_such_data' found in this system"); + } + + #[test] + fn custom_data_validation() { + let mut system = Arc::new(System::new( + "Angstrom".into(), + type_tensor(&[1, 6, 8]), + positions_tensor(3, "f32"), + cell_tensor(10.0, "f32"), + pbc_tensor(&[true, true, true]), + ).unwrap()); + for name in ["types", "type", "Positions", "position", "CELL", "neighbors", "neighbor", "pair", "pairs", "Types", "POSITIONS", "Cell", "Neighbors"] { + let data = valid_custom_data("f32"); + let err = system.add_custom_data(name.to_string(), data, false).unwrap_err(); + assert_eq!(err.to_string(), format!("invalid parameter: custom data can not be named '{}'", name)); + } + + let err = system.add_custom_data("my_data", valid_custom_data("f32"), false).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: 'my_data' is not a standard quantity name; custom quantity names must use '::'"); + + let keys = Labels::empty(vec!["key"]); + let empty = TensorMap::new(keys, vec![]).unwrap(); + let err = system.add_custom_data("test::empty", empty, false).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: custom data 'test::empty' has no blocks"); + + let dtype_mismatch = valid_custom_data("f64"); + let err = system.add_custom_data("test::dtype", dtype_mismatch, false).unwrap_err(); + assert_eq!(err.to_string(), "invalid parameter: invalid dtype for quantity 'test::dtype': expected f32, got f64"); + } + + #[test] + fn system_clone() { + let system = test_system("f32"); + let cloned = system.try_clone().unwrap(); + + // same metadata + assert_eq!(cloned.length_unit(), system.length_unit()); + assert_eq!(cloned.size(), system.size()); + assert_eq!(cloned.device(), system.device()); + assert_eq!(cloned.dtype(), system.dtype()); + + // same pairs + assert_eq!(cloned.known_pairs().len(), system.known_pairs().len()); + + // same custom data + assert_eq!(cloned.known_custom_data(), system.known_custom_data()); + + // the data is independent — modifying one doesn't affect the other. + assert_ne!( + system.positions().raw.data as usize, + cloned.positions().raw.data as usize, + ); + assert_ne!( + system.cell().raw.data as usize, + cloned.cell().raw.data as usize, + ); + assert_ne!( + system.types().raw.data as usize, + cloned.types().raw.data as usize, + ); + } + + #[test] + fn system_convert_units() { + let system = test_system("f64"); + let ref_system = system.try_clone().unwrap(); + + let requested_inputs = [Quantity { + name: QuantityName::new("custom::data/name".into()).unwrap(), + unit: "kJ/mol".into(), + description: None, + gradients: vec![], + sample_kind: SampleKind::System, + }]; + + // convert from Angstrom to nanometer, and eV to kJ/mol + let converted = system.convert_units("nanometer", &requested_inputs).unwrap(); + + // length unit should be updated + assert_eq!(converted.length_unit(), "nanometer"); + + // positions should be scaled: Angstrom -> nanometer (factor 0.1) + let positions: ArrayViewD = converted.positions().try_into().unwrap(); + let ref_positions: ArrayViewD = ref_system.positions().try_into().unwrap(); + assert_relative_eq!(positions, ref_positions.to_owned() * 0.1, max_relative = 1e-6); + + // cell should be scaled + let cell: ArrayViewD = converted.cell().try_into().unwrap(); + let ref_cell: ArrayViewD = ref_system.cell().try_into().unwrap(); + assert_relative_eq!(cell, ref_cell.to_owned() * 0.1, max_relative = 1e-6); + + // pair list values should be scaled + let options = PairListOptions { + cutoff: 3.5, + full_list: true, + strict: false, + requestors: vec![], + }; + let pairs = converted.get_pairs(&options).unwrap(); + let pairs_values = pairs.values().to_ndarray_lock::().read().unwrap(); + let pairs_ref = ref_system.get_pairs(&options).unwrap(); + let pairs_values_ref = pairs_ref.values().to_ndarray_lock::().read().unwrap().clone(); + assert_relative_eq!(*pairs_values, pairs_values_ref * 0.1, max_relative = 1e-6); + + // custom data should have unit info updated + let converted_data = converted.get_custom_data("custom::data/name").unwrap(); + assert_eq!(converted_data.get_info("unit"), Some("kJ/mol")); + assert_eq!(converted_data.keys().count(), 1); + + // custom data values should be scaled: eV -> kJ/mol + let expected_factor = crate::unit_conversion_factor("eV", "kJ/mol").unwrap(); + + let ref_data = ref_system.get_custom_data("custom::data/name").unwrap(); + let ref_block = ref_data.block_by_id(0); + let ref_values = ref_block.values().to_ndarray_lock::().read().unwrap().clone(); + + let block = converted_data.block_by_id(0); + let values = block.values().to_ndarray_lock::().read().unwrap(); + assert_relative_eq!(*values, ref_values * expected_factor, max_relative = 1e-6); + } +} diff --git a/metatomic-core/src/units.rs b/metatomic-core/src/units.rs new file mode 100644 index 000000000..4d5a908e8 --- /dev/null +++ b/metatomic-core/src/units.rs @@ -0,0 +1,746 @@ +use crate::Error; + +use std::sync::LazyLock; +use std::collections::HashMap; +use std::fmt; +use std::ops::{Add, Sub}; + +/// Physical dimension vector with named integer exponents: +/// [Length, Time, Mass, Electric Current, Temperature] +/// +/// Note: quantity of substance (mole) is intentionally not included, since we +/// want `kJ/mol` and `eV` to have the same dimension. +#[derive(Debug, Clone, PartialEq, Eq)] +struct Dimension { + length: i32, + time: i32, + mass: i32, + electric_current: i32, + temperature: i32, +} + +impl Dimension { + /// Dimensionless — all exponents are zero. + const NONE: Dimension = Dimension { length: 0, time: 0, mass: 0, electric_current: 0, temperature: 0 }; + + /// Length dimension + const LENGTH: Dimension = Dimension { length: 1, time: 0, mass: 0, electric_current: 0, temperature: 0 }; + /// Time dimension + const TIME: Dimension = Dimension { length: 0, time: 1, mass: 0, electric_current: 0, temperature: 0 }; + /// Mass dimension + const MASS: Dimension = Dimension { length: 0, time: 0, mass: 1, electric_current: 0, temperature: 0 }; + /// Electric charge dimension (current × time) + const CHARGE: Dimension = Dimension { length: 0, time: 1, mass: 0, electric_current: 1, temperature: 0 }; + /// Temperature dimension + const TEMPERATURE: Dimension = Dimension { length: 0, time: 0, mass: 0, electric_current: 0, temperature: 1 }; + + /// Energy dimension: L² T⁻² M¹ + const ENERGY: Dimension = Dimension { length: 2, time: -2, mass: 1, electric_current: 0, temperature: 0 }; + /// Pressure dimension: L⁻¹ T⁻² M¹ + const PRESSURE: Dimension = Dimension { length: -1, time: -2, mass: 1, electric_current: 0, temperature: 0 }; + /// Electric dipole dimension: L¹ T¹ I¹ + const ELECTRIC_DIPOLE: Dimension = Dimension { length: 1, time: 1, mass: 0, electric_current: 1, temperature: 0 }; + + fn pow(&self, p: f64) -> Dimension { + Dimension { + length: round_if_integer(f64::from(self.length) * p), + time: round_if_integer(f64::from(self.time) * p), + mass: round_if_integer(f64::from(self.mass) * p), + electric_current: round_if_integer(f64::from(self.electric_current) * p), + temperature: round_if_integer(f64::from(self.temperature) * p), + } + } +} + +impl Add<&Dimension> for &Dimension { + type Output = Dimension; + + fn add(self, other: &Dimension) -> Dimension { + Dimension { + length: self.length + other.length, + time: self.time + other.time, + mass: self.mass + other.mass, + electric_current: self.electric_current + other.electric_current, + temperature: self.temperature + other.temperature, + } + } +} + +impl Sub<&Dimension> for &Dimension { + type Output = Dimension; + + fn sub(self, other: &Dimension) -> Dimension { + Dimension { + length: self.length - other.length, + time: self.time - other.time, + mass: self.mass - other.mass, + electric_current: self.electric_current - other.electric_current, + temperature: self.temperature - other.temperature, + } + } +} + +#[allow(clippy::cast_possible_truncation)] +fn round_if_integer(v: f64) -> i32 { + let rounded = v.round(); + assert!((v - rounded).abs() <= 1e-10, "non-integer dimension exponent {} is not supported", v); + return rounded as i32; +} + +impl fmt::Display for Dimension { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + use fmt::Write; + let mut first = true; + f.write_char('[')?; + + for (name, v) in [ + ("L", self.length), + ("T", self.time), + ("M", self.mass), + ("I", self.electric_current), + ("Θ", self.temperature), + ] { + if v == 0 { + continue; + } + + if !first { + f.write_char(' ')?; + } + first = false; + + f.write_str(name)?; + + if v != 1 && v != -1 { + write!(f, "^{}", v)?; + } + + if v == -1 { + f.write_str("^-1")?; + } + } + + if first { + f.write_str("dimensionless")?; + } + f.write_char(']')?; + + Ok(()) + } +} + +/// A parsed unit value: SI conversion factor and physical dimension. +#[derive(Debug, Clone)] +struct UnitValue { + factor: f64, + dim: Dimension, +} + +/// All base units with SI factors and dimensions. +/// Factors are expressed in SI base units (m, s, kg, C, K). +/// Case-insensitive lookup: names are lowercased before searching. +static BASE_UNITS: LazyLock> = LazyLock::new(|| { + let mut map = HashMap::new(); + + // --- Temperature --- + map.insert("kelvin", UnitValue { factor: 1.0, dim: Dimension::TEMPERATURE }); + map.insert("k", UnitValue { factor: 1.0, dim: Dimension::TEMPERATURE }); + + // --- Length --- + map.insert("angstrom", UnitValue { factor: 1e-10, dim: Dimension::LENGTH }); + map.insert("a", UnitValue { factor: 1e-10, dim: Dimension::LENGTH }); + map.insert("bohr", UnitValue { factor: 5.2917721054482e-11, dim: Dimension::LENGTH }); + map.insert("nm", UnitValue { factor: 1e-9, dim: Dimension::LENGTH }); + map.insert("nanometer", UnitValue { factor: 1e-9, dim: Dimension::LENGTH }); + map.insert("meter", UnitValue { factor: 1.0, dim: Dimension::LENGTH }); + map.insert("m", UnitValue { factor: 1.0, dim: Dimension::LENGTH }); + map.insert("cm", UnitValue { factor: 1e-2, dim: Dimension::LENGTH }); + map.insert("centimeter", UnitValue { factor: 1e-2, dim: Dimension::LENGTH }); + map.insert("mm", UnitValue { factor: 1e-3, dim: Dimension::LENGTH }); + map.insert("millimeter", UnitValue { factor: 1e-3, dim: Dimension::LENGTH }); + map.insert("um", UnitValue { factor: 1e-6, dim: Dimension::LENGTH }); + map.insert("µm", UnitValue { factor: 1e-6, dim: Dimension::LENGTH }); + map.insert("micrometer", UnitValue { factor: 1e-6, dim: Dimension::LENGTH }); + + // --- Energy --- + map.insert("electronvolt", UnitValue { factor: 1.602176634e-19, dim: Dimension::ENERGY }); + map.insert("ev", UnitValue { factor: 1.602176634e-19, dim: Dimension::ENERGY }); + map.insert("mev", UnitValue { factor: 1.602176634e-19 * 1e-3, dim: Dimension::ENERGY }); + map.insert("hartree", UnitValue { factor: 4.359744722206048e-18, dim: Dimension::ENERGY }); + map.insert("ry", UnitValue { factor: 2.179872361103024e-18, dim: Dimension::ENERGY }); + map.insert("rydberg", UnitValue { factor: 2.179872361103024e-18, dim: Dimension::ENERGY }); + map.insert("joule", UnitValue { factor: 1.0, dim: Dimension::ENERGY }); + map.insert("j", UnitValue { factor: 1.0, dim: Dimension::ENERGY }); + map.insert("kcal", UnitValue { factor: 4184.0, dim: Dimension::ENERGY }); + map.insert("kj", UnitValue { factor: 1000.0, dim: Dimension::ENERGY }); + + // --- Time --- + map.insert("s", UnitValue { factor: 1.0, dim: Dimension::TIME }); + map.insert("second", UnitValue { factor: 1.0, dim: Dimension::TIME }); + map.insert("ms", UnitValue { factor: 1e-3, dim: Dimension::TIME }); + map.insert("millisecond", UnitValue { factor: 1e-3, dim: Dimension::TIME }); + map.insert("us", UnitValue { factor: 1e-6, dim: Dimension::TIME }); + map.insert("µs", UnitValue { factor: 1e-6, dim: Dimension::TIME }); + map.insert("microsecond", UnitValue { factor: 1e-6, dim: Dimension::TIME }); + map.insert("ns", UnitValue { factor: 1e-9, dim: Dimension::TIME }); + map.insert("nanosecond", UnitValue { factor: 1e-9, dim: Dimension::TIME }); + map.insert("ps", UnitValue { factor: 1e-12, dim: Dimension::TIME }); + map.insert("picosecond", UnitValue { factor: 1e-12, dim: Dimension::TIME }); + map.insert("fs", UnitValue { factor: 1e-15, dim: Dimension::TIME }); + map.insert("femtosecond", UnitValue { factor: 1e-15, dim: Dimension::TIME }); + + // --- Mass --- + map.insert("u", UnitValue { factor: 1.6605390689252e-27, dim: Dimension::MASS }); + map.insert("dalton", UnitValue { factor: 1.6605390689252e-27, dim: Dimension::MASS }); + map.insert("kg", UnitValue { factor: 1.0, dim: Dimension::MASS }); + map.insert("kilogram", UnitValue { factor: 1.0, dim: Dimension::MASS }); + map.insert("g", UnitValue { factor: 1e-3, dim: Dimension::MASS }); + map.insert("gram", UnitValue { factor: 1e-3, dim: Dimension::MASS }); + map.insert("electron_mass", UnitValue { factor: 9.109383713928e-31, dim: Dimension::MASS }); + map.insert("m_e", UnitValue { factor: 9.109383713928e-31, dim: Dimension::MASS }); + + // --- Charge --- + map.insert("e", UnitValue { factor: 1.602176634e-19, dim: Dimension::CHARGE }); + map.insert("coulomb", UnitValue { factor: 1.0, dim: Dimension::CHARGE }); + map.insert("c", UnitValue { factor: 1.0, dim: Dimension::CHARGE }); + + // --- Pressure --- + map.insert("pa", UnitValue { factor: 1.0, dim: Dimension::PRESSURE }); + map.insert("pascal", UnitValue { factor: 1.0, dim: Dimension::PRESSURE }); + map.insert("kpa", UnitValue { factor: 1e3, dim: Dimension::PRESSURE }); + map.insert("kilopascal", UnitValue { factor: 1e3, dim: Dimension::PRESSURE }); + map.insert("mpa", UnitValue { factor: 1e6, dim: Dimension::PRESSURE }); + map.insert("megapascal", UnitValue { factor: 1e6, dim: Dimension::PRESSURE }); + map.insert("gpa", UnitValue { factor: 1e9, dim: Dimension::PRESSURE }); + map.insert("gigapascal", UnitValue { factor: 1e9, dim: Dimension::PRESSURE }); + map.insert("bar", UnitValue { factor: 100000.0, dim: Dimension::PRESSURE }); + map.insert("atm", UnitValue { factor: 101325.0, dim: Dimension::PRESSURE }); + + // --- Electric dipole moment --- + map.insert("debye", UnitValue { factor: 1.0 / 299792458.0 * 1e-21, dim: Dimension::ELECTRIC_DIPOLE }); + map.insert("d", UnitValue { factor: 1.0 / 299792458.0 * 1e-21, dim: Dimension::ELECTRIC_DIPOLE }); + + // --- Dimensionless --- + map.insert("mol", UnitValue { factor: 6.02214076e23, dim: Dimension::NONE }); + + // --- Derived --- + map.insert("hbar", UnitValue { + factor: 1.0545718176462e-34, + dim: Dimension { length: 2, time: -1, mass: 1, electric_current: 0, temperature: 0 }, + }); + + map +}); + +// ---- Tokenizer ---- + +#[derive(Debug, Clone)] +enum Token { + LParen, + RParen, + Mul, + Div, + Pow, + Value(String), +} + +impl Token { + fn precedence(&self) -> i32 { + match self { + Token::LParen | Token::RParen => 0, + Token::Mul | Token::Div => 10, + Token::Pow => 20, + Token::Value(_) => -1, + } + } +} + +impl fmt::Display for Token { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Token::LParen => write!(f, "("), + Token::RParen => write!(f, ")"), + Token::Mul => write!(f, "*"), + Token::Div => write!(f, "/"), + Token::Pow => write!(f, "^"), + Token::Value(v) => write!(f, "{}", v), + } + } +} + +fn tokenize(unit: &str) -> Vec { + let mut tokens = Vec::new(); + let mut current = String::new(); + + for c in unit.chars() { + if c == '*' || c == '/' || c == '^' || c == '(' || c == ')' { + if !current.is_empty() { + tokens.push(Token::Value(current.clone())); + current.clear(); + } + let t = match c { + '*' => Token::Mul, + '/' => Token::Div, + '^' => Token::Pow, + '(' => Token::LParen, + ')' => Token::RParen, + _ => unreachable!(), + }; + tokens.push(t); + } else if !c.is_whitespace() { + current.push(c); + } + } + + if !current.is_empty() { + tokens.push(Token::Value(current)); + } + + tokens +} + +// ---- Shunting-Yard ---- + +/// Convert infix tokens to [Reverse Polish Notation] (RPN) using the +/// [Shunting-Yard] algorithm. +/// +/// RPN (also called postfix notation) writes operators after their operands, +/// e.g. `kJ / mol` becomes `kJ mol /`. This removes the need for parentheses +/// and precedence rules, making the expression easy to evaluate with a stack. +/// +/// All operators are treated as left-associative. +/// +/// [Reverse Polish Notation]: https://en.wikipedia.org/wiki/Reverse_Polish_notation +/// [Shunting-Yard]: https://en.wikipedia.org/wiki/Shunting-yard_algorithm +fn shunting_yard(tokens: &[Token]) -> Result, Error> { + let mut output: Vec = Vec::new(); + let mut operators: Vec = Vec::new(); + + for token in tokens { + match token { + Token::Value(_) => { + output.push(token.clone()); + } + Token::Mul | Token::Div | Token::Pow => { + while let Some(top) = operators.last() { + if token.precedence() <= top.precedence() { + output.push(operators.pop().unwrap()); + } else { + break; + } + } + operators.push(token.clone()); + } + Token::LParen => { + operators.push(token.clone()); + } + Token::RParen => { + while let Some(top) = operators.last() { + if matches!(top, Token::LParen) { + break; + } + output.push(operators.pop().unwrap()); + } + if operators.is_empty() || !matches!(operators.last(), Some(Token::LParen)) { + return Err(Error::InvalidParameter( + "unit expression has unbalanced parentheses".into(), + )); + } + operators.pop(); // discard LParen + } + } + } + + while let Some(top) = operators.pop() { + if matches!(top, Token::LParen | Token::RParen) { + return Err(Error::InvalidParameter( + "unit expression has unbalanced parentheses".into(), + )); + } + output.push(top); + } + + Ok(output) +} + +// ---- AST evaluator ---- + +struct UnitExpr { + val: UnitExprData, +} + +enum UnitExprData { + Val(UnitValue, String), + Mul(Box, Box), + Div(Box, Box), + Pow(Box, Box), +} + +impl fmt::Display for UnitExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match &self.val { + UnitExprData::Val(_, name) => f.write_str(name), + UnitExprData::Mul(lhs, rhs) => { + write!(f, "({} * {})", lhs, rhs) + } + UnitExprData::Div(lhs, rhs) => { + write!(f, "({} / {})", lhs, rhs) + } + UnitExprData::Pow(base, exponent) => { + write!(f, "({} ^ {})", base, exponent) + } + } + } +} + +impl UnitExpr { + fn eval(&self) -> Result { + match &self.val { + UnitExprData::Val(v, _) => Ok(v.clone()), + UnitExprData::Mul(lhs, rhs) => { + let l = lhs.eval()?; + let r = rhs.eval()?; + let result_factor = l.factor * r.factor; + if !result_factor.is_finite() { + return Err(Error::InvalidParameter(format!( + "unit conversion factor overflows: multiplication result is infinite \ + or NaN for '{}'", + self + ))); + } + Ok(UnitValue { + factor: result_factor, + dim: &l.dim + &r.dim, + }) + } + UnitExprData::Div(lhs, rhs) => { + let l = lhs.eval()?; + let r = rhs.eval()?; + let result_factor = l.factor / r.factor; + if !result_factor.is_finite() { + return Err(Error::InvalidParameter(format!( + "unit conversion factor overflows: division result is infinite \ + or NaN for '{}'", + self + ))); + } + Ok(UnitValue { + factor: result_factor, + dim: &l.dim - &r.dim, + }) + } + UnitExprData::Pow(base, exponent) => { + let b = base.eval()?; + let e = exponent.eval()?; + + if e.dim != Dimension::NONE { + return Err(Error::InvalidParameter(format!( + "exponent in unit expression must be dimensionless, got dimension {} \ + for exponent '{}'", + e.dim, + exponent + ))); + } + let result_factor = b.factor.powf(e.factor); + if !result_factor.is_finite() { + return Err(Error::InvalidParameter(format!( + "unit conversion factor overflows: exponentiation result is infinite \ + or NaN for '{}'", + self + ))); + } + Ok(UnitValue { + factor: result_factor, + dim: b.dim.pow(e.factor), + }) + } + } + } +} + +/// Read one expression from the [RPN] stream (recursive, pops from the back). +/// +/// RPN arranges expressions as `lhs rhs op`, so `rhs` is on top of the stack +/// and must be popped first. For example `kJ mol /` pops `mol` (rhs) then +/// `kJ` (lhs) to build `Div(lhs=kJ, rhs=mol)`. +/// +/// [RPN]: https://en.wikipedia.org/wiki/Reverse_Polish_notation +fn read_expr(stream: &mut Vec) -> Result { + let token = stream.pop().ok_or_else(|| { + Error::InvalidParameter("malformed unit expression: missing a value".into()) + })?; + + match token { + Token::Value(v) => { + let lower = v.to_lowercase(); + if let Some(uv) = BASE_UNITS.get(lower.as_str()) { + return Ok(UnitExpr { + val: UnitExprData::Val(uv.clone(), v), + }); + } + if let Ok(val) = v.parse::() { + return Ok(UnitExpr { + val: UnitExprData::Val(UnitValue { factor: val, dim: Dimension::NONE }, v), + }); + } + Err(Error::InvalidParameter(format!("unknown unit '{}'", v))) + } + // RPN: lhs rhs Mul — pop rhs first, then lhs + Token::Mul => { + let rhs = read_expr(stream)?; + let lhs = read_expr(stream)?; + Ok(UnitExpr { + val: UnitExprData::Mul(Box::new(lhs), Box::new(rhs)), + }) + } + // RPN: lhs rhs Div — pop rhs first, then lhs + Token::Div => { + let rhs = read_expr(stream)?; + let lhs = read_expr(stream)?; + Ok(UnitExpr { + val: UnitExprData::Div(Box::new(lhs), Box::new(rhs)), + }) + } + // RPN: base exponent Pow — pop exponent first, then base + Token::Pow => { + let exponent = read_expr(stream)?; + let base = read_expr(stream)?; + Ok(UnitExpr { + val: UnitExprData::Pow(Box::new(base), Box::new(exponent)), + }) + } + _ => Err(Error::InvalidParameter(format!( + "unexpected symbol in unit expression: '{}'", + token + ))), + } +} + +/// Parse a unit expression string and return the evaluated `UnitValue`. +fn parse_unit_expression(unit: &str) -> Result { + if unit.is_empty() { + return Ok(UnitValue { factor: 1.0, dim: Dimension::NONE }); + } + + let tokens = tokenize(unit); + if tokens.is_empty() { + return Ok(UnitValue { factor: 1.0, dim: Dimension::NONE }); + } + + let mut rpn = shunting_yard(&tokens)?; + let ast = read_expr(&mut rpn)?; + + if !rpn.is_empty() { + let remaining: Vec = rpn.iter().map(|t| t.to_string()).collect(); + return Err(Error::InvalidParameter(format!( + "malformed unit expression: leftover input '{}'", + remaining.join(" ") + ))); + } + + ast.eval() +} + +/// Get the multiplicative conversion factor to use to convert from +/// `from_unit` to `to_unit`. Both units are parsed as expressions (e.g. +/// "kJ/mol/A^2", "(eV*u)^(1/2)") and their dimensions must match. +/// +/// Unit expressions are built from base units combined with `*`, `/`, `^`, +/// and parentheses. Unit lookup is case-insensitive, and whitespace is +/// ignored. For example: +/// +/// - `"kJ/mol"` -- energy per mole +/// - `"eV/Angstrom^3"` -- pressure +/// - `"(eV*u)^(1/2)"` -- momentum (fractional powers) +/// - `"Hartree/Bohr"` -- force in atomic units +pub fn unit_conversion_factor(from_unit: &str, to_unit: &str) -> Result { + if from_unit.is_empty() || to_unit.is_empty() { + return Ok(1.0); + } + + let from = parse_unit_expression(from_unit)?; + let to = parse_unit_expression(to_unit)?; + + if from.dim != to.dim { + return Err(Error::InvalidParameter(format!( + "dimension mismatch in unit conversion: '{}' has dimension {} but '{}' has dimension {}", + from_unit, + from.dim, + to_unit, + to.dim + ))); + } + + Ok(from.factor / to.factor) +} + + +/// Check if a unit expression is valid and has the same dimension as the reference unit. +pub fn validate_unit(unit: &str, reference_unit: &str, context: Option<&str>) -> Result<(), Error> { + let unit_value = parse_unit_expression(unit)?; + let reference_value = parse_unit_expression(reference_unit)?; + + if unit_value.dim != reference_value.dim { + return Err(Error::InvalidParameter(format!( + "dimension mismatch{}: '{}' has dimension {} but expected dimension {}", + context.map_or_else(String::new, |c| format!(" in {}", c)), + unit, + unit_value.dim, + reference_value.dim + ))); + } + + Ok(()) +} + + +#[cfg(test)] +#[allow(clippy::float_cmp)] +mod tests { + use super::*; + + #[test] + fn test_tokenize_simple() { + let tokens = tokenize("eV"); + assert_eq!(tokens.len(), 1); + assert!(matches!(&tokens[0], Token::Value(v) if v == "eV")); + } + + #[test] + fn test_tokenize_operators() { + let tokens = tokenize("kJ/mol/A^2"); + let types: Vec = tokens.iter().map(|t| t.to_string()).collect(); + assert_eq!(types, vec!["kJ", "/", "mol", "/", "A", "^", "2"]); + } + + #[test] + fn test_tokenize_parens() { + let tokens = tokenize("(eV*u)^(1/2)"); + let types: Vec = tokens.iter().map(|t| t.to_string()).collect(); + assert_eq!(types, vec!["(", "eV", "*", "u", ")", "^", "(", "1", "/", "2", ")"]); + } + + #[test] + fn test_tokenize_whitespace() { + let tokens = tokenize(" kJ / mol "); + let types: Vec = tokens.iter().map(|t| t.to_string()).collect(); + assert_eq!(types, vec!["kJ", "/", "mol"]); + } + + #[test] + fn test_shunting_yard() { + let tokens = tokenize("kJ/mol"); + let rpn = shunting_yard(&tokens).unwrap(); + let types: Vec = rpn.iter().map(|t| t.to_string()).collect(); + assert_eq!(types, vec!["kJ", "mol", "/"]); + + let tokens = tokenize("kJ/mol/A^2"); + let rpn = shunting_yard(&tokens).unwrap(); + let types: Vec = rpn.iter().map(|t| t.to_string()).collect(); + assert_eq!(types, vec!["kJ", "mol", "/", "A", "2", "^", "/"]); + } + + #[test] + fn test_parens_mismatch() { + let tokens = tokenize("("); + let err = shunting_yard(&tokens).expect_err("expected error"); + assert_eq!( + err.to_string(), + "invalid parameter: unit expression has unbalanced parentheses" + ); + + let tokens = tokenize("(eV*u"); + let err = shunting_yard(&tokens).expect_err("expected error"); + assert_eq!( + err.to_string(), + "invalid parameter: unit expression has unbalanced parentheses" + ); + } + + #[test] + fn test_simple_conversion() { + let factor = unit_conversion_factor("eV", "eV").unwrap(); + assert_eq!(factor, 1.0); + + let factor = unit_conversion_factor("m", "A").unwrap(); + assert!((factor - 1e10).abs() < 1e-5); + + let factor = unit_conversion_factor("eV", "kJ").unwrap(); + assert!((factor - 1.602176634e-22).abs() < 1e-30); + } + + #[test] + fn test_dimension_mismatch() { + let err = unit_conversion_factor("eV", "m").expect_err("expected error"); + assert_eq!( + err.to_string(), + "invalid parameter: dimension mismatch in unit conversion: \ + 'eV' has dimension [L^2 T^-2 M] but 'm' has dimension [L]" + ); + } + + #[test] + fn test_empty_units() { + let factor = unit_conversion_factor("", "").unwrap(); + assert_eq!(factor, 1.0); + + let factor = unit_conversion_factor("eV", "").unwrap(); + assert_eq!(factor, 1.0); + } + + #[test] + fn test_compound_units() { + let from = unit_conversion_factor("kJ/mol", "eV").unwrap(); + assert!((from - 0.010364269656262174).abs() < 1e-15); + + let from = unit_conversion_factor("eV/A^3", "GPa").unwrap(); + assert!((from - 160.21766339999996).abs() < 1e-12); + } + + #[test] + fn test_case_insensitive() { + let f1 = unit_conversion_factor("eV", "eV").unwrap(); + let f2 = unit_conversion_factor("EV", "eV").unwrap(); + assert_eq!(f1, f2); + + let factor = unit_conversion_factor("eV", "MeV").unwrap(); + assert!((factor - 1000.0).abs() < 1e-12); + } + + #[test] + fn test_unknown_unit() { + let err = unit_conversion_factor("foo", "eV").expect_err("expected error"); + assert_eq!(err.to_string(), "invalid parameter: unknown unit 'foo'"); + } + + #[test] + fn test_numeric_literal() { + let factor = unit_conversion_factor("2", "1").unwrap(); + assert_eq!(factor, 2.0); + } + + #[test] + fn test_fractional_power() { + let err = unit_conversion_factor("(eV*u)^(1/2)", "eV*u").expect_err("expected error"); + assert_eq!( + err.to_string(), + "invalid parameter: dimension mismatch in unit conversion: \ + '(eV*u)^(1/2)' has dimension [L T^-1 M] but 'eV*u' has dimension [L^2 T^-2 M^2]" + ); + + let factor = unit_conversion_factor("(eV*u)^(1/2)", "(eV*u)^(1/2)").unwrap(); + assert_eq!(factor, 1.0); + } + + #[test] + fn test_dimension_to_string() { + assert_eq!(Dimension::NONE.to_string(), "[dimensionless]"); + assert_eq!(Dimension::LENGTH.to_string(), "[L]"); + assert_eq!(Dimension::ENERGY.to_string(), "[L^2 T^-2 M]"); + assert_eq!(Dimension::PRESSURE.to_string(), "[L^-1 T^-2 M]"); + assert_eq!(Dimension::TEMPERATURE.to_string(), "[Θ]"); + + let velocity = Dimension { length: 1, time: -1, mass: 0, electric_current: 0, temperature: 0 }; + assert_eq!(velocity.to_string(), "[L T^-1]"); + } +} diff --git a/metatomic-core/src/utils.rs b/metatomic-core/src/utils.rs new file mode 100644 index 000000000..aa7e97704 --- /dev/null +++ b/metatomic-core/src/utils.rs @@ -0,0 +1,148 @@ +use metatensor::{TensorBlock, TensorMap}; + +use crate::kernels; +use crate::Error; + +/// Scale all values and gradients in a `TensorMap` by `factor`. +/// +/// If a block's values or gradients are marked as read-only in DLPack, a copy +/// is made via `mts_array_t.copy` and the copy is scaled; otherwise the data is +/// scaled in place. This is a no-op when `factor == 1.0`. +pub(crate) fn scale_tensormap(tensor: TensorMap, factor: f64) -> Result { + #[allow(clippy::float_cmp)] + if factor == 1.0 { + return Ok(tensor); + } + + let dlpack_version = dlpk::sys::DLPackVersion::current(); + + // Check if all blocks' values and gradients are writable + let mut all_writable = true; + for block in tensor.blocks() { + let device = block.values().device()?; + let dlpack = block.values().as_dlpack(device, None, dlpack_version)?; + if dlpack.is_read_only() { + all_writable = false; + break; + } + for (_, gradient) in block.gradients() { + let grad_dlpack = gradient.values().as_dlpack(device, None, dlpack_version)?; + if grad_dlpack.is_read_only() { + all_writable = false; + break; + } + } + if !all_writable { + break; + } + } + + if all_writable { + // Scale in place + let mut tensor = tensor; + for mut block in tensor.blocks_mut() { + let device = block.values().device()?; + let mut dlpack = block.values_mut().as_dlpack(device, None, dlpack_version)?; + kernels::scale_inplace(dlpack.as_mut(), factor)?; + + for (_, mut gradient) in block.gradients_mut() { + let mut dlpack = gradient.values_mut().as_dlpack(device, None, dlpack_version)?; + kernels::scale_inplace(dlpack.as_mut(), factor)?; + } + } + Ok(tensor) + } else { + // At least one block is read-only: build new blocks with copied + scaled data + let keys = tensor.keys().clone(); + let mut new_blocks = Vec::new(); + + for block in tensor.blocks() { + // copy values, then scale the copy in place + let device = block.values().device()?; + let values_copy = block.values().copy(device)?; + let mut dlpack = values_copy.as_dlpack(device, None, dlpack_version)?; + assert!(!dlpack.is_read_only(), "copy of the value is still read only"); + kernels::scale_inplace(dlpack.as_mut(), factor)?; + + let samples = block.samples(); + let components = block.components(); + let properties = block.properties(); + let mut new_block = TensorBlock::new( + values_copy, + &samples, + &components, + &properties, + )?; + + // copy and scale all gradients + for (parameter, gradient) in block.gradients() { + let grad_copy = gradient.values().copy(device)?; + let mut dlpack = grad_copy.as_dlpack(device, None, dlpack_version)?; + assert!(!dlpack.is_read_only(), "copy of the gradients is still read only"); + kernels::scale_inplace(dlpack.as_mut(), factor)?; + + let grad_samples = gradient.samples(); + let grad_components = gradient.components(); + let grad_properties = gradient.properties(); + let new_gradient = TensorBlock::new( + grad_copy, + &grad_samples, + &grad_components, + &grad_properties, + )?; + + new_block.add_gradient(parameter, new_gradient)?; + } + + new_blocks.push(new_block); + } + + TensorMap::new(keys, new_blocks).map_err(Error::from) + } +} + + + +#[cfg(test)] +mod tests { + use super::*; + use metatensor::Labels; + + #[test] + fn test_scale_tensormap() { + // build a TensorMap with one block containing values and a gradient + let samples = Labels::new(["system"], [[0i32]]); + let components = vec![]; + let properties = Labels::new(["energy"], [[0i32]]); + let values = ndarray::ArrayD::::from_shape_vec(vec![1, 1], vec![1.0]).unwrap(); + let mut block = TensorBlock::new(values, &samples, &components, &properties).unwrap(); + + // add a gradient w.r.t. positions + let grad_samples = Labels::new(["sample"], [[0i32]]); + let grad_components = vec![Labels::new(["xyz"], [[0i32], [1i32], [2]])]; + let grad_properties = Labels::new(["energy"], [[0i32]]); + let grad_values = ndarray::ArrayD::::from_shape_vec(vec![1, 3, 1], vec![2.0, 4.0, 6.0]).unwrap(); + let gradient = TensorBlock::new(grad_values, &grad_samples, &grad_components, &grad_properties).unwrap(); + block.add_gradient("positions", gradient).unwrap(); + + let keys = Labels::new(["_"], [[0i32]]); + let tensor_map = TensorMap::new(keys, vec![block]).unwrap(); + + // scale by 2.5 + let scaled = scale_tensormap(tensor_map, 2.5).unwrap(); + + // check values + let block = scaled.block_by_id(0); + let device = block.values().device().unwrap(); + let dlpack = block.values().as_dlpack(device, None, dlpk::sys::DLPackVersion::current()).unwrap(); + let values: ndarray::ArrayViewD = dlpack.as_ref().try_into().unwrap(); + assert_eq!(values, ndarray::arr2(&[[2.5_f32]]).into_dyn()); + + // check gradient + let gradient = block.gradient("positions").unwrap(); + let device = gradient.values().device().unwrap(); + let dlpack = gradient.values().as_dlpack(device, None, dlpk::sys::DLPackVersion::current()).unwrap(); + let grad_values: ndarray::ArrayViewD = dlpack.as_ref().try_into().unwrap(); + assert_eq!(grad_values, ndarray::arr1(&[5.0_f32, 10.0, 15.0]).to_shape(vec![1, 3, 1]).unwrap()); + } +} diff --git a/metatomic-core/tests/CMakeLists.txt b/metatomic-core/tests/CMakeLists.txt new file mode 100644 index 000000000..ac716ebb4 --- /dev/null +++ b/metatomic-core/tests/CMakeLists.txt @@ -0,0 +1,131 @@ +cmake_minimum_required(VERSION 3.22) +project(metatomic-tests) + +if (${CMAKE_CURRENT_SOURCE_DIR} STREQUAL ${CMAKE_SOURCE_DIR}) + if("${CMAKE_BUILD_TYPE}" STREQUAL "" AND "${CMAKE_CONFIGURATION_TYPES}" STREQUAL "") + message(STATUS "Setting build type to 'release' as none was specified.") + set(CMAKE_BUILD_TYPE "release" + CACHE STRING + "Choose the type of build, options are: debug or release" + FORCE) + set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS release debug) + endif() +endif() + +if (MINGW) + # CI can't find libsdc++, so we statically link it + set(CMAKE_EXE_LINKER_FLAGS "-static-libstdc++") +endif() + +add_subdirectory(../ metatomic) +get_target_property(METATOMIC_IMPORTED_LOCATION metatomic::shared IMPORTED_LOCATION) +get_filename_component(METATOMIC_DIR ${METATOMIC_IMPORTED_LOCATION} DIRECTORY) + +add_subdirectory(external) + +find_program(VALGRIND valgrind) +if (VALGRIND) + if (NOT "$ENV{METATOMIC_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" + "--gen-suppressions=all" + ) + endif() +else() + set(TEST_COMMAND "") +endif() + +if (CMAKE_CXX_COMPILER_ID MATCHES "Clang") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Weverything") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-c++98-compat") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-c++98-compat-pedantic") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-weak-vtables") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-float-equal") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-prototypes") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-shadow-uncaptured-local") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-padded") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unsafe-buffer-usage") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-poison-system-directories") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-allocator-wrappers") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-covered-switch-default") +endif() + + +enable_testing() +add_subdirectory(test-plugins) + +if (TARGET metatensor::shared) + get_target_property(METATENSOR_LOCATION metatensor::shared IMPORTED_LOCATION) + get_filename_component(METATENSOR_DIR ${METATENSOR_LOCATION} DIRECTORY) +elseif (TARGET metatensor) + get_target_property(METATENSOR_LOCATION metatensor LOCATION) + get_filename_component(METATENSOR_DIR ${METATENSOR_LOCATION} DIRECTORY) +else() + set(METATENSOR_DIR "") +endif() + +function(metatomic_add_test source target) + add_executable(${target} ${source}) + target_link_libraries(${target} metatomic) + + set_target_properties(${target} PROPERTIES + # make sure we can dlsym the function generated by MTA_REGISTER_PLUGIN + # on linux when loading from the same executable directly + ENABLE_EXPORTS ON + + # Ensure that the binaries find the right shared library. + # + # Without this, when configuring with cmake before the library is built, + # cmake does not find the library on the filesystem and does not add the + # RPATH to executables linking to it + BUILD_RPATH ${METATOMIC_DIR} + NO_SYSTEM_FROM_IMPORTED ON + ) + + target_compile_definitions(${target} PRIVATE PLUGIN_DIR="$") + + add_test( + NAME ${target} + COMMAND ${TEST_COMMAND} $ + ) + + if(WIN32) + # We need to set the path to allow access to metatomic.dll and + # metatensor.dll. This does a similar job to the BUILD_RPATH above. + STRING(REPLACE ";" "\\;" PATH_STRING "$ENV{PATH}") + set_tests_properties(${target} PROPERTIES + ENVIRONMENT "PATH=${PATH_STRING}\;$\;${METATENSOR_DIR}" + ) + endif() +endfunction() + +file(GLOB ALL_TESTS *.cpp) +foreach(_file_ ${ALL_TESTS}) + get_filename_component(_name_ ${_file_} NAME_WE) + metatomic_add_test(${_file_} ${_name_}) + target_link_libraries(${_name_} catch) +endforeach() + +file(GLOB ALL_CPP_TESTS cxx/*.cpp) +foreach(_file_ ${ALL_CPP_TESTS}) + get_filename_component(_name_ ${_file_} NAME_WE) + metatomic_add_test(${_file_} "cxx-${_name_}") + target_link_libraries("cxx-${_name_}" catch) +endforeach() + +file(GLOB ALL_EXAMPLES ${CMAKE_CURRENT_SOURCE_DIR}/../../examples/c/*.c) +foreach(_file_ ${ALL_EXAMPLES}) + get_filename_component(_name_ ${_file_} NAME_WE) + metatomic_add_test(${_file_} "example-${_name_}") + # keep assert in the examples + target_compile_options("example-${_name_}" PRIVATE -UNDEBUG) +endforeach() + +file(GLOB ALL_CXX_EXAMPLES ${CMAKE_CURRENT_SOURCE_DIR}/../../examples/cxx/*.cpp) +foreach(_file_ ${ALL_CXX_EXAMPLES}) + get_filename_component(_name_ ${_file_} NAME_WE) + metatomic_add_test(${_file_} "example-cxx-${_name_}") + target_compile_options("example-cxx-${_name_}" PRIVATE -UNDEBUG) +endforeach() diff --git a/metatomic-core/tests/c-model.cpp b/metatomic-core/tests/c-model.cpp new file mode 100644 index 000000000..9c4418bf9 --- /dev/null +++ b/metatomic-core/tests/c-model.cpp @@ -0,0 +1,252 @@ +#include + +#include + +#include "metatomic.h" +#include "tensor_utils.hpp" + +#include + + +struct SimpleModelData { + double scale; +}; + +static mta_status_t unload_impl(void* model_data) { + delete static_cast(model_data); + return MTA_SUCCESS; +} + +static mta_status_t metadata_impl(const void* model_data, mta_string_t* metadata_json) { + (void) model_data; + + *metadata_json = mta_string_create(R"({ + "name": "test C model", + "description": "small model used as a C API example", + "authors": [], + "references": { + "model": [], + "implementation": [], + "architecture": [] + } + })"); + return MTA_SUCCESS; +} + +static mta_status_t capabilities_impl(const void* model_data, mta_string_t* capabilities_json) { + (void) model_data; + + *capabilities_json = mta_string_create(R"({ + "type": "metatomic_model_capabilities", + "outputs": [{ + "type": "metatomic_quantity", + "name": "energy", + "unit": "eV", + "gradients": [], + "sample_kind": "system" + }], + "atomic_types": [1, 6, 8], + "interaction_range": 4.5, + "length_unit": "nm", + "supported_devices": ["cpu"], + "dtype": "float32" + })"); + return MTA_SUCCESS; +} + +static mta_status_t requested_pair_lists_impl( + const void* model_data, + mta_string_t* pair_options_json +) { + (void) model_data; + *pair_options_json = mta_string_create("[]"); + return MTA_SUCCESS; +} + +static mta_status_t requested_inputs_impl( + const void* model_data, + mta_string_t* requested_inputs_json +) { + (void) model_data; + *requested_inputs_json = mta_string_create("[]"); + return MTA_SUCCESS; +} + + +mts_tensormap_t* scalar_tensormap(double value) { + auto values = std::make_unique>( + std::vector{1, 1}, + std::vector{static_cast(value)} + ); + + auto array = metatensor::DataArrayBase::to_mts_array(std::move(values)); + + auto samples = metatensor::Labels({"system"}, {{0}}); + auto properties = metatensor::Labels({"energy"}, {{0}}); + + auto* block = mts_block( + std::move(array).release(), + samples.as_mts_labels_t(), + nullptr, + 0, + properties.as_mts_labels_t() + ); + if (block == nullptr) { + return nullptr; + } + + auto keys = metatensor::Labels({"_"}, {{0}}); + auto blocks = std::vector{block}; + return mts_tensormap(keys.as_mts_labels_t(), blocks.data(), blocks.size()); +} + +static mta_status_t execute_inner_impl( + void* model_data, + const mta_system_t* const* systems, + uintptr_t systems_count, + const mts_labels_t* selected_atoms, + const char* requested_outputs_json, + mts_tensormap_t** outputs, + uintptr_t outputs_count +) { + (void)systems; + (void)systems_count; + (void)selected_atoms; + (void)requested_outputs_json; + + auto* data = static_cast(model_data); + + for (uintptr_t i = 0; i < outputs_count; i++) { + outputs[i] = scalar_tensormap(data->scale); + } + + return MTA_SUCCESS; +} + +static mta_status_t load_model_impl( + const char* load_from, + const char* options_json, + mta_model_t* model +) { + (void)options_json; + assert(model != nullptr); + + if (std::strcmp(load_from, "test-c-model") != 0) { + return MTA_UNSUPPORTED_MODEL_ERROR; + } + + model->data = new SimpleModelData{2.0}; + model->unload = unload_impl; + model->metadata = metadata_impl; + model->capabilities = capabilities_impl; + model->requested_pair_lists = requested_pair_lists_impl; + model->requested_inputs = requested_inputs_impl; + model->execute_inner = execute_inner_impl; + + return MTA_SUCCESS; +} + +TEST_CASE("simple C model can be registered and loaded through the C API") { + static auto PLUGIN = mta_plugin_t { + MTA_ABI_VERSION, + "test-c-plugin", + load_model_impl, + }; + mta_register_plugin(PLUGIN); + + auto model = mta_model_t{}; + auto status = mta_load_model("test-c-model", "{}", "test-c-plugin", &model); + REQUIRE(status == MTA_SUCCESS); + + CHECK(model.data != nullptr); + CHECK(model.unload != nullptr); + CHECK(model.metadata != nullptr); + CHECK(model.capabilities != nullptr); + CHECK(model.requested_pair_lists != nullptr); + CHECK(model.requested_inputs != nullptr); + CHECK(model.execute_inner != nullptr); + + mta_string_t metadata = nullptr; + status = model.metadata(model.data, &metadata); + REQUIRE(status == MTA_SUCCESS); + + CHECK(metadata != nullptr); + auto metadata_str = std::string(mta_string_view(metadata)); + mta_string_free(metadata); + + CHECK(metadata_str.find("\"name\": \"test C model\"") != std::string::npos); + + + mta_string_t pair_lists = nullptr; + status = model.requested_pair_lists(model.data, &pair_lists); + REQUIRE(status == MTA_SUCCESS); + + CHECK(pair_lists != nullptr); + CHECK(std::strcmp(mta_string_view(pair_lists), "[]") == 0); + mta_string_free(pair_lists); + + REQUIRE(model.unload(model.data) == MTA_SUCCESS); +} + +TEST_CASE("execute a model through the C API") { + static auto PLUGIN = mta_plugin_t { + MTA_ABI_VERSION, + "test-c-plugin", + load_model_impl, + }; + mta_register_plugin(PLUGIN); + + auto model = mta_model_t{}; + auto status = mta_load_model("test-c-model", "{}", "test-c-plugin", &model); + REQUIRE(status == MTA_SUCCESS); + + // create a water molecule system (O, H, H) with types matching capabilities [1, 6, 8] + auto types_array = std::make_unique>( + std::vector{3}, std::vector{8, 1, 1} + ); + auto types_mts = metatensor::DataArrayBase::to_mts_array(std::move(types_array)); + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + auto* types_dlpack = types_mts.as_dlpack(cpu, nullptr, version); + + mta_system_t* system = nullptr; + status = mta_system_create( + "nm", + types_dlpack, + positions_tensor(3), + cell_tensor(), + pbc_tensor(), + &system + ); + REQUIRE(status == MTA_SUCCESS); + REQUIRE(system != nullptr); + + // run the model + const char* requested_outputs = R"([{ + "type": "metatomic_quantity", + "name": "energy", + "unit": "eV", + "gradients": [], + "sample_kind": "system" + }])"; + + mts_tensormap_t* output = nullptr; + status = mta_execute_model( + model, + &system, + 1, + nullptr, + requested_outputs, + true, + &output, + 1 + ); + CHECK(status == MTA_SUCCESS); + + if (output != nullptr) { + mts_tensormap_free(output); + } + + mta_system_free(system); + REQUIRE(model.unload(model.data) == MTA_SUCCESS); +} diff --git a/metatomic-core/tests/check-cxx-install.rs b/metatomic-core/tests/check-cxx-install.rs new file mode 100644 index 000000000..6baa5b4e1 --- /dev/null +++ b/metatomic-core/tests/check-cxx-install.rs @@ -0,0 +1,66 @@ +use std::path::PathBuf; +use std::sync::Mutex; + +mod utils; + +lazy_static::lazy_static! { + // Make sure only one of the tests below run at the time, since they both + // try to modify the same files + static ref LOCK: Mutex<()> = Mutex::new(()); +} + + +/// Check that metatomic can be built and installed with cmake, and that the +/// installed version can be used from another cmake project with `find_package` +#[test] +fn check_cxx_install() { + let _guard = match LOCK.lock() { + Ok(guard) => guard, + Err(_) => { + panic!("another test failed, stopping") + } + }; + + const CARGO_TARGET_TMPDIR: &str = env!("CARGO_TARGET_TMPDIR"); + + let mut build_dir = PathBuf::from(CARGO_TARGET_TMPDIR); + build_dir.push("cxx-install"); + build_dir.push("cmake-find-package"); + std::fs::create_dir_all(&build_dir).expect("failed to create build dir"); + + // ====================================================================== // + // install dependencies with pip + let deps_dir = build_dir.join("deps"); + let virtualenv_dir = deps_dir.join("virtualenv"); + std::fs::create_dir_all(&virtualenv_dir).expect("failed to create virtualenv dir"); + let python_exe = utils::create_python_venv(virtualenv_dir); + let metatensor_cmake_prefix = utils::setup_metatensor_pip(&python_exe); + + // ====================================================================== // + // build and install metatomic with cmake + let metatomic_dep = deps_dir.join("metatomic-core"); + let source_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + + let cmake_args = vec![ + format!("-DCMAKE_PREFIX_PATH={}", metatensor_cmake_prefix.display()), + ]; + let metatomic_cmake_prefix = utils::setup_metatomic_cmake(&source_dir, &metatomic_dep, cmake_args); + + // ====================================================================== // + // try to use the installed metatomic from cmake + let mut tests_source_dir = source_dir; + tests_source_dir.extend(["tests", "cmake-project"]); + + // configure cmake for the test cmake project + let mut cmake_config = utils::cmake_config(&tests_source_dir, &build_dir); + cmake_config.arg(format!("-DCMAKE_PREFIX_PATH={};{}", metatensor_cmake_prefix.display(), metatomic_cmake_prefix.display())); + utils::run_command(cmake_config, "cmake configuration"); + + // build the code, linking to metatomic + let cmake_build = utils::cmake_build(&build_dir); + utils::run_command(cmake_build, "cmake build"); + + // run the executables + let ctest = utils::ctest(&build_dir); + utils::run_command(ctest, "ctest"); +} diff --git a/metatomic-core/tests/cmake-project/CMakeLists.txt b/metatomic-core/tests/cmake-project/CMakeLists.txt new file mode 100644 index 000000000..2b04acfa4 --- /dev/null +++ b/metatomic-core/tests/cmake-project/CMakeLists.txt @@ -0,0 +1,84 @@ +cmake_minimum_required(VERSION 3.22) + +message(STATUS "Running with CMake version ${CMAKE_VERSION}") + +project(metatomic-test-cmake-project C CXX) + +option(USE_CMAKE_SUBDIRECTORY OFF) + +if (MINGW) + # CI can't find libsdc++, so we statically link it + set(CMAKE_EXE_LINKER_FLAGS "-static-libstdc++") +endif() + + +if (USE_CMAKE_SUBDIRECTORY) + message(STATUS "Using metatomic with add_subdirectory") + # build metatomic as part of this project + add_subdirectory(../../ metatomic) + + # load metatomic from the build path + set(CMAKE_BUILD_RPATH "$") +else() + message(STATUS "Using metatomic with find_package") + # If building a dev version, we also need to update the REQUIRED_METATOMIC_VERSION + # in the same way we update the metatomic-torch version + include(../../cmake/dev-versions.cmake) + set(REQUIRED_METATOMIC_VERSION "0.1.0") + create_development_version("${REQUIRED_METATOMIC_VERSION}" METATOMIC_CORE_FULL_VERSION "metatomic-core-v") + string(REGEX REPLACE "([0-9]*)\\.([0-9]*).*" "\\1.\\2" REQUIRED_METATOMIC_VERSION ${METATOMIC_CORE_FULL_VERSION}) + + find_package(metatomic ${REQUIRED_METATOMIC_VERSION} REQUIRED) + + if(TARGET metatomic::shared) + get_target_property(mta_build_version metatomic::shared BUILD_VERSION) + if (NOT ${mta_build_version} STREQUAL ${METATOMIC_CORE_FULL_VERSION}) + message(FATAL_ERROR "Invalid BUILD_VERSION for metatomic::shared, expected ${METATOMIC_CORE_FULL_VERSION} but got ${mta_build_version}") + endif() + endif() + + if(TARGET metatomic::static) + get_target_property(mta_build_version metatomic::static BUILD_VERSION) + if (NOT ${mta_build_version} STREQUAL ${METATOMIC_CORE_FULL_VERSION}) + message(FATAL_ERROR "Invalid BUILD_VERSION for metatomic::static, expected ${METATOMIC_CORE_FULL_VERSION} but got ${mta_build_version}") + endif() + endif() +endif() + +enable_testing() + + +if(TARGET metatomic::shared) + add_executable(c-main src/main.c) + target_link_libraries(c-main metatomic::shared) + + add_executable(cxx-main src/main.cpp) + target_link_libraries(cxx-main metatomic::shared) + + add_test(NAME c-main COMMAND c-main) + add_test(NAME cxx-main COMMAND cxx-main) + + if(WIN32) + # We need to set the path to allow access to metatomic.dll + STRING(REPLACE ";" "\\;" PATH_STRING "$ENV{PATH}") + set_tests_properties(c-main PROPERTIES + ENVIRONMENT "PATH=${PATH_STRING}\;$" + ) + + set_tests_properties(cxx-main PROPERTIES + ENVIRONMENT "PATH=${PATH_STRING}\;$" + ) + endif() +endif() + + +if(TARGET metatomic::static) + add_executable(c-main-static src/main.c) + target_link_libraries(c-main-static metatomic::static) + + add_executable(cxx-main-static src/main.cpp) + target_link_libraries(cxx-main-static metatomic::static) + + add_test(NAME c-main-static COMMAND c-main-static) + add_test(NAME cxx-main-static COMMAND cxx-main-static) +endif() diff --git a/metatomic-core/tests/cmake-project/README.md b/metatomic-core/tests/cmake-project/README.md new file mode 100644 index 000000000..70a687bf0 --- /dev/null +++ b/metatomic-core/tests/cmake-project/README.md @@ -0,0 +1,3 @@ +# Sample CMake project using metatomic + +This is a basic cmake project linking to metatomic from C and C++ code. diff --git a/metatomic-core/tests/cmake-project/src/main.c b/metatomic-core/tests/cmake-project/src/main.c new file mode 100644 index 000000000..dcad0f764 --- /dev/null +++ b/metatomic-core/tests/cmake-project/src/main.c @@ -0,0 +1,8 @@ +#include + +#include + +int main(void) { + printf("Metatomic version: %s\n", mta_version()); + return 0; +} diff --git a/metatomic-core/tests/cmake-project/src/main.cpp b/metatomic-core/tests/cmake-project/src/main.cpp new file mode 100644 index 000000000..04ec152b6 --- /dev/null +++ b/metatomic-core/tests/cmake-project/src/main.cpp @@ -0,0 +1,9 @@ +#include + +#include + + +int main() { + std::cout << "Metatomic version: " << mta_version() << std::endl; + return 0; +} diff --git a/metatomic-core/tests/cxx/helpers.hpp b/metatomic-core/tests/cxx/helpers.hpp new file mode 100644 index 000000000..9df99d472 --- /dev/null +++ b/metatomic-core/tests/cxx/helpers.hpp @@ -0,0 +1,126 @@ +#pragma once + +#include +#include + +#include +#include "metatomic.hpp" + +/// Build a `types` tensor for `n_atoms` atoms. +/// +/// The tensor has dtype `T` and shape `(n_atoms,)`. +template +inline metatomic::DLPackTensor types_tensor(size_t n_atoms) { + auto type_data = std::vector(); + type_data.reserve(n_atoms); + for (size_t i=0; i(i * 3 + 1)); + } + + auto array = std::make_unique>( + std::vector{n_atoms}, std::move(type_data) + ); + auto mts = metatensor::DataArrayBase::to_mts_array(std::move(array)); + + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + return metatomic::DLPackTensor(mts.as_dlpack(cpu, nullptr, version)); +} + +/// Build a `positions` tensor for `n_atoms` atoms. +/// +/// The tensor has dtype `T` and shape `(n_atoms, 3)`. +template +inline metatomic::DLPackTensor positions_tensor(size_t n_atoms) { + auto position_data = std::vector(); + position_data.reserve(n_atoms * 3); + for (size_t i=0; i(i * 3 + 1)); + position_data.push_back(static_cast(i * 3 + 2)); + position_data.push_back(static_cast(i * 3 + 3)); + } + + auto array = std::make_unique>( + std::vector{n_atoms, 3}, std::move(position_data) + ); + auto mts = metatensor::DataArrayBase::to_mts_array(std::move(array)); + + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + return metatomic::DLPackTensor(mts.as_dlpack(cpu, nullptr, version)); +} + +/// Build a `cell` tensor. +/// +/// The tensor has dtype `T` and shape `(3, 3)`. The `y` row is zero to +/// match the non-periodic `y` direction in `pbc`. +template +inline metatomic::DLPackTensor cell_tensor() { + auto array = std::make_unique>( + std::vector{3, 3}, + std::vector{ + T(10.0), T(0.0), T(0.0), + T(0.0), T(0.0), T(0.0), + T(0.0), T(0.0), T(10.0), + } + ); + auto mts = metatensor::DataArrayBase::to_mts_array(std::move(array)); + + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + return metatomic::DLPackTensor(mts.as_dlpack(cpu, nullptr, version)); +} + +/// Build a `pbc` tensor. +/// +/// The tensor has dtype bool and shape `(3,)`. +inline metatomic::DLPackTensor pbc_tensor() { + // `SimpleDataArray` does not compile (`std::vector` has no + // `data()` method), so we use `uint8_t` and patch the dtype code to + // `kDLBool`. + auto array = std::make_unique>( + std::vector{3}, std::vector{1, 0, 1} + ); + auto mts = metatensor::DataArrayBase::to_mts_array(std::move(array)); + + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + auto* tensor = mts.as_dlpack(cpu, nullptr, version); + tensor->dl_tensor.dtype.code = DLDataTypeCode::kDLBool; + + return metatomic::DLPackTensor(tensor); +} + +/// Build a simple `System` with `n_atoms` atoms. +inline metatomic::System test_system(size_t n_atoms = 4) { + return metatomic::System( + "nm", + types_tensor(n_atoms), + positions_tensor(n_atoms), + cell_tensor(), + pbc_tensor() + ); +} + +/// Build a `TensorMap` holding a single scalar value. +/// +/// @param value scalar value to store +/// @param property name of the single property in the returned tensor map +inline metatensor::TensorMap scalar_tensor(double value, const std::string& property) { + auto values = std::make_unique>( + std::vector{1, 1}, std::vector{value} + ); + + auto samples = metatensor::Labels({"system"}, {{0}}); + auto properties = metatensor::Labels({property}, {{0}}); + + auto block = metatensor::TensorBlock( + std::move(values), samples, {}, properties + ); + + auto blocks = std::vector(); + blocks.push_back(std::move(block)); + + auto keys = metatensor::Labels({"_"}, {{0}}); + return metatensor::TensorMap(keys, std::move(blocks)); +} diff --git a/metatomic-core/tests/cxx/metadata.cpp b/metatomic-core/tests/cxx/metadata.cpp new file mode 100644 index 000000000..ec5590fed --- /dev/null +++ b/metatomic-core/tests/cxx/metadata.cpp @@ -0,0 +1,871 @@ +#include + +#include "metatomic.hpp" + +TEST_CASE("JSON serialization C++ API") { + SECTION("PairListOptions"){ + double cutoff = 3.0; + std::string cutoff_hex = "0x4008000000000000"; + + SECTION("Builder construction") { + auto p1 = metatomic::PairListOptions::builder() + .cutoff(cutoff) + .full_list(true) + .strict(false) + .add_requestor("model1") + .add_requestor("model2") + .build(); + + nlohmann::json j = p1; + + CHECK(j["cutoff"] == cutoff_hex); + CHECK(j["full_list"] == true); + CHECK(j["strict"] == false); + CHECK(j["requestors"].is_array()); + CHECK(j["requestors"].size() == 2); + CHECK(j["requestors"][0] == "model1"); + CHECK(j["requestors"][1] == "model2"); + + auto p2 = j.get(); + CHECK(p2.cutoff() == Approx(cutoff)); + CHECK(p2.full_list() == true); + CHECK(p2.strict() == false); + CHECK(p2.requestors().size() == 2); + CHECK(p2.requestors()[0] == "model1"); + CHECK(p2.requestors()[1] == "model2"); + } + + SECTION("Builder with requestors set as a list") { + auto p1 = metatomic::PairListOptions::builder() + .cutoff(cutoff) + .full_list(true) + .strict(false) + .requestors({"model1", "model2"}) + .build(); + + nlohmann::json j = p1; + + CHECK(j["cutoff"] == cutoff_hex); + CHECK(j["full_list"] == true); + CHECK(j["strict"] == false); + CHECK(j["requestors"].is_array()); + CHECK(j["requestors"].size() == 2); + CHECK(j["requestors"][0] == "model1"); + CHECK(j["requestors"][1] == "model2"); + + auto p2 = j.get(); + CHECK(p2.cutoff() == Approx(cutoff)); + CHECK(p2.full_list() == true); + CHECK(p2.strict() == false); + CHECK(p2.requestors().size() == 2); + CHECK(p2.requestors()[0] == "model1"); + CHECK(p2.requestors()[1] == "model2"); + } + + SECTION("build() validates completeness") { + CHECK_THROWS_WITH( + metatomic::PairListOptions::builder().build(), + Catch::Matchers::StartsWith("cutoff must be set before building PairListOptions") + ); + + CHECK_THROWS_WITH( + metatomic::PairListOptions::builder().full_list(true).build(), + Catch::Matchers::StartsWith("cutoff must be set before building PairListOptions") + ); + + CHECK_THROWS_WITH( + metatomic::PairListOptions::builder().cutoff(cutoff).build(), + Catch::Matchers::StartsWith("full_list must be set before building PairListOptions") + ); + + CHECK_THROWS_WITH( + metatomic::PairListOptions::builder().cutoff(-1.0), + Catch::Matchers::StartsWith("cutoff must be a finite positive number") + ); + } + + SECTION("add_requestor ignores empty strings and duplicates") { + auto p1 = metatomic::PairListOptions::builder() + .cutoff(cutoff) + .full_list(true) + .add_requestor("model1") + .add_requestor("") + .add_requestor("model2") + .add_requestor("model1") + .build(); + + auto requestors = p1.requestors(); + CHECK(requestors.size() == 2); + CHECK(requestors[0] == "model1"); + CHECK(requestors[1] == "model2"); + + nlohmann::json j = p1; + CHECK(j["requestors"].size() == 2); + CHECK(j["requestors"][0] == "model1"); + CHECK(j["requestors"][1] == "model2"); + } + + } + + SECTION("References") { + SECTION("Builder construction") { + auto r1 = metatomic::ModelMetadata::References::builder() + .model({"model ref 1", "model ref 2"}) + .architecture({"architecture ref 1"}) + .implementation({"implementation ref 1", "implementation ref 2"}) + .build(); + + nlohmann::json j = r1; + + CHECK(j["model"].is_array()); + CHECK(j["model"].size() == 2); + CHECK(j["model"][0] == "model ref 1"); + CHECK(j["model"][1] == "model ref 2"); + + CHECK(j["architecture"].is_array()); + CHECK(j["architecture"].size() == 1); + CHECK(j["architecture"][0] == "architecture ref 1"); + + CHECK(j["implementation"].is_array()); + CHECK(j["implementation"].size() == 2); + CHECK(j["implementation"][0] == "implementation ref 1"); + CHECK(j["implementation"][1] == "implementation ref 2"); + + auto r2 = j.get(); + CHECK(r2.model()[0] == "model ref 1"); + CHECK(r2.model()[1] == "model ref 2"); + CHECK(r2.architecture().size() == 1); + CHECK(r2.architecture()[0] == "architecture ref 1"); + CHECK(r2.implementation().size() == 2); + CHECK(r2.implementation()[0] == "implementation ref 1"); + CHECK(r2.implementation()[1] == "implementation ref 2"); + } + + SECTION("Builder with no setters succeeds") { + auto r1 = metatomic::ModelMetadata::References::builder().build(); + + CHECK(r1.model().empty()); + CHECK(r1.architecture().empty()); + CHECK(r1.implementation().empty()); + } + + SECTION("Builder accumulates references") { + auto r1 = metatomic::ModelMetadata::References::builder() + .add_model("model ref 1") + .add_model("model ref 2") + .add_architecture("architecture ref 1") + .add_implementation("implementation ref 1") + .add_implementation("implementation ref 2") + .build(); + + CHECK(r1.model().size() == 2); + CHECK(r1.model()[0] == "model ref 1"); + CHECK(r1.model()[1] == "model ref 2"); + CHECK(r1.architecture().size() == 1); + CHECK(r1.architecture()[0] == "architecture ref 1"); + CHECK(r1.implementation().size() == 2); + CHECK(r1.implementation()[1] == "implementation ref 2"); + } + } + + SECTION("ModelMetadata") { + auto create_example = []() { + return metatomic::ModelMetadata::builder() + .name("test-model") + .authors({"Alice", "Bob"}) + .description("A test model") + .references(metatomic::ModelMetadata::References::builder() + .model({"doi:10.1234/test"}) + .architecture({"doi:10.1234/arch"}) + .implementation({"https://github.com/test"}) + .build()) + .extra(std::map{ + {"key1", "value1"}, + {"key2", "value2"} + }) + .build(); + }; + + auto create_example_with_setters = []() { + return metatomic::ModelMetadata::builder() + .name("test-model") + .add_author("Alice") + .add_author("Bob") + .description("A test model") + .add_reference("model", "doi:10.1234/test") + .add_reference("architecture", "doi:10.1234/arch") + .add_reference("implementation", "https://github.com/test") + .add_extra("key1", "value1") + .add_extra("key2", "value2") + .build(); + }; + + SECTION("JSON roundtrip conversion with builder") { + auto m1 = create_example(); + nlohmann::json j = m1; + + CHECK(j["type"] == "metatomic_model_metadata"); + CHECK(j["name"] == "test-model"); + CHECK(j["authors"].is_array()); + CHECK(j["authors"].size() == 2); + CHECK(j["authors"][0] == "Alice"); + CHECK(j["authors"][1] == "Bob"); + CHECK(j["description"] == "A test model"); + CHECK(j["references"]["model"][0] == "doi:10.1234/test"); + CHECK(j["references"]["architecture"][0] == "doi:10.1234/arch"); + CHECK(j["references"]["implementation"][0] == "https://github.com/test"); + CHECK(j["extra"]["key1"] == "value1"); + CHECK(j["extra"]["key2"] == "value2"); + + auto m2 = j.get(); + CHECK(m2.name() == m1.name()); + CHECK(m2.authors() == m1.authors()); + CHECK(m2.description() == m1.description()); + CHECK(m2.references().model() == m1.references().model()); + CHECK(m2.references().architecture() == m1.references().architecture()); + CHECK(m2.references().implementation() == m1.references().implementation()); + CHECK(m2.extra() == m1.extra()); + } + + SECTION("JSON roundtrip conversion with builder (accumulated fields)") { + auto m1 = create_example_with_setters(); + nlohmann::json j = m1; + + CHECK(j["type"] == "metatomic_model_metadata"); + CHECK(j["name"] == "test-model"); + CHECK(j["authors"].is_array()); + CHECK(j["authors"].size() == 2); + CHECK(j["authors"][0] == "Alice"); + CHECK(j["authors"][1] == "Bob"); + CHECK(j["description"] == "A test model"); + CHECK(j["references"]["model"][0] == "doi:10.1234/test"); + CHECK(j["references"]["architecture"][0] == "doi:10.1234/arch"); + CHECK(j["references"]["implementation"][0] == "https://github.com/test"); + CHECK(j["extra"]["key1"] == "value1"); + CHECK(j["extra"]["key2"] == "value2"); + + auto m2 = j.get(); + CHECK(m2.name() == m1.name()); + CHECK(m2.authors() == m1.authors()); + CHECK(m2.description() == m1.description()); + CHECK(m2.references().model() == m1.references().model()); + CHECK(m2.references().architecture() == m1.references().architecture()); + CHECK(m2.references().implementation() == m1.references().implementation()); + CHECK(m2.extra() == m1.extra()); + } + + SECTION("Invalid JSON data") { + auto m1 = create_example(); + nlohmann::json j = m1; + + CHECK_THROWS_WITH( + nlohmann::json("not an object").get(), + Catch::Matchers::StartsWith("invalid JSON data for ModelMetadata, expected an object") + ); + + { + auto j_copy = j; + j_copy["type"] = "something-else"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'type' in JSON for ModelMetadata must be 'metatomic_model_metadata'") + ); + } + + { + auto j_copy = j; + j_copy.erase("name"); + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'name' in JSON for ModelMetadata must be a string") + ); + } + + { + auto j_copy = j; + j_copy["name"] = 42; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'name' in JSON for ModelMetadata must be a string") + ); + } + + { + auto j_copy = j; + j_copy["authors"] = "Alice"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'authors' in JSON for ModelMetadata must be an array") + ); + } + + { + auto j_copy = j; + j_copy["authors"] = {"Alice", 42}; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'authors' in JSON for ModelMetadata must be an array of strings") + ); + } + + { + auto j_copy = j; + j_copy.erase("description"); + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'description' in JSON for ModelMetadata must be a string") + ); + } + + { + auto j_copy = j; + j_copy["extra"] = "not-an-object"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'extra' in JSON for ModelMetadata must be an object") + ); + } + + { + auto j_copy = j; + j_copy["extra"] = {{"key", 42}}; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'extra' in JSON for ModelMetadata must be an object with string values") + ); + } + + { + auto j_copy = j; + j_copy["references"] = "not-an-object"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("invalid JSON data for references in ModelMetadata, expected an object") + ); + } + } + + SECTION("Model metadata formatting") { + auto m1 = create_example(); + std::string output = m1.print(); + std::string expected = + "This is the test-model model\n" + "============================\n" + "\n" + "A test model\n" + "\n" + "Model authors\n" + "-------------\n" + "\n" + "- Alice\n" + "- Bob\n" + "\n" + "Model references\n" + "----------------\n" + "\n" + "Please cite the following references when using this model:\n" + "- about this specific model:\n" + " * doi:10.1234/test\n" + "- about the architecture of this model:\n" + " * doi:10.1234/arch\n" + "- about the implementation of this model:\n" + " * https://github.com/test\n"; + + CHECK(output == expected); + } + + SECTION("Builder with no setters succeeds") { + auto m1 = metatomic::ModelMetadata::builder().build(); + + CHECK(m1.name().empty()); + CHECK(m1.authors().empty()); + CHECK(m1.description().empty()); + CHECK(m1.references().model().empty()); + CHECK(m1.references().architecture().empty()); + CHECK(m1.references().implementation().empty()); + CHECK(m1.extra().empty()); + } + + SECTION("add_reference validates section") { + CHECK_THROWS_WITH( + metatomic::ModelMetadata::builder().add_reference("invalid", "ref"), + Catch::Matchers::StartsWith("reference section must be 'model', 'architecture', or 'implementation', got 'invalid'") + ); + } + } + + SECTION("DType") { + SECTION("JSON roundtrip conversion") { + auto dtype1 = metatomic::ModelCapabilities::DType::Float32; + nlohmann::json j = dtype1; + CHECK(j == "float32"); + auto dtype2 = j.get(); + CHECK(dtype2 == metatomic::ModelCapabilities::DType::Float32); + + auto dtype3 = metatomic::ModelCapabilities::DType::Float64; + nlohmann::json j2 = dtype3; + CHECK(j2 == "float64"); + auto dtype4 = j2.get(); + CHECK(dtype4 == metatomic::ModelCapabilities::DType::Float64); + } + + SECTION("Invalid JSON data") { + CHECK_THROWS_WITH( + nlohmann::json(42).get(), + Catch::Matchers::StartsWith("dtype in JSON for ModelCapabilities must be a string") + ); + + CHECK_THROWS_WITH( + nlohmann::json("float16").get(), + Catch::Matchers::StartsWith("invalid string for dtype in JSON for ModelCapabilities, expected 'float32' or 'float64'") + ); + } + } + + SECTION("Quantity") { + SECTION("JSON roundtrip conversion with description") { + auto q1 = metatomic::Quantity::builder() + .name("energy") + .unit("eV") + .sample_kind(metatomic::SampleKind::System) + .description("total energy of the system") + .gradients({metatomic::Gradients::Positions}) + .build(); + + nlohmann::json j = q1; + + CHECK(j["type"] == "metatomic_quantity"); + CHECK(j["name"] == "energy"); + CHECK(j["unit"] == "eV"); + CHECK(j["description"] == "total energy of the system"); + CHECK(j["gradients"].is_array()); + CHECK(j["gradients"].size() == 1); + CHECK(j["gradients"][0] == "positions"); + CHECK(j["sample_kind"] == "system"); + + auto q2 = j.get(); + CHECK(q2.name() == q1.name()); + CHECK(q2.unit() == q1.unit()); + CHECK(q2.description() == q1.description()); + CHECK(q2.gradients() == q1.gradients()); + CHECK(q2.sample_kind() == q1.sample_kind()); + } + + SECTION("JSON roundtrip conversion without description") { + auto q1 = metatomic::Quantity::builder() + .name("charge") + .unit("e") + .sample_kind(metatomic::SampleKind::Atom) + .build(); + + nlohmann::json j = q1; + + CHECK(j["type"] == "metatomic_quantity"); + CHECK(j["name"] == "charge"); + CHECK(j["unit"] == "e"); + CHECK(!j.contains("description")); + CHECK(j["gradients"].is_array()); + CHECK(j["gradients"].size() == 0); + CHECK(j["sample_kind"] == "atom"); + + auto q2 = j.get(); + CHECK(q2.name() == q1.name()); + CHECK(q2.unit() == q1.unit()); + CHECK(q2.description().empty()); + CHECK(q2.gradients().empty()); + CHECK(q2.sample_kind() == q1.sample_kind()); + } + + SECTION("build() validates completeness") { + CHECK_THROWS_WITH( + metatomic::Quantity::builder().build(), + Catch::Matchers::StartsWith("name must be set before building Quantity") + ); + + CHECK_THROWS_WITH( + metatomic::Quantity::builder().name("energy").build(), + Catch::Matchers::StartsWith("unit must be set before building Quantity") + ); + + CHECK_THROWS_WITH( + metatomic::Quantity::builder().name("energy").unit("eV").build(), + Catch::Matchers::StartsWith("sample_kind must be set before building Quantity") + ); + } + + SECTION("add gradients") { + auto q1 = metatomic::Quantity::builder() + .name("energy") + .unit("eV") + .sample_kind(metatomic::SampleKind::System) + .add_gradient(metatomic::Gradients::Positions) + .add_gradient(metatomic::Gradients::Strain) + .build(); + + CHECK(q1.gradients().size() == 2); + CHECK(q1.gradients()[0] == metatomic::Gradients::Positions); + CHECK(q1.gradients()[1] == metatomic::Gradients::Strain); + + nlohmann::json j = q1; + CHECK(j["gradients"].size() == 2); + CHECK(j["gradients"][0] == "positions"); + CHECK(j["gradients"][1] == "strain"); + } + + SECTION("Empty description is treated as no description") { + nlohmann::json j = { + {"type", "metatomic_quantity"}, + {"name", "charge"}, + {"unit", "e"}, + {"description", ""}, + {"gradients", nlohmann::json::array()}, + {"sample_kind", "atom"} + }; + + auto q = j.get(); + CHECK(q.name() == "charge"); + CHECK(q.unit() == "e"); + CHECK(q.description().empty()); + CHECK(q.gradients().empty()); + CHECK(q.sample_kind() == metatomic::SampleKind::Atom); + } + + SECTION("Invalid JSON data") { + CHECK_THROWS_WITH( + nlohmann::json("not an object").get(), + Catch::Matchers::StartsWith("invalid JSON data for Quantity, expected an object") + ); + + { + nlohmann::json j = {{"type", "wrong-type"}}; + CHECK_THROWS_WITH( + j.get(), + Catch::Matchers::StartsWith("'type' in JSON for Quantity must be 'metatomic_quantity'") + ); + } + + { + nlohmann::json j = { + {"type", "metatomic_quantity"}, + {"name", 42} + }; + CHECK_THROWS_WITH( + j.get(), + Catch::Matchers::StartsWith("'name' in JSON for Quantity must be a string") + ); + } + + { + nlohmann::json j = { + {"type", "metatomic_quantity"}, + {"name", "energy"}, + {"unit", "eV"}, + {"gradients", "positions"} + }; + CHECK_THROWS_WITH( + j.get(), + Catch::Matchers::StartsWith("'gradients' in JSON for Quantity must be an array") + ); + } + + { + nlohmann::json j = { + {"type", "metatomic_quantity"}, + {"name", "energy"}, + {"unit", "eV"}, + {"gradients", {"positions"}}, + {"sample_kind", "unknown"} + }; + CHECK_THROWS_WITH( + j.get(), + Catch::Matchers::StartsWith("'sample_kind' in JSON for Quantity must be 'atom', 'system' or 'atom_pair', got 'unknown'") + ); + } + } + } + + SECTION("ModelCapabilities") { + auto make_quantity = [](const std::string& name, const std::string& unit, + metatomic::SampleKind sample_kind, + const std::string& description = "", + std::vector gradients = {}) { + return metatomic::Quantity::builder() + .name(name) + .unit(unit) + .sample_kind(sample_kind) + .description(description) + .gradients(std::move(gradients)) + .build(); + }; + + auto create_example = [&]() { + std::vector outputs = { + make_quantity("energy", "eV", metatomic::SampleKind::System, + "total energy", {metatomic::Gradients::Positions}), + make_quantity("charge", "e", metatomic::SampleKind::Atom), + }; + + return metatomic::ModelCapabilities::builder() + .atomic_types({1, 6, 8}) + .interaction_range(5.0) + .length_unit("Angstrom") + .supported_devices({metatomic::ModelCapabilities::Device::CPU, + metatomic::ModelCapabilities::Device::CUDA}) + .dtype(metatomic::ModelCapabilities::DType::Float32) + .outputs(std::move(outputs)) + .build(); + }; + + auto create_example_with_setters = [&]() { + std::vector outputs = { + make_quantity("energy", "eV", metatomic::SampleKind::System, + "total energy", {metatomic::Gradients::Positions}), + make_quantity("charge", "e", metatomic::SampleKind::Atom), + }; + + return metatomic::ModelCapabilities::builder() + .atomic_types({1, 6, 8}) + .interaction_range(5.0) + .length_unit("Angstrom") + .supported_devices({metatomic::ModelCapabilities::Device::CPU, + metatomic::ModelCapabilities::Device::CUDA}) + .dtype(metatomic::ModelCapabilities::DType::Float32) + .outputs(std::move(outputs)) + .build(); + }; + + SECTION("JSON roundtrip conversion with builder") { + auto c1 = create_example(); + nlohmann::json j = c1; + + CHECK(j["type"] == "metatomic_model_capabilities"); + CHECK(j["outputs"].is_array()); + CHECK(j["outputs"].size() == 2); + CHECK(j["outputs"][0]["name"] == "energy"); + CHECK(j["outputs"][1]["name"] == "charge"); + CHECK(j["atomic_types"].is_array()); + CHECK(j["atomic_types"].size() == 3); + CHECK(j["atomic_types"][0] == 1); + CHECK(j["atomic_types"][1] == 6); + CHECK(j["atomic_types"][2] == 8); + CHECK(j["interaction_range"] == Approx(5.0)); + CHECK(j["length_unit"] == "Angstrom"); + CHECK(j["supported_devices"].is_array()); + CHECK(j["supported_devices"].size() == 2); + CHECK(j["supported_devices"][0] == "cpu"); + CHECK(j["supported_devices"][1] == "cuda"); + CHECK(j["dtype"] == "float32"); + + auto c2 = j.get(); + CHECK(c2.outputs().size() == c1.outputs().size()); + CHECK(c2.outputs()[0].name() == c1.outputs()[0].name()); + CHECK(c2.outputs()[1].name() == c1.outputs()[1].name()); + CHECK(c2.atomic_types() == c1.atomic_types()); + CHECK(c2.interaction_range() == Approx(c1.interaction_range())); + CHECK(c2.length_unit() == c1.length_unit()); + CHECK(c2.supported_devices() == c1.supported_devices()); + CHECK(c2.dtype() == c1.dtype()); + } + + SECTION("JSON roundtrip conversion with builder (accumulated fields)") { + auto c1 = create_example_with_setters(); + nlohmann::json j = c1; + + CHECK(j["type"] == "metatomic_model_capabilities"); + CHECK(j["outputs"].is_array()); + CHECK(j["outputs"].size() == 2); + CHECK(j["outputs"][0]["name"] == "energy"); + CHECK(j["outputs"][1]["name"] == "charge"); + CHECK(j["atomic_types"].is_array()); + CHECK(j["atomic_types"].size() == 3); + CHECK(j["atomic_types"][0] == 1); + CHECK(j["atomic_types"][1] == 6); + CHECK(j["atomic_types"][2] == 8); + CHECK(j["interaction_range"] == Approx(5.0)); + CHECK(j["length_unit"] == "Angstrom"); + CHECK(j["supported_devices"].is_array()); + CHECK(j["supported_devices"].size() == 2); + CHECK(j["supported_devices"][0] == "cpu"); + CHECK(j["supported_devices"][1] == "cuda"); + CHECK(j["dtype"] == "float32"); + + auto c2 = j.get(); + CHECK(c2.outputs().size() == c1.outputs().size()); + CHECK(c2.outputs()[0].name() == c1.outputs()[0].name()); + CHECK(c2.outputs()[1].name() == c1.outputs()[1].name()); + CHECK(c2.atomic_types() == c1.atomic_types()); + CHECK(c2.interaction_range() == Approx(c1.interaction_range())); + CHECK(c2.length_unit() == c1.length_unit()); + CHECK(c2.supported_devices() == c1.supported_devices()); + CHECK(c2.dtype() == c1.dtype()); + } + + SECTION("build() validates completeness") { + CHECK_THROWS_WITH( + metatomic::ModelCapabilities::builder().build(), + Catch::Matchers::StartsWith("atomic_types must be set before building ModelCapabilities") + ); + + CHECK_THROWS_WITH( + metatomic::ModelCapabilities::builder().atomic_types({1}).build(), + Catch::Matchers::StartsWith("interaction_range must be set before building ModelCapabilities") + ); + + CHECK_THROWS_WITH( + metatomic::ModelCapabilities::builder() + .atomic_types({1}).interaction_range(5.0).build(), + Catch::Matchers::StartsWith("length_unit must be set before building ModelCapabilities") + ); + + CHECK_THROWS_WITH( + metatomic::ModelCapabilities::builder() + .atomic_types({1}).interaction_range(5.0).length_unit("Angstrom").build(), + Catch::Matchers::StartsWith("supported_devices must be set before building ModelCapabilities") + ); + + CHECK_THROWS_WITH( + metatomic::ModelCapabilities::builder() + .atomic_types({1}).interaction_range(5.0).length_unit("Angstrom") + .supported_devices({metatomic::ModelCapabilities::Device::CPU}).build(), + Catch::Matchers::StartsWith("dtype must be set before building ModelCapabilities") + ); + + CHECK_THROWS_WITH( + metatomic::ModelCapabilities::builder().interaction_range(-1.0), + Catch::Matchers::StartsWith("interaction_range must be non-negative") + ); + } + + SECTION("add outputs, atomic types, and supported devices") { + auto c1 = metatomic::ModelCapabilities::builder() + .interaction_range(5.0) + .length_unit("Angstrom") + .dtype(metatomic::ModelCapabilities::DType::Float32) + .add_output(make_quantity("energy", "eV", metatomic::SampleKind::System, + "total energy", {metatomic::Gradients::Positions})) + .add_output(make_quantity("charge", "e", metatomic::SampleKind::Atom)) + .add_atomic_type(1) + .add_atomic_type(6) + .add_atomic_type(8) + .add_supported_device(metatomic::ModelCapabilities::Device::CPU) + .add_supported_device(metatomic::ModelCapabilities::Device::CUDA) + .build(); + + CHECK(c1.outputs().size() == 2); + CHECK(c1.outputs()[0].name() == "energy"); + CHECK(c1.outputs()[1].name() == "charge"); + CHECK(c1.atomic_types().size() == 3); + CHECK(c1.atomic_types()[0] == 1); + CHECK(c1.atomic_types()[1] == 6); + CHECK(c1.atomic_types()[2] == 8); + CHECK(c1.supported_devices().size() == 2); + CHECK(c1.supported_devices()[0] == metatomic::ModelCapabilities::Device::CPU); + CHECK(c1.supported_devices()[1] == metatomic::ModelCapabilities::Device::CUDA); + + nlohmann::json j = c1; + CHECK(j["outputs"].size() == 2); + CHECK(j["atomic_types"].size() == 3); + CHECK(j["supported_devices"].size() == 2); + } + + SECTION("Invalid JSON data") { + auto c1 = create_example(); + nlohmann::json j = c1; + + CHECK_THROWS_WITH( + nlohmann::json("not an object").get(), + Catch::Matchers::StartsWith("invalid JSON data for ModelCapabilities, expected an object") + ); + + { + auto j_copy = j; + j_copy["type"] = "something-else"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'type' in JSON for ModelCapabilities must be 'metatomic_model_capabilities'") + ); + } + + { + auto j_copy = j; + j_copy["outputs"] = "energy"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'outputs' in JSON for ModelCapabilities must be an array") + ); + } + + { + auto j_copy = j; + j_copy["atomic_types"] = "1"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'atomic_types' in JSON for ModelCapabilities must be an array") + ); + } + + { + auto j_copy = j; + j_copy["atomic_types"] = {1, "x"}; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'atomic_types' in JSON for ModelCapabilities must be an array of integers") + ); + } + + { + auto j_copy = j; + j_copy.erase("interaction_range"); + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'interaction_range' in JSON for ModelCapabilities must be a number") + ); + } + + { + auto j_copy = j; + j_copy["interaction_range"] = -1.0; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'interaction_range' in JSON for ModelCapabilities must be non-negative") + ); + } + + { + auto j_copy = j; + j_copy["length_unit"] = "eV"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("invalid parameter: dimension mismatch") + ); + } + + { + auto j_copy = j; + j_copy["supported_devices"] = "cpu"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("'supported_devices' in JSON for ModelCapabilities must be an array") + ); + } + + { + auto j_copy = j; + j_copy["supported_devices"] = {"cpu", "wat"}; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("invalid string for device in JSON for ModelCapabilities, expected 'cpu', 'cuda', 'rocm', or 'metal'") + ); + } + + { + auto j_copy = j; + j_copy["dtype"] = "float16"; + CHECK_THROWS_WITH( + j_copy.get(), + Catch::Matchers::StartsWith("invalid string for dtype in JSON for ModelCapabilities, expected 'float32' or 'float64'") + ); + } + } + } +} diff --git a/metatomic-core/tests/cxx/misc.cpp b/metatomic-core/tests/cxx/misc.cpp new file mode 100644 index 000000000..ed08aaf40 --- /dev/null +++ b/metatomic-core/tests/cxx/misc.cpp @@ -0,0 +1,64 @@ +#include + +#include "metatomic.hpp" + + +TEST_CASE("unit conversion factor") { + // same unit -> factor = 1.0 + auto factor = metatomic::unit_conversion_factor("m", "m"); + CHECK(factor == 1.0); + + // kJ/mol -> eV + factor = metatomic::unit_conversion_factor("kJ/mol", "eV"); + CHECK(factor == Approx(0.010364269656262174).epsilon(1e-15)); + + REQUIRE_THROWS_WITH( + metatomic::unit_conversion_factor("m", "kg"), + "invalid parameter: dimension mismatch in unit conversion: " + "'m' has dimension [L] but 'kg' has dimension [M]" + ); +} + + +TEST_CASE("metatdata formatting") { + std::string json =R"({ + "type": "metatomic_model_metadata", + "name": "name", + "description": "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation.", + "authors": ["Short author", "Some extremely long author that will take more than one line in the printed output"], + "references": { + "architecture": ["ref-2", "ref-3"], + "model": ["a very long reference that will take more than one line in the printed output"], + "implementation": [] + }, + "extra": {} +})"; + + const auto* expected = R"(This is the name model +====================== + +Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor +incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis +nostrud exercitation. + +Model authors +------------- + +- Short author +- Some extremely long author that will take more than one line in the printed + output + +Model references +---------------- + +Please cite the following references when using this model: +- about this specific model: + * a very long reference that will take more than one line in the printed + output +- about the architecture of this model: + * ref-2 + * ref-3 +)"; + + CHECK(metatomic::format_metadata(json) == expected); +} diff --git a/metatomic-core/tests/cxx/model.cpp b/metatomic-core/tests/cxx/model.cpp new file mode 100644 index 000000000..971f078db --- /dev/null +++ b/metatomic-core/tests/cxx/model.cpp @@ -0,0 +1,379 @@ +#include +#include +#include +#include +#include + +#include +#include + +#include "metatomic.hpp" +#include "helpers.hpp" + + +class SimpleCppModel: public metatomic::BaseModel { +public: + explicit SimpleCppModel(double scale): scale_(scale) {} + + metatomic::ModelCapabilities capabilities() const override final { + return metatomic::ModelCapabilities::builder() + .atomic_types({1, 4, 7, 10}) + .interaction_range(4.5) + .length_unit("nm") + .supported_devices({metatomic::ModelCapabilities::Device::CPU}) + .dtype(metatomic::ModelCapabilities::DType::Float32) + .add_output(metatomic::Quantity::builder() + .name("energy") + .unit("eV") + .sample_kind(metatomic::SampleKind::System) + .build()) + .add_output(metatomic::Quantity::builder() + .name("custom::output") + .unit("eV") + .sample_kind(metatomic::SampleKind::Atom) + .build()) + .build(); + } + + metatomic::ModelMetadata metadata() const override final { + return metatomic::ModelMetadata::builder() + .name("simple C++ model") + .description("test model for BaseModel") + .build(); + } + + std::vector requested_pair_lists() const final { + return {}; + } + + std::vector requested_inputs() const final { + return {}; + } + + std::vector execute_inner( + const std::vector& systems, + const metatensor::Labels* selected_atoms, + const std::vector& requested_outputs + ) final { + std::vector outputs; + outputs.reserve(requested_outputs.size()); + + size_t atom_count = 0; + if (selected_atoms != nullptr) { + atom_count = selected_atoms->count(); + } else { + for (const auto& system: systems) { + atom_count += system.size(); + } + } + + for (const auto& output: requested_outputs) { + if (output.name() != "energy") { + throw metatomic::Error("unknown output: " + output.name()); + } + + double energy = scale_ * static_cast(atom_count); + outputs.push_back(scalar_tensor(energy, output.name())); + } + + return outputs; + } + +private: + double scale_; +}; + +TEST_CASE("BaseModel") { + auto model = std::make_unique(2.5); + + auto capabilities = model->capabilities(); + CHECK(capabilities.atomic_types().size() == 4); + + const auto& outputs = capabilities.outputs(); + CHECK(outputs.size() == 2); + CHECK(outputs[0].name() == "energy"); + CHECK(outputs[1].name() == "custom::output"); + CHECK(outputs[1].sample_kind() == metatomic::SampleKind::Atom); + + auto system = test_system(4); + auto systems = std::vector(); + systems.push_back(std::move(system)); + + // NOTE: we call execute_inner directly only for testing + // in practice, the model should be executed through the `mta_execute_model` function + auto requested_outputs = std::vector{outputs[0]}; + auto results = model->execute_inner(systems, nullptr, requested_outputs); + + REQUIRE(results.size() == 1); + CHECK(results[0].keys().count() == 1); + + auto block = results[0].block_by_id(0); + auto values = block.values(); + REQUIRE(values.data() != nullptr); + CHECK(values.data()[0] == Approx(10.0)); +} + + +TEST_CASE("Wrap mta_model_t with ExternalModel") { + auto raw_model = metatomic::BaseModel::to_mta_model( + std::make_unique(3.0) + ); + auto model = metatomic::ExternalModel(raw_model); + + auto outputs = model.capabilities().outputs(); + CHECK(outputs.size() == 2); + CHECK(outputs[0].name() == "energy"); + CHECK(outputs[1].name() == "custom::output"); + + auto system = test_system(4); + std::vector systems; + systems.push_back(std::move(system)); + + // Request only "energy" output + // The "custom::output" errors out + auto out = metatomic::execute_model( + model, systems, std::nullopt, std::vector{outputs[0]}, false + ); + REQUIRE(out.size() == 1); + + // 4 atoms * scale 3.0 = 12.0 + auto block = out[0].block_by_id(0); + auto values = block.values(); + REQUIRE(values.data() != nullptr); + CHECK(values.data()[0] == Approx(12.0)); +} + + +TEST_CASE("ExternalModel move semantics") { + auto raw = metatomic::BaseModel::to_mta_model( + std::make_unique(1.0) + ); + auto model = metatomic::ExternalModel(raw); + CHECK(model.as_mta_model_t() != nullptr); + + auto moved = std::move(model); + CHECK(moved.as_mta_model_t() != nullptr); + CHECK(moved.capabilities().outputs().size() == 2); +} + + +TEST_CASE("ExternalModel release transfers ownership") { + auto raw = metatomic::BaseModel::to_mta_model( + std::make_unique(2.0) + ); + auto model = metatomic::ExternalModel(raw); + + // release the raw model back to the caller; the ExternalModel is empty + // and will not call unload on destruction + auto released = model.release(); + CHECK(released.unload != nullptr); + + // re-wrap the released model to verify it is still valid + auto wrapped = metatomic::ExternalModel(released); + + auto outputs = wrapped.capabilities().outputs(); + CHECK(outputs.size() == 2); + CHECK(outputs[0].name() == "energy"); + CHECK(outputs[1].name() == "custom::output"); + + auto system = test_system(4); + std::vector systems; + systems.push_back(std::move(system)); + + // Request only "energy" output + // The "custom::output" errors out + auto out = metatomic::execute_model( + wrapped, systems, std::nullopt, std::vector{outputs[0]}, false + ); + REQUIRE(out.size() == 1); + + // 4 atoms * scale 2.0 = 8.0 + auto block = out[0].block_by_id(0); + auto values = block.values(); + REQUIRE(values.data() != nullptr); + CHECK(values.data()[0] == Approx(8.0)); +} + + +TEST_CASE("to_mta_model for ExternalModel") { + auto inner_raw = metatomic::BaseModel::to_mta_model( + std::make_unique(3.0) + ); + void* inner_data = inner_raw.data; + auto external = std::make_unique(inner_raw); + auto outer_raw = metatomic::BaseModel::to_mta_model(std::move(external)); + + // `to_mta_model` short-circuits for `ExternalModel` + // The raw model's data pointer should be the same as the inner model's data pointer. + CHECK(outer_raw.data == inner_data); + + // The raw model's callbacks must all be set by `to_mta_model`. + CHECK(outer_raw.capabilities != nullptr); + CHECK(outer_raw.metadata != nullptr); + CHECK(outer_raw.requested_pair_lists != nullptr); + CHECK(outer_raw.requested_inputs != nullptr); + CHECK(outer_raw.execute_inner != nullptr); + CHECK(outer_raw.unload != nullptr); + + // Wrap the raw model back in an ExternalModel to test through the C++ interface. + auto model = metatomic::ExternalModel(outer_raw); + + auto capabilities = model.capabilities(); + CHECK(capabilities.length_unit() == "nm"); + + auto metadata = model.metadata(); + CHECK(metadata.name() == "simple C++ model"); + + const auto& outputs = capabilities.outputs(); + REQUIRE(outputs.size() == 2); + CHECK(outputs[0].name() == "energy"); + CHECK(outputs[1].name() == "custom::output"); + CHECK(outputs[1].sample_kind() == metatomic::SampleKind::Atom); + + CHECK(model.requested_pair_lists().empty()); + CHECK(model.requested_inputs().empty()); +} + + +TEST_CASE("execute_model with a BaseModel") { + auto model = SimpleCppModel(2.5); + + auto outputs = model.capabilities().outputs(); + auto requested_outputs = std::vector{outputs[0]}; + + std::vector systems; + systems.push_back(test_system(4)); + + auto out = metatomic::execute_model( + model, systems, std::nullopt, requested_outputs, false + ); + REQUIRE(out.size() == 1); + + // 4 atoms * scale 2.5 = 10.0 + auto values = out[0].block_by_id(0).values(); + REQUIRE(values.data() != nullptr); + CHECK(values.data()[0] == Approx(10.0)); + + // The `model` should still be valid after execution of `execute_model` + CHECK(model.capabilities().length_unit() == "nm"); + + // executing again should give the same result + auto again = metatomic::execute_model( + model, systems, std::nullopt, requested_outputs, false + ); + REQUIRE(again.size() == 1); + + auto again_values = again[0].block_by_id(0).values(); + REQUIRE(again_values.data() != nullptr); + CHECK(again_values.data()[0] == Approx(10.0)); +} + + +TEST_CASE("execute_model with an ExternalModel") { + auto model = metatomic::ExternalModel(metatomic::BaseModel::to_mta_model( + std::make_unique(3.0) + )); + + auto outputs = model.capabilities().outputs(); + auto requested_outputs = std::vector{outputs[0]}; + + std::vector systems; + systems.push_back(test_system(4)); + + // execute model twice to make sure the model remains valid + for (int i = 0; i < 2; i++) { + auto out = metatomic::execute_model( + model, systems, std::nullopt, requested_outputs, false + ); + REQUIRE(out.size() == 1); + + // 4 atoms * scale 3.0 = 12.0 + auto values = out[0].block_by_id(0).values(); + REQUIRE(values.data() != nullptr); + CHECK(values.data()[0] == Approx(12.0)); + } + + // mta_model_t is still owned by the ExternalModel + CHECK(model.as_mta_model_t()->unload != nullptr); + CHECK(model.metadata().name() == "simple C++ model"); +} + + +TEST_CASE("mta_model_view does not take ownership") { + auto model = SimpleCppModel(1.0); + + auto model_view = metatomic::BaseModel::mta_model_view(model); + CHECK(model_view.data == static_cast(&model)); + CHECK(model_view.unload == nullptr); + + CHECK(model_view.capabilities != nullptr); + CHECK(model_view.metadata != nullptr); + CHECK(model_view.requested_pair_lists != nullptr); + CHECK(model_view.requested_inputs != nullptr); + CHECK(model_view.execute_inner != nullptr); + + // borrowing an ExternalModel gives back its own callbacks, without `unload` + auto external = metatomic::ExternalModel(metatomic::BaseModel::to_mta_model( + std::make_unique(1.0) + )); + auto* raw = external.as_mta_model_t(); + + auto model_view_external = metatomic::BaseModel::mta_model_view(external); + CHECK(model_view_external.data == raw->data); + CHECK(model_view_external.execute_inner == raw->execute_inner); + CHECK(model_view_external.unload == nullptr); + CHECK(raw->unload != nullptr); +} + +class ThrowingModel: public metatomic::BaseModel { +public: + [[noreturn]] metatomic::ModelCapabilities capabilities() const final { + throw std::out_of_range("ThrowingCppModel: intentional failure in capabilities"); + } + + metatomic::ModelMetadata metadata() const final { + return metatomic::ModelMetadata::builder() + .name("throwing C++ model") + .build(); + } + + std::vector requested_pair_lists() const final { + return {}; + } + + std::vector requested_inputs() const final { + return {}; + } + + std::vector execute_inner( + const std::vector&, + const metatensor::Labels*, + const std::vector& + ) final { + return {}; + } +}; + + +TEST_CASE("C++ exception") { + auto raw = metatomic::BaseModel::to_mta_model( + std::make_unique() + ); + + // An exception thrown by a C++ model is reported as `MTA_MODEL_ERROR` + mta_string_t capabilities_json = nullptr; + auto status = raw.capabilities(raw.data, &capabilities_json); + CHECK(status == MTA_MODEL_ERROR); + CHECK(status != MTA_UNSUPPORTED_MODEL_ERROR); + CHECK(capabilities_json == nullptr); + + const char* message = nullptr; + const char* origin = nullptr; + mta_last_error(&message, &origin, nullptr); + CHECK(std::string(origin) == "C++ exception"); + CHECK(std::string(message) == "ThrowingCppModel: intentional failure in capabilities"); + + // Going back through the C++ API rethrows the original exception + auto model = metatomic::ExternalModel(raw); + CHECK_THROWS_AS(model.capabilities(), std::out_of_range); +} diff --git a/metatomic-core/tests/cxx/plugins.cpp b/metatomic-core/tests/cxx/plugins.cpp new file mode 100644 index 000000000..fe67ade2e --- /dev/null +++ b/metatomic-core/tests/cxx/plugins.cpp @@ -0,0 +1,70 @@ +#include + +#include + +#include "metatomic.hpp" + + +TEST_CASE("Load C plugins") { + metatomic::load_plugin(PLUGIN_DIR "/test-c-plugin.so"); + + REQUIRE_THROWS_WITH( + metatomic::load_model("some_model", "{}", "test-c-plugin"), + "invalid parameter: failed to load model from 'some_model': plugin 'test-c-plugin' could not load the model" + ); + + REQUIRE_THROWS_WITH( + metatomic::load_model("some_model"), + "invalid parameter: failed to load model from 'some_model': tried the " + "following plugins, but none could load the model: test-c-plugin" + ); + + REQUIRE_THROWS_WITH( + metatomic::load_plugin(PLUGIN_DIR "/bad-abi-plugin.so"), + "invalid parameter: can not register plugin 'bad-abi-plugin': " + "plugin ABI version is 2, but metatomic expects 1" + ); +} + + +TEST_CASE("Load C++ plugins") { + metatomic::load_plugin(PLUGIN_DIR "/test-cxx-plugin.so"); + + auto model = metatomic::load_model("test-cxx-model", "{}", "test-cxx-plugin"); + + auto metadata = model.metadata(); + CHECK(metadata.name() == "simple C++ plugin model"); + + auto capabilities = model.capabilities(); + CHECK(capabilities.length_unit() == "nm"); + REQUIRE(capabilities.outputs().size() == 1); + CHECK(capabilities.outputs()[0].name() == "energy"); + + REQUIRE_THROWS_WITH( + metatomic::load_model("unknown", "{}", "test-cxx-plugin"), + "invalid parameter: failed to load model from 'unknown': plugin " + "'test-cxx-plugin' could not load the model" + ); + + // an exception thrown by the plugin is reported as `MTA_MODEL_ERROR`, + // the actual exception makes it back to the caller + REQUIRE_THROWS_WITH( + metatomic::load_model("throws", "{}", "test-cxx-plugin"), + "load_model_cxx: intentional failure for 'throws'" + ); + CHECK_THROWS_AS( + metatomic::load_model("throws", "{}", "test-cxx-plugin"), + metatomic::Error + ); + + CHECK_THROWS_AS( + metatomic::load_model("throws-std", "{}", "test-cxx-plugin"), + std::out_of_range + ); + + REQUIRE_THROWS_WITH( + metatomic::load_model("unknown"), + Catch::Contains("tried the following plugins, but none could load the model") + && Catch::Contains("test-cxx-plugin") + ); +} diff --git a/metatomic-core/tests/cxx/system.cpp b/metatomic-core/tests/cxx/system.cpp new file mode 100644 index 000000000..4e65c7385 --- /dev/null +++ b/metatomic-core/tests/cxx/system.cpp @@ -0,0 +1,295 @@ +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include "metatomic.hpp" +#include "helpers.hpp" + +static metatensor::TensorBlock pair_block() { + auto samples = metatensor::Labels( + {"first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"}, + {{0, 1, 0, 0, 0}} + ); + auto components = std::vector{ + metatensor::Labels({"xyz"}, {{0}, {1}, {2}}) + }; + auto properties = metatensor::Labels({"distance"}, {{0}}); + + auto values = std::make_unique>( + std::vector{1, 3, 1}, std::vector{1.5F, 2.5F, 3.5F} + ); + + return metatensor::TensorBlock(std::move(values), samples, components, properties); +} + +static metatensor::TensorMap custom_data() { + auto keys = metatensor::Labels({"key"}, {{0}}); + + auto samples = metatensor::Labels({"sample"}, {{0}}); + auto properties = metatensor::Labels({"property"}, {{0}}); + auto values = std::make_unique>( + std::vector{1, 1}, std::vector{42.0F} + ); + auto block = metatensor::TensorBlock(std::move(values), samples, {}, properties); + + auto blocks = std::vector(); + blocks.push_back(std::move(block)); + return metatensor::TensorMap(keys, std::move(blocks)); +} + +// Helper function to check that two DLPack tensors have the same shape, strides, dtype, and data. +template +static void check_tensors(const DLManagedTensorVersioned* expected, const DLManagedTensorVersioned* loaded) { + REQUIRE((expected != nullptr && loaded != nullptr)); + CHECK(loaded->dl_tensor.device.device_type == expected->dl_tensor.device.device_type); + CHECK(loaded->dl_tensor.device.device_id == expected->dl_tensor.device.device_id); + REQUIRE(loaded->dl_tensor.ndim == expected->dl_tensor.ndim); + CHECK(loaded->dl_tensor.dtype.code == expected->dl_tensor.dtype.code); + CHECK(loaded->dl_tensor.dtype.bits == expected->dl_tensor.dtype.bits); + CHECK(loaded->dl_tensor.dtype.lanes == expected->dl_tensor.dtype.lanes); + + for (int64_t i = 0; i < expected->dl_tensor.ndim; i++) { + CHECK(loaded->dl_tensor.shape[i] == expected->dl_tensor.shape[i]); + CHECK(loaded->dl_tensor.strides[i] == expected->dl_tensor.strides[i]); + } + + CHECK(metatensor::details::vector_from_dlpack(expected->dl_tensor) == metatensor::details::vector_from_dlpack(loaded->dl_tensor)); +} + +// Helper function to check that two Systems are equal +static void check_systems(const metatomic::System& system, const metatomic::System& loaded) { + CHECK(loaded.size() == system.size()); + CHECK(loaded.length_unit() == system.length_unit()); + + check_tensors(system.types().as_dlpack(), loaded.types().as_dlpack()); + check_tensors(system.positions().as_dlpack(), loaded.positions().as_dlpack()); + check_tensors(system.cell().as_dlpack(), loaded.cell().as_dlpack()); + check_tensors(system.pbc().as_dlpack(), loaded.pbc().as_dlpack()); +} + + +TEST_CASE("System basics") { + auto system = test_system(4); + + CHECK(system.size() == 4); + CHECK(system.length_unit() == "nm"); +} + +TEST_CASE("System construction errors") { + // wrong dtype for `types` (float instead of int32) + REQUIRE_THROWS_WITH( + metatomic::System( + "Angstrom", + types_tensor(3), + positions_tensor(3), + cell_tensor(), + pbc_tensor() + ), + "invalid parameter: `types` must be a tensor of 32-bit integers" + ); +} + +TEST_CASE("System data") { + auto system = test_system(4); + + SECTION("types") { + auto types = system.types(); + REQUIRE(static_cast(types)); + CHECK(types->dl_tensor.ndim == 1); + CHECK(types->dl_tensor.shape[0] == 4); + CHECK(types->dl_tensor.dtype.code == kDLInt); + CHECK(types->dl_tensor.dtype.bits == 32); + + auto* data = reinterpret_cast( + static_cast(types->dl_tensor.data) + types->dl_tensor.byte_offset + ); + CHECK(data[0] == 1); + CHECK(data[3] == 10); + } + + SECTION("positions") { + auto positions = system.positions(); + REQUIRE(static_cast(positions)); + CHECK(positions->dl_tensor.ndim == 2); + CHECK(positions->dl_tensor.shape[0] == 4); + CHECK(positions->dl_tensor.shape[1] == 3); + CHECK(positions->dl_tensor.dtype.code == kDLFloat); + + auto* data = reinterpret_cast( + static_cast(positions->dl_tensor.data) + positions->dl_tensor.byte_offset + ); + CHECK(data[0] == 1.0F); + CHECK(data[9] == 10.0F); + } + + SECTION("cell") { + auto cell = system.cell(); + REQUIRE(static_cast(cell)); + CHECK(cell->dl_tensor.ndim == 2); + CHECK(cell->dl_tensor.shape[0] == 3); + CHECK(cell->dl_tensor.shape[1] == 3); + } + + SECTION("pbc") { + auto pbc = system.pbc(); + REQUIRE(static_cast(pbc)); + CHECK(pbc->dl_tensor.ndim == 1); + CHECK(pbc->dl_tensor.shape[0] == 3); + CHECK(pbc->dl_tensor.dtype.code == kDLBool); + + auto* data = reinterpret_cast( + static_cast(pbc->dl_tensor.data) + pbc->dl_tensor.byte_offset + ); + CHECK(data[0] == true); + CHECK(data[1] == false); + CHECK(data[2] == true); + } +} + +TEST_CASE("System pairs") { + auto system = test_system(4); + + auto options = metatomic::PairListOptions::builder() + .cutoff(1.0) + .full_list(true) + .strict(false) + .add_requestor("test") + .build(); + + system.add_pairs(options, pair_block()); + + const auto* options_json = R"({ + "type": "metatomic_pair_list_options", + "cutoff": "0x40364ccccccccccd", + "full_list": false, + "strict": true, + "requestors": [""] + })"; + + system.add_pairs(options_json, pair_block()); + + auto pairs = system.pairs(options); + CHECK(pairs.samples().count() == 1); + CHECK(pairs.properties().size() == 1); + + auto known = system.known_pairs(); + CHECK(known.size() == 2); + CHECK(known[0].cutoff() == 1.0); + CHECK(known[0].full_list() == true); + CHECK(known[0].strict() == false); + CHECK(known[0].requestors().size() == 1); + CHECK(known[0].requestors()[0] == "test"); + + CHECK(known[1].cutoff() == 22.3); + CHECK(known[1].full_list() == false); + CHECK(known[1].strict() == true); + CHECK(known[1].requestors().size() == 0); +} + +TEST_CASE("System custom data") { + auto system = test_system(4); + + system.add_custom_data("test::my_data", custom_data()); + + auto data = system.custom_data("test::my_data"); + CHECK(data.keys().count() == 1); + + // retrieving unknown data throws + REQUIRE_THROWS(system.custom_data("test::no_such_data")); + + system.add_custom_data("test::other_data", custom_data()); + auto names = system.known_custom_data(); + std::sort(names.begin(), names.end()); + CHECK(names.size() == 2); + CHECK(names[0] == "test::my_data"); + CHECK(names[1] == "test::other_data"); +} + +TEST_CASE("System ownership") { + SECTION("move") { + auto system = test_system(4); + auto* ptr = system.as_mta_system_t(); + + auto moved = std::move(system); + CHECK(moved.as_mta_system_t() == ptr); + CHECK(moved.size() == 4); + } + + SECTION("release / unsafe_from_ptr round-trip") { + auto system = test_system(4); + auto* raw = system.release(); + REQUIRE(raw != nullptr); + + auto owned = metatomic::System::unsafe_from_ptr(raw); + CHECK(owned.size() == 4); + } + + SECTION("unsafe_view_from_ptr does not free") { + auto system = test_system(4); + + { + auto view = metatomic::System::unsafe_view_from_ptr(system.as_mta_system_t()); + CHECK(view.size() == 4); + } + + // the original system is still usable after the view is destroyed + CHECK(system.size() == 4); + } +} + +TEST_CASE("System serialization") { + SECTION("save and load a file") { + auto system = test_system(4); + const auto path = (std::filesystem::temp_directory_path() / "metatomic-test-system.mta").string(); + + struct FileCleanup { + const std::string& path; + + ~FileCleanup() { + std::remove(path.c_str()); + } + } cleanup{path}; + + metatomic::io::save(path, system); + auto loaded = metatomic::io::load(path); + + check_systems(system, loaded); + } + + SECTION("load a legacy file") { + auto path = std::filesystem::path(__FILE__).parent_path().parent_path() / "data" / "legacy.mta"; + auto system = metatomic::io::load(path.string()); + + CHECK(system.as_mta_system_t() != nullptr); + CHECK(system.size() == 4); + CHECK(system.length_unit().empty()); + + auto types = metatensor::details::vector_from_dlpack(system.types()->dl_tensor); + CHECK((types == std::vector{1, 6, 7, 8})); + + auto positions = metatensor::details::vector_from_dlpack(system.positions()->dl_tensor); + CHECK((positions == std::vector{ + 0.0, 0.0, 0.0, + 1.0, 2.0, 3.0, + 4.0, 5.0, 6.0, + 7.0, 8.0, 9.0, + })); + } + + SECTION("save and load an in-memory buffer") { + auto system = test_system(4); + + auto buffer = metatomic::io::save_buffer>(system); + REQUIRE_FALSE(buffer.empty()); + + auto loaded = metatomic::io::load_buffer(buffer); + check_systems(system, loaded); + } +} diff --git a/metatomic-core/tests/data/legacy.mta b/metatomic-core/tests/data/legacy.mta new file mode 100644 index 000000000..1eee677e7 Binary files /dev/null and b/metatomic-core/tests/data/legacy.mta differ diff --git a/metatomic-torch/tests/external/.gitattributes b/metatomic-core/tests/external/.gitattributes similarity index 100% rename from metatomic-torch/tests/external/.gitattributes rename to metatomic-core/tests/external/.gitattributes diff --git a/metatomic-torch/tests/external/CMakeLists.txt b/metatomic-core/tests/external/CMakeLists.txt similarity index 100% rename from metatomic-torch/tests/external/CMakeLists.txt rename to metatomic-core/tests/external/CMakeLists.txt diff --git a/metatomic-torch/tests/external/catch/catch.cpp b/metatomic-core/tests/external/catch/catch.cpp similarity index 100% rename from metatomic-torch/tests/external/catch/catch.cpp rename to metatomic-core/tests/external/catch/catch.cpp diff --git a/metatomic-torch/tests/external/catch/catch.hpp b/metatomic-core/tests/external/catch/catch.hpp similarity index 100% rename from metatomic-torch/tests/external/catch/catch.hpp rename to metatomic-core/tests/external/catch/catch.hpp diff --git a/metatomic-core/tests/misc.cpp b/metatomic-core/tests/misc.cpp new file mode 100644 index 000000000..8b3b7656f --- /dev/null +++ b/metatomic-core/tests/misc.cpp @@ -0,0 +1,121 @@ +#include + +#include + +#include "metatomic.h" + + +TEST_CASE("Version macros") { + CHECK(std::string(METATOMIC_VERSION) == mta_version()); + + auto version = std::to_string(METATOMIC_VERSION_MAJOR) + "." + + std::to_string(METATOMIC_VERSION_MINOR) + "." + + std::to_string(METATOMIC_VERSION_PATCH); + + // METATOMIC_VERSION should start with `x.y.z` + CHECK(std::string(METATOMIC_VERSION).find(version) == 0); +} + +TEST_CASE("mta_string_t") { + auto* str = mta_string_create("hello"); + REQUIRE(str != nullptr); + + const char* view = mta_string_view(str); + CHECK(std::strlen(view) == 5); + CHECK(std::string(view) == "hello"); + mta_string_free(str); + + // empty string + str = mta_string_create(""); + REQUIRE(str != nullptr); + CHECK(std::string(mta_string_view(str)) == ""); + mta_string_free(str); + + // special characters + str = mta_string_create("a\nb\tc\xFFºµ"); + REQUIRE(str != nullptr); + CHECK(std::string(mta_string_view(str)) == std::string("a\nb\tc\xFFºµ")); + mta_string_free(str); + + // long string + std::string long_str(10000, 'x'); + str = mta_string_create(long_str.c_str()); + REQUIRE(str != nullptr); + CHECK(std::string(mta_string_view(str)) == long_str); + mta_string_free(str); + + // free on a null pointer should work + mta_string_free(nullptr); +} + +TEST_CASE("unit conversion factor") { + double factor = 0.0; + + // same unit -> factor = 1.0 + auto status = mta_unit_conversion_factor("m", "m", &factor); + REQUIRE(status == MTA_SUCCESS); + CHECK(factor == 1.0); + + // kJ/mol -> eV + CHECK(mta_unit_conversion_factor("kJ/mol", "eV", &factor) == MTA_SUCCESS); + CHECK(factor == Approx(0.010364269656262174).epsilon(1e-15)); + + // dimension mismatch -> error + status = mta_unit_conversion_factor("m", "kg", &factor); + REQUIRE(status != MTA_SUCCESS); + + const char* error_msg = nullptr; + mta_last_error(&error_msg, nullptr, nullptr); + CHECK(std::string(error_msg) == + "invalid parameter: dimension mismatch in unit conversion: " + "'m' has dimension [L] but 'kg' has dimension [M]" + ); +} + +TEST_CASE("metatdata formatting") { + std::string json =R"({ + "type": "metatomic_model_metadata", + "name": "name", + "description": "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation.", + "authors": ["Short author", "Some extremely long author that will take more than one line in the printed output"], + "references": { + "architecture": ["ref-2", "ref-3"], + "model": ["a very long reference that will take more than one line in the printed output"], + "implementation": [] + }, + "extra": {} +})"; + + const auto* expected = R"(This is the name model +====================== + +Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor +incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis +nostrud exercitation. + +Model authors +------------- + +- Short author +- Some extremely long author that will take more than one line in the printed + output + +Model references +---------------- + +Please cite the following references when using this model: +- about this specific model: + * a very long reference that will take more than one line in the printed + output +- about the architecture of this model: + * ref-2 + * ref-3 +)"; + + auto* mta_string = mta_string_create(""); + REQUIRE(mta_string != nullptr); + auto status = mta_format_metadata(json.c_str(), &mta_string); + REQUIRE(status == MTA_SUCCESS); + CHECK(std::string(mta_string_view(mta_string)) == expected); + mta_string_free(mta_string); +} diff --git a/metatomic-core/tests/plugins.cpp b/metatomic-core/tests/plugins.cpp new file mode 100644 index 000000000..70a2dd14f --- /dev/null +++ b/metatomic-core/tests/plugins.cpp @@ -0,0 +1,49 @@ +#include + +#include "metatomic.h" + + +TEST_CASE("Load plugins") { + auto status = mta_load_plugin(PLUGIN_DIR "/test-c-plugin.so"); + CHECK(status == MTA_SUCCESS); + + const char* error_message; + const char* error_origin; + + struct mta_model_t model; + status = mta_load_model("some_model", "{}", "test-c-plugin", &model); + CHECK(status == MTA_INVALID_PARAMETER_ERROR); + + status = mta_last_error(&error_message, &error_origin, nullptr); + REQUIRE(status == MTA_SUCCESS); + + CHECK(std::string(error_origin) == "metatomic-core"); + CHECK(std::string(error_message) == ( + "invalid parameter: failed to load model from 'some_model': plugin 'test-c-plugin' could not load the model" + )); + + status = mta_load_model("some_model", "{}", nullptr, &model); + CHECK(status == MTA_INVALID_PARAMETER_ERROR); + + status = mta_last_error(&error_message, &error_origin, nullptr); + REQUIRE(status == MTA_SUCCESS); + + CHECK(std::string(error_origin) == "metatomic-core"); + CHECK(std::string(error_message) == ( + "invalid parameter: failed to load model from 'some_model': tried the " + "following plugins, but none could load the model: test-c-plugin" + )); + + + status = mta_load_plugin(PLUGIN_DIR "/bad-abi-plugin.so"); + CHECK(status == MTA_INVALID_PARAMETER_ERROR); + + status = mta_last_error(&error_message, &error_origin, nullptr); + REQUIRE(status == MTA_SUCCESS); + + CHECK(std::string(error_origin) == "metatomic-core"); + CHECK(std::string(error_message) == ( + "invalid parameter: can not register plugin 'bad-abi-plugin': " + "plugin ABI version is 2, but metatomic expects 1" + )); +} diff --git a/metatomic-core/tests/run-cxx-tests.rs b/metatomic-core/tests/run-cxx-tests.rs new file mode 100644 index 000000000..0d3b48d9d --- /dev/null +++ b/metatomic-core/tests/run-cxx-tests.rs @@ -0,0 +1,40 @@ +use std::path::PathBuf; + +mod utils; + +#[test] +fn run_cxx_tests() { + const CARGO_TARGET_TMPDIR: &str = env!("CARGO_TARGET_TMPDIR"); + + let mut build_dir = PathBuf::from(CARGO_TARGET_TMPDIR); + build_dir.push("cxx-tests"); + std::fs::create_dir_all(&build_dir).expect("failed to create build dir"); + + // ====================================================================== // + // setup dependencies for the torch tests + let deps_dir = build_dir.join("deps"); + let virtualenv_dir = deps_dir.join("virtualenv"); + std::fs::create_dir_all(&virtualenv_dir).expect("failed to create virtualenv dir"); + let python_exe = utils::create_python_venv(virtualenv_dir); + let metatensor_cmake_prefix = utils::setup_metatensor_pip(&python_exe); + + // ====================================================================== // + // build the metatomic C++ tests and run them + + let mut source_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + source_dir.push("tests"); + + // configure cmake for the tests + let mut cmake_config = utils::cmake_config(&source_dir, &build_dir); + cmake_config.arg("-DCMAKE_EXPORT_COMPILE_COMMANDS=ON"); + cmake_config.arg(format!("-DCMAKE_PREFIX_PATH={}", metatensor_cmake_prefix.display())); + utils::run_command(cmake_config, "cmake configuration"); + + // build the tests + let cmake_build = utils::cmake_build(&build_dir); + utils::run_command(cmake_build, "cmake build"); + + // run the tests + let ctest = utils::ctest(&build_dir); + utils::run_command(ctest, "ctest"); +} diff --git a/metatomic-core/tests/system.cpp b/metatomic-core/tests/system.cpp new file mode 100644 index 000000000..1ccb6a7d3 --- /dev/null +++ b/metatomic-core/tests/system.cpp @@ -0,0 +1,698 @@ +#include +#include +#include +#include +#include +#include + +#include + +#include +#include "metatomic.h" +#include "tensor_utils.hpp" + +static mts_block_t* pair_block() { + auto samples = metatensor::Labels( + {"first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"}, + {{0, 1, 0, 0, 0}} + ); + + auto components = metatensor::Labels({"xyz"}, {{0}, {1}, {2}}); + std::vector components_list = { + components.as_mts_labels_t() + }; + + auto properties = metatensor::Labels({"distance"}, {{0}}); + + auto values = std::make_unique>( + std::vector{1, 3, 1}, + std::vector{1.5F, 2.5F, 3.5F} + ); + auto values_mts = metatensor::DataArrayBase::to_mts_array(std::move(values)); + + auto* block = mts_block( + std::move(values_mts).release(), + samples.as_mts_labels_t(), + components_list.data(), + components_list.size(), + properties.as_mts_labels_t() + ); + REQUIRE(block != nullptr); + return block; +} + +static mts_tensormap_t* custom_data() { + auto keys = metatensor::Labels({"key"}, {{0}}); + auto samples = metatensor::Labels({"sample"}, {{0}}); + auto properties = metatensor::Labels({"property"}, {{0}}); + + auto values = std::make_unique>( + std::vector{1, 1}, + std::vector{42.0F} + ); + auto values_mts = metatensor::DataArrayBase::to_mts_array(std::move(values)); + + auto* block = mts_block( + std::move(values_mts).release(), + samples.as_mts_labels_t(), + nullptr, + 0, + properties.as_mts_labels_t() + ); + REQUIRE(block != nullptr); + + std::vector blocks = {block}; + auto* tensormap = mts_tensormap( + keys.as_mts_labels_t(), + blocks.data(), + blocks.size() + ); + REQUIRE(tensormap != nullptr); + + return tensormap; +} + +TEST_CASE("system") { + SECTION("create and free") { + mta_system_t* system_f32 = nullptr; + auto status = mta_system_create( + "nm", + types_tensor(4), + positions_tensor(4), + cell_tensor(), + pbc_tensor(), + &system_f32 + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(system_f32 != nullptr); + + status = mta_system_free(system_f32); + CHECK(status == MTA_SUCCESS); + + mta_system_t* system_f64 = nullptr; + status = mta_system_create( + "nm", + types_tensor(4), + positions_tensor(4), + cell_tensor(), + pbc_tensor(), + &system_f64 + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(system_f64 != nullptr); + + status = mta_system_free(system_f64); + CHECK(status == MTA_SUCCESS); + + // free on null pointer is fine + status = mta_system_free(nullptr); + REQUIRE(status == MTA_SUCCESS); + } + + SECTION("errors") { + mta_system_t* system = nullptr; + + // wrong dtype for types (float instead of int32) + auto status = mta_system_create( + "Angstrom", + types_tensor(3), + positions_tensor(3), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status != MTA_SUCCESS); + CHECK(system == nullptr); + + const char* message = nullptr; + mta_last_error(&message, nullptr, nullptr); + CHECK(std::string(message) == "invalid parameter: `types` must be a tensor of 32-bit integers"); + + // wrong dtype for positions (int32 instead of float) + status = mta_system_create( + "Angstrom", + types_tensor(3), + positions_tensor(3), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status != MTA_SUCCESS); + CHECK(system == nullptr); + + mta_last_error(&message, nullptr, nullptr); + CHECK(std::string(message) == "invalid parameter: `positions` must be a tensor of 32 or 64-bit floating point data"); + + // wrong dtype for cell (int32 instead of float) + status = mta_system_create( + "Angstrom", + types_tensor(3), + positions_tensor(3), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status != MTA_SUCCESS); + CHECK(system == nullptr); + + mta_last_error(&message, nullptr, nullptr); + CHECK(std::string(message) == "invalid parameter: `cell` must have the same dtype as `positions`, got i32 and f32"); + + // wrong dtype for pbc (float instead of bool) + auto* float_pbc = pbc_tensor(); + float_pbc->dl_tensor.dtype.code = kDLFloat; + status = mta_system_create( + "Angstrom", + types_tensor(3), + positions_tensor(3), + cell_tensor(), + float_pbc, + &system + ); + CHECK(status != MTA_SUCCESS); + CHECK(system == nullptr); + + mta_last_error(&message, nullptr, nullptr); + CHECK(std::string(message) == "invalid parameter: `pbc` must be a tensor of booleans"); + + // mismatched positions/type shapes + status = mta_system_create( + "Angstrom", + types_tensor(3), + positions_tensor(5), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status != MTA_SUCCESS); + CHECK(system == nullptr); + + mta_last_error(&message, nullptr, nullptr); + CHECK(std::string(message) == "invalid parameter: `positions` must be a (n_atoms x 3) tensor, got a tensor with shape [5, 3]"); + + + // wrong cell shape + auto* cell = cell_tensor(); + cell->dl_tensor.shape[0] = 9; + cell->dl_tensor.shape[1] = 1; + status = mta_system_create( + "Angstrom", + types_tensor(3), + positions_tensor(3), + cell, + pbc_tensor(), + &system + ); + CHECK(status != MTA_SUCCESS); + CHECK(system == nullptr); + + mta_last_error(&message, nullptr, nullptr); + CHECK(std::string(message) == "invalid parameter: `cell` must be a (3 x 3) tensor, got a tensor with shape [9, 1]"); + + + // wrong pbc shape + auto* pbc = pbc_tensor(); + pbc->dl_tensor.shape[0] = 2; + status = mta_system_create( + "Angstrom", + types_tensor(3), + positions_tensor(3), + cell_tensor(), + pbc, + &system + ); + CHECK(status != MTA_SUCCESS); + CHECK(system == nullptr); + + mta_last_error(&message, nullptr, nullptr); + CHECK(std::string(message) == "invalid parameter: `pbc` must contain 3 entries, got a tensor with shape [2]"); + } + + SECTION("size") { + mta_system_t* system = nullptr; + auto status = mta_system_create( + "nm", + types_tensor(4), + positions_tensor(4), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(system != nullptr); + + uintptr_t size = 0; + status = mta_system_size(system, &size); + CHECK(status == MTA_SUCCESS); + CHECK(size == 4); + + status = mta_system_free(system); + CHECK(status == MTA_SUCCESS); + } + + SECTION("length unit") { + mta_system_t* system = nullptr; + auto status = mta_system_create( + "nm", + types_tensor(4), + positions_tensor(4), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(system != nullptr); + + mta_string_t unit = nullptr; + status = mta_system_get_length_unit(system, &unit); + CHECK(status == MTA_SUCCESS); + CHECK(std::string(mta_string_view(unit)) == "nm"); + mta_string_free(unit); + + status = mta_system_free(system); + CHECK(status == MTA_SUCCESS); + } +} + +TEST_CASE("system data") { + mta_system_t* system = nullptr; + auto status = mta_system_create( + "nm", + types_tensor(4), + positions_tensor(4), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(system != nullptr); + + DLManagedTensorVersioned* data = nullptr; + + SECTION("types") { + status = mta_system_get_data( + system, MTA_SYSTEM_DATA_TYPES, &data + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(data != nullptr); + + CHECK(data->dl_tensor.ndim == 1); + CHECK(data->dl_tensor.shape[0] == 4); + CHECK(data->dl_tensor.dtype.code == kDLInt); + CHECK(data->dl_tensor.dtype.bits == 32); + + auto* types = reinterpret_cast(static_cast(data->dl_tensor.data) + data->dl_tensor.byte_offset); + CHECK(types[0] == 1); + CHECK(types[1] == 4); + CHECK(types[2] == 7); + CHECK(types[3] == 10); + } + + SECTION("positions") { + status = mta_system_get_data( + system, MTA_SYSTEM_DATA_POSITIONS, &data + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(data != nullptr); + + CHECK(data->dl_tensor.ndim == 2); + CHECK(data->dl_tensor.shape[0] == 4); + CHECK(data->dl_tensor.shape[1] == 3); + CHECK(data->dl_tensor.dtype.code == kDLFloat); + CHECK(data->dl_tensor.dtype.bits == 32); + + auto* positions = reinterpret_cast(static_cast(data->dl_tensor.data) + data->dl_tensor.byte_offset); + CHECK(positions[0] == 1.0F); + CHECK(positions[3] == 4.0F); + CHECK(positions[6] == 7.0F); + CHECK(positions[9] == 10.0F); + } + + SECTION("cell") { + status = mta_system_get_data( + system, MTA_SYSTEM_DATA_CELL, &data + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(data != nullptr); + + CHECK(data->dl_tensor.ndim == 2); + CHECK(data->dl_tensor.shape[0] == 3); + CHECK(data->dl_tensor.shape[1] == 3); + CHECK(data->dl_tensor.dtype.code == kDLFloat); + CHECK(data->dl_tensor.dtype.bits == 32); + + auto* cell = reinterpret_cast(static_cast(data->dl_tensor.data) + data->dl_tensor.byte_offset); + CHECK(cell[0] == 10.0F); + CHECK(cell[4] == 0.0F); + CHECK(cell[8] == 10.0F); + } + + SECTION("pbc") { + status = mta_system_get_data( + system, MTA_SYSTEM_DATA_PBC, &data + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(data != nullptr); + + CHECK(data->dl_tensor.ndim == 1); + CHECK(data->dl_tensor.shape[0] == 3); + CHECK(data->dl_tensor.dtype.code == kDLBool); + CHECK(data->dl_tensor.dtype.bits == 8); + + auto* pbc = reinterpret_cast(static_cast(data->dl_tensor.data) + data->dl_tensor.byte_offset); + CHECK(pbc[0] == true); + CHECK(pbc[1] == false); + CHECK(pbc[2] == true); + } + + data->deleter(data); + + status = mta_system_free(system); + CHECK(status == MTA_SUCCESS); +} + + +TEST_CASE("system pairs") { + mta_system_t* system = nullptr; + auto status = mta_system_create( + "nm", + types_tensor(4), + positions_tensor(4), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(system != nullptr); + + const auto* options_json = R"({ + "type": "metatomic_pair_list_options", + "cutoff": "0x00001000", + "full_list": true, + "strict": false, + "requestors": ["test"] + })"; + + auto* pairs = pair_block(); + status = mta_system_add_pairs(system, options_json, pairs); + CHECK(status == MTA_SUCCESS); + + const mts_block_t* recovered_pairs = nullptr; + status = mta_system_get_pairs(system, options_json, &recovered_pairs); + CHECK(status == MTA_SUCCESS); + // we get the same pointer back + CHECK(static_cast(recovered_pairs) == static_cast(pairs)); + + // Add a second block with different options + const auto* other_json = R"({ + "type": "metatomic_pair_list_options", + "cutoff": "0x00001000", + "full_list": true, + "strict": true, + "requestors": [] + })"; + + pairs = pair_block(); + status = mta_system_add_pairs(system, other_json, pairs); + CHECK(status == MTA_SUCCESS); + + // Check known pairs contains both + mta_string_t known = nullptr; + status = mta_system_known_pairs(system, &known); + CHECK(status == MTA_SUCCESS); + REQUIRE(known != nullptr); + + auto known_str = std::string(mta_string_view(known)); + mta_string_free(known); + + auto first = known_str.find("metatomic_pair_list_options"); + CHECK(first != std::string::npos); + known_str = known_str.substr(first + 1); + auto second = known_str.find("metatomic_pair_list_options"); + CHECK(second != std::string::npos); + + mta_system_free(system); +} + +TEST_CASE("system custom data") { + mta_system_t* system = nullptr; + auto status = mta_system_create( + "Angstrom", + types_tensor(4), + positions_tensor(4), + cell_tensor(), + pbc_tensor(), + &system + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(system != nullptr); + + auto* data = custom_data(); + status = mta_system_add_custom_data(system, "test::my_data", data); + CHECK(status == MTA_SUCCESS); + + const mts_tensormap_t* retrieved = nullptr; + status = mta_system_get_custom_data( + system, "test::my_data", &retrieved + ); + CHECK(status == MTA_SUCCESS); + CHECK(retrieved != nullptr); + CHECK(static_cast(retrieved) == static_cast(data)); + + retrieved = nullptr; + status = mta_system_get_custom_data( + system, "test::no_such_data", &retrieved + ); + CHECK(status != MTA_SUCCESS); + CHECK(retrieved == nullptr); + + data = custom_data(); + status = mta_system_add_custom_data(system, "test::other_data", data); + CHECK(status == MTA_SUCCESS); + + mta_string_t names = nullptr; + status = mta_system_known_custom_data(system, &names); + CHECK(status == MTA_SUCCESS); + CHECK(names != nullptr); + + auto names_str = std::string(mta_string_view(names)); + CHECK(names_str.find("test::my_data") != std::string::npos); + CHECK(names_str.find("test::other_data") != std::string::npos); + + mta_system_free(system); +} + +/// Build a system containing all kinds of data (basic data, pairs, and custom +/// data) for use in serialization round-trip tests. +static mta_system_t* full_test_system() { + mta_system_t* system = nullptr; + auto status = mta_system_create( + "Angstrom", + types_tensor(4), + positions_tensor(4), + cell_tensor(), + pbc_tensor(), + &system + ); + REQUIRE(status == MTA_SUCCESS); + REQUIRE(system != nullptr); + + const auto* pairs_options_json = R"({ + "type": "metatomic_pair_list_options", + "cutoff": "0x00001000", + "full_list": true, + "strict": false, + "requestors": ["test"] + })"; + status = mta_system_add_pairs(system, pairs_options_json, pair_block()); + CHECK(status == MTA_SUCCESS); + + status = mta_system_add_custom_data(system, "test::my_data", custom_data()); + CHECK(status == MTA_SUCCESS); + + return system; +} + +/// Check that the given system contains the data expected from +/// `full_test_system`, independently of how it was loaded back. +static void check_full_system_data(const mta_system_t* system) { + uintptr_t size = 0; + CHECK(mta_system_size(system, &size) == MTA_SUCCESS); + CHECK(size == 4); + + mta_string_t unit = nullptr; + CHECK(mta_system_get_length_unit(system, &unit) == MTA_SUCCESS); + CHECK(std::string(mta_string_view(unit)) == "Angstrom"); + mta_string_free(unit); + + DLManagedTensorVersioned* data = nullptr; + + // types + CHECK(mta_system_get_data(system, MTA_SYSTEM_DATA_TYPES, &data) == MTA_SUCCESS); + REQUIRE(data != nullptr); + CHECK(data->dl_tensor.ndim == 1); + CHECK(data->dl_tensor.shape[0] == 4); + CHECK(data->dl_tensor.dtype.code == kDLInt); + CHECK(data->dl_tensor.dtype.bits == 32); + { + auto* types = reinterpret_cast( + static_cast(data->dl_tensor.data) + data->dl_tensor.byte_offset + ); + CHECK(types[0] == 1); + CHECK(types[1] == 4); + CHECK(types[2] == 7); + CHECK(types[3] == 10); + } + data->deleter(data); + + // positions + CHECK(mta_system_get_data(system, MTA_SYSTEM_DATA_POSITIONS, &data) == MTA_SUCCESS); + REQUIRE(data != nullptr); + CHECK(data->dl_tensor.ndim == 2); + CHECK(data->dl_tensor.shape[0] == 4); + CHECK(data->dl_tensor.shape[1] == 3); + CHECK(data->dl_tensor.dtype.code == kDLFloat); + CHECK(data->dl_tensor.dtype.bits == 32); + { + auto* positions = reinterpret_cast( + static_cast(data->dl_tensor.data) + data->dl_tensor.byte_offset + ); + CHECK(positions[0] == 1.0F); + CHECK(positions[3] == 4.0F); + CHECK(positions[6] == 7.0F); + CHECK(positions[9] == 10.0F); + } + data->deleter(data); + + // cell + CHECK(mta_system_get_data(system, MTA_SYSTEM_DATA_CELL, &data) == MTA_SUCCESS); + REQUIRE(data != nullptr); + CHECK(data->dl_tensor.ndim == 2); + CHECK(data->dl_tensor.shape[0] == 3); + CHECK(data->dl_tensor.shape[1] == 3); + CHECK(data->dl_tensor.dtype.code == kDLFloat); + CHECK(data->dl_tensor.dtype.bits == 32); + { + auto* cell = reinterpret_cast( + static_cast(data->dl_tensor.data) + data->dl_tensor.byte_offset + ); + CHECK(cell[0] == 10.0F); + CHECK(cell[4] == 0.0F); + CHECK(cell[8] == 10.0F); + } + data->deleter(data); + + // pbc + CHECK(mta_system_get_data(system, MTA_SYSTEM_DATA_PBC, &data) == MTA_SUCCESS); + REQUIRE(data != nullptr); + CHECK(data->dl_tensor.ndim == 1); + CHECK(data->dl_tensor.shape[0] == 3); + CHECK(data->dl_tensor.dtype.code == kDLBool); + CHECK(data->dl_tensor.dtype.bits == 8); + { + auto* pbc = reinterpret_cast( + static_cast(data->dl_tensor.data) + data->dl_tensor.byte_offset + ); + CHECK(pbc[0] == 1); + CHECK(pbc[1] == 0); + CHECK(pbc[2] == 1); + } + data->deleter(data); + + // known pairs survive the round-trip + mta_string_t known = nullptr; + CHECK(mta_system_known_pairs(system, &known) == MTA_SUCCESS); + REQUIRE(known != nullptr); + { + auto known_str = std::string(mta_string_view(known)); + CHECK(known_str.find("metatomic_pair_list_options") != std::string::npos); + CHECK(known_str.find("\"full_list\":true") != std::string::npos); + } + mta_string_free(known); + + // the pairs block can be retrieved + const auto* pairs_options_json = R"({ + "type": "metatomic_pair_list_options", + "cutoff": "0x00001000", + "full_list": true, + "strict": false, + "requestors": [] + })"; + const mts_block_t* pairs = nullptr; + CHECK(mta_system_get_pairs(system, pairs_options_json, &pairs) == MTA_SUCCESS); + CHECK(pairs != nullptr); + + // custom data survives the round-trip + mta_string_t names = nullptr; + CHECK(mta_system_known_custom_data(system, &names) == MTA_SUCCESS); + REQUIRE(names != nullptr); + { + auto names_str = std::string(mta_string_view(names)); + CHECK(names_str.find("test::my_data") != std::string::npos); + } + mta_string_free(names); + + const mts_tensormap_t* retrieved = nullptr; + CHECK(mta_system_get_custom_data(system, "test::my_data", &retrieved) == MTA_SUCCESS); + CHECK(retrieved != nullptr); +} + +/// `mts_realloc_buffer_t` callback backed by a `std::vector`. +static uint8_t* vector_realloc(void* user_data, uint8_t* /*ptr*/, uintptr_t new_size) { + auto* buffer = static_cast*>(user_data); + buffer->resize(new_size, 0); + return buffer->data(); +} + +TEST_CASE("system serialization") { + SECTION("save and load to a file") { + auto* system = full_test_system(); + + auto path = (std::filesystem::temp_directory_path() / "metatomic-test-system.mta").string(); + + CHECK(mta_save(path.c_str(), system) == MTA_SUCCESS); + + mta_system_t* loaded = nullptr; + auto status = mta_load( + path.c_str(), + metatensor::details::default_create_array, + &loaded + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(loaded != nullptr); + + check_full_system_data(loaded); + + CHECK(mta_system_free(loaded) == MTA_SUCCESS); + CHECK(mta_system_free(system) == MTA_SUCCESS); + std::remove(path.c_str()); + } + + SECTION("save and load to an in-memory buffer") { + auto* system = full_test_system(); + + std::vector buffer; + uint8_t* ptr = buffer.data(); + uintptr_t size = buffer.size(); + + auto status = mta_save_buffer( + &ptr, &size, &buffer, vector_realloc, system + ); + CHECK(status == MTA_SUCCESS); + buffer.resize(size); + + mta_system_t* loaded = nullptr; + status = mta_load_buffer( + buffer.data(), buffer.size(), + metatensor::details::default_create_array, + &loaded + ); + CHECK(status == MTA_SUCCESS); + REQUIRE(loaded != nullptr); + + check_full_system_data(loaded); + + CHECK(mta_system_free(loaded) == MTA_SUCCESS); + CHECK(mta_system_free(system) == MTA_SUCCESS); + } +} diff --git a/metatomic-core/tests/tensor_utils.hpp b/metatomic-core/tests/tensor_utils.hpp new file mode 100644 index 000000000..875ff24b7 --- /dev/null +++ b/metatomic-core/tests/tensor_utils.hpp @@ -0,0 +1,77 @@ +#ifndef METATOMIC_TESTS_UTILS_HPP +#define METATOMIC_TESTS_UTILS_HPP + +#include +#include + +#include + +/// Helpers to create DLPack tensors for System creation through the C API. +/// These return raw `DLManagedTensorVersioned*` pointers that are transferred +/// to `mta_system_create` (which takes ownership). + +template +inline DLManagedTensorVersioned* types_tensor(size_t n_atoms) { + std::vector type_data; + type_data.reserve(n_atoms); + for (size_t i = 0; i < n_atoms; i++) { + type_data.push_back(static_cast(i * 3 + 1)); + } + auto array = std::make_unique>( + std::vector{n_atoms}, + std::move(type_data) + ); + auto mts = metatensor::DataArrayBase::to_mts_array(std::move(array)); + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + return mts.as_dlpack(cpu, nullptr, version); +} + +template +inline DLManagedTensorVersioned* positions_tensor(size_t n_atoms) { + std::vector position_data; + position_data.reserve(n_atoms * 3); + for (size_t i = 0; i < n_atoms; i++) { + position_data.push_back(static_cast(i * 3 + 1)); + position_data.push_back(static_cast(i * 3 + 2)); + position_data.push_back(static_cast(i * 3 + 3)); + } + auto array = std::make_unique>( + std::vector{n_atoms, 3}, + std::move(position_data) + ); + auto mts = metatensor::DataArrayBase::to_mts_array(std::move(array)); + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + return mts.as_dlpack(cpu, nullptr, version); +} + +template +inline DLManagedTensorVersioned* cell_tensor() { + auto array = std::make_unique>( + std::vector{3, 3}, + std::vector{ + T(10.0), T(0.0), T(0.0), + T(0.0), T(0.0), T(0.0), + T(0.0), T(0.0), T(10.0), + } + ); + auto mts = metatensor::DataArrayBase::to_mts_array(std::move(array)); + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + return mts.as_dlpack(cpu, nullptr, version); +} + +inline DLManagedTensorVersioned* pbc_tensor() { + std::vector pbc_data = {1, 0, 1}; + auto array = std::make_unique>( + std::vector{3}, + std::move(pbc_data) + ); + auto mts = metatensor::DataArrayBase::to_mts_array(std::move(array)); + DLDevice cpu = {kDLCPU, 0}; + DLPackVersion version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION}; + return mts.as_dlpack(cpu, nullptr, version); +} + +#endif // METATOMIC_TESTS_UTILS_HPP diff --git a/metatomic-core/tests/test-plugins/CMakeLists.txt b/metatomic-core/tests/test-plugins/CMakeLists.txt new file mode 100644 index 000000000..9ebeea89d --- /dev/null +++ b/metatomic-core/tests/test-plugins/CMakeLists.txt @@ -0,0 +1,23 @@ +# create test plugins with a consistent name across platforms +# .so suffix is used for all the platforms + +add_library(test-c-plugin SHARED plugin.c) +target_link_libraries(test-c-plugin metatomic) +set_target_properties(test-c-plugin PROPERTIES + PREFIX "" + SUFFIX ".so" +) + +add_library(bad-abi-plugin SHARED bad-abi.c) +target_link_libraries(bad-abi-plugin metatomic) +set_target_properties(bad-abi-plugin PROPERTIES + PREFIX "" + SUFFIX ".so" +) + +add_library(test-cxx-plugin SHARED cxx-plugin.cpp) +target_link_libraries(test-cxx-plugin metatomic) +set_target_properties(test-cxx-plugin PROPERTIES + PREFIX "" + SUFFIX ".so" +) diff --git a/metatomic-core/tests/test-plugins/bad-abi.c b/metatomic-core/tests/test-plugins/bad-abi.c new file mode 100644 index 000000000..ddde49263 --- /dev/null +++ b/metatomic-core/tests/test-plugins/bad-abi.c @@ -0,0 +1,17 @@ +#include + + +static mta_status_t load_model(const char *load_from, const char *options_json, struct mta_model_t *model) { + // This plugin can not load any model + return MTA_UNSUPPORTED_MODEL_ERROR; +} + + +MTA_REGISTER_PLUGIN(register_plugin, { + mta_plugin_t plugin = { + .abi_version = MTA_ABI_VERSION + 1, // incompatible ABI version + .name = "bad-abi-plugin", + .load_model = load_model, + }; + return register_plugin(plugin); +}); diff --git a/metatomic-core/tests/test-plugins/cxx-plugin.cpp b/metatomic-core/tests/test-plugins/cxx-plugin.cpp new file mode 100644 index 000000000..5c1c12e77 --- /dev/null +++ b/metatomic-core/tests/test-plugins/cxx-plugin.cpp @@ -0,0 +1,78 @@ +#include +#include +#include +#include +#include + +#include +#include "metatomic.hpp" + + +class SimpleModel: public metatomic::BaseModel { +public: + metatomic::ModelCapabilities capabilities() const final { + return metatomic::ModelCapabilities::builder() + .atomic_types({1, 6, 8}) + .interaction_range(4.5) + .length_unit("nm") + .supported_devices({metatomic::ModelCapabilities::Device::CPU}) + .dtype(metatomic::ModelCapabilities::DType::Float32) + .add_output(metatomic::Quantity::builder() + .name("energy") + .unit("eV") + .sample_kind(metatomic::SampleKind::System) + .build()) + .build(); + } + + metatomic::ModelMetadata metadata() const final { + return metatomic::ModelMetadata::builder() + .name("simple C++ plugin model") + .description("test model for MTA_REGISTER_CXX_PLUGIN") + .build(); + } + + std::vector requested_pair_lists() const final { + return {}; + } + + std::vector requested_inputs() const final { + return {}; + } + + std::vector execute_inner( + const std::vector&, + const metatensor::Labels*, + const std::vector& + ) final { + return {}; + } +}; + + +std::unique_ptr load_model_cxx( + const std::string& load_from, + const std::map& options +) { + (void)options; + + if (load_from == "throws") { + throw metatomic::Error("load_model_cxx: intentional failure for '" + load_from + "'"); + } + + if (load_from == "throws-std") { + // exceptions that are not `metatomic::Error` should also make it back + // to the caller with their original type + throw std::out_of_range("load_model_cxx: intentional failure for '" + load_from + "'"); + } + + if (load_from != "test-cxx-model") { + // this plugin can not load this model + return nullptr; + } + + return std::make_unique(); +} + + +MTA_REGISTER_CXX_PLUGIN("test-cxx-plugin", load_model_cxx); diff --git a/metatomic-core/tests/test-plugins/plugin.c b/metatomic-core/tests/test-plugins/plugin.c new file mode 100644 index 000000000..2f203e655 --- /dev/null +++ b/metatomic-core/tests/test-plugins/plugin.c @@ -0,0 +1,17 @@ +#include + + +static mta_status_t load_model(const char *load_from, const char *options_json, struct mta_model_t *model) { + // This plugin can not load any model + return MTA_UNSUPPORTED_MODEL_ERROR; +} + + +MTA_REGISTER_PLUGIN(register_plugin, { + mta_plugin_t plugin = { + .abi_version = MTA_ABI_VERSION, + .name = "test-c-plugin", + .load_model = load_model, + }; + return register_plugin(plugin); +}); diff --git a/metatomic-core/tests/utils/mod.rs b/metatomic-core/tests/utils/mod.rs new file mode 100644 index 000000000..ff2ae89ff --- /dev/null +++ b/metatomic-core/tests/utils/mod.rs @@ -0,0 +1,463 @@ +#![allow(dead_code)] +#![allow(clippy::needless_return)] + +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; +use std::process::{Command, Stdio}; + +fn build_type() -> &'static str { + // assume that debug assertion means that we are building the code in + // debug mode, even if that could be not true in some cases + if cfg!(debug_assertions) { + "debug" + } else { + "release" + } +} + +fn append_flags(existing: Option, extra: &str) -> String { + match existing { + Some(flags) if !flags.trim().is_empty() => format!("{flags} {extra}"), + _ => extra.into(), + } +} + +pub fn cmake_config(source_dir: &Path, build_dir: &Path) -> Command { + let cmake = which::which("cmake").expect("could not find cmake"); + + let mut cmake_config = Command::new(cmake); + cmake_config.current_dir(build_dir); + cmake_config.arg(source_dir); + cmake_config.arg("--no-warn-unused-cli"); + cmake_config.arg(format!("-DCMAKE_BUILD_TYPE={}", build_type())); + + // the cargo executable currently running + let cargo_exe = std::env::var("CARGO").expect("CARGO env var is not set"); + cmake_config.arg(format!("-DCARGO_EXE={}", cargo_exe)); + + if std::env::var_os("CARGO_LLVM_COV").is_some() { + let coverage_compile_flags = "-fprofile-instr-generate -fcoverage-mapping"; + let coverage_link_flags = "-fprofile-instr-generate"; + + let c_flags = append_flags(std::env::var("CFLAGS").ok(), coverage_compile_flags); + let cxx_flags = append_flags(std::env::var("CXXFLAGS").ok(), coverage_compile_flags); + let exe_linker_flags = + append_flags(std::env::var("LDFLAGS").ok(), coverage_link_flags); + + cmake_config.arg(format!("-DCMAKE_C_FLAGS={c_flags}")); + cmake_config.arg(format!("-DCMAKE_CXX_FLAGS={cxx_flags}")); + cmake_config.arg(format!("-DCMAKE_EXE_LINKER_FLAGS={exe_linker_flags}")); + cmake_config.arg(format!("-DCMAKE_SHARED_LINKER_FLAGS={exe_linker_flags}")); + } + + return cmake_config; +} + +pub fn cmake_build(build_dir: &Path) -> Command { + let cmake = which::which("cmake").expect("could not find cmake"); + + let mut cmake_build = Command::new(cmake); + cmake_build.current_dir(build_dir); + cmake_build.arg("--build"); + cmake_build.arg("."); + cmake_build.arg("--parallel"); + cmake_build.arg("--config"); + cmake_build.arg(build_type()); + + return cmake_build; +} + + +pub fn ctest(build_dir: &Path) -> Command { + let ctest = which::which("ctest").expect("could not find ctest"); + + let mut ctest = Command::new(ctest); + ctest.current_dir(build_dir); + ctest.arg("--output-on-failure"); + ctest.arg("--build-config"); + ctest.arg(build_type()); + + return ctest +} + +/// Find the path to the uv binary, or None if not present +fn find_uv() -> Option { + which::which("uv").ok() +} + +/// Find the path to the `python`or `python3` binary on the user system +fn find_python() -> PathBuf { + if let Ok(python) = which::which("python") { + let output = Command::new(&python) + .arg("-c") + .arg("import sys; print(sys.version_info.major)") + .output() + .expect("could not run python"); + + if output.status.success() { + let stdout = String::from_utf8_lossy(&output.stdout); + + if stdout.trim() == "3" { + // we found Python 3 + return python; + } + } + } + + // try python3 + let python = which::which("python3").expect("failed to run `which python3`"); + let output = Command::new(&python) + .arg("-c") + .arg("import sys; print(sys.version_info.major)") + .output() + .expect("could not run python"); + + if output.status.success() { + let stdout = String::from_utf8_lossy(&output.stdout); + if stdout.trim() == "3" { + // we found Python 3 + return python; + } + } + + panic!("could not find Python 3") +} + +/// Helper: get python executable path inside a venv +fn python_in_venv(venv_dir: &Path) -> PathBuf { + let mut python = venv_dir.to_path_buf(); + if cfg!(target_os = "windows") { + python.extend(["Scripts", "python.exe"]); + } else { + python.extend(["bin", "python"]); + } + python +} + +/// Create a Python virtualenv using uv if available, else fallback to +/// `python -m venv`, and return the path to the python executable in the venv +pub fn create_python_venv(build_dir: PathBuf) -> PathBuf { + if let Some(uv_bin) = find_uv() { + let mut cmd = Command::new(&uv_bin); + cmd.arg("venv"); + cmd.arg("--allow-existing"); + cmd.arg(&build_dir); + + run_command(cmd, "uv venv creation"); + } else { + let mut cmd = Command::new(find_python()); + cmd.arg("-m"); + cmd.arg("venv"); + cmd.arg("--upgrade-deps"); + cmd.arg(&build_dir); + + run_command(cmd, "python to create virtualenv with `venv`"); + } + + python_in_venv(&build_dir) +} + +#[derive(Default)] +pub struct PipInstallOptions { + pub upgrade: bool, + pub no_deps: bool, + pub no_build_isolation: bool, +} + +/// Install a package with pip (uses uv if present, else falls back to python) +fn pip_install( + python: &Path, + packages: &[&str], + options: PipInstallOptions, +) { + if let Some(uv_bin) = find_uv() { + let mut cmd = Command::new(&uv_bin); + cmd.arg("pip").arg("install").arg("--python").arg(python); + + // follow the same behavior as pip when there are multiple indexes + cmd.arg("--index-strategy"); + cmd.arg("unsafe-best-match"); + + if options.upgrade { + cmd.arg("--upgrade"); + } + if options.no_deps { + cmd.arg("--no-deps"); + } + if options.no_build_isolation { + cmd.arg("--no-build-isolation"); + // uv doesn't support --check-build-dependencies + } + + for package in packages { + cmd.arg(package); + } + + run_command(cmd, "uv pip install"); + } else { + let mut cmd = Command::new(python); + cmd.arg("-m").arg("pip").arg("install"); + if options.upgrade { + cmd.arg("--upgrade"); + } + if options.no_deps { + cmd.arg("--no-deps"); + } + if options.no_build_isolation { + // If pip, add both supported options + cmd.arg("--no-build-isolation"); + cmd.arg("--check-build-dependencies"); + } + + for package in packages { + cmd.arg(package); + } + + run_command(cmd, "pip install"); + } +} + +/// Download PyTorch in a Python virtualenv, and return the +/// CMAKE_PREFIX_PATH for the corresponding libtorch +pub fn setup_torch_pip(python: &Path) -> PathBuf { + let torch_version = std::env::var("METATOMIC_TESTS_TORCH_VERSION").unwrap_or("2.13".into()); + pip_install( + python, + &[&format!("torch=={}.*", torch_version)], + PipInstallOptions { upgrade: true, no_deps: false, no_build_isolation: false } + ); + + let mut cmd = Command::new(python); + cmd.arg("-c"); + cmd.arg("import torch; print(torch.utils.cmake_prefix_path)"); + + let output = run_command(cmd, "python to get torch cmake prefix"); + + let stdout = String::from_utf8_lossy(&output.stdout); + let prefix = PathBuf::from(stdout.trim()); + if !prefix.exists() { + panic!("'torch.utils.cmake_prefix' at '{}' does not exist", prefix.display()); + } + + return prefix; +} + +/// Install metatensor in a Python virtualenv with pip, and return the +/// CMAKE_PREFIX_PATH for the installed libmetatensor. +pub fn setup_metatensor_pip(python: &Path) -> PathBuf { + pip_install(python, &["metatensor-core >=0.2.4,<0.3"], PipInstallOptions::default()); + + let mut cmd = Command::new(python); + cmd.arg("-c"); + cmd.arg("import metatensor; print(metatensor.utils.cmake_prefix_path)"); + + let output = run_command(cmd, "python to get metatensor cmake prefix"); + + let stdout = String::from_utf8_lossy(&output.stdout); + let prefix = PathBuf::from(stdout.trim()); + if !prefix.exists() { + panic!("'metatensor.utils.cmake_prefix' at '{}' does not exist", prefix.display()); + } + + return prefix; +} + +/// Install metatensor-torch in a Python virtualenv with pip, and return the +/// CMAKE_PREFIX_PATH for the installed libmetatensor_torch. +pub fn setup_metatensor_torch_pip(python: &Path) -> PathBuf { + pip_install(python, &["metatensor-torch >=0.10.0,<0.11"], PipInstallOptions::default()); + + let mut cmd = Command::new(python); + cmd.arg("-c"); + cmd.arg("import metatensor.torch; print(metatensor.torch.utils.cmake_prefix_path)"); + + let output = run_command(cmd, "python to get metatensor_torch cmake prefix"); + + let stdout = String::from_utf8_lossy(&output.stdout); + let prefix = PathBuf::from(stdout.trim()); + if !prefix.exists() { + panic!("'metatensor.torch.utils.cmake_prefix' at '{}' does not exist", prefix.display()); + } + + return prefix; +} + +/// Build metatomic-torch located in `source_dir` inside `build_dir`, and return +/// the installation prefix. +pub fn setup_metatomic_torch_cmake(source_dir: &Path, build_dir: &Path, cmake_args: Vec) -> PathBuf { + std::fs::create_dir_all(build_dir).expect("failed to create metatomic build dir"); + + // configure cmake for metatomic-torch + let mut cmake_config = cmake_config(source_dir, build_dir); + + let install_prefix = build_dir.join("usr"); + cmake_config.arg(format!("-DCMAKE_INSTALL_PREFIX={}", install_prefix.display())); + + // Add any additional cmake arguments + for arg in cmake_args { + cmake_config.arg(arg); + } + + run_command(cmake_config, "cmake configuration for metatomic_torch"); + + // build and install metatomic-torch + let mut cmake_build = cmake_build(build_dir); + cmake_build.arg("--target"); + cmake_build.arg("install"); + + run_command(cmake_build, "cmake build for metatomic_torch"); + + install_prefix +} + +/// Build metatomic-core located in `source_dir` inside `build_dir`, and return +/// the installation prefix +pub fn setup_metatomic_cmake(source_dir: &Path, build_dir: &Path, cmake_args: Vec) -> PathBuf { + std::fs::create_dir_all(build_dir).expect("failed to create metatomic build dir"); + + // configure cmake for metatomic + let mut cmake_config = cmake_config(source_dir, build_dir); + + let install_prefix = build_dir.join("usr"); + cmake_config.arg(format!("-DCMAKE_INSTALL_PREFIX={}", install_prefix.display())); + + // Add any additional cmake arguments + for arg in cmake_args { + cmake_config.arg(arg); + } + + run_command(cmake_config, "cmake configuration for metatomic"); + + // build and install metatomic + let mut cmake_build = cmake_build(build_dir); + cmake_build.arg("--target"); + cmake_build.arg("install"); + + run_command(cmake_build, "cmake build for metatomic"); + + install_prefix +} + +/// Install metatomic-core in a Python virtualenv with pip, and return the +/// CMAKE_PREFIX_PATH for the installed libmetatomic. +pub fn setup_metatomic_core_pip(python: &Path, source_dir: &Path) -> PathBuf { + // build dependencies + pip_install( + python, + &["cmake", "packaging >=26", "setuptools >=77"], + PipInstallOptions::default() + ); + // runtime dependencies which are not just metatensor and metatensor-torch + pip_install(python, &["wigners"], PipInstallOptions::default()); + + pip_install( + python, + &[&source_dir.display().to_string()], + PipInstallOptions { + upgrade: true, + no_deps: true, + no_build_isolation: true + } + ); + + // let mut cmd = Command::new(python); + // cmd.arg("-c"); + // cmd.arg("import metatomic; print(metatomic.utils.cmake_prefix_path)"); + + // let output = run_command(cmd, "python to get metatomic cmake prefix"); + + // let stdout = String::from_utf8_lossy(&output.stdout); + // let prefix = PathBuf::from(stdout.trim()); + // if !prefix.exists() { + // panic!("'metatomic.utils.cmake_prefix' at '{}' does not exist", prefix.display()); + // } + + // return prefix; + return PathBuf::new(); +} + + +/// Install metatomic-torch in a Python virtualenv with pip, and return the +/// CMAKE_PREFIX_PATH for the installed libmetatomic_torch. +pub fn setup_metatomic_torch_pip(python: &Path, source_dir: &Path) -> PathBuf { + pip_install( + python, + &[&source_dir.display().to_string()], + PipInstallOptions { + upgrade: true, + no_deps: true, + no_build_isolation: true + } + ); + + let mut cmd = Command::new(python); + cmd.arg("-c"); + cmd.arg("import metatomic.torch; print(metatomic.torch.utils.cmake_prefix_path)"); + + let output = run_command(cmd, "python to get metatomic_torch cmake prefix"); + + let stdout = String::from_utf8_lossy(&output.stdout); + let prefix = PathBuf::from(stdout.trim()); + if !prefix.exists() { + panic!("'metatomic.torch.utils.cmake_prefix' at '{}' does not exist", prefix.display()); + } + + return prefix; +} + +pub fn run_command(mut command: Command, context: &str) -> std::process::Output { + write!(std::io::stdout().lock(), "\n\n[Running] {:?}\n\n", command).unwrap(); + + let mut child = command + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn().unwrap_or_else(|_| panic!("failed to spawn {}", context)); + + let mut child_stdout = child.stdout.take().expect("missing stdout"); + let mut child_stderr = child.stderr.take().expect("missing stderr"); + + let out_handle = std::thread::spawn(move || -> std::io::Result> { + let mut buf = [0u8; 8192]; + let mut captured = Vec::new(); + let mut sink = std::io::stdout().lock(); + loop { + let n = child_stdout.read(&mut buf)?; + if n == 0 { + break; + } + sink.write_all(&buf[..n])?; + sink.flush()?; + captured.extend_from_slice(&buf[..n]); + } + Ok(captured) + }); + + let err_handle = std::thread::spawn(move || -> std::io::Result> { + let mut buf = [0u8; 8192]; + let mut captured = Vec::new(); + let mut sink = std::io::stderr().lock(); + loop { + let n = child_stderr.read(&mut buf)?; + if n == 0 { + break; + } + sink.write_all(&buf[..n])?; + sink.flush()?; + captured.extend_from_slice(&buf[..n]); + } + Ok(captured) + }); + + let status = child.wait().unwrap_or_else(|_| panic!("failed to run {}", context)); + let stdout = String::from_utf8_lossy(&out_handle.join().unwrap().unwrap()).into_owned(); + let stderr = String::from_utf8_lossy(&err_handle.join().unwrap().unwrap()).into_owned(); + + if !status.success() { + panic!( + "{} failed, status: {}\nstderr:\n\n{}\nstdout:\n\n{}\n", + context, status, stderr, stdout + ); + } + + return std::process::Output { status, stdout: stdout.into_bytes(), stderr: stderr.into_bytes() }; +} diff --git a/metatomic-torch/Cargo.toml b/metatomic-torch/Cargo.toml new file mode 100644 index 000000000..6387cd4db --- /dev/null +++ b/metatomic-torch/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "metatomic-torch" +version = "0.0.0" +edition = "2024" +publish = false +rust-version = "1.88" + +[lib] +path = "lib.rs" + +[dev-dependencies] +lazy_static = "1" +which = "8" diff --git a/metatomic-torch/lib.rs b/metatomic-torch/lib.rs new file mode 100644 index 000000000..59bc69bb6 --- /dev/null +++ b/metatomic-torch/lib.rs @@ -0,0 +1 @@ +// empty lib.rs, this crate only exists to run TorchScript C++ tests with cargo diff --git a/metatomic-torch/tests/CMakeLists.txt b/metatomic-torch/tests/CMakeLists.txt index 89a3db0f2..7d6257a0d 100644 --- a/metatomic-torch/tests/CMakeLists.txt +++ b/metatomic-torch/tests/CMakeLists.txt @@ -1,4 +1,5 @@ -add_subdirectory(external) +# re-use catch from metatomic-core C++ tests +add_subdirectory(../../metatomic-core/tests/external external) # make sure we compile catch with the flags that torch requires. In particular, # torch sets -D_GLIBCXX_USE_CXX11_ABI=0 on Linux, which changes some of the @@ -14,9 +15,11 @@ if (VALGRIND) "--leak-check=full" "--show-leak-kinds=definite,indirect,possible" "--track-origins=yes" "--gen-suppressions=all" "--suppressions=${CMAKE_CURRENT_SOURCE_DIR}/valgrind.supp" ) + set(USING_VALGRIND ON) endif() else() set(TEST_COMMAND "") + set(USING_VALGRIND OFF) endif() @@ -46,7 +49,9 @@ foreach(_file_ ${ALL_TESTS}) ) # stop tests if they run for more than 30s - set_tests_properties(torch-${_name_} PROPERTIES TIMEOUT 30) + if (NOT USING_VALGRIND) + set_tests_properties(torch-${_name_} PROPERTIES TIMEOUT 30) + endif() if(WIN32) # We need to set the path to allow access to torch.dll diff --git a/metatomic-torch/tests/check-torch-install.rs b/metatomic-torch/tests/check-torch-install.rs new file mode 100644 index 000000000..14e85628a --- /dev/null +++ b/metatomic-torch/tests/check-torch-install.rs @@ -0,0 +1,216 @@ +use std::path::PathBuf; +use std::sync::Mutex; + +mod utils; + +lazy_static::lazy_static! { + // Make sure only one of the tests below run at the time, since they both + // try to modify the same files + static ref LOCK: Mutex<()> = Mutex::new(()); +} + +/// Check that metatomic-torch can be built and installed with cmake, and that +/// the installed version can be used from another cmake project with +/// `find_package` +#[test] +fn check_torch_install() { + let _guard = match LOCK.lock() { + Ok(guard) => guard, + Err(_) => { + panic!("another test failed, stopping") + } + }; + + const CARGO_TARGET_TMPDIR: &str = env!("CARGO_TARGET_TMPDIR"); + let cargo_manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + + let mut build_dir = PathBuf::from(CARGO_TARGET_TMPDIR); + build_dir.push("torch-install"); + build_dir.push("cmake-find-package"); + std::fs::create_dir_all(&build_dir).expect("failed to create build dir"); + + // ====================================================================== // + // install dependencies with pip + let deps_dir = build_dir.join("deps"); + + let torch_dep = deps_dir.join("virtualenv"); + std::fs::create_dir_all(&torch_dep).expect("failed to create virtualenv dir"); + let python = utils::create_python_venv(torch_dep); + let pytorch_cmake_prefix = utils::setup_torch_pip(&python); + let metatensor_cmake_prefix = utils::setup_metatensor_pip(&python); + let metatensor_torch_cmake_prefix = utils::setup_metatensor_torch_pip(&python); + + // ====================================================================== // + // build and install metatomic-torch with cmake + let metatomic_torch_dep = deps_dir.join("metatomic-torch"); + + let cmake_options = vec![ + format!( + "-DCMAKE_PREFIX_PATH={};{};{}", + pytorch_cmake_prefix.display(), + metatensor_cmake_prefix.display(), + metatensor_torch_cmake_prefix.display() + ), + // The two properties below handle the RPATH for metatomic_torch, + // setting it in such a way that we can always load libmetatensor.so and + // libtorch.so from the location they are found at when compiling + // metatomic-torch. See + // https://gitlab.kitware.com/cmake/community/-/wikis/doc/cmake/RPATH-handling + // for more information on CMake RPATH handling + "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON".into(), + "-DCMAKE_INSTALL_RPATH_USE_LINK_PATH=ON".into(), + ]; + + let install_prefix = utils::setup_metatomic_torch_cmake( + &cargo_manifest_dir, + &metatomic_torch_dep, + cmake_options, + ); + + // ====================================================================== // + // try to use the installed metatomic-torch from cmake + let mut source_dir = PathBuf::from(&cargo_manifest_dir); + source_dir.extend(["tests", "cmake-project"]); + + // configure cmake for the test cmake project + let mut cmake_config = utils::cmake_config(&source_dir, &build_dir); + cmake_config.arg(format!( + "-DCMAKE_PREFIX_PATH={};{};{};{}", + metatensor_cmake_prefix.display(), + pytorch_cmake_prefix.display(), + metatensor_torch_cmake_prefix.display(), + install_prefix.display(), + )); + + utils::run_command(cmake_config, "cmake configuration"); + + // build the code, linking to metatomic-torch + let cmake_build = utils::cmake_build(&build_dir); + utils::run_command(cmake_build, "cmake build"); + + // run the executables + let ctest = utils::ctest(&build_dir); + utils::run_command(ctest, "ctest"); +} + +/// Same as above, but using metatomic-torch from the Python wheel, +/// instead of building it from source with cmake. +#[test] +fn check_python_install() { + let _guard = match LOCK.lock() { + Ok(guard) => guard, + Err(_) => { + panic!("another test failed, stopping") + } + }; + + const CARGO_TARGET_TMPDIR: &str = env!("CARGO_TARGET_TMPDIR"); + + let mut build_dir = PathBuf::from(CARGO_TARGET_TMPDIR); + build_dir.push("torch-install"); + build_dir.push("python-wheels"); + std::fs::create_dir_all(&build_dir).expect("failed to create build dir"); + + // ====================================================================== // + // install dependencies with pip + let mut venv_dir = build_dir.clone(); + venv_dir.push("virtualenv"); + + let python_exe = utils::create_python_venv(venv_dir); + + let cargo_manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + let pytorch_cmake_prefix = utils::setup_torch_pip(&python_exe); + let metatensor_cmake_prefix = utils::setup_metatensor_pip(&python_exe); + let metatensor_torch_cmake_prefix = utils::setup_metatensor_torch_pip(&python_exe); + + // ====================================================================== // + // build and install metatomic and metatomic-torch with pip + let mta_core_source_dir = cargo_manifest_dir.parent().unwrap().join("python").join("metatomic_core"); + let metatomic_core_cmake_prefix = utils::setup_metatomic_core_pip(&python_exe, &mta_core_source_dir); + + let mta_torch_source_dir = cargo_manifest_dir.parent().unwrap().join("python").join("metatomic_torch"); + let metatomic_torch_cmake_prefix = utils::setup_metatomic_torch_pip(&python_exe, &mta_torch_source_dir); + + // ====================================================================== // + // try to use the installed metatomic-torch from cmake + let mut source_dir = PathBuf::from(&cargo_manifest_dir); + source_dir.extend(["tests", "cmake-project"]); + + // configure cmake for the test cmake project + let mut cmake_config = utils::cmake_config(&source_dir, &build_dir); + cmake_config.arg(format!( + "-DCMAKE_PREFIX_PATH={};{};{};{};{}", + pytorch_cmake_prefix.display(), + metatensor_cmake_prefix.display(), + metatensor_torch_cmake_prefix.display(), + metatomic_core_cmake_prefix.display(), + metatomic_torch_cmake_prefix.display(), + )); + + utils::run_command(cmake_config, "cmake configuration"); + + // build the code, linking to metatomic-torch + let cmake_build = utils::cmake_build(&build_dir); + utils::run_command(cmake_build, "cmake build"); + + // run the executables + let ctest = utils::ctest(&build_dir); + utils::run_command(ctest, "ctest"); +} + +/// Same test as above, but building metatomic-torch in the same +/// CMake project (i.e. using add_subdirectory instead of find_package) +#[test] +fn check_cmake_subdirectory() { + let _guard = match LOCK.lock() { + Ok(guard) => guard, + Err(_) => { + panic!("another test failed, stopping") + } + }; + + const CARGO_TARGET_TMPDIR: &str = env!("CARGO_TARGET_TMPDIR"); + + // install torch + let mut build_dir = PathBuf::from(CARGO_TARGET_TMPDIR); + build_dir.push("torch-install"); + build_dir.push("cmake-subdirectory"); + std::fs::create_dir_all(&build_dir).expect("failed to create build dir"); + + // ====================================================================== // + // install dependencies with pip + let deps_dir = build_dir.join("deps"); + + let virtualenv_dir = deps_dir.join("virtualenv"); + std::fs::create_dir_all(&virtualenv_dir).expect("failed to create virtualenv dir"); + let python = utils::create_python_venv(virtualenv_dir); + let pytorch_cmake_prefix = utils::setup_torch_pip(&python); + let metatensor_cmake_prefix = utils::setup_metatensor_pip(&python); + let metatensor_torch_cmake_prefix = utils::setup_metatensor_torch_pip(&python); + + // ====================================================================== // + // build metatomic-torch with cmake, using add_subdirectory + let cargo_manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + let mut source_dir = PathBuf::from(&cargo_manifest_dir); + source_dir.extend(["tests", "cmake-project"]); + + // configure cmake for the test cmake project + let mut cmake_config = utils::cmake_config(&source_dir, &build_dir); + cmake_config.arg(format!( + "-DCMAKE_PREFIX_PATH={};{};{}", + pytorch_cmake_prefix.display(), + metatensor_cmake_prefix.display(), + metatensor_torch_cmake_prefix.display() + )); + cmake_config.arg("-DUSE_CMAKE_SUBDIRECTORY=ON"); + + utils::run_command(cmake_config, "cmake configuration"); + + // build the code, linking to metatomic-torch + let cmake_build = utils::cmake_build(&build_dir); + utils::run_command(cmake_build, "cmake build"); + + // run the executables + let ctest = utils::ctest(&build_dir); + utils::run_command(ctest, "ctest"); +} diff --git a/metatomic-torch/tests/run-torch-tests.rs b/metatomic-torch/tests/run-torch-tests.rs new file mode 100644 index 000000000..93772f0a6 --- /dev/null +++ b/metatomic-torch/tests/run-torch-tests.rs @@ -0,0 +1,47 @@ +use std::path::PathBuf; + +mod utils; + +#[test] +fn run_torch_tests() { + const CARGO_TARGET_TMPDIR: &str = env!("CARGO_TARGET_TMPDIR"); + let cargo_manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + + // ====================================================================== // + // setup dependencies for the torch tests + + let mut build_dir = PathBuf::from(CARGO_TARGET_TMPDIR); + build_dir.push("torch-tests"); + let deps_dir = build_dir.join("deps"); + + let torch_dep = deps_dir.join("virtualenv"); + std::fs::create_dir_all(&torch_dep).expect("failed to create virtualenv dir"); + let python_exe = utils::create_python_venv(torch_dep); + let pytorch_cmake_prefix = utils::setup_torch_pip(&python_exe); + let metatensor_cmake_prefix = utils::setup_metatensor_pip(&python_exe); + let metatensor_torch_cmake_prefix = utils::setup_metatensor_torch_pip(&python_exe); + + // ====================================================================== // + // build the metatomic-torch C++ tests and run them + let source_dir = cargo_manifest_dir; + + // configure cmake for the tests + let mut cmake_config = utils::cmake_config(&source_dir, &build_dir); + cmake_config.arg("-DMETATOMIC_TORCH_TESTS=ON"); + cmake_config.arg(format!( + "-DCMAKE_PREFIX_PATH={};{};{}", + pytorch_cmake_prefix.display(), + metatensor_cmake_prefix.display(), + metatensor_torch_cmake_prefix.display() + )); + + utils::run_command(cmake_config, "cmake configuration"); + + // build the tests + let cmake_build = utils::cmake_build(&build_dir); + utils::run_command(cmake_build, "cmake build"); + + // run the tests + let ctest = utils::ctest(&build_dir); + utils::run_command(ctest, "ctest"); +} diff --git a/metatomic-torch/tests/utils/mod.rs b/metatomic-torch/tests/utils/mod.rs new file mode 120000 index 000000000..20b8b0094 --- /dev/null +++ b/metatomic-torch/tests/utils/mod.rs @@ -0,0 +1 @@ +../../../metatomic-core/tests/utils/mod.rs \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 88dc392b9..9deadf6c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ name = "metatomic" version = "0.1.0" dynamic = ["authors", "optional-dependencies"] -requires-python = ">=3.10" +requires-python = ">=3.11" readme = "README.md" license = "BSD-3-Clause" @@ -46,7 +46,10 @@ packages = [] [tool.coverage.paths] # path remapping for coverage. `coverage` will map path matching the second # string to the first string. -torch = ["python/metatensor_torch/", ".tox/*/lib/python*/site-packages/"] +core = ["python/metatomic_core/src/", ".tox/*/lib/python*/site-packages/"] +torch = ["python/metatomic_torch/src/", ".tox/*/lib/python*/site-packages/"] +ase = ["python/metatomic_ase/src/", ".tox/*/lib/python*/site-packages/"] +torchsim = ["python/metatomic_torchsim/src/", ".tox/*/lib/python*/site-packages/"] [tool.coverage.report] show_missing = true @@ -63,13 +66,19 @@ filterwarnings = [ "ignore:ast.NameConstant is deprecated and will be removed in Python 3.14:DeprecationWarning", # TorchScript deprecation warnings "ignore:`torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`:DeprecationWarning", + "ignore:`torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`:DeprecationWarning", "ignore:`torch.jit.save` is deprecated. Please switch to `torch.export`:DeprecationWarning", - "ignore:.*vesin.metatomic was only tested with metatomic.torch >=0.1.3,<0.2.*:UserWarning", "ignore:`torch.jit.load` is deprecated. Please switch to `torch.export`.:DeprecationWarning", "ignore:`torch.jit.script` is not supported in Python 3.14+:DeprecationWarning", + "ignore:`torch.jit.script_method` is not supported in Python 3.14+:DeprecationWarning", "ignore:`torch.jit.save` is not supported in Python 3.14+:DeprecationWarning", - # deprecation warning from warp/nvalchemi + # vesin and metatomic warning + "ignore:.*vesin.metatomic was only tested with metatomic.torch >=0.1.3,<0.2.*:UserWarning", + # Warnings from warp (dependency of nvalchemi) + "ignore:.*Structure will use memory layout compatible with MSVC:DeprecationWarning", "ignore:warp.config.quiet is deprecated:DeprecationWarning", + # Warning in old version of torchsim + "ignore:`vesin.torch` is deprecated and will be removed in a future release:DeprecationWarning" ] ### ======================================================================== ### @@ -95,6 +104,8 @@ docstring-code-format = true [tool.uv.pip] reinstall-package = [ - "metatomic-torch", - "metatomic-torchsim", + "metatomic_core", + "metatomic_torch", + "metatomic_torchsim", + "metatomic_ase", ] diff --git a/python/Cargo.toml b/python/Cargo.toml new file mode 100644 index 000000000..3546f0179 --- /dev/null +++ b/python/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "metatomic-python" +version = "0.0.0" +edition = "2024" +publish = false +rust-version = "1.88" + +[lib] +path = "lib.rs" + +[dev-dependencies] +which = "8" diff --git a/python/examples/README.rst b/python/examples/README.rst deleted file mode 100644 index 318815652..000000000 --- a/python/examples/README.rst +++ /dev/null @@ -1,4 +0,0 @@ -.. _atomistic-tutorials: - -Tutorials -========= diff --git a/python/lib.rs b/python/lib.rs new file mode 100644 index 000000000..5ef74bad8 --- /dev/null +++ b/python/lib.rs @@ -0,0 +1 @@ +// empty lib.rs, this crate only exists to run Python tests with cargo diff --git a/python/metatomic_ase/pyproject.toml b/python/metatomic_ase/pyproject.toml index 184462e69..1e83306b4 100644 --- a/python/metatomic_ase/pyproject.toml +++ b/python/metatomic_ase/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "metatomic-ase" dynamic = ["version", "authors", "dependencies"] -requires-python = ">=3.10" +requires-python = ">=3.11" readme = "README.md" license = "BSD-3-Clause" diff --git a/python/metatomic_ase/setup.py b/python/metatomic_ase/setup.py index c9123f408..9c4b1f606 100644 --- a/python/metatomic_ase/setup.py +++ b/python/metatomic_ase/setup.py @@ -1,4 +1,5 @@ import os +import pathlib import subprocess import sys @@ -8,8 +9,8 @@ from setuptools.command.sdist import sdist -ROOT = os.path.realpath(os.path.dirname(__file__)) -METATOMIC_TORCH = os.path.realpath(os.path.join(ROOT, "..", "metatomic_torch")) +ROOT = pathlib.Path(__file__).parent.resolve() +METATOMIC_TORCH = (ROOT / ".." / "metatomic_torch").resolve() METATOMIC_ASE_VERSION = "0.1.3" @@ -53,15 +54,15 @@ def git_version_info(): """ TAG_PREFIX = "metatomic-ase-v" - if os.path.exists("git_version_info"): + if (ROOT / "git_version_info").exists(): # we are building from a sdist, without git available, but the git # version was recorded in the `git_version_info` file - with open("git_version_info") as fd: + with open(ROOT / "git_version_info") as fd: n_commits = int(fd.readline().strip()) git_hash = fd.readline().strip() else: - script = os.path.join(ROOT, "..", "..", "scripts", "git-version-info.py") - assert os.path.exists(script) + script = (ROOT / ".." / ".." / "scripts" / "git-version-info.py").resolve() + assert script.exists() output = subprocess.run( [sys.executable, script, TAG_PREFIX], @@ -76,12 +77,15 @@ def git_version_info(): f"stdout: {output.stdout}\n" f"stderr: {output.stderr}\n" ) - elif output.stderr: + if output.stderr: print(output.stderr, file=sys.stderr) + + lines = output.stdout.splitlines() + if len(lines) < 2: + # the script gave up early (cf. `warn_and_exit`) n_commits = 0 git_hash = "" else: - lines = output.stdout.splitlines() n_commits = int(lines[0].strip()) git_hash = lines[1].strip() @@ -127,19 +131,19 @@ def create_version_number(version): # when packaging a sdist for release, we should never use local dependencies METATOMIC_NO_LOCAL_DEPS = os.environ.get("METATOMIC_NO_LOCAL_DEPS", "0") == "1" - if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_TORCH): + if not METATOMIC_NO_LOCAL_DEPS and METATOMIC_TORCH.exists(): # we are building from a git checkout or full repo archive - install_requires.append(f"metatomic-torch @ file://{METATOMIC_TORCH}") + install_requires.append(f"metatomic-torch @ {METATOMIC_TORCH.as_uri()}") else: # we are building from a sdist/installing from a wheel install_requires.append("metatomic-torch >=0.1.12,<0.2") - with open(os.path.join(ROOT, "AUTHORS")) as fd: + with open(ROOT / "AUTHORS") as fd: authors = fd.read().splitlines() if authors[0].startswith(".."): # handle "raw" symlink files (on Windows or from full repo tarball) - with open(os.path.join(ROOT, authors[0])) as fd: + with open(ROOT / authors[0]) as fd: authors = fd.read().splitlines() setup( diff --git a/python/metatomic_core/AUTHORS b/python/metatomic_core/AUTHORS new file mode 120000 index 000000000..f04b7e8a2 --- /dev/null +++ b/python/metatomic_core/AUTHORS @@ -0,0 +1 @@ +../../AUTHORS \ No newline at end of file diff --git a/python/metatomic_core/CMakeLists.txt b/python/metatomic_core/CMakeLists.txt new file mode 100644 index 000000000..02b9f407d --- /dev/null +++ b/python/metatomic_core/CMakeLists.txt @@ -0,0 +1,62 @@ +# This file allow the python module in metatomic-core to either use an +# externally-provided version of the shared metatomic library; or to build the +# code from source and bundle the shared library inside the wheel. +# +# The first case is used when distributing the code in conda (since we have a +# separate libmetatomic package), the second one is used everywhere else (for +# local development builds and for the PyPI distribution). + +cmake_minimum_required(VERSION 3.22) +project(metatomic-python NONE) + +option(METATOMIC_CORE_PYTHON_USE_EXTERNAL_LIB "Force the usage of an external version of metatomic-core" OFF) +set(METATOMIC_CORE_SOURCE_DIR "" CACHE PATH "Path to the sources of metatomic-core") + +file(REMOVE ${CMAKE_INSTALL_PREFIX}/_external.py) + +set(REQUIRED_METATOMIC_VERSION "0.1.0") +if(${METATOMIC_CORE_PYTHON_USE_EXTERNAL_LIB}) + # when building a source checkout, update version to include git information + # this will not apply when building a sdist + if (EXISTS ${CMAKE_SOURCE_DIR}/../../metatomic-core/cmake/dev-versions.cmake) + include(${CMAKE_SOURCE_DIR}/../../metatomic-core/cmake/dev-versions.cmake) + create_development_version("${REQUIRED_METATOMIC_VERSION}" REQUIRED_METATOMIC_VERSION "metatomic-core-v") + # strip any -dev/-rc suffix on the version since find_package does not support it + string(REGEX REPLACE "([0-9]*)\\.([0-9]*)\\.([0-9]*).*" "\\1.\\2.\\3" REQUIRED_METATOMIC_VERSION ${REQUIRED_METATOMIC_VERSION}) + endif() + + find_package(metatomic ${REQUIRED_METATOMIC_VERSION} REQUIRED) + + get_target_property(METATOMIC_LOCATION metatomic::shared LOCATION) + message(STATUS "Using external metatomic-core v${metatomic_VERSION} at ${METATOMIC_LOCATION}") + + # Get the prefix to use as cmake_prefix_path when trying to load this + # version of the library again + get_filename_component(METATOMIC_PREFIX "${METATOMIC_LOCATION}" DIRECTORY) + get_filename_component(METATOMIC_PREFIX "${METATOMIC_PREFIX}" DIRECTORY) + + file(WRITE ${CMAKE_INSTALL_PREFIX}/_external.py + "EXTERNAL_METATOMIC_PATH = \"${METATOMIC_LOCATION}\"\n\n" + ) + file(APPEND ${CMAKE_INSTALL_PREFIX}/_external.py + "EXTERNAL_METATOMIC_PREFIX = \"${METATOMIC_PREFIX}\"\n" + ) + + install(CODE "message(STATUS \"nothing to install\")") +else() + if ("${METATOMIC_CORE_SOURCE_DIR}" STREQUAL "") + message(FATAL_ERROR + "Missing METATOMIC_CORE_SOURCE_DIR, please specify where to \ + find the source code for metatomic-core" + ) + endif() + + message(STATUS "Using internal metatomic-core from ${METATOMIC_CORE_SOURCE_DIR}") + + set(BUILD_SHARED_LIBS ON) + set(METATOMIC_INSTALL_BOTH_STATIC_SHARED OFF) + # strip dynamic library for smaller wheels to download/install + set(EXTRA_RUST_FLAGS "-Cstrip=symbols") + + add_subdirectory("${METATOMIC_CORE_SOURCE_DIR}" metatomic-core) +endif() diff --git a/python/metatomic_core/LICENSE b/python/metatomic_core/LICENSE new file mode 120000 index 000000000..30cff7403 --- /dev/null +++ b/python/metatomic_core/LICENSE @@ -0,0 +1 @@ +../../LICENSE \ No newline at end of file diff --git a/python/metatomic_core/MANIFEST.in b/python/metatomic_core/MANIFEST.in new file mode 100644 index 000000000..17a63224a --- /dev/null +++ b/python/metatomic_core/MANIFEST.in @@ -0,0 +1,8 @@ +include pyproject.toml +include CMakeLists.txt +include AUTHORS +include LICENSE + +include git_version_info + +include metatomic-core-cxx-*.tar.gz diff --git a/python/metatomic_core/pyproject.toml b/python/metatomic_core/pyproject.toml new file mode 100644 index 000000000..6f6725d58 --- /dev/null +++ b/python/metatomic_core/pyproject.toml @@ -0,0 +1,54 @@ +[project] +name = "metatomic-core" +dynamic = ["version", "authors", "dependencies"] +requires-python = ">=3.11" + +# readme = "TODO" +license = "BSD-3-Clause" +description = "Interface between atomistic machine learning models and simulation tools" + +keywords = ["machine learning", "molecular modeling"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "Operating System :: POSIX", + "Operating System :: MacOS :: MacOS X", + "Operating System :: Microsoft :: Windows", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Bio-Informatics", + "Topic :: Scientific/Engineering :: Chemistry", + "Topic :: Scientific/Engineering :: Physics", + "Topic :: Software Development :: Libraries", + "Topic :: Software Development :: Libraries :: Python Modules", +] + +[project.urls] +homepage = "https://docs.metatensor.org/metatomic/" +documentation = "https://docs.metatensor.org/metatomic/" +repository = "https://github.com/metatensor/metatomic" +# changelog = "TODO" + +### ======================================================================== ### +[build-system] +requires = [ + "setuptools >=77", + "packaging >=26", + "cmake", + "metatensor-core >=0.2.4,<0.3", +] + +build-backend = "setuptools.build_meta" + + +[tool.setuptools] +zip-safe = false + +### ======================================================================== ### +[tool.pytest.ini_options] +python_files = ["*.py"] +testpaths = ["tests"] +filterwarnings = [ + "error", +] diff --git a/python/metatomic_core/setup.py b/python/metatomic_core/setup.py new file mode 100644 index 000000000..09fba7da8 --- /dev/null +++ b/python/metatomic_core/setup.py @@ -0,0 +1,326 @@ +import glob +import os +import pathlib +import subprocess +import sys +import tomllib + +import packaging.version +from setuptools import Extension, setup +from setuptools.command.bdist_egg import bdist_egg +from setuptools.command.bdist_wheel import bdist_wheel +from setuptools.command.build_ext import build_ext +from setuptools.command.sdist import sdist + + +ROOT = pathlib.Path(__file__).parent.resolve() + +METATOMIC_BUILD_TYPE = os.environ.get("METATOMIC_BUILD_TYPE", "release") +if METATOMIC_BUILD_TYPE not in ["debug", "release"]: + raise Exception( + f"invalid build type passed: '{METATOMIC_BUILD_TYPE}', " + "expected 'debug' or 'release'" + ) + +# the root path to the _native_ source code of metatomic (Rust source, but built with +# cmake) +METATOMIC_CORE_SRC = os.path.join(ROOT, "..", "..", "metatomic-core") + + +class universal_wheel(bdist_wheel): + # When building the wheel, the `wheel` package assumes that if we have a + # binary extension then we are linking to `libpython.so`; and thus the wheel + # is only usable with a single python version. This is not the case for + # here, and the wheel will be compatible with any Python >=3. This is + # tracked in https://github.com/pypa/wheel/issues/185, but until then we + # manually override the wheel tag. + def get_tag(self): + tag = bdist_wheel.get_tag(self) + # tag[2:] contains the os/arch tags, we want to keep them + return ("py3", "none") + tag[2:] + + +class cmake_ext(build_ext): + """ + Build the native library using cmake + """ + + def finalize_options(self): + if self.editable_mode: + raise RuntimeError( + "metatomic-core does not support editable installation yet" + ) + return super().finalize_options() + + def run(self): + import metatensor + + source_dir = ROOT + build_dir = os.path.join(ROOT, "build", "cmake-build") + install_dir = os.path.join(os.path.realpath(self.build_lib), "metatomic") + + os.makedirs(build_dir, exist_ok=True) + + use_external_lib = os.environ.get( + "METATOMIC_CORE_PYTHON_USE_EXTERNAL_LIB", "OFF" + ) + + cmake_options = [ + "-DCMAKE_VERBOSE_MAKEFILE=ON", + f"-DCMAKE_INSTALL_PREFIX={install_dir}", + f"-DMETATOMIC_CORE_SOURCE_DIR={METATOMIC_CORE_SRC}", + "-DCMAKE_INSTALL_LIBDIR=lib", + f"-DCMAKE_BUILD_TYPE={METATOMIC_BUILD_TYPE}", + f"-DMETATOMIC_CORE_PYTHON_USE_EXTERNAL_LIB={use_external_lib}", + f"-DCMAKE_PREFIX_PATH={metatensor.utils.cmake_prefix_path}", + ] + + if "CARGO" in os.environ: + cmake_options.append(f"-DCARGO_EXE={os.environ['CARGO']}") + + # Handle cross-compilation by detecting cibuildwheels environnement + # variables + if sys.platform.startswith("darwin"): + # ARCHFLAGS is set by cibuildwheels + ARCHFLAGS = os.environ.get("ARCHFLAGS") + if ARCHFLAGS is not None: + archs = filter( + lambda u: bool(u), + ARCHFLAGS.strip().split("-arch "), + ) + archs = list(archs) + assert len(archs) == 1 + arch = archs[0].strip() + + if arch == "x86_64": + cmake_options.append("-DRUST_BUILD_TARGET=x86_64-apple-darwin") + elif arch == "arm64": + cmake_options.append("-DRUST_BUILD_TARGET=aarch64-apple-darwin") + else: + raise ValueError(f"unknown arch: {arch}") + + elif sys.platform.startswith("linux"): + # we set RUST_BUILD_TARGET in our custom docker image + RUST_BUILD_TARGET = os.environ.get("RUST_BUILD_TARGET") + if RUST_BUILD_TARGET is not None: + cmake_options.append(f"-DRUST_BUILD_TARGET={RUST_BUILD_TARGET}") + + elif sys.platform.startswith("win32"): + # CARGO_BUILD_TARGET is set by cibuildwheels + CARGO_BUILD_TARGET = os.environ.get("CARGO_BUILD_TARGET") + if CARGO_BUILD_TARGET is not None: + cmake_options.append(f"-DRUST_BUILD_TARGET={CARGO_BUILD_TARGET}") + + else: + raise ValueError(f"unknown platform: {sys.platform}") + + subprocess.run( + ["cmake", source_dir, *cmake_options], + cwd=build_dir, + check=True, + ) + subprocess.run( + ["cmake", "--build", build_dir, "--parallel", "--target", "install"], + check=True, + ) + + +class bdist_egg_disabled(bdist_egg): + """Disabled version of bdist_egg + + Prevents setup.py install performing setuptools' default easy_install, + which it should never ever do. + """ + + def run(self): + sys.exit( + "Aborting implicit building of eggs.\nUse `pip install .` or " + "`python -m build --wheel . && pip install dist/metatomic_torch-*.whl` " + "to install from source." + ) + + +class sdist_generate_data(sdist): + """ + Create a sdist with an additional generated files: + - `git_version_info` + """ + + def run(self): + n_commits, git_hash = git_version_info() + with open("git_version_info", "w") as fd: + fd.write(f"{n_commits}\n{git_hash}\n") + + generate_cxx_tar() + + # run original sdist + super().run() + + os.unlink("git_version_info") + for path in glob.glob("metatomic-core-cxx-*.tar.gz"): + os.unlink(path) + + +def generate_cxx_tar(): + script = os.path.join(ROOT, "..", "..", "scripts", "package-core.sh") + assert os.path.exists(script) + + try: + output = subprocess.run( + ["bash", "--version"], + stderr=subprocess.PIPE, + stdout=subprocess.PIPE, + encoding="utf8", + ) + except Exception as e: + raise RuntimeError("could not run `bash`, is it installed?") from e + + output = subprocess.run( + ["bash", script, os.getcwd()], + stderr=subprocess.PIPE, + stdout=subprocess.PIPE, + encoding="utf8", + ) + if output.returncode != 0: + stderr = output.stderr + stdout = output.stdout + raise RuntimeError( + "failed to collect C++ sources for Python sdist\n" + f"stdout:\n {stdout}\n\nstderr:\n {stderr}" + ) + + +def get_rust_version(): + # read version from Cargo.toml + with open(os.path.join(METATOMIC_CORE_SRC, "Cargo.toml"), "rb") as fd: + cargo_toml = tomllib.load(fd) + return cargo_toml["package"]["version"] + + +def git_version_info(): + """ + If git is available and we are building from a checkout, get the number of commits + since the last tag & full hash of the code. Otherwise, this always returns (0, ""). + """ + TAG_PREFIX = "metatomic-core-v" + + if (ROOT / "git_version_info").exists(): + # we are building from a sdist, without git available, but the git + # version was recorded in the `git_version_info` file + with open(ROOT / "git_version_info") as fd: + n_commits = int(fd.readline().strip()) + git_hash = fd.readline().strip() + else: + script = (ROOT / ".." / ".." / "scripts" / "git-version-info.py").resolve() + assert script.exists() + + output = subprocess.run( + [sys.executable, script, TAG_PREFIX], + stderr=subprocess.PIPE, + stdout=subprocess.PIPE, + encoding="utf8", + ) + + if output.returncode != 0: + raise Exception( + "failed to get git version info.\n" + f"stdout: {output.stdout}\n" + f"stderr: {output.stderr}\n" + ) + if output.stderr: + print(output.stderr, file=sys.stderr) + + lines = output.stdout.splitlines() + if len(lines) < 2: + # the script gave up early (cf. `warn_and_exit`), and only printed + # the number of commits + n_commits = 0 + git_hash = "" + else: + n_commits = int(lines[0].strip()) + git_hash = lines[1].strip() + + return n_commits, git_hash + + +def create_version_number(version): + version = packaging.version.parse(version) + + n_commits, git_hash = git_version_info() + + if n_commits != 0: + # if we have commits since the last tag, this mean we are in a pre-release of + # the next version. So we increase either the minor version number or the + # release candidate number (if we are closing up on a release) + if version.pre is not None: + assert version.pre[0] == "rc" + pre = ("rc", version.pre[1] + 1) + release = version.release + else: + major, minor, _ = version.release + release = (major, minor + 1, 0) + pre = None + + version = version.__replace__( + release=release, + pre=pre, + dev=n_commits, + local=git_hash, + ) + + return str(version) + + +if __name__ == "__main__": + if not os.path.exists(METATOMIC_CORE_SRC): + # we are building from a sdist, which should include metatomic-core Rust + # sources as a tarball + tarballs = glob.glob(os.path.join(ROOT, "metatomic-core-*.tar.gz")) + + if not len(tarballs) == 1: + raise RuntimeError( + "expected a single 'metatomic-core-*.tar.gz' file containing " + "metatomic-core Rust sources. remove all files and re-run " + "scripts/package-core.sh" + ) + + METATOMIC_CORE_SRC = os.path.realpath(tarballs[0]) + subprocess.run( + ["cmake", "-E", "tar", "xf", METATOMIC_CORE_SRC], + cwd=ROOT, + check=True, + ) + + METATOMIC_CORE_SRC = ".".join(METATOMIC_CORE_SRC.split(".")[:-2]) + + with open(ROOT / "AUTHORS") as fd: + authors = fd.read().splitlines() + + if authors[0].startswith(".."): + # handle "raw" symlink files (on Windows or from full repo tarball) + with open(ROOT / authors[0]) as fd: + authors = fd.read().splitlines() + + install_requires = [ + "metatensor-core >=0.2.2,<0.3", + ] + + setup( + version=create_version_number(get_rust_version()), + author=", ".join(authors), + ext_modules=[Extension(name="metatomic", sources=[])], + install_requires=install_requires, + cmdclass={ + "build_ext": cmake_ext, + "bdist_egg": bdist_egg if "bdist_egg" in sys.argv else bdist_egg_disabled, + "bdist_wheel": universal_wheel, + "sdist": sdist_generate_data, + }, + package_data={ + "metatomic-core": [ + "metatomic/lib/*", + "metatomic/include/*", + "metatomic/include/metatomic/*", + ] + }, + ) diff --git a/python/metatomic_core/src/metatomic/__init__.py b/python/metatomic_core/src/metatomic/__init__.py new file mode 100644 index 000000000..4b345959f --- /dev/null +++ b/python/metatomic_core/src/metatomic/__init__.py @@ -0,0 +1,16 @@ +from . import utils # noqa: F401 +from ._capabilities import ModelCapabilities +from ._metadata import ModelMetadata, References +from ._quantity import Quantity +from ._status import MetatomicError +from ._system import PairListOptions +from ._version import __version__ # noqa: F401 + + +# pretend the classes are defined in the top-level module for better error messages +MetatomicError.__module__ = __name__ +ModelCapabilities.__module__ = __name__ +ModelMetadata.__module__ = __name__ +PairListOptions.__module__ = __name__ +Quantity.__module__ = __name__ +References.__module__ = __name__ diff --git a/python/metatomic_core/src/metatomic/_c_api.py b/python/metatomic_core/src/metatomic/_c_api.py new file mode 100644 index 000000000..f3aa94053 --- /dev/null +++ b/python/metatomic_core/src/metatomic/_c_api.py @@ -0,0 +1,307 @@ +# fmt: off +# flake8: noqa +""" +This file declares the C-API corresponding to metatomic.h, in a way compatible +with the ctypes Python module. + +This file is automatically generated by `scripts/update-declarations.py`, +do not edit it manually! +""" + +import ctypes +import platform +from ctypes import CFUNCTYPE, POINTER + +from ctypes_dlpack import DLDataType, DLDevice, DLManagedTensorVersioned, DLPackVersion +from metatensor._c_api import ( + mts_labels_t, + mts_block_t, + mts_tensormap_t, + mts_realloc_buffer_t, + mts_create_array_callback_t, +) + + +class _EnumType(type(ctypes.c_int32)): + def __new__(metacls, name, bases, namespace): + if "_members_" not in namespace: + members = {} + for key, value in namespace.items(): + if not key.startswith("_"): + members[key] = value + namespace["_members_"] = members + else: + members = namespace["_members_"] + + namespace["_reverse_map_"] = {v: k for k, v in members.items()} + return type(ctypes.c_int32).__new__(metacls, name, bases, namespace) + + def __repr__(self): + return f"" + + +class _Enum(ctypes.c_int32, metaclass=_EnumType): + _members_ = {} + + def __repr__(self): + value_name = self._reverse_map_.get(self.value, str(self.value)) + return f"{self.__class__.__name__}.{value_name}" + + def __eq__(self, other): + if isinstance(other, int): + return self.value == other + if type(self) is type(other): + return self.value == other.value + return NotImplemented + + def __hash__(self): + return hash(self.value) + + +arch = platform.architecture()[0] +if arch == "32bit": + c_uintptr_t = ctypes.c_uint32 +elif arch == "64bit": + c_uintptr_t = ctypes.c_uint64 + + + +class mta_status_t(_Enum): + MTA_SUCCESS = 0 + MTA_INVALID_PARAMETER_ERROR = 1 + MTA_IO_ERROR = 2 + MTA_MEMORY_ERROR = 3 + MTA_SERIALIZATION_ERROR = 4 + MTA_DLPACK_ERROR = 5 + MTA_METATENSOR_ERROR = 6 + MTA_UNSUPPORTED_MODEL_ERROR = 7 + MTA_MODEL_ERROR = 8 + MTA_INTERNAL_ERROR = 255 + + +class mta_system_data_kind(_Enum): + MTA_SYSTEM_DATA_TYPES = 0 + MTA_SYSTEM_DATA_POSITIONS = 1 + MTA_SYSTEM_DATA_CELL = 2 + MTA_SYSTEM_DATA_PBC = 3 + + +class mta_opaque_string_t(ctypes.Structure): + pass + + +class mta_system_t(ctypes.Structure): + pass + + +class mta_model_t(ctypes.Structure): + pass + + +class mta_plugin_t(ctypes.Structure): + pass + + +mta_string_t = POINTER(mta_opaque_string_t) + + +mta_model_t._fields_ = [ + ("data", ctypes.c_void_p), + ("unload", CFUNCTYPE(mta_status_t, ctypes.c_void_p)), + ("capabilities", CFUNCTYPE(mta_status_t, ctypes.c_void_p, POINTER(mta_string_t))), + ("metadata", CFUNCTYPE(mta_status_t, ctypes.c_void_p, POINTER(mta_string_t))), + ("requested_pair_lists", CFUNCTYPE(mta_status_t, ctypes.c_void_p, POINTER(mta_string_t))), + ("requested_inputs", CFUNCTYPE(mta_status_t, ctypes.c_void_p, POINTER(mta_string_t))), + ("execute_inner", CFUNCTYPE(mta_status_t, ctypes.c_void_p, POINTER(POINTER(mta_system_t)), c_uintptr_t, POINTER(mts_labels_t), ctypes.c_char_p, POINTER(POINTER(mts_tensormap_t)), c_uintptr_t)), +] + +mta_plugin_t._fields_ = [ + ("abi_version", ctypes.c_int32), + ("name", ctypes.c_char_p), + ("load_model", CFUNCTYPE(mta_status_t, ctypes.c_char_p, ctypes.c_char_p, POINTER(mta_model_t))), +] + + +def setup_functions(lib): + from ._status import check_status + + lib.mta_last_error.argtypes = [ + POINTER(ctypes.c_char_p), + POINTER(ctypes.c_char_p), + POINTER(POINTER(None)), + ] + lib.mta_last_error.restype = mta_status_t + + lib.mta_set_last_error.argtypes = [ + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_void_p, + CFUNCTYPE(None, ctypes.c_void_p), + ] + lib.mta_set_last_error.restype = check_status + + lib.mta_version.argtypes = [ + ] + lib.mta_version.restype = ctypes.c_char_p + + lib.mta_string_create.argtypes = [ + ctypes.c_char_p, + ] + lib.mta_string_create.restype = mta_string_t + + lib.mta_string_free.argtypes = [ + mta_string_t, + ] + lib.mta_string_free.restype = None + + lib.mta_string_view.argtypes = [ + mta_string_t, + ] + lib.mta_string_view.restype = ctypes.c_char_p + + lib.mta_unit_conversion_factor.argtypes = [ + ctypes.c_char_p, + ctypes.c_char_p, + POINTER(ctypes.c_double), + ] + lib.mta_unit_conversion_factor.restype = check_status + + lib.mta_system_create.argtypes = [ + ctypes.c_char_p, + POINTER(DLManagedTensorVersioned), + POINTER(DLManagedTensorVersioned), + POINTER(DLManagedTensorVersioned), + POINTER(DLManagedTensorVersioned), + POINTER(POINTER(mta_system_t)), + ] + lib.mta_system_create.restype = check_status + + lib.mta_system_free.argtypes = [ + POINTER(mta_system_t), + ] + lib.mta_system_free.restype = check_status + + lib.mta_system_size.argtypes = [ + POINTER(mta_system_t), + POINTER(c_uintptr_t), + ] + lib.mta_system_size.restype = check_status + + lib.mta_system_get_data.argtypes = [ + POINTER(mta_system_t), + mta_system_data_kind, + POINTER(POINTER(DLManagedTensorVersioned)), + ] + lib.mta_system_get_data.restype = check_status + + lib.mta_system_get_length_unit.argtypes = [ + POINTER(mta_system_t), + POINTER(mta_string_t), + ] + lib.mta_system_get_length_unit.restype = check_status + + lib.mta_system_add_pairs.argtypes = [ + POINTER(mta_system_t), + ctypes.c_char_p, + POINTER(mts_block_t), + ] + lib.mta_system_add_pairs.restype = check_status + + lib.mta_system_get_pairs.argtypes = [ + POINTER(mta_system_t), + ctypes.c_char_p, + POINTER(POINTER(mts_block_t)), + ] + lib.mta_system_get_pairs.restype = check_status + + lib.mta_system_known_pairs.argtypes = [ + POINTER(mta_system_t), + POINTER(mta_string_t), + ] + lib.mta_system_known_pairs.restype = check_status + + lib.mta_system_add_custom_data.argtypes = [ + POINTER(mta_system_t), + ctypes.c_char_p, + POINTER(mts_tensormap_t), + ] + lib.mta_system_add_custom_data.restype = check_status + + lib.mta_system_get_custom_data.argtypes = [ + POINTER(mta_system_t), + ctypes.c_char_p, + POINTER(POINTER(mts_tensormap_t)), + ] + lib.mta_system_get_custom_data.restype = check_status + + lib.mta_system_known_custom_data.argtypes = [ + POINTER(mta_system_t), + POINTER(mta_string_t), + ] + lib.mta_system_known_custom_data.restype = check_status + + lib.mta_execute_model.argtypes = [ + mta_model_t, + POINTER(POINTER(mta_system_t)), + c_uintptr_t, + POINTER(mts_labels_t), + ctypes.c_char_p, + ctypes.c_bool, + POINTER(POINTER(mts_tensormap_t)), + c_uintptr_t, + ] + lib.mta_execute_model.restype = check_status + + lib.mta_format_metadata.argtypes = [ + ctypes.c_char_p, + POINTER(mta_string_t), + ] + lib.mta_format_metadata.restype = check_status + + lib.mta_register_plugin.argtypes = [ + mta_plugin_t, + ] + lib.mta_register_plugin.restype = check_status + + lib.mta_load_plugin.argtypes = [ + ctypes.c_char_p, + ] + lib.mta_load_plugin.restype = check_status + + lib.mta_load_model.argtypes = [ + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_char_p, + POINTER(mta_model_t), + ] + lib.mta_load_model.restype = check_status + + lib.mta_save.argtypes = [ + ctypes.c_char_p, + POINTER(mta_system_t), + ] + lib.mta_save.restype = check_status + + lib.mta_save_buffer.argtypes = [ + POINTER(ctypes.c_char_p), + POINTER(c_uintptr_t), + ctypes.c_void_p, + mts_realloc_buffer_t, + POINTER(mta_system_t), + ] + lib.mta_save_buffer.restype = check_status + + lib.mta_load.argtypes = [ + ctypes.c_char_p, + mts_create_array_callback_t, + POINTER(POINTER(mta_system_t)), + ] + lib.mta_load.restype = check_status + + lib.mta_load_buffer.argtypes = [ + ctypes.c_char_p, + c_uintptr_t, + mts_create_array_callback_t, + POINTER(POINTER(mta_system_t)), + ] + lib.mta_load_buffer.restype = check_status diff --git a/python/metatomic_core/src/metatomic/_c_lib.py b/python/metatomic_core/src/metatomic/_c_lib.py new file mode 100644 index 000000000..f1f082247 --- /dev/null +++ b/python/metatomic_core/src/metatomic/_c_lib.py @@ -0,0 +1,182 @@ +import ctypes +import os +import re +import sys +from collections import namedtuple +from ctypes import cdll, wintypes + +from metatensor._c_lib import _get_library as _get_mts_library + +from ._c_api import setup_functions +from ._version import __version__ + + +_HERE = os.path.realpath(os.path.dirname(__file__)) + +Version = namedtuple("Version", ["major", "minor", "patch"]) + + +def parse_version(version): + match = re.match(r"(\d+)\.(\d+)\.(\d+).*", version) + if match: + return Version(*map(int, match.groups())) + else: + raise ValueError("Invalid version string format") + + +def _compatible_versions(actual, minimal): + actual = parse_version(actual) + minimal = parse_version(minimal) + + # Different major version are not compatible + if actual.major != minimal.major: + return False + + # If the major version is 0, different minor version are not compatible + if actual.major == 0 and actual.minor != minimal.minor: + return False + + return True + + +class LibraryFinder: + def __init__(self): + self._cached_dll = None + + def __call__(self): + if self._cached_dll is None: + # make sure to load metatensor first, since we want to resolve symbols from + # there + _get_mts_library() + + # if the library is already loaded in the current process, use this one + # instead of loading a second, independent copy of it + dll = _already_loaded(_lib_name()) + if dll is None: + path = _lib_path() + dll = cdll.LoadLibrary(path) + else: + path = "" + + self._cached_dll = dll + setup_functions(self._cached_dll) + + version = self._cached_dll.mta_version().decode("utf8") + if not _compatible_versions(version, __version__): + self._cached_dll = None + raise RuntimeError( + f"wrong version for libmetatomic, we want v{__version__}, " + f"but we got v{version} @ '{path}'" + ) + + return self._cached_dll + + +def _lib_name(): + """Name of the metatomic shared library on the current platform""" + if sys.platform.startswith("darwin"): + return "libmetatomic.dylib" + elif sys.platform.startswith("linux"): + return "libmetatomic.so" + elif sys.platform.startswith("win"): + return "metatomic.dll" + else: + raise ImportError("Unknown platform. Please edit this file") + + +def _already_loaded(name): + """ + Check if the library with the given ``name`` is already loaded in the current + process, and return the corresponding ``CDLL`` if it is. This makes sure we share + the global state of the library (error buffers, registered data origins, ...) with + whoever loaded it first, instead of using a second, independent copy of the library. + + Returns ``None`` if the library is not already loaded. + """ + if sys.platform.startswith("win"): + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + kernel32.GetModuleHandleW.restype = wintypes.HMODULE + kernel32.GetModuleHandleW.argtypes = [wintypes.LPCWSTR] + + handle = kernel32.GetModuleHandleW(name) + if not handle: + return None + + return ctypes.CDLL(name, handle=handle) + else: + # RTLD_NOLOAD gives us a handle if the library is already loaded, and fails + # instead of loading it otherwise. + RTLD_NOLOAD = getattr(os, "RTLD_NOLOAD", None) + if RTLD_NOLOAD is None: + return None + + try: + return ctypes.CDLL(name, mode=RTLD_NOLOAD | os.RTLD_LOCAL) + except OSError: + return None + + +def _lib_path(): + try: + # check if we are using an externally-provided version of the shared library + from ._external import EXTERNAL_METATOMIC_PATH + + return EXTERNAL_METATOMIC_PATH + except ImportError: + pass + + # otherwise load from the local installation + windows = sys.platform.startswith("win") + if windows: + path = os.path.join(_HERE, "bin", _lib_name()) + else: + path = os.path.join(_HERE, "lib", _lib_name()) + + if os.path.isfile(path): + if windows: + _check_dll(path) + return path + + raise ImportError("Could not find metatomic shared library at " + path) + + +def _check_dll(path): + """Check if the DLL at ``path`` matches the architecture of Python""" + import platform + import struct + + IMAGE_FILE_MACHINE_I386 = 332 + IMAGE_FILE_MACHINE_AMD64 = 34404 + IMAGE_FILE_MACHINE_ARM64 = 43620 + + machine = None + with open(path, "rb") as fd: + header = fd.read(2).decode(encoding="utf-8", errors="strict") + if header != "MZ": + raise ImportError(path + " is not a DLL") + else: + fd.seek(60) + header = fd.read(4) + header_offset = struct.unpack(" list[Quantity]: + """ + Outputs this model can provide. + + During a specific run, a model might be asked to only compute a subset of these + outputs. + + The returned list is a copy, assign to this property to change the outputs. + """ + return list(self._outputs) + + @outputs.setter + def outputs(self, value: Sequence[Quantity]): + outputs = [] + for output in value: + if not isinstance(output, Quantity): + raise ValueError(f"outputs must be Quantity, got {type(output)}") + outputs.append(output) + + self._outputs = outputs + + @property + def atomic_types(self) -> list[int]: + """ + Atomic types this model supports. + + The meaning of the integers in this list is up to the model, and is not required + to be the atomic numbers. + + The returned list is a copy, assign to this property to change the atomic types. + """ + return list(self._atomic_types) + + @atomic_types.setter + def atomic_types(self, value: Sequence[int]): + atomic_types = [] + for atomic_type in value: + if isinstance(atomic_type, bool) or not isinstance(atomic_type, int): + raise ValueError( + f"atomic types must be integers, got {type(atomic_type)}" + ) + atomic_types.append(atomic_type) + + self._atomic_types = atomic_types + + @property + def interaction_range(self) -> float: + """ + Interaction range of the model, in the length unit of the model. + + This is the maximum distance between two atoms for which the model's output can + depend on their relative position. For a short range model, this is the same as + the largest pair list cutoff; for a message passing model, this is the cutoff of + one environment times the number of message passing steps; and for an explicit + long range model, this should be set to infinity (``float("inf")``). + """ + return self._interaction_range + + @interaction_range.setter + def interaction_range(self, value: float): + interaction_range = float(value) + + if math.isnan(interaction_range) or interaction_range < 0.0: + raise ValueError("interaction_range must be non-negative") + + self._interaction_range = interaction_range + + @property + def length_unit(self) -> str: + """ + Length unit used by the model for its inputs, e.g. ``"angstrom"`` or + ``"nanometer"``. + + This applies to the :py:attr:`interaction_range`, any cutoff in pair lists, the + system positions, cell and pair lists given to the model. + """ + return self._length_unit + + @length_unit.setter + def length_unit(self, value: str): + if not isinstance(value, str): + raise ValueError(f"length_unit must be a string, got {type(value)}") + + self._length_unit = value + + @property + def supported_devices(self) -> list[str]: + """ + Devices on which this model can run. + + The devices should be ordered by preference: the first entry in this list should + be the best device for this model, and so on. + + The returned list is a copy, assign to this property to change the supported + devices. + """ + return list(self._supported_devices) + + @supported_devices.setter + def supported_devices(self, value: Sequence[str]): + for device in value: + if not isinstance(device, str): + raise ValueError( + f"devices must be a list of strings, got {type(device)}" + ) + if device not in _VALID_DEVICES: + raise ValueError( + f"device must be one of {list(_VALID_DEVICES)}, got {device}" + ) + + self._supported_devices = list(value) + + @property + def dtype(self) -> str: + """ + Data type of this model, used for all its inputs and outputs. + + The model is free to use a different data type for its internal computations. + """ + return self._dtype + + @dtype.setter + def dtype(self, value: str): + if not isinstance(value, str): + raise ValueError(f"dtype must be a string, got {type(value)}") + if value not in _VALID_DTYPES: + raise ValueError(f"dtype must be one of {list(_VALID_DTYPES)}, got {value}") + + self._dtype = value + + def find_output(self, request: Quantity) -> Optional[Quantity]: + """ + Find the output matching the name and sample kind of ``request``, or ``None`` if + this model does not declare such an output. + + :param request: the quantity to look for + """ + for output in self._outputs: + if ( + output.name == request.name + and output.sample_kind == request.sample_kind + ): + return output + + return None + + def to_dict(self) -> dict: + """ + Convert this object to a JSON-compatible dictionary, following the + :ref:`documented format `. + """ + return { + "type": "metatomic_model_capabilities", + "outputs": [output.to_dict() for output in self._outputs], + "atomic_types": self.atomic_types, + "interaction_range": self.interaction_range, + "length_unit": self.length_unit, + "supported_devices": self._supported_devices, + "dtype": self.dtype, + } + + @classmethod + def from_dict(cls, data: dict) -> "ModelCapabilities": + """ + Create a :py:class:`ModelCapabilities` from a JSON-compatible dictionary, + following the :ref:`documented format `. + + :param data: dictionary containing the data, typically obtained by parsing JSON + with :py:func:`json.loads` + :raises ValueError: if the data does not match the expected format + """ + if not isinstance(data, dict): + raise ValueError( + "invalid JSON data for ModelCapabilities, expected an object" + ) + + valid_keys = set( + [ + "type", + "outputs", + "atomic_types", + "interaction_range", + "length_unit", + "supported_devices", + "dtype", + ] + ) + for key in data.keys(): + if key not in valid_keys: + raise ValueError( + f"unexpected key '{key}' in JSON for ModelCapabilities" + ) + + if data.get("type") != "metatomic_model_capabilities": + raise ValueError( + "'type' in JSON for ModelCapabilities must be " + "'metatomic_model_capabilities'" + ) + + if not isinstance(data.get("outputs"), list): + raise ValueError("'outputs' in JSON for ModelCapabilities must be an array") + outputs = [Quantity.from_dict(output) for output in data["outputs"]] + + if not isinstance(data.get("atomic_types"), list): + raise ValueError( + "'atomic_types' in JSON for ModelCapabilities must be an array" + ) + + for atomic_type in data["atomic_types"]: + if isinstance(atomic_type, bool) or not isinstance(atomic_type, int): + raise ValueError( + "'atomic_types' in JSON for ModelCapabilities must be an " + "array of integers" + ) + + interaction_range = data.get("interaction_range") + if isinstance(interaction_range, bool) or not isinstance( + interaction_range, (int, float) + ): + raise ValueError( + "'interaction_range' in JSON for ModelCapabilities must be a number" + ) + if interaction_range < 0.0: + raise ValueError( + "'interaction_range' in JSON for ModelCapabilities must be non-negative" + ) + + if not isinstance(data.get("length_unit"), str): + raise ValueError( + "'length_unit' in JSON for ModelCapabilities must be a string" + ) + + if not isinstance(data.get("supported_devices"), list): + raise ValueError( + "'supported_devices' in JSON for ModelCapabilities must be an array" + ) + + if not isinstance(data.get("dtype"), str): + raise ValueError("dtype in JSON for ModelCapabilities must be a string") + + return cls( + atomic_types=data["atomic_types"], + interaction_range=interaction_range, + length_unit=data["length_unit"], + supported_devices=data["supported_devices"], + dtype=data["dtype"], + outputs=outputs, + ) + + def __repr__(self) -> str: + return ( + f"ModelCapabilities(outputs={self._outputs!r}, " + f"atomic_types={self._atomic_types}, " + f"interaction_range={self._interaction_range}, " + f"length_unit='{self._length_unit}', " + f"supported_devices={[str(d) for d in self._supported_devices]}, " + f"dtype='{self._dtype}')" + ) + + def __eq__(self, other) -> bool: + if not isinstance(other, ModelCapabilities): + return NotImplemented + + return ( + self._outputs == other._outputs + and self._atomic_types == other._atomic_types + and self._interaction_range == other._interaction_range + and self._length_unit == other._length_unit + and self._supported_devices == other._supported_devices + and self._dtype == other._dtype + ) diff --git a/python/metatomic_core/src/metatomic/_metadata.py b/python/metatomic_core/src/metatomic/_metadata.py new file mode 100644 index 000000000..48f636e20 --- /dev/null +++ b/python/metatomic_core/src/metatomic/_metadata.py @@ -0,0 +1,428 @@ +import ctypes +import json +from collections.abc import Mapping, Sequence +from typing import Optional, Union + +from ._c_api import mta_string_t + + +def _format_metadata(metadata: dict) -> str: + """ + Call ``mta_format_metadata`` to render a JSON-serialized + :py:class:`ModelMetadata` as human-readable text. + + The formatting is done by the shared library to make sure all the languages + supported by metatomic produce exactly the same output. + """ + from ._c_lib import _get_library + + lib = _get_library() + + printed = mta_string_t() + lib.mta_format_metadata(json.dumps(metadata).encode("utf8"), ctypes.byref(printed)) + try: + return lib.mta_string_view(printed).decode("utf8") + finally: + lib.mta_string_free(printed) + + +def _check_string_list(values, context: str) -> list[str]: + """ + Check that ``values`` is a list of strings, and return it as a new + :py:class:`list`. ``context`` is used to build the error messages. + """ + if not isinstance(values, list): + raise ValueError(f"{context} must be an array") + + for value in values: + if not isinstance(value, str): + raise ValueError(f"{context} must be an array of strings") + + return list(values) + + +class References: + """ + References for a model, divided into three categories: references about the model as + a whole, references about the architecture of the model, and references about the + implementation of the model. + + Each category is a list of strings, which can be DOIs, URLs, or any other format the + model author finds useful. + """ + + def __init__( + self, + *, + model: Optional[Sequence[str]] = None, + architecture: Optional[Sequence[str]] = None, + implementation: Optional[Sequence[str]] = None, + ): + """ + :param model: references about the model as a whole, e.g. a paper describing the + model or a website presenting it + :param architecture: references about the architecture of the model, e.g. papers + describing the mathematical form of the model + :param implementation: references about the implementation of the model, e.g. a + link to the source code repository or a paper describing the software + """ + self.model = [] if model is None else model + self.architecture = [] if architecture is None else architecture + self.implementation = [] if implementation is None else implementation + + @staticmethod + def _check_section(values: Sequence[str], section: str) -> list[str]: + result = [] + for value in values: + if not isinstance(value, str): + raise ValueError( + f"reference must be a string (in '{section}' section), " + f"got {type(value)}" + ) + if value == "": + raise ValueError( + f"reference can not be empty string (in '{section}' section)" + ) + result.append(value) + + return result + + @property + def model(self) -> list[str]: + """ + References about the model as a whole, e.g. a paper describing the model or a + website presenting it. + + The returned list is a copy, use :py:meth:`add` to register a new reference. + """ + return list(self._model) + + @model.setter + def model(self, value: Sequence[str]): + self._model = self._check_section(value, "model") + + @property + def architecture(self) -> list[str]: + """ + References about the architecture of the model, e.g. papers describing the + mathematical form of the model. + + The returned list is a copy, use :py:meth:`add` to register a new reference. + """ + return list(self._architecture) + + @architecture.setter + def architecture(self, value: Sequence[str]): + self._architecture = self._check_section(value, "architecture") + + @property + def implementation(self) -> list[str]: + """ + References about the implementation of the model, e.g. a link to the source code + repository or a paper describing the software. + + The returned list is a copy, use :py:meth:`add` to register a new reference. + """ + return list(self._implementation) + + @implementation.setter + def implementation(self, value: Sequence[str]): + self._implementation = self._check_section(value, "implementation") + + def add(self, section: str, reference: str): + """ + Add ``reference`` to the given ``section``. + + :param section: one of ``"model"``, ``"architecture"``, or ``"implementation"`` + :param reference: the reference to add + """ + if section not in ["model", "architecture", "implementation"]: + raise ValueError( + "reference section must be 'model', 'architecture', or " + f"'implementation', got '{section}'" + ) + + checked = self._check_section([reference], section) + getattr(self, f"_{section}").extend(checked) + + def to_dict(self) -> dict: + """Convert this object to a JSON-compatible dictionary""" + return { + "model": self.model, + "architecture": self.architecture, + "implementation": self.implementation, + } + + @classmethod + def from_dict(cls, data: dict) -> "References": + """ + Create a :py:class:`References` from a JSON-compatible dictionary. + + :param data: dictionary containing the data, typically obtained by parsing JSON + with :py:func:`json.loads` + :raises ValueError: if the data does not match the expected format + """ + if not isinstance(data, dict): + raise ValueError( + "invalid JSON data for references in ModelMetadata, expected an object" + ) + + valid_keys = set(["model", "architecture", "implementation"]) + for key in data.keys(): + if key not in valid_keys: + raise ValueError(f"unexpected key '{key}' in JSON for references") + + return cls( + model=_check_string_list( + data.get("model"), "'model' in references of ModelMetadata" + ), + architecture=_check_string_list( + data.get("architecture"), + "'architecture' in references of ModelMetadata", + ), + implementation=_check_string_list( + data.get("implementation"), + "'implementation' in references of ModelMetadata", + ), + ) + + def __repr__(self) -> str: + return ( + f"References(model={self._model}, architecture={self._architecture}, " + f"implementation={self._implementation})" + ) + + def __eq__(self, other) -> bool: + if not isinstance(other, References): + return NotImplemented + + return ( + self._model == other._model + and self._architecture == other._architecture + and self._implementation == other._implementation + ) + + +class ModelMetadata: + """ + Metadata about a model: who created it, what it does, and which references should be + cited when using it. + + This class implements ``__str__``, so the metadata can be pretty-printed with + ``print(metadata)``. + """ + + def __init__( + self, + *, + name: str = "", + description: str = "", + authors: Optional[Sequence[str]] = None, + references: Union["References", Mapping[str, Sequence[str]], None] = None, + extra: Optional[Mapping[str, str]] = None, + ): + """ + :param name: name of the model, e.g. ``"MyCoolModel v1.2"`` + :param description: free-text description of the model + :param authors: authors of the model, e.g. ``["Alice Smith", "Bob Johnson + "]`` + :param references: references for the model that should be cited when using it, + either a :py:class:`References` or a dictionary with ``"model"``, + ``"architecture"`` and ``"implementation"`` keys + :param extra: any other key-value pairs the model author wants to include in the + metadata. This can be used for any purpose. + """ + self.name = name + self.description = description + self.authors = [] if authors is None else authors + self.references = References() if references is None else references + self.extra = {} if extra is None else extra + + @property + def name(self) -> str: + """Name of the model, e.g. ``"MyCoolModel v1.2"``""" + return self._name + + @name.setter + def name(self, value: str): + if not isinstance(value, str): + raise ValueError(f"name must be a string, got {type(value)}") + self._name = value + + @property + def description(self) -> str: + """Free-text description of the model""" + return self._description + + @description.setter + def description(self, value: str): + if not isinstance(value, str): + raise ValueError(f"description must be a string, got {type(value)}") + self._description = value + + @property + def authors(self) -> list[str]: + """ + Authors of the model, e.g. ``["Alice Smith", "Bob Johnson + "]``. + + The returned list is a copy, use :py:meth:`add_author` to register a new author. + """ + return list(self._authors) + + @authors.setter + def authors(self, value: Sequence[str]): + self._authors = [] + for author in value: + self.add_author(author) + + def add_author(self, author: str): + """ + Add ``author`` to the list of authors of this model. + + :param author: name of the author, optionally followed by an email address + between angle brackets + """ + if not isinstance(author, str): + raise ValueError(f"author must be a string, got {type(author)}") + + if author == "": + raise ValueError("author can not be empty string in ModelMetadata") + + self._authors.append(author) + + @property + def references(self) -> References: + """References for the model that should be cited when using it""" + return self._references + + @references.setter + def references(self, value: Union[References, Mapping[str, Sequence[str]]]): + if isinstance(value, References): + self._references = value + elif isinstance(value, Mapping): + unknown = set(value.keys()) - {"model", "architecture", "implementation"} + if unknown: + raise ValueError( + "reference section must be 'model', 'architecture', or " + f"'implementation', got '{sorted(unknown)[0]}'" + ) + self._references = References(**value) + else: + raise ValueError( + f"references must be a References or a dict, got {type(value)}" + ) + + @property + def extra(self) -> dict[str, str]: + """ + Any other key-value pairs the model author wants to include in the metadata. + This can be used for any purpose. + + The returned dictionary is a copy, assign to this property to change the extra + metadata. + """ + return dict(self._extra) + + @extra.setter + def extra(self, value: Mapping[str, str]): + extra = {} + for key, entry in value.items(): + if not isinstance(key, str): + raise ValueError(f"extra keys must be strings, got {type(key)}") + if not isinstance(entry, str): + raise ValueError(f"extra values must be strings, got {type(entry)}") + extra[key] = entry + + self._extra = extra + + def to_dict(self) -> dict: + """ + Convert this object to a JSON-compatible dictionary, following the + :ref:`documented format `. + """ + return { + "type": "metatomic_model_metadata", + "name": self.name, + "authors": self.authors, + "description": self.description, + "references": self.references.to_dict(), + "extra": self.extra, + } + + @classmethod + def from_dict(cls, data: dict) -> "ModelMetadata": + """ + Create a :py:class:`ModelMetadata` from a JSON-compatible dictionary, following + the :ref:`documented format `. + + :param data: dictionary containing the data, typically obtained by parsing JSON + with :py:func:`json.loads` + :raises ValueError: if the data does not match the expected format + """ + if not isinstance(data, dict): + raise ValueError("invalid JSON data for ModelMetadata, expected an object") + + valid_keys = set( + ["type", "name", "authors", "description", "references", "extra"] + ) + for key in data.keys(): + if key not in valid_keys: + raise ValueError(f"unexpected key '{key}' in JSON for ModelMetadata") + + if data.get("type") != "metatomic_model_metadata": + raise ValueError( + "'type' in JSON for ModelMetadata must be 'metatomic_model_metadata'" + ) + + if not isinstance(data.get("name"), str): + raise ValueError("'name' in JSON for ModelMetadata must be a string") + + authors = _check_string_list( + data.get("authors"), "'authors' in JSON for ModelMetadata" + ) + + if not isinstance(data.get("description"), str): + raise ValueError("'description' in JSON for ModelMetadata must be a string") + + references = References.from_dict(data.get("references")) + + if not isinstance(data.get("extra"), dict): + raise ValueError("'extra' in JSON for ModelMetadata must be an object") + + for value in data["extra"].values(): + if not isinstance(value, str): + raise ValueError( + "'extra' in JSON for ModelMetadata must be an object with " + "string values" + ) + + return cls( + name=data["name"], + description=data["description"], + authors=authors, + references=references, + extra=data["extra"], + ) + + def __str__(self) -> str: + return _format_metadata(self.to_dict()) + + def __repr__(self) -> str: + return ( + f"ModelMetadata(name='{self._name}', description='{self._description}', " + f"authors={self._authors}, references={self._references!r}, " + f"extra={self._extra})" + ) + + def __eq__(self, other) -> bool: + if not isinstance(other, ModelMetadata): + return NotImplemented + + return ( + self._name == other._name + and self._description == other._description + and self._authors == other._authors + and self._references == other._references + and self._extra == other._extra + ) diff --git a/python/metatomic_core/src/metatomic/_quantity.py b/python/metatomic_core/src/metatomic/_quantity.py new file mode 100644 index 000000000..633bcf1f8 --- /dev/null +++ b/python/metatomic_core/src/metatomic/_quantity.py @@ -0,0 +1,220 @@ +from collections.abc import Sequence +from typing import Optional + + +_VALID_SAMPLE_KINDS = ["system", "atom", "atom_pair"] +_VALID_GRADIENTS = ["positions", "strain"] + + +class Quantity: + """ + A physical quantity that a model can take as input or produce as output. + """ + + def __init__( + self, + *, + name: str, + unit: str, + sample_kind: str, + description: Optional[str] = None, + gradients: Optional[Sequence[str]] = None, + ): + """ + :param name: name of the quantity, either one of the :ref:`standard names + ` or a custom name of the form + ``::[/]``. + :param unit: unit of the quantity + :param sample_kind: kind of samples this quantity is associated with + :param description: optional description of this quantity, especially useful + when a model defines multiple variants of the same quantity + :param gradients: list of gradients stored explicitly in the + :py:class:`TensorMap ` for this quantity + """ + self.name = name + self.unit = unit + self.sample_kind = sample_kind + self.description = description + self.gradients = [] if gradients is None else gradients + + @property + def name(self) -> str: + """ + Name of this quantity, this can be one of the :ref:`standard names + ` or a custom name of the form + ``::[/]``. + + This is not validated here: an invalid name will be rejected by the shared + library when the quantity is sent to it. + """ + return self._name + + @name.setter + def name(self, value: str): + if not isinstance(value, str): + raise ValueError(f"name must be a string, got {type(value)}") + self._name = value + + @property + def unit(self) -> str: + """Unit of this quantity""" + return self._unit + + @unit.setter + def unit(self, value: str): + if not isinstance(value, str): + raise ValueError(f"unit must be a string, got {type(value)}") + self._unit = value + + @property + def sample_kind(self) -> str: + """ + Kind of samples this quantity is associated with. This is can be one of the + following: + + - ``system`` for per-system/global quantity + - ``atom`` for per-atom quantity + - ``atom_pair`` for quantities defined over a pair of atoms + """ + return self._sample_kind + + @sample_kind.setter + def sample_kind(self, value: str): + if not isinstance(value, str): + raise ValueError(f"sample_kind must be a string, got {type(value)}") + + if value not in _VALID_SAMPLE_KINDS: + raise ValueError( + f"sample_kind must be one of {list(_VALID_SAMPLE_KINDS)}, got {value}" + ) + + self._sample_kind = value + + @property + def description(self) -> Optional[str]: + """ + Description of this quantity, used to provide more details about it, especially + when a model defines multiple variants of the same quantity. + """ + return self._description + + @description.setter + def description(self, value: Optional[str]): + if value is None or value == "": + # an empty description is the same as no description at all + self._description = None + elif isinstance(value, str): + self._description = value + else: + raise ValueError(f"description must be a string, got {type(value)}") + + @property + def gradients(self) -> list[str]: + """ + List of gradients stored explicitly in the :py:class:`TensorMap + ` for this quantity. + + Gradients can be one of the following: + + - ``positions`` for gradients with respect to positions, e.g. the forces + - ``strain`` for gradients with respect to the strain, e.g. the stress tensor + + The returned list is a copy, assign to this property to change the gradients. + """ + return list(self._gradients) + + @gradients.setter + def gradients(self, value: Sequence[str]): + for gradient in value: + if not isinstance(gradient, str): + raise ValueError( + f"gradients must be a list of strings, got {type(gradient)}" + ) + if gradient not in _VALID_GRADIENTS: + raise ValueError( + f"gradient must be one of {list(_VALID_GRADIENTS)}, got {gradient}" + ) + + self._gradients = list(value) + + def to_dict(self) -> dict: + """ + Convert this object to a JSON-compatible dictionary, following the + :ref:`documented format `. + """ + result = { + "type": "metatomic_quantity", + "name": self.name, + "unit": self.unit, + "sample_kind": self.sample_kind, + "gradients": self._gradients, + } + + if self.description is not None: + result["description"] = self.description + + return result + + @classmethod + def from_dict(cls, data: dict) -> "Quantity": + """ + Create a :py:class:`Quantity` from a JSON-compatible dictionary, following the + :ref:`documented format `. + + :param data: dictionary containing the data, typically obtained by parsing JSON + """ + if not isinstance(data, dict): + raise ValueError("invalid JSON data for Quantity, expected an object") + + valid_keys = set( + ["type", "name", "unit", "sample_kind", "description", "gradients"] + ) + for key in data.keys(): + if key not in valid_keys: + raise ValueError(f"unexpected key '{key}' in JSON for Quantity") + + if data.get("type") != "metatomic_quantity": + raise ValueError("'type' in JSON for Quantity must be 'metatomic_quantity'") + + if not isinstance(data.get("name"), str): + raise ValueError("'name' in JSON for Quantity must be a string") + + if not isinstance(data.get("unit"), str): + raise ValueError("'unit' in JSON for Quantity must be a string") + + description = data.get("description") + if description is not None and not isinstance(description, str): + raise ValueError("'description' in JSON for Quantity must be a string") + + if not isinstance(data.get("gradients"), list): + raise ValueError("'gradients' in JSON for Quantity must be an array") + + if not isinstance(data.get("sample_kind"), str): + raise ValueError("'sample_kind' in JSON for Quantity must be a string") + + return cls( + name=data["name"], + unit=data["unit"], + sample_kind=data["sample_kind"], + description=description, + gradients=data["gradients"], + ) + + def __repr__(self) -> str: + return ( + f"Quantity(name='{self._name}', unit='{self._unit}', " + f"sample_kind='{self._sample_kind}', description={self._description!r}, " + f"gradients={[str(g) for g in self._gradients]})" + ) + + def __eq__(self, other) -> bool: + if not isinstance(other, Quantity): + return NotImplemented + + return ( + self._name == other._name + and self._unit == other._unit + and self._sample_kind == other._sample_kind + and self._description == other._description + and self._gradients == other._gradients + ) diff --git a/python/metatomic_core/src/metatomic/_status.py b/python/metatomic_core/src/metatomic/_status.py new file mode 100644 index 000000000..47ccd9464 --- /dev/null +++ b/python/metatomic_core/src/metatomic/_status.py @@ -0,0 +1,103 @@ +import ctypes +import sys +from typing import Optional + +from ._c_api import mta_status_t + + +class MetatomicError(Exception): + """This class is used to throw exceptions for all errors in metatomic.""" + + def __init__(self, message, status=None): + super(Exception, self).__init__(message) + + self.message: str = message + """error message for this exception""" + + self.status: Optional[int] = status + """status code for this exception""" + + +def check_status(status): + if status == mta_status_t.MTA_SUCCESS: + return + else: + raise _get_exception(status) + + +def check_pointer(pointer): + if not pointer: + raise _get_exception() + + +def _delete_exception(exception): + # decrement the reference count of the exception + exception_ptr = ctypes.cast(exception, ctypes.py_object) + ctypes.pythonapi.Py_DecRef(exception_ptr) + + +_DELETE_EXCEPTION = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(_delete_exception) + + +def save_exception(e): + """ + Save the given exception in meatomic's thread-local storage, so that it can be + retrieved later with `_get_exception()`. + """ + from ._c_lib import _get_library + + lib = _get_library() + + # increment the reference count of the exception + exception_ptr = ctypes.py_object(e) + ctypes.pythonapi.Py_IncRef(exception_ptr) + + try: + lib.mta_set_last_error( + ctypes.c_char_p(str(e).encode("utf8")), + ctypes.c_char_p(b"Python exception"), + ctypes.c_void_p.from_buffer(exception_ptr), + _DELETE_EXCEPTION, + ) + except Exception as err: + # if we failed to save the exception, we are in a very bad state, but we should + # still try to report the original error message if possible. + print( + "INTERNAL ERROR: unable to save last error after Python callback failure", + file=sys.stderr, + ) + print( + f"original error was: {e}, error while saving was {err}", + file=sys.stderr, + ) + ctypes.pythonapi.Py_DecRef(exception_ptr) + + +def _get_exception(status=None): + """ + Get the last error from libmetatensor that happened on the current thread. + + If the last error was caused by a Python exception, this returns the exception as + is, otherwise it returns a new MetatensorError with the last error message. + """ + from ._c_lib import _get_library + + lib = _get_library() + message = ctypes.c_char_p() + origin = ctypes.c_char_p() + user_data = ctypes.c_void_p() + status = lib.mta_last_error( + ctypes.byref(message), ctypes.byref(origin), ctypes.byref(user_data) + ) + + if status != mta_status_t.MTA_SUCCESS: + return MetatomicError( + "INTERNAL ERROR: failed to get the last error", status=status + ) + + if origin.value == b"Python exception" and user_data.value is not None: + # This error was caused by a Python exception, so we re-raise it here + # (the exception is stored in the user_data pointer) + return ctypes.cast(user_data, ctypes.py_object).value + + return MetatomicError(message.value.decode("utf8"), status=status) diff --git a/python/metatomic_core/src/metatomic/_system.py b/python/metatomic_core/src/metatomic/_system.py new file mode 100644 index 000000000..380ac5217 --- /dev/null +++ b/python/metatomic_core/src/metatomic/_system.py @@ -0,0 +1,313 @@ +import ctypes +import json +import math +import re +import struct +from collections.abc import Sequence +from typing import Optional + +from ._c_api import mta_string_t + + +_HEX_NUMBER = re.compile(r"(0[xX])?[0-9a-fA-F]+") + + +def _hex_from_cutoff(value: float) -> str: + """ + Get the hexadecimal representation of the bit pattern of ``value``. + + Storing floating point values as their bit pattern makes the JSON + round-trip exact, without relying on the precision of the decimal + representation. + """ + bits = struct.unpack(" float: + """ + Inverse of :py:func:`_hex_from_cutoff`, reading a ``f64`` from the hexadecimal + representation of its bit pattern. + + ``context`` is used to build the error message if ``value`` is not a valid + hexadecimal string. + """ + if isinstance(value, str) and _HEX_NUMBER.fullmatch(value) is not None: + bits = int(value, 16) + else: + bits = 2**64 + + if bits >= 2**64: + raise ValueError( + "'cutoff' in JSON for PairListOptions must be a hex-encoded string, " + f"got '{value}'" + ) + + return struct.unpack(" str: + """ + Call ``mta_format_metadata`` to render a JSON-serialized + :py:class:`ModelMetadata` as human-readable text. + + The formatting is done by the shared library to make sure all the languages + supported by metatomic produce exactly the same output. + """ + from ._c_lib import _get_library + + lib = _get_library() + + printed = mta_string_t() + lib.mta_format_metadata(json.dumps(metadata).encode("utf8"), ctypes.byref(printed)) + try: + return lib.mta_string_view(printed).decode("utf8") + finally: + lib.mta_string_free(printed) + + +def _check_string_list(values, context: str) -> list[str]: + """ + Check that ``values`` is a list of strings, and return it as a new + :py:class:`list`. ``context`` is used to build the error messages. + """ + if not isinstance(values, list): + raise ValueError(f"{context} must be an array") + + for value in values: + if not isinstance(value, str): + raise ValueError(f"{context} must be an array of strings") + + return list(values) + + +### ================================================================================ ### + + +class PairListOptions: + """ + Options for the calculation of a pair list (also known as a neighbor list). + + A model declares the pair lists it needs with these options, and the engine running + the model is then responsible for computing matching pair lists and attaching them + to the systems given to the model. + """ + + def __init__( + self, + *, + cutoff: float, + full_list: bool, + strict: bool = True, + requestors: Optional[Sequence[str]] = None, + ): + """ + :param cutoff: spherical cutoff radius for this pair list, in the length unit of + the model + :param full_list: should the pair list be a full list (containing both the pair + ``i -> j`` and ``j -> i``) or a half list (containing only ``i -> j``) + :param strict: does the list only contain pairs within the cutoff (``True``) or + can it also contain pairs slightly beyond the cutoff (``False``) + :param requestors: list of strings describing who requested this pair list. More + requestors can be added later with :py:meth:`add_requestor`. + """ + self.cutoff = cutoff + self.full_list = full_list + self.strict = strict + self.requestors = [] if requestors is None else requestors + + @property + def cutoff(self) -> float: + """ + Spherical cutoff radius for this pair list, in the length unit of the model. + """ + return self._cutoff + + @cutoff.setter + def cutoff(self, value: float): + cutoff = float(value) + + if not math.isfinite(cutoff) or cutoff <= 0.0: + raise ValueError("cutoff must be a finite positive number") + + self._cutoff = cutoff + + @property + def full_list(self) -> bool: + """ + Should the pair list be a full list (containing both the pair ``i -> j`` and ``j + -> i``) or a half list (containing only ``i -> j``)? + """ + return self._full_list + + @full_list.setter + def full_list(self, value: bool): + self._full_list = bool(value) + + @property + def strict(self) -> bool: + """ + Does this list only contain pairs within the cutoff (``True``), or can it also + contain pairs slightly beyond the cutoff (``False``)? + """ + return self._strict + + @strict.setter + def strict(self, value: bool): + self._strict = bool(value) + + @property + def requestors(self) -> list[str]: + """ + List of strings describing who requested this pair list. + + The returned list is a copy, use :py:meth:`add_requestor` to register a new + requestor. + """ + return list(self._requestors) + + @requestors.setter + def requestors(self, value: Sequence[str]): + self._requestors = [] + for requestor in value: + self.add_requestor(requestor) + + def add_requestor(self, requestor: str): + """ + Add ``requestor`` to the list of entities requesting this pair list. Empty + strings and duplicates are ignored. + + :param requestor: string describing who is requesting this pair list + """ + requestor = str(requestor) + if requestor != "" and requestor not in self._requestors: + self._requestors.append(requestor) + + def to_dict(self) -> dict: + """ + Convert this object to a JSON-compatible dictionary, following the + :ref:`documented format `. + """ + return { + "type": "metatomic_pair_list_options", + # store the bit pattern so the float round-trips exactly + "cutoff": _hex_from_cutoff(self.cutoff), + "full_list": self.full_list, + "strict": self.strict, + "requestors": self.requestors, + } + + @classmethod + def from_dict(cls, data: dict) -> "PairListOptions": + """ + Create a :py:class:`PairListOptions` from a JSON-compatible dictionary, + following the :ref:`documented format `. + + :param data: dictionary containing the data, typically obtained by parsing JSON + with :py:func:`json.loads` + :raises ValueError: if the data does not match the expected format + """ + if not isinstance(data, dict): + raise ValueError( + "invalid JSON data for PairListOptions, expected an object" + ) + + valid_keys = set(["type", "cutoff", "full_list", "strict", "requestors"]) + for key in data.keys(): + if key not in valid_keys: + raise ValueError(f"unexpected key '{key}' in JSON for PairListOptions") + + if data.get("type") != "metatomic_pair_list_options": + raise ValueError( + "'type' in JSON for PairListOptions must be " + "'metatomic_pair_list_options'" + ) + + cutoff = _cutoff_from_hex(data.get("cutoff")) + if not math.isfinite(cutoff) or cutoff <= 0.0: + raise ValueError( + "'cutoff' in JSON for PairListOptions must be a finite positive number" + ) + + if not isinstance(data.get("full_list"), bool): + raise ValueError( + "'full_list' in JSON for PairListOptions must be a boolean" + ) + + if not isinstance(data.get("strict"), bool): + raise ValueError("'strict' in JSON for PairListOptions must be a boolean") + + requestors = [] + if "requestors" in data: + requestors = _check_string_list( + data["requestors"], "'requestors' in JSON for PairListOptions" + ) + + return cls( + cutoff=cutoff, + full_list=data["full_list"], + strict=data["strict"], + requestors=requestors, + ) + + def __repr__(self) -> str: + return ( + f"PairListOptions(cutoff={self._cutoff}, full_list={self._full_list}, " + f"strict={self._strict})" + ) + + def _comparison_key(self): + # the list of requestors is intentionally left out: two requests with + # the same parameters can be fulfilled by the same pair list, whoever + # asked for them + return (self._cutoff, self._full_list, self._strict) + + # `PairListOptions` are compared by cutoff first, then `full_list` and + # finally `strict`; the list of requestors is ignored everywhere. + + def __eq__(self, other) -> bool: + """ + Check if two :py:class:`PairListOptions` are equal. + + The list of requestors is ignored when checking for equality. + """ + if not isinstance(other, PairListOptions): + return NotImplemented + return self._comparison_key() == other._comparison_key() + + def __ne__(self, other) -> bool: + """ + Check if two :py:class:`PairListOptions` are different. + + The list of requestors is ignored when checking for equality. + """ + if not isinstance(other, PairListOptions): + return NotImplemented + return self._comparison_key() != other._comparison_key() + + def __lt__(self, other) -> bool: + """Check if this pair list sorts before ``other``""" + if not isinstance(other, PairListOptions): + return NotImplemented + return self._comparison_key() < other._comparison_key() + + def __le__(self, other) -> bool: + """Check if this pair list sorts before ``other`` or is equal to it""" + if not isinstance(other, PairListOptions): + return NotImplemented + return self._comparison_key() <= other._comparison_key() + + def __gt__(self, other) -> bool: + """Check if this pair list sorts after ``other``""" + if not isinstance(other, PairListOptions): + return NotImplemented + return self._comparison_key() > other._comparison_key() + + def __ge__(self, other) -> bool: + """Check if this pair list sorts after ``other`` or is equal to it""" + if not isinstance(other, PairListOptions): + return NotImplemented + return self._comparison_key() >= other._comparison_key() + + def __hash__(self) -> int: + return hash(self._comparison_key()) diff --git a/python/metatomic_core/src/metatomic/_version.py b/python/metatomic_core/src/metatomic/_version.py new file mode 100644 index 000000000..4220b3fa1 --- /dev/null +++ b/python/metatomic_core/src/metatomic/_version.py @@ -0,0 +1,4 @@ +import importlib.metadata + + +__version__ = importlib.metadata.version("metatomic-core") diff --git a/python/metatomic_core/src/metatomic/torch.py b/python/metatomic_core/src/metatomic/torch.py new file mode 100644 index 000000000..060e7bccf --- /dev/null +++ b/python/metatomic_core/src/metatomic/torch.py @@ -0,0 +1,14 @@ +import sys + + +try: + import metatomic_torch +except ImportError as e: + raise ImportError( + "metatomic-torch is required to use the metatomic.torch module. " + "Please install it with `pip install metatomic-torch` or using " + "your favorite Python package manager." + ) from e + +# metatomic.torch is registered as an alias in metatomic_torch's __init__.py +assert sys.modules["metatomic.torch"] is metatomic_torch diff --git a/python/metatomic_core/src/metatomic/utils.py b/python/metatomic_core/src/metatomic/utils.py new file mode 100644 index 000000000..48a870235 --- /dev/null +++ b/python/metatomic_core/src/metatomic/utils.py @@ -0,0 +1,16 @@ +import os + + +try: + from ._external import EXTERNAL_METATOMIC_PREFIX + + cmake_prefix_path = EXTERNAL_METATOMIC_PREFIX + """ + Path containing the CMake configuration files for the underlying C library + """ + +except ImportError: + cmake_prefix_path = os.path.join(os.path.dirname(__file__), "lib", "cmake") + """ + Path containing the CMake configuration files for the underlying C library + """ diff --git a/python/metatomic_core/tests/capabilities.py b/python/metatomic_core/tests/capabilities.py new file mode 100644 index 000000000..be3b7e77a --- /dev/null +++ b/python/metatomic_core/tests/capabilities.py @@ -0,0 +1,191 @@ +import math +import re + +import pytest + +from metatomic import ModelCapabilities, Quantity + + +@pytest.fixture +def capabilities(): + return ModelCapabilities( + outputs=[ + Quantity( + name="energy", + unit="eV", + sample_kind="system", + description="total energy", + gradients=["positions"], + ), + Quantity( + name="custom::charge/with_variant", + unit="e", + sample_kind="atom", + ), + ], + atomic_types=[1, 6, 8], + interaction_range=5.0, + length_unit="Angstrom", + supported_devices=["cpu", "cuda"], + dtype="float32", + ) + + +def test_model_capabilities(capabilities): + assert len(capabilities.outputs) == 2 + assert capabilities.atomic_types == [1, 6, 8] + assert capabilities.interaction_range == 5.0 + assert capabilities.length_unit == "Angstrom" + assert capabilities.supported_devices == ["cpu", "cuda"] + assert capabilities.dtype == "float32" + + # explicit long range models can use an infinite interaction range + capabilities.interaction_range = math.inf + assert capabilities.interaction_range == math.inf + + # the returned lists are copies + capabilities.atomic_types.append(16) + assert capabilities.atomic_types == [1, 6, 8] + + capabilities.supported_devices.append("metal") + assert capabilities.supported_devices == ["cpu", "cuda"] + + capabilities.outputs.append(capabilities.outputs[0]) + assert len(capabilities.outputs) == 2 + + +def test_model_capabilities_errors(): + def capabilities(**kwargs): + parameters = { + "atomic_types": [1], + "interaction_range": 5.0, + "length_unit": "angstrom", + "supported_devices": ["cpu"], + "dtype": "float32", + } + parameters.update(kwargs) + return ModelCapabilities(**parameters) + + with pytest.raises(ValueError, match="interaction_range must be non-negative"): + capabilities(interaction_range=-1.0) + + with pytest.raises(ValueError, match="atomic types must be integers"): + capabilities(atomic_types=["1"]) + + with pytest.raises(ValueError, match="length_unit must be a string"): + capabilities(length_unit=42) + + message = "device must be one of ['cpu', 'cuda', 'rocm', 'metal'], got wat" + with pytest.raises(ValueError, match=re.escape(message)): + capabilities(supported_devices=["cpu", "wat"]) + + message = "dtype must be one of ['float32', 'float64'], got float16" + with pytest.raises(ValueError, match=re.escape(message)): + capabilities(dtype="float16") + + with pytest.raises(ValueError, match="outputs must be Quantity"): + capabilities(outputs=["energy"]) + + +def test_model_capabilities_roundtrip(capabilities): + data = capabilities.to_dict() + + assert data == { + "type": "metatomic_model_capabilities", + "outputs": [ + { + "type": "metatomic_quantity", + "name": "energy", + "unit": "eV", + "sample_kind": "system", + "gradients": ["positions"], + "description": "total energy", + }, + { + "type": "metatomic_quantity", + "name": "custom::charge/with_variant", + "unit": "e", + "sample_kind": "atom", + "gradients": [], + }, + ], + "atomic_types": [1, 6, 8], + "interaction_range": 5.0, + "length_unit": "Angstrom", + "supported_devices": ["cpu", "cuda"], + "dtype": "float32", + } + + parsed = ModelCapabilities.from_dict(data) + assert parsed == capabilities + + assert parsed.outputs[0].name == "energy" + assert parsed.outputs[1].name == "custom::charge/with_variant" + + +def test_model_capabilities_from_dict_errors(capabilities): + def corrupted(**kwargs): + data = capabilities.to_dict() + data.update(kwargs) + return data + + def without(key): + data = capabilities.to_dict() + del data[key] + return data + + cases = [ + ( + "not an object", + "invalid JSON data for ModelCapabilities, expected an object", + ), + ( + corrupted(type="something-else"), + "'type' in JSON for ModelCapabilities must be " + "'metatomic_model_capabilities'", + ), + ( + corrupted(outputs="energy"), + "'outputs' in JSON for ModelCapabilities must be an array", + ), + ( + corrupted(atomic_types="1"), + "'atomic_types' in JSON for ModelCapabilities must be an array", + ), + ( + corrupted(atomic_types=[1, "x"]), + "'atomic_types' in JSON for ModelCapabilities must be an array of integers", + ), + ( + without("interaction_range"), + "'interaction_range' in JSON for ModelCapabilities must be a number", + ), + ( + corrupted(interaction_range=-1.0), + "'interaction_range' in JSON for ModelCapabilities must be non-negative", + ), + ( + without("length_unit"), + "'length_unit' in JSON for ModelCapabilities must be a string", + ), + ( + corrupted(supported_devices="cpu"), + "'supported_devices' in JSON for ModelCapabilities must be an array", + ), + ( + corrupted(supported_devices=["cpu", "wat"]), + "device must be one of ['cpu', 'cuda', 'rocm', 'metal'], got wat", + ), + ( + without("dtype"), + "dtype in JSON for ModelCapabilities must be a string", + ), + ( + corrupted(dtype="float16"), + "dtype must be one of ['float32', 'float64'], got float16", + ), + ] + + for data, message in cases: + with pytest.raises(ValueError, match=re.escape(message)): + ModelCapabilities.from_dict(data) diff --git a/python/metatomic_core/tests/metadata.py b/python/metatomic_core/tests/metadata.py new file mode 100644 index 000000000..d2c8f9151 --- /dev/null +++ b/python/metatomic_core/tests/metadata.py @@ -0,0 +1,220 @@ +import pytest + +from metatomic import ModelMetadata, References + + +def test_references(): + references = References() + assert references.model == [] + assert references.architecture == [] + assert references.implementation == [] + + references.add("model", "doi:10.1234/test") + references.add("architecture", "doi:10.1234/arch") + references.add("implementation", "https://github.com/test") + + assert references.model == ["doi:10.1234/test"] + assert references.architecture == ["doi:10.1234/arch"] + assert references.implementation == ["https://github.com/test"] + + # the returned lists are copies + references.model.append("doi:10.1234/other") + assert references.model == ["doi:10.1234/test"] + + message = ( + "reference section must be 'model', 'architecture', or 'implementation', " + "got 'wrong'" + ) + with pytest.raises(ValueError, match=message): + references.add("wrong", "doi:10.1234/test") + + message = "reference can not be empty string \\(in 'model' section\\)" + with pytest.raises(ValueError, match=message): + references.add("model", "") + + with pytest.raises(ValueError, match=message): + References(model=["doi:10.1234/test", ""]) + + +@pytest.fixture +def metadata(): + return ModelMetadata( + name="test-model", + description="A test model", + authors=["Alice", "Bob "], + references=References( + model=["doi:10.1234/test"], + architecture=["doi:10.1234/arch"], + implementation=["https://github.com/test"], + ), + extra={"key1": "value1", "key2": "value2"}, + ) + + +def test_model_metadata(metadata): + assert metadata.name == "test-model" + assert metadata.description == "A test model" + assert metadata.authors == ["Alice", "Bob "] + assert metadata.references.model == ["doi:10.1234/test"] + assert metadata.extra == {"key1": "value1", "key2": "value2"} + + metadata.add_author("Charlie") + assert metadata.authors == ["Alice", "Bob ", "Charlie"] + + # the references can also be given as a plain dict + metadata = ModelMetadata(references={"model": ["doi:10.1234/test"]}) + assert metadata.references.model == ["doi:10.1234/test"] + assert metadata.references.architecture == [] + + defaults = ModelMetadata() + assert defaults.name == "" + assert defaults.description == "" + assert defaults.authors == [] + assert defaults.references == References() + assert defaults.extra == {} + + +def test_model_metadata_errors(): + with pytest.raises(ValueError, match="author can not be empty string"): + ModelMetadata(authors=["Alice", ""]) + + with pytest.raises(ValueError, match="name must be a string"): + ModelMetadata(name=42) + + with pytest.raises(ValueError, match="description must be a string"): + ModelMetadata(description=42) + + with pytest.raises(ValueError, match="extra values must be strings"): + ModelMetadata(extra={"key": 42}) + + with pytest.raises(ValueError, match="extra keys must be strings"): + ModelMetadata(extra={42: "value"}) + + message = ( + "reference section must be 'model', 'architecture', or 'implementation', " + "got 'wrong'" + ) + with pytest.raises(ValueError, match=message): + ModelMetadata(references={"wrong": []}) + + message = "references must be a References or a dict" + with pytest.raises(ValueError, match=message): + ModelMetadata(references=42) + + +def test_model_metadata_roundtrip(metadata): + data = metadata.to_dict() + + assert data == { + "type": "metatomic_model_metadata", + "name": "test-model", + "authors": ["Alice", "Bob "], + "description": "A test model", + "references": { + "model": ["doi:10.1234/test"], + "architecture": ["doi:10.1234/arch"], + "implementation": ["https://github.com/test"], + }, + "extra": {"key1": "value1", "key2": "value2"}, + } + + assert ModelMetadata.from_dict(data) == metadata + + +def test_model_metadata_from_dict_errors(metadata): + def corrupted(**kwargs): + data = metadata.to_dict() + data.update(kwargs) + return data + + def without(key): + data = metadata.to_dict() + del data[key] + return data + + cases = [ + ("not an object", "invalid JSON data for ModelMetadata, expected an object"), + ( + corrupted(type="something-else"), + "'type' in JSON for ModelMetadata must be 'metatomic_model_metadata'", + ), + (without("name"), "'name' in JSON for ModelMetadata must be a string"), + (corrupted(name=42), "'name' in JSON for ModelMetadata must be a string"), + ( + corrupted(authors="Alice"), + "'authors' in JSON for ModelMetadata must be an array", + ), + ( + corrupted(authors=["Alice", 42]), + "'authors' in JSON for ModelMetadata must be an array of strings", + ), + ( + without("description"), + "'description' in JSON for ModelMetadata must be a string", + ), + ( + corrupted(extra="not-an-object"), + "'extra' in JSON for ModelMetadata must be an object", + ), + ( + corrupted(extra={"key": 42}), + "'extra' in JSON for ModelMetadata must be an object with string values", + ), + ( + corrupted(references="not-an-object"), + "invalid JSON data for references in ModelMetadata, expected an object", + ), + ( + without("references"), + "invalid JSON data for references in ModelMetadata, expected an object", + ), + ( + corrupted(references={"model": "doi:10.1234/test"}), + "'model' in references of ModelMetadata must be an array", + ), + ( + corrupted( + references={"model": [42], "architecture": [], "implementation": []} + ), + "'model' in references of ModelMetadata must be an array of strings", + ), + ( + corrupted(references={"model": [], "implementation": []}), + "'architecture' in references of ModelMetadata must be an array", + ), + ] + + for data, message in cases: + with pytest.raises(ValueError, match=message): + ModelMetadata.from_dict(data) + + +def test_model_metadata_printing(metadata): + expected = """This is the test-model model +============================ + +A test model + +Model authors +------------- + +- Alice +- Bob + +Model references +---------------- + +Please cite the following references when using this model: +- about this specific model: + * doi:10.1234/test +- about the architecture of this model: + * doi:10.1234/arch +- about the implementation of this model: + * https://github.com/test +""" + assert str(metadata) == expected + + expected = """This is an unnamed model +======================== +""" + assert str(ModelMetadata()) == expected diff --git a/python/metatomic_core/tests/quantity.py b/python/metatomic_core/tests/quantity.py new file mode 100644 index 000000000..db9657d19 --- /dev/null +++ b/python/metatomic_core/tests/quantity.py @@ -0,0 +1,121 @@ +import re + +import pytest + +from metatomic import Quantity + + +@pytest.fixture +def quantity(): + return Quantity( + name="energy", + unit="eV", + sample_kind="atom", + description="total energy of the system", + gradients=["positions"], + ) + + +def test_quantity(quantity): + assert quantity.name == "energy" + assert quantity.unit == "eV" + assert quantity.sample_kind == "atom" + assert quantity.description == "total energy of the system" + assert quantity.gradients == ["positions"] + + # an empty description is the same as no description + quantity.description = "" + assert quantity.description is None + + # the returned list of gradients is a copy + quantity.gradients.append("strain") + assert quantity.gradients == ["positions"] + + +def test_quantity_errors(): + with pytest.raises(ValueError, match="name must be a string"): + Quantity(name=42, unit="eV", sample_kind="atom") + + message = "sample_kind must be one of ['system', 'atom', 'atom_pair'], got foo" + with pytest.raises(ValueError, match=re.escape(message)): + Quantity(name="energy", unit="eV", sample_kind="foo") + + message = "sample_kind must be a string, got " + with pytest.raises(ValueError, match=message): + Quantity(name="energy", unit="eV", sample_kind=42) + + message = "gradient must be one of ['positions', 'strain'], got foo" + with pytest.raises(ValueError, match=re.escape(message)): + Quantity( + name="energy", unit="eV", sample_kind="atom", gradients=["positions", "foo"] + ) + + with pytest.raises(ValueError, match="unit must be a string"): + Quantity(name="energy", unit=42, sample_kind="atom") + + with pytest.raises(ValueError, match="description must be a string"): + Quantity(name="energy", unit="eV", sample_kind="atom", description=42) + + +def test_quantity_roundtrip(quantity): + data = quantity.to_dict() + + assert data == { + "type": "metatomic_quantity", + "name": "energy", + "unit": "eV", + "sample_kind": "atom", + "gradients": ["positions"], + "description": "total energy of the system", + } + + assert Quantity.from_dict(data) == quantity + + # `description` is left out when it is not set + quantity.description = None + assert "description" not in quantity.to_dict() + assert Quantity.from_dict(quantity.to_dict()) == quantity + + +def test_quantity_from_dict_errors(quantity): + def corrupted(**kwargs): + data = quantity.to_dict() + data.update(kwargs) + return data + + def without(key): + data = quantity.to_dict() + del data[key] + return data + + cases = [ + ("not an object", "invalid JSON data for Quantity, expected an object"), + ( + corrupted(type="something-else"), + "'type' in JSON for Quantity must be 'metatomic_quantity'", + ), + (without("name"), "'name' in JSON for Quantity must be a string"), + (without("unit"), "'unit' in JSON for Quantity must be a string"), + (corrupted(description=42), "'description' in JSON for Quantity must be a"), + (without("gradients"), "'gradients' in JSON for Quantity must be an array"), + ( + corrupted(gradients="positions"), + "'gradients' in JSON for Quantity must be an array", + ), + ( + corrupted(gradients=["positions", "foo"]), + "gradient must be one of ['positions', 'strain'], got foo", + ), + ( + without("sample_kind"), + "'sample_kind' in JSON for Quantity must be a string", + ), + ( + corrupted(sample_kind="foo"), + "sample_kind must be one of ['system', 'atom', 'atom_pair'], got foo", + ), + ] + + for data, message in cases: + with pytest.raises(ValueError, match=re.escape(message)): + Quantity.from_dict(data) diff --git a/python/metatomic_core/tests/system.py b/python/metatomic_core/tests/system.py new file mode 100644 index 000000000..c44522cc5 --- /dev/null +++ b/python/metatomic_core/tests/system.py @@ -0,0 +1,210 @@ +import copy +import math +import operator + +import pytest + +from metatomic import PairListOptions + + +### ================================================================================ ### +### PairListOptions ### +### ================================================================================ ### + + +@pytest.fixture +def pair_options(): + return PairListOptions( + cutoff=3.5, + full_list=True, + strict=False, + requestors=["nl-1", "nl-2"], + ) + + +def test_pair_options(pair_options): + assert pair_options.strict is False + assert pair_options.requestors == ["nl-1", "nl-2"] + + # setters & getters + pair_options.cutoff = 1.0 + pair_options.full_list = False + pair_options.strict = True + pair_options.requestors = ["foo"] + assert pair_options.cutoff == 1.0 + assert pair_options.full_list is False + assert pair_options.strict is True + assert pair_options.requestors == ["foo"] + + defaults = PairListOptions(cutoff=3.5, full_list=True) + assert defaults.cutoff == 3.5 + assert defaults.full_list is True + # `strict` defaults to True + assert defaults.strict is True + assert defaults.requestors == [] + + +def test_pair_options_invalid_cutoff(): + for cutoff in [0.0, -1.0, math.inf, -math.inf, math.nan]: + message = "cutoff must be a finite positive number" + with pytest.raises(ValueError, match=message): + PairListOptions(cutoff=cutoff, full_list=True) + + +def test_pair_options_requestors(pair_options): + options = PairListOptions(cutoff=3.5, full_list=True) + + options.add_requestor("nl-1") + options.add_requestor("nl-2") + # empty strings and duplicates are ignored, first-seen order is preserved + options.add_requestor("nl-1") + options.add_requestor("") + assert options.requestors == ["nl-1", "nl-2"] + + # the returned list is a copy + options.requestors.append("nl-3") + assert options.requestors == ["nl-1", "nl-2"] + + options.requestors = ["a", "", "b", "a"] + assert options.requestors == ["a", "b"] + + # handle duplicate/empty strings in from_dict + data = pair_options.to_dict() + data["requestors"] = ["a", "", "b", "a"] + + parsed = PairListOptions.from_dict(data) + assert parsed.requestors == ["a", "b"] + + +def test_pair_options_comparison(pair_options): + # the requestors are ignored when comparing + other = copy.deepcopy(pair_options) + other.add_requestor("nl-3") + assert pair_options == other + assert hash(pair_options) == hash(other) + + other = copy.deepcopy(pair_options) + other.cutoff = 4.0 + assert pair_options != other + assert pair_options < other + assert pair_options <= other + assert other > pair_options + assert other >= pair_options + + other = copy.deepcopy(pair_options) + other.strict = True + assert pair_options != other + assert pair_options < other + + # a pair list compares equal to (and neither before nor after) itself + same = copy.deepcopy(pair_options) + assert pair_options <= same + assert pair_options >= same + assert not pair_options < same + assert not pair_options > same + + assert pair_options != "not a PairListOptions" + for op in [operator.lt, operator.le, operator.gt, operator.ge]: + with pytest.raises(TypeError): + op(pair_options, "not a PairListOptions") + + unsorted = [ + PairListOptions(cutoff=4.0, full_list=False), + PairListOptions(cutoff=1.0, full_list=True), + PairListOptions(cutoff=1.0, full_list=False), + ] + assert sorted(unsorted) == [unsorted[2], unsorted[1], unsorted[0]] + + +def test_pair_options_roundtrip(pair_options): + data = pair_options.to_dict() + + assert data == { + "type": "metatomic_pair_list_options", + "cutoff": "0x400c000000000000", + "full_list": True, + "strict": False, + "requestors": ["nl-1", "nl-2"], + } + + parsed = PairListOptions.from_dict(data) + assert parsed == pair_options + assert parsed.requestors == pair_options.requestors + + +def test_pair_options_cutoff_keeps_full_precision(): + options = PairListOptions(cutoff=1.0 / 3.0, full_list=True) + parsed = PairListOptions.from_dict(options.to_dict()) + assert parsed.cutoff == options.cutoff + + +def test_pair_options_from_dict_errors(pair_options): + def corrupted(**kwargs): + data = pair_options.to_dict() + data.update(kwargs) + return data + + def without(key): + data = pair_options.to_dict() + del data[key] + return data + + # each case corrupts exactly one field of an otherwise valid object + cases = [ + ( + "not an object", + "invalid JSON data for PairListOptions, expected an object", + ), + ( + corrupted(type="something-else"), + "'type' in JSON for PairListOptions must be 'metatomic_pair_list_options'", + ), + ( + without("cutoff"), + "'cutoff' in JSON for PairListOptions must be a hex-encoded string", + ), + ( + corrupted(cutoff="not-hex"), + "'cutoff' in JSON for PairListOptions must be a hex-encoded string", + ), + ( + corrupted(cutoff=3.5), + "'cutoff' in JSON for PairListOptions must be a hex-encoded string", + ), + ( + corrupted(cutoff="0x7ff8000000000000"), # NaN + "'cutoff' in JSON for PairListOptions must be a finite positive number", + ), + ( + corrupted(cutoff="0x7ff0000000000000"), # +inf + "'cutoff' in JSON for PairListOptions must be a finite positive number", + ), + ( + corrupted(cutoff="0xbff0000000000000"), # -1.0 + "'cutoff' in JSON for PairListOptions must be a finite positive number", + ), + ( + corrupted(cutoff="0x0"), # 0.0 + "'cutoff' in JSON for PairListOptions must be a finite positive number", + ), + ( + corrupted(full_list="yes"), + "'full_list' in JSON for PairListOptions must be a boolean", + ), + ( + without("strict"), + "'strict' in JSON for PairListOptions must be a boolean", + ), + ( + corrupted(requestors="nl-1"), + "'requestors' in JSON for PairListOptions must be an array", + ), + ( + corrupted(requestors=["nl-1", 42]), + "'requestors' in JSON for PairListOptions must be an array of strings", + ), + ] + + for data, message in cases: + with pytest.raises(ValueError, match=message): + PairListOptions.from_dict(data) diff --git a/python/metatomic_core/tests/utils.py b/python/metatomic_core/tests/utils.py new file mode 100644 index 000000000..bbbc21153 --- /dev/null +++ b/python/metatomic_core/tests/utils.py @@ -0,0 +1,15 @@ +import os + +import metatomic as mta + + +def test_cmake_prefix_path(): + assert os.path.exists(mta.utils.cmake_prefix_path) + + +def test_library_loading(): + # temporary test to be removed as soon as some other test actually load the library + import metatomic._c_lib # noqa: F401 + + lib = mta._c_lib._get_library() + assert lib.mta_version().decode("utf8").replace("-", ".") == mta.__version__ diff --git a/python/metatomic_torch/CMakeLists.txt b/python/metatomic_torch/CMakeLists.txt index 3578cd11f..74702d3ac 100644 --- a/python/metatomic_torch/CMakeLists.txt +++ b/python/metatomic_torch/CMakeLists.txt @@ -63,6 +63,9 @@ else() add_subdirectory("${METATOMIC_TORCH_SOURCE_DIR}" metatomic-torch) + if (CMAKE_VERSION VERSION_LESS "3.25") + set(LINUX $) + endif() if (LINUX OR APPLE) if (LINUX) @@ -74,12 +77,12 @@ else() set(metatomic_install_rpath "${CMAKE_INSTALL_RPATH}") # when loading the libraries from a Python installation: - # - $ORIGIN/../../../../torch/lib is where libtorch.so will be - # - $ORIGIN/../../../../metatensor/lib is where libmetatensor.so will be - # - $ORIGIN/../../../../metatensor/torch/torch-x.y/lib is where libmetatensor_torch.so will be - set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../../torch/lib") - set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../../metatensor/lib") - set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../../metatensor/torch/torch-${Torch_VERSION_MAJOR}.${Torch_VERSION_MINOR}/lib") + # - $ORIGIN/../../../torch/lib is where libtorch.so will be + # - $ORIGIN/../../../metatensor/lib is where libmetatensor.so will be + # - $ORIGIN/../../../metatensor_torch/torch-${Torch_VERSION_MAJOR}.${Torch_VERSION_MINOR}/lib is where libmetatensor_torch.so will be + set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../torch/lib") + set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../metatensor/lib") + set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../metatensor_torch/torch-${Torch_VERSION_MAJOR}.${Torch_VERSION_MINOR}/lib") set_target_properties( metatomic_torch PROPERTIES INSTALL_RPATH "${metatomic_install_rpath}" diff --git a/python/metatomic_torch/MANIFEST.in b/python/metatomic_torch/MANIFEST.in index 5f3ae9425..eb7359c7c 100644 --- a/python/metatomic_torch/MANIFEST.in +++ b/python/metatomic_torch/MANIFEST.in @@ -5,7 +5,7 @@ include LICENSE include git_version_info -include metatomic-torch-*.tar.gz +include metatomic-torch-cxx-*.tar.gz recursive-include build-backend *.py diff --git a/python/metatomic_torch/README.rst b/python/metatomic_torch/README.rst index f06f2b8af..994fda75e 100644 --- a/python/metatomic_torch/README.rst +++ b/python/metatomic_torch/README.rst @@ -1,4 +1,4 @@ -metatensor-torch -================ +metatomic-torch +=============== -This package contains the TorchScript bindings to the core API of metatensor. +This package contains the TorchScript bindings to the core API of metatomic. diff --git a/python/metatomic_torch/build-backend/backend.py b/python/metatomic_torch/build-backend/backend.py index c762d91e6..be0389a2c 100644 --- a/python/metatomic_torch/build-backend/backend.py +++ b/python/metatomic_torch/build-backend/backend.py @@ -1,11 +1,24 @@ # This is a custom Python build backend wrapping setuptool's to only depend on # torch/metatensor-torch when building the wheel and not the sdist import os +import pathlib from setuptools import build_meta -ROOT = os.path.realpath(os.path.dirname(__file__)) +ROOT = pathlib.Path(__file__).parent.resolve() + +METATOMIC_CORE = (ROOT / ".." / ".." / "metatomic_core").resolve() +METATOMIC_NO_LOCAL_DEPS = os.environ.get("METATOMIC_NO_LOCAL_DEPS", "0") == "1" + + +if not METATOMIC_NO_LOCAL_DEPS and METATOMIC_CORE.exists(): + # we are building from a git checkout + METATOMIC_CORE_DEP = f"metatomic-core @ {METATOMIC_CORE.as_uri()}" +else: + # we are building from a sdist + METATOMIC_CORE_DEP = "metatomic-core >=0.1.0,<0.2" + FORCED_TORCH_VERSION = os.environ.get("METATOMIC_TORCH_BUILD_WITH_TORCH_VERSION") if FORCED_TORCH_VERSION is not None: @@ -27,7 +40,7 @@ # Special dependencies to build the wheels def get_requires_for_build_wheel(config_settings=None): defaults = build_meta.get_requires_for_build_wheel(config_settings) - return defaults + [TORCH_DEP] + return defaults + [TORCH_DEP, METATOMIC_CORE_DEP] def build_editable(wheel_directory, config_settings=None, metadata_directory=None): diff --git a/python/metatomic_torch/metatomic/torch/__init__.py b/python/metatomic_torch/metatomic/torch/__init__.py deleted file mode 100644 index 06a9ae9c5..000000000 --- a/python/metatomic_torch/metatomic/torch/__init__.py +++ /dev/null @@ -1,70 +0,0 @@ -import os -from typing import TYPE_CHECKING - -import torch - -from ._c_lib import _load_library -from .version import __version__ # noqa: F401 - - -if os.environ.get("METATOMIC_IMPORT_FOR_SPHINX", "0") != "0" or TYPE_CHECKING: - from .documentation import ( - ModelCapabilities, - ModelEvaluationOptions, - ModelMetadata, - ModelOutput, - NeighborListOptions, - System, - check_atomistic_model, - load_model_extensions, - pick_device, - pick_output, - read_model_metadata, - register_autograd_neighbors, - unit_conversion_factor, - unit_dimension_for_quantity, - ) - - _check_quantities = None - -else: - _load_library() - - System = torch.classes.metatomic.System - NeighborListOptions = torch.classes.metatomic.NeighborListOptions - - ModelOutput = torch.classes.metatomic.ModelOutput - ModelEvaluationOptions = torch.classes.metatomic.ModelEvaluationOptions - ModelCapabilities = torch.classes.metatomic.ModelCapabilities - ModelMetadata = torch.classes.metatomic.ModelMetadata - - read_model_metadata = torch.ops.metatomic.read_model_metadata - load_model_extensions = torch.ops.metatomic.load_model_extensions - check_atomistic_model = torch.ops.metatomic.check_atomistic_model - _check_quantities = torch.ops.metatomic._check_quantities - - register_autograd_neighbors = torch.ops.metatomic.register_autograd_neighbors - - unit_conversion_factor = torch.ops.metatomic.unit_conversion_factor - unit_dimension_for_quantity = torch.ops.metatomic.unit_dimension_for_quantity - - pick_device = torch.ops.metatomic.pick_device - pick_output = torch.ops.metatomic.pick_output - -from . import ( # noqa: F401 - ase_calculator, - o3, -) -from .model import ( # noqa: F401 - AtomisticModel, - ModelInterface, - is_atomistic_model, - load_atomistic_model, -) -from .serialization import ( # noqa: F401 - load_system, - load_system_buffer, - save, - save_buffer, -) -from .systems_to_torch import systems_to_torch # noqa: F401 diff --git a/python/metatomic_torch/pyproject.toml b/python/metatomic_torch/pyproject.toml index 2d3c34368..1bfb2d679 100644 --- a/python/metatomic_torch/pyproject.toml +++ b/python/metatomic_torch/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "metatomic-torch" dynamic = ["version", "authors", "dependencies"] -requires-python = ">=3.10" +requires-python = ">=3.11" readme = "README.rst" license = "BSD-3-Clause" @@ -48,10 +48,6 @@ backend-path = ["build-backend"] [tool.setuptools] zip-safe = false -[tool.setuptools.packages.find] -include = ["metatomic*"] -namespaces = true - ### ======================================================================== ### [tool.pytest.ini_options] python_files = ["*.py"] diff --git a/python/metatomic_torch/setup.py b/python/metatomic_torch/setup.py index bfc81072c..bbf11bd6f 100644 --- a/python/metatomic_torch/setup.py +++ b/python/metatomic_torch/setup.py @@ -1,5 +1,6 @@ import glob import os +import pathlib import subprocess import sys @@ -12,7 +13,7 @@ from setuptools.command.sdist import sdist -ROOT = os.path.realpath(os.path.dirname(__file__)) +ROOT = pathlib.Path(__file__).parent.resolve() METATOMIC_BUILD_TYPE = os.environ.get("METATOMIC_BUILD_TYPE", "release") if METATOMIC_BUILD_TYPE not in ["debug", "release"]: @@ -21,10 +22,9 @@ "expected 'debug' or 'release'" ) -METATOMIC_TORCH_SRC = os.path.realpath( - os.path.join(ROOT, "..", "..", "metatomic-torch") -) -METATOMIC_ASE = os.path.realpath(os.path.join(ROOT, "..", "metatomic_ase")) +METATOMIC_TORCH_SRC = (ROOT / ".." / ".." / "metatomic-torch").resolve() +METATOMIC_CORE = (ROOT / ".." / "metatomic_core").resolve() +METATOMIC_ASE = (ROOT / ".." / "metatomic_ase").resolve() class universal_wheel(bdist_wheel): @@ -49,10 +49,10 @@ def run(self): import torch source_dir = ROOT - build_dir = os.path.join(ROOT, "build", "cmake-build") - install_dir = os.path.join(os.path.realpath(self.build_lib), "metatomic/torch") + build_dir = ROOT / "build" / "cmake-build" + install_dir = pathlib.Path(self.build_lib).resolve() / "metatomic_torch" - os.makedirs(build_dir, exist_ok=True) + build_dir.mkdir(parents=True, exist_ok=True) # Tell CMake where to find metatensor, metatensor_torch, and torch cmake_prefix_path = [ @@ -65,9 +65,7 @@ def run(self): # compile the code. This allows having multiple version of this shared library # inside the wheel; and dynamically pick the right one. torch_major, torch_minor, *_ = torch.__version__.split(".") - cmake_install_prefix = os.path.join( - install_dir, f"torch-{torch_major}.{torch_minor}" - ) + cmake_install_prefix = install_dir / f"torch-{torch_major}.{torch_minor}" use_external_lib = os.environ.get( "METATOMIC_TORCH_PYTHON_USE_EXTERNAL_LIB", "OFF" @@ -141,8 +139,8 @@ def run(self): def generate_cxx_tar(): - script = os.path.join(ROOT, "..", "..", "scripts", "package-torch.sh") - assert os.path.exists(script) + script = (ROOT / ".." / ".." / "scripts" / "package-torch.sh").resolve() + assert script.exists() try: output = subprocess.run( @@ -179,15 +177,15 @@ def git_version_info(): """ TAG_PREFIX = "metatomic-torch-v" - if os.path.exists("git_version_info"): + if (ROOT / "git_version_info").exists(): # we are building from a sdist, without git available, but the git # version was recorded in the `git_version_info` file - with open("git_version_info") as fd: + with open(ROOT / "git_version_info") as fd: n_commits = int(fd.readline().strip()) git_hash = fd.readline().strip() else: - script = os.path.join(ROOT, "..", "..", "scripts", "git-version-info.py") - assert os.path.exists(script) + script = (ROOT / ".." / ".." / "scripts" / "git-version-info.py").resolve() + assert script.exists() output = subprocess.run( [sys.executable, script, TAG_PREFIX], @@ -202,12 +200,15 @@ def git_version_info(): f"stdout: {output.stdout}\n" f"stderr: {output.stderr}\n" ) - elif output.stderr: + if output.stderr: print(output.stderr, file=sys.stderr) + + lines = output.stdout.splitlines() + if len(lines) < 2: + # the script gave up early (cf. `warn_and_exit`) n_commits = 0 git_hash = "" else: - lines = output.stdout.splitlines() n_commits = int(lines[0].strip()) git_hash = lines[1].strip() @@ -274,10 +275,10 @@ def create_version_number(version): # End of Windows/MKL/PIP hack - if not os.path.exists(METATOMIC_TORCH_SRC): + if not METATOMIC_TORCH_SRC.exists(): # we are building from a sdist, which should include metatomic-torch C++ # sources as a tarball - tarballs = glob.glob(os.path.join(ROOT, "metatomic-torch-cxx-*.tar.gz")) + tarballs = glob.glob(ROOT / "metatomic-torch-cxx-*.tar.gz") if not len(tarballs) == 1: raise RuntimeError( @@ -285,7 +286,7 @@ def create_version_number(version): "metatomic-torch C++ sources" ) - METATOMIC_TORCH_SRC = os.path.realpath(tarballs[0]) + METATOMIC_TORCH_SRC = pathlib.Path(tarballs[0]).resolve() subprocess.run( ["cmake", "-E", "tar", "xf", METATOMIC_TORCH_SRC], cwd=ROOT, @@ -294,15 +295,15 @@ def create_version_number(version): METATOMIC_TORCH_SRC = ".".join(METATOMIC_TORCH_SRC.split(".")[:-2]) - with open(os.path.join(METATOMIC_TORCH_SRC, "VERSION")) as fd: + with open(METATOMIC_TORCH_SRC / "VERSION") as fd: METATOMIC_TORCH_VERSION = fd.read().strip() - with open(os.path.join(ROOT, "AUTHORS")) as fd: + with open(ROOT / "AUTHORS") as fd: authors = fd.read().splitlines() if authors[0].startswith(".."): # handle "raw" symlink files (on Windows or from full repo tarball) - with open(os.path.join(ROOT, authors[0])) as fd: + with open(ROOT / authors[0]) as fd: authors = fd.read().splitlines() try: @@ -326,11 +327,14 @@ def create_version_number(version): # when packaging a sdist for release, we should never use local dependencies METATOMIC_NO_LOCAL_DEPS = os.environ.get("METATOMIC_NO_LOCAL_DEPS", "0") == "1" - if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_ASE): + if not METATOMIC_NO_LOCAL_DEPS and METATOMIC_CORE.exists(): + assert METATOMIC_ASE.exists() # we are building from a git checkout or full repo archive - install_requires.append(f"metatomic-ase @ file://{METATOMIC_ASE}") + install_requires.append(f"metatomic-core @ {METATOMIC_CORE.as_uri()}") + install_requires.append(f"metatomic-ase @ {METATOMIC_ASE.as_uri()}") else: # we are building from a sdist/installing from a wheel + install_requires.append("metatomic-core >=0.1.0,<0.2.0") install_requires.append("metatomic-ase >=0.1.1,<0.2.0") setup( diff --git a/python/metatomic_torch/src/metatomic_torch/__init__.py b/python/metatomic_torch/src/metatomic_torch/__init__.py new file mode 100644 index 000000000..d867277f6 --- /dev/null +++ b/python/metatomic_torch/src/metatomic_torch/__init__.py @@ -0,0 +1,131 @@ +import importlib.abc +import importlib.util +import os +import sys +from typing import TYPE_CHECKING + +import torch + +import metatomic + + +sys.modules["metatomic.torch"] = sys.modules[__name__] +if not hasattr(metatomic, "torch"): + metatomic.torch = sys.modules[__name__] + + +class _MetatomicTorchAliasLoader(importlib.abc.Loader): + def __init__(self, canonical): + self._canonical = canonical + + def create_module(self, spec): + return importlib.import_module(self._canonical) + + def exec_module(self, module): + pass + + +class _MetatomicTorchAliasFinder(importlib.abc.MetaPathFinder): + """Resolve ``metatomic.torch.`` to the same module as ``metatomic_torch.``. + + ``metatomic.torch`` is an alias of the ``metatomic_torch`` package, but Python + otherwise loads submodules under the two names as distinct module objects (it + names a submodule after the import path, not after the package's ``__name__``). + This finder makes both namespaces refer to the exact same module object, so e.g. + ``metatomic.torch.AtomisticModel`` and ``metatomic_torch.model.AtomisticModel`` + are identical rather than two independent copies of the same class. + """ + + def find_spec(self, fullname, path, target=None): + if not fullname.startswith("metatomic.torch."): + return None + + canonical = "metatomic_torch" + fullname[len("metatomic.torch") :] + try: + canonical_spec = importlib.util.find_spec(canonical) + except ModuleNotFoundError: + return None + + if canonical_spec is None: + return None + + spec = importlib.util.spec_from_loader( + fullname, + _MetatomicTorchAliasLoader(canonical), + origin=canonical_spec.origin, + ) + if canonical_spec.submodule_search_locations: + spec.submodule_search_locations = list( + canonical_spec.submodule_search_locations + ) + return spec + + +sys.meta_path.insert(0, _MetatomicTorchAliasFinder()) + + +from ._c_lib import _load_library # noqa: E402 +from .version import __version__ # noqa: F401, E402 + + +if os.environ.get("METATOMIC_IMPORT_FOR_SPHINX", "0") != "0" or TYPE_CHECKING: + from .documentation import ( + ModelCapabilities, + ModelEvaluationOptions, + ModelMetadata, + ModelOutput, + NeighborListOptions, + System, + check_atomistic_model, + load_model_extensions, + pick_device, + pick_output, + read_model_metadata, + register_autograd_neighbors, + unit_conversion_factor, + unit_dimension_for_quantity, + ) + + _check_quantities = None + +else: + _load_library() + + System = torch.classes.metatomic.System + NeighborListOptions = torch.classes.metatomic.NeighborListOptions + + ModelOutput = torch.classes.metatomic.ModelOutput + ModelEvaluationOptions = torch.classes.metatomic.ModelEvaluationOptions + ModelCapabilities = torch.classes.metatomic.ModelCapabilities + ModelMetadata = torch.classes.metatomic.ModelMetadata + + read_model_metadata = torch.ops.metatomic.read_model_metadata + load_model_extensions = torch.ops.metatomic.load_model_extensions + check_atomistic_model = torch.ops.metatomic.check_atomistic_model + _check_quantities = torch.ops.metatomic._check_quantities + + register_autograd_neighbors = torch.ops.metatomic.register_autograd_neighbors + + unit_conversion_factor = torch.ops.metatomic.unit_conversion_factor + unit_dimension_for_quantity = torch.ops.metatomic.unit_dimension_for_quantity + + pick_device = torch.ops.metatomic.pick_device + pick_output = torch.ops.metatomic.pick_output + +from . import ( # noqa: F401, E402 + ase_calculator, + o3, +) +from .model import ( # noqa: F401, E402 + AtomisticModel, + ModelInterface, + is_atomistic_model, + load_atomistic_model, +) +from .serialization import ( # noqa: F401, E402 + load_system, + load_system_buffer, + save, + save_buffer, +) +from .systems_to_torch import systems_to_torch # noqa: F401, E402 diff --git a/python/metatomic_torch/metatomic/torch/_c_lib.py b/python/metatomic_torch/src/metatomic_torch/_c_lib.py similarity index 93% rename from python/metatomic_torch/metatomic/torch/_c_lib.py rename to python/metatomic_torch/src/metatomic_torch/_c_lib.py index c80d07ae6..23a246708 100644 --- a/python/metatomic_torch/metatomic/torch/_c_lib.py +++ b/python/metatomic_torch/src/metatomic_torch/_c_lib.py @@ -62,11 +62,16 @@ def _lib_path(): "is not ABI compatible" ) else: - all_versions = ", ".join(map(lambda version: f"v{version}", existing_versions)) + found_versions = "we found builds for torch" + ", ".join( + map(lambda version: f"v{version}", existing_versions) + ) + + if not found_versions: + found_versions = "we found no build" + raise ImportError( f"Trying to load metatomic-torch with torch v{torch.__version__}, " - f"we found builds for torch {all_versions}; which are not ABI compatible.\n" - "You can try to re-install from source with " + f"{found_versions} in {_HERE}.\nYou can try to re-install from source with " "`pip install metatomic-torch --no-binary=metatomic-torch`" ) diff --git a/python/metatomic_torch/metatomic/torch/_extensions.py b/python/metatomic_torch/src/metatomic_torch/_extensions.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/_extensions.py rename to python/metatomic_torch/src/metatomic_torch/_extensions.py diff --git a/python/metatomic_torch/metatomic/torch/_quantities.py b/python/metatomic_torch/src/metatomic_torch/_quantities.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/_quantities.py rename to python/metatomic_torch/src/metatomic_torch/_quantities.py diff --git a/python/metatomic_torch/metatomic/torch/ase_calculator.py b/python/metatomic_torch/src/metatomic_torch/ase_calculator.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/ase_calculator.py rename to python/metatomic_torch/src/metatomic_torch/ase_calculator.py diff --git a/python/metatomic_torch/metatomic/torch/data/dftd3_parameters.npz b/python/metatomic_torch/src/metatomic_torch/data/dftd3_parameters.npz similarity index 100% rename from python/metatomic_torch/metatomic/torch/data/dftd3_parameters.npz rename to python/metatomic_torch/src/metatomic_torch/data/dftd3_parameters.npz diff --git a/python/metatomic_torch/metatomic/torch/dftd3.py b/python/metatomic_torch/src/metatomic_torch/dftd3.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/dftd3.py rename to python/metatomic_torch/src/metatomic_torch/dftd3.py diff --git a/python/metatomic_torch/metatomic/torch/documentation.py b/python/metatomic_torch/src/metatomic_torch/documentation.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/documentation.py rename to python/metatomic_torch/src/metatomic_torch/documentation.py diff --git a/python/metatomic_torch/metatomic/torch/heat_flux.py b/python/metatomic_torch/src/metatomic_torch/heat_flux.py similarity index 99% rename from python/metatomic_torch/metatomic/torch/heat_flux.py rename to python/metatomic_torch/src/metatomic_torch/heat_flux.py index 6e103ad93..167149b06 100644 --- a/python/metatomic_torch/metatomic/torch/heat_flux.py +++ b/python/metatomic_torch/src/metatomic_torch/heat_flux.py @@ -5,13 +5,13 @@ from vesin.metatomic import NeighborList from . import ( + AtomisticModel, ModelCapabilities, ModelOutput, NeighborListOptions, System, unit_conversion_factor, ) -from .model import AtomisticModel def _wrap_positions(positions: torch.Tensor, cell: torch.Tensor) -> torch.Tensor: diff --git a/python/metatomic_torch/metatomic/torch/model.py b/python/metatomic_torch/src/metatomic_torch/model.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/model.py rename to python/metatomic_torch/src/metatomic_torch/model.py diff --git a/python/metatomic_torch/metatomic/torch/o3/__init__.py b/python/metatomic_torch/src/metatomic_torch/o3/__init__.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/o3/__init__.py rename to python/metatomic_torch/src/metatomic_torch/o3/__init__.py diff --git a/python/metatomic_torch/metatomic/torch/o3/_decompose.py b/python/metatomic_torch/src/metatomic_torch/o3/_decompose.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/o3/_decompose.py rename to python/metatomic_torch/src/metatomic_torch/o3/_decompose.py diff --git a/python/metatomic_torch/metatomic/torch/o3/_projections.py b/python/metatomic_torch/src/metatomic_torch/o3/_projections.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/o3/_projections.py rename to python/metatomic_torch/src/metatomic_torch/o3/_projections.py diff --git a/python/metatomic_torch/metatomic/torch/o3/_quadrature.py b/python/metatomic_torch/src/metatomic_torch/o3/_quadrature.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/o3/_quadrature.py rename to python/metatomic_torch/src/metatomic_torch/o3/_quadrature.py diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/src/metatomic_torch/o3/_symmetrized.py similarity index 99% rename from python/metatomic_torch/metatomic/torch/o3/_symmetrized.py rename to python/metatomic_torch/src/metatomic_torch/o3/_symmetrized.py index c56b5476c..ac68ac2e0 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/src/metatomic_torch/o3/_symmetrized.py @@ -669,7 +669,7 @@ def wrap( :param batch_size: number of transformed Systems evaluated in one model call """ if not isinstance(model, AtomisticModel): - raise TypeError("model must be an AtomisticModel") + raise TypeError(f"model must be an AtomisticModel, got '{type(model)}'") capabilities = model.capabilities() diff --git a/python/metatomic_torch/src/metatomic_torch/o3/_tranformations.py b/python/metatomic_torch/src/metatomic_torch/o3/_tranformations.py new file mode 100644 index 000000000..c0a4736b1 --- /dev/null +++ b/python/metatomic_torch/src/metatomic_torch/o3/_tranformations.py @@ -0,0 +1,979 @@ +""" +Rotate systems and tensor maps under O(3) transformations, routing rows of +multi-system tensors by their ``"system"`` sample label. +""" + +from numbers import Integral + +import torch +from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap + +from .. import System, register_autograd_neighbors +from ._wigner import build_wigner_D_cache + + +_INTEGER_DTYPES = ( + torch.uint8, + torch.uint16, + torch.uint32, + torch.uint64, + torch.int8, + torch.int16, + torch.int32, + torch.int64, +) + + +def _validate_nonnegative_integer(name: str, value: int) -> int: + """Validate a non-negative integer and return it as a Python int.""" + if torch.jit.is_scripting(): + integer_value = value + else: + if isinstance(value, bool) or not isinstance(value, Integral): + raise TypeError( + f"{name} must be a non-negative integer, got {type(value).__name__}." + ) + integer_value = int(value) + if integer_value < 0: + raise ValueError(f"{name} must be a non-negative integer, got {integer_value}.") + + return integer_value + + +def _spherical_parity_factor( + ell: int, + sigma: int, + is_improper: bool, +) -> int: + """Return ``sigma * (-1) ** ell`` for an improper transformation, else ``1``.""" + if torch.jit.is_scripting(): + integer_sigma = sigma + else: + if isinstance(sigma, bool) or not isinstance(sigma, Integral): + raise TypeError(f"sigma must be an integer, got {type(sigma).__name__}.") + integer_sigma = int(sigma) + if integer_sigma not in (-1, 1): + raise ValueError(f"sigma must be either -1 or +1, got {integer_sigma}.") + + if is_improper: + return integer_sigma * int((-1) ** ell) + + return 1 + + +def _validate_system_ids( + systems: list[System], + transformations: list["O3Transformation"], + system_ids: list[int] | torch.Tensor | None, + *, + expected_device: torch.device | None, +) -> torch.Tensor: + """Check and normalize the ``system_ids`` argument of ``transform_tensor``. + + ``system_ids[i]`` is the value in a block's ``"system"`` sample column that + selects ``transformations[i]``. This checks that systems and transformations + pair up one-to-one and that there is one distinct integer id per system, + returning the ids as a ``torch.long`` tensor (``0..n_systems - 1`` when + ``system_ids`` is ``None``). + """ + n_systems = len(systems) + n_transformations = len(transformations) + if n_systems != n_transformations: + raise ValueError( + "Expected one transformation per system, but got " + f"len(systems)={n_systems} and " + f"len(transformations)={n_transformations}." + ) + + if system_ids is None: + return torch.arange(n_systems, dtype=torch.long, device=expected_device) + + if isinstance(system_ids, torch.Tensor): + if system_ids.ndim != 1: + raise ValueError( + "system_ids must be one-dimensional, but got a tensor with shape " + f"{tuple(system_ids.shape)}." + ) + if system_ids.dtype not in _INTEGER_DTYPES: + raise ValueError( + "system_ids must contain integers, but got a tensor with dtype " + f"{system_ids.dtype}." + ) + if expected_device is not None and system_ids.device != expected_device: + raise ValueError( + f"system_ids are on device {system_ids.device}, but the values to " + f"transform are on device {expected_device}." + ) + validated_ids = system_ids.to(dtype=torch.long) + else: + python_ids: list[int] = [] + for system_id in system_ids: + if isinstance(system_id, bool) or not isinstance(system_id, Integral): + raise ValueError("system_ids must contain integers.") + python_ids.append(int(system_id)) + validated_ids = torch.tensor( + python_ids, + dtype=torch.long, + device=expected_device, + ) + + if len(validated_ids) != n_systems: + raise ValueError( + "system_ids must contain exactly one entry per system, but got " + f"len(system_ids)={len(validated_ids)} and len(systems)={n_systems}." + ) + if torch.unique(validated_ids).numel() != n_systems: + raise ValueError( + "system_ids must contain one distinct entry per system, but got " + f"{validated_ids.tolist()}." + ) + + return validated_ids + + +def _validate_transformations_dtype_device( + transformations: list["O3Transformation"], + *, + expected_dtype: torch.dtype, + expected_device: torch.device, +) -> None: + """Check that every transformation has the expected dtype and device.""" + for index, transformation in enumerate(transformations): + if ( + transformation.dtype != expected_dtype + or transformation.device != expected_device + ): + raise ValueError( + f"Transformation at index {index} has dtype/device " + f"({transformation.dtype}, {transformation.device}), differing from " + f"the values to transform ({expected_dtype}, {expected_device})." + ) + + +class O3Transformation: + """ + A single O(3) transformation, represented by a (3, 3) rotation or improper-rotation + matrix. + + The constructor stores a copy of ``matrix``. + """ + + def __init__(self, matrix: torch.Tensor, max_angular_momentum: int): + """ + :param matrix: (3, 3) rotation or improper-rotation matrix + :param max_angular_momentum: non-negative maximum angular momentum for + which Wigner-D matrices are available + """ + max_angular_momentum = _validate_nonnegative_integer( + "max_angular_momentum", max_angular_momentum + ) + + if matrix.shape != (3, 3): + raise ValueError( + f"Transformation has shape {tuple(matrix.shape)}; expected (3, 3)." + ) + + identity = torch.eye(3, device=matrix.device, dtype=matrix.dtype) + if not torch.allclose(matrix @ matrix.T, identity, atol=1e-5): + raise ValueError( + "Transformation is not orthogonal (R @ R.T deviates from I)." + ) + + # Keep an independent copy so modifying the input tensor later cannot make + # the matrix disagree with the cached parity and Wigner-D matrices. + self._matrix = matrix.clone() + self._max_angular_momentum = max_angular_momentum + self._is_improper = bool(torch.det(self._matrix) < 0) + + self._wigner_D_cache: dict[int, torch.Tensor] | None = None + + @classmethod + def _create_no_checks( + cls, + matrix: torch.Tensor, + max_angular_momentum: int, + *, + is_improper: bool, + ) -> "O3Transformation": + """Create a transformation after validation in ``random_transformations``. + + The random factory validates its arguments and matrices before calling this + method. This avoids repeating the public constructor's checks, matrix copy, + and determinant calculation for every matrix. ``is_improper`` must match + ``matrix``. + """ + transformation = cls.__new__(cls) + transformation._matrix = matrix + transformation._max_angular_momentum = max_angular_momentum + transformation._is_improper = is_improper + transformation._wigner_D_cache = None + return transformation + + def _ensure_wigner_D_cache(self) -> dict[int, torch.Tensor]: + """Ensure that the Wigner-D cache has been built and return it.""" + if self._wigner_D_cache is None: + self._wigner_D_cache = build_wigner_D_cache( + self._max_angular_momentum, + self._matrix, + device=self._matrix.device, + dtype=self._matrix.dtype, + ) + + return self._wigner_D_cache + + def _wigner_D_cache_entry(self, ell: int) -> torch.Tensor: + """Return the internal cache entry for ``ell`` without copying it.""" + ell = self._validate_ell_range(ell) + + D = self._ensure_wigner_D_cache().get(ell) + if D is None: + raise ValueError(f"Wigner-D matrix for ell={ell} not found in cache.") + + return D + + @property + def matrix(self) -> torch.Tensor: + """The (3, 3) rotation or improper-rotation matrix.""" + return self._matrix + + @property + def dtype(self) -> torch.dtype: + """The dtype of the transformation matrix.""" + return self._matrix.dtype + + @property + def device(self) -> torch.device: + """The device of the transformation matrix.""" + return self._matrix.device + + @property + def is_improper(self) -> bool: + """Whether this transformation is improper, with negative determinant.""" + return self._is_improper + + def transform_cartesian(self, vectors: torch.Tensor) -> torch.Tensor: + """Apply the transformation to Cartesian vectors. + + :param vectors: (..., 3) tensor of Cartesian vectors + :return: (..., 3) tensor of transformed vectors + """ + return vectors @ self._matrix.T + + def _validate_ell_range(self, ell: int) -> int: + """Check that ``ell`` is an integer in ``[0, max_angular_momentum]``.""" + ell = _validate_nonnegative_integer("ell", ell) + + if ell > self._max_angular_momentum: + raise ValueError( + f"ell={ell} exceeds max_angular_momentum={self._max_angular_momentum}." + ) + + return ell + + def transform_spherical( + self, values: torch.Tensor, ell: int, sigma: int + ) -> torch.Tensor: + """Apply the transformation to spherical values. + + :param values: (..., 2*ell+1) tensor of spherical values + :param ell: angular momentum in ``[0, max_angular_momentum]`` + :param sigma: ``+1`` for a proper spherical representation or ``-1`` for + a pseudo one. Under an improper transformation, the representation + acquires the factor ``sigma * (-1) ** ell``. + :return: (..., 2*ell+1) tensor of transformed spherical values + """ + ell = self._validate_ell_range(ell) + parity_factor = _spherical_parity_factor( + ell, + sigma, + is_improper=self.is_improper, + ) + + D = self._wigner_D_cache_entry(ell) + transformed = values @ D.T + if parity_factor != 1: + transformed = transformed * parity_factor + + return transformed + + def wigner_D_matrix(self, ell: int) -> torch.Tensor: + """Return the proper-part Wigner-D matrix for ``ell``. + + For an improper transformation, :meth:`transform_spherical` applies the + inversion-parity factor separately. + + :param ell: angular momentum in ``[0, max_angular_momentum]`` + :return: (2*ell+1, 2*ell+1) Wigner-D matrix + """ + return self._wigner_D_cache_entry(ell) + + +def random_transformations( + n: int, + max_angular_momentum: int = 0, + *, + device: torch.device, + dtype: torch.dtype, + include_inversions: bool = False, + generator: torch.Generator | None = None, +) -> list[O3Transformation]: + """Sample ``n`` transformations uniformly from SO(3), or from O(3) when + inversions are included. + + Rotations are sampled from the Haar measure on SO(3) via random unit quaternions. + When ``include_inversions`` is ``True``, each matrix is independently negated with + probability 0.5, giving a uniform distribution over the full O(3) group. + + :param n: non-negative number of transformations to generate + :param max_angular_momentum: non-negative maximum angular momentum for + Wigner-D matrices + :param device: target device for the output tensors + :param dtype: target dtype for the output tensors; must be + :attr:`torch.float32` or :attr:`torch.float64` + :param include_inversions: if ``True``, sample from O(3) instead of SO(3) + :param generator: optional :class:`torch.Generator` for reproducible sampling; when + ``None`` the global RNG is used + :return: list of ``n`` :class:`O3Transformation` objects + """ + n = _validate_nonnegative_integer("n", n) + max_angular_momentum = _validate_nonnegative_integer( + "max_angular_momentum", max_angular_momentum + ) + + if dtype not in (torch.float32, torch.float64): + raise ValueError(f"dtype must be torch.float32 or torch.float64, got {dtype}.") + + q = torch.randn(n, 4, device=device, dtype=dtype, generator=generator) + q = q / q.norm(dim=1, keepdim=True) + w, x, y, z = q.unbind(1) + # Quaternion to rotation matrix (standard formula) + R = torch.stack( + [ + 1 - 2 * (y * y + z * z), + 2 * (x * y - w * z), + 2 * (x * z + w * y), + 2 * (x * y + w * z), + 1 - 2 * (x * x + z * z), + 2 * (y * z - w * x), + 2 * (x * z - w * y), + 2 * (y * z + w * x), + 1 - 2 * (x * x + y * y), + ], + dim=1, + ).reshape(n, 3, 3) + + matrices_are_improper = [False] * n + + if include_inversions: + signs = torch.randint(0, 2, (n,), device=device, generator=generator) * 2 - 1 + R = R * signs.to(dtype=dtype).reshape(n, 1, 1) + matrices_are_improper = (signs < 0).tolist() + + identity = torch.eye( + 3, + device=R.device, + dtype=R.dtype, + ).expand(n, 3, 3) + if not torch.allclose( + R @ R.transpose(-1, -2), + identity, + atol=1e-5, + ): + raise ValueError("Generated transformations are not orthogonal.") + + return [ + O3Transformation._create_no_checks( + matrix, + max_angular_momentum, + is_improper=is_improper, + ) + for matrix, is_improper in zip( + R.unbind(0), + matrices_are_improper, + strict=True, + ) + ] + + +def _value_row_indices_by_system( + block: TensorBlock, + system_ids: torch.Tensor, +) -> list[torch.Tensor]: + """Return value-row indices in ``system_ids`` order, or all rows for one system.""" + if len(system_ids) == 1: + return [torch.arange(block.values.shape[0], device=block.values.device)] + + if "system" not in block.samples.names: + raise ValueError( + "Rotational augmentation expects output samples to include a 'system' " + "dimension when transforming multiple systems." + ) + system_labels = block.samples.column("system").to(dtype=torch.long) + unique_labels = torch.unique(system_labels) + labels_are_known = torch.isin(unique_labels, system_ids) + if not labels_are_known.all(): + unknown_labels = unique_labels[~labels_are_known] + raise ValueError( + f"Block samples contain system labels {unknown_labels.tolist()} that are " + f"not in system_ids={system_ids.tolist()}. Every sample must be " + f"assigned to a system in the transformation." + ) + return [ + torch.nonzero(system_labels == system_id, as_tuple=False).reshape(-1) + for system_id in system_ids + ] + + +def _gradient_row_indices_by_system( + grad_block: TensorBlock, + parent_block: TensorBlock, + system_ids: torch.Tensor, +) -> list[torch.Tensor]: + """Group gradient rows by the system of their referenced value row.""" + if len(system_ids) == 1: + return [ + torch.arange(grad_block.values.shape[0], device=grad_block.values.device) + ] + + if "system" not in parent_block.samples.names: + raise ValueError( + "Rotational augmentation expects the values samples to include a 'system' " + "dimension when transforming gradients of multiple systems." + ) + + parent_system_labels = parent_block.samples.column("system").to(dtype=torch.long) + parent_value_rows = grad_block.samples.column("sample").to(dtype=torch.long) + gradient_system_labels = parent_system_labels[parent_value_rows] + + return [ + torch.nonzero( + gradient_system_labels == system_id, + as_tuple=False, + ).reshape(-1) + for system_id in system_ids + ] + + +def transform_system(system: System, transformation: O3Transformation) -> System: + """Apply an O(3) transformation to a single System. + + Positions, cell vectors, neighbor-list displacements, and custom data following + :ref:`o3-conventions` are transformed. Atomic types and periodic-boundary flags + are preserved. + + :param system: input system + :param transformation: O(3) transformation to apply, matching + ``system.positions`` in dtype and device + :return: new System with transformed geometry + """ + if ( + system.positions.dtype != transformation.dtype + or system.positions.device != transformation.device + ): + raise ValueError( + f"System has positions with dtype/device " + f"({system.positions.dtype}, {system.positions.device}) differing " + f"from the transformations ({transformation.dtype}, " + f"{transformation.device})." + ) + + new_system = System( + positions=transformation.transform_cartesian(system.positions), + types=system.types, + cell=transformation.transform_cartesian(system.cell), + pbc=system.pbc, + ) + + for data_name in system.known_data(): + data = system.get_data(data_name) + new_system.add_data( + data_name, transform_tensor(data, [system], [transformation]) + ) + + for options in system.known_neighbor_lists(): + neighbors = system.get_neighbor_list(options) + # neighbor vectors are stored as (N, 3, 1); squeeze/unsqueeze around the matmul + # Detach the input graph before registering the rotated values below. + neighbors_values = neighbors.values.detach().squeeze(-1) + new_values = transformation.transform_cartesian(neighbors_values) + rotated_neighbors = TensorBlock( + values=new_values.unsqueeze(-1), + samples=neighbors.samples, + components=neighbors.components, + properties=neighbors.properties, + ) + register_autograd_neighbors(new_system, rotated_neighbors) + new_system.add_neighbor_list(options, rotated_neighbors) + + return new_system + + +def _contract_component_axes( + values: torch.Tensor, + matrices: list[torch.Tensor], +) -> torch.Tensor: + """Rotate each component axis of ``values`` by its matrix. + + ``values`` has shape ``(n_rows, d_1, ..., d_k, n_properties)`` and ``matrices[j]`` + (shape ``(d_j, d_j)``) is contracted with component axis ``j`` as + ``out[..., A, ...] = sum_a matrices[j][A, a] * values[..., a, ...]``. + + :param values: values tensor of a value or gradient block + :param matrices: one rotation matrix per component axis (empty for scalars) + :return: rotated values, same shape as the input + """ + # Reserve einsum indices for all ten component axes supported by Metatomic. + _EINSUM_IN = "abcdefghjk" + _EINSUM_OUT = "ABCDEFGHIJ" + + if len(matrices) == 0: + return values + n_axes = len(matrices) + if n_axes > len(_EINSUM_IN): + raise ValueError(f"can not transform a tensor with {n_axes} component axes") + in_subscript = "i" + _EINSUM_IN[:n_axes] + "p" + out_subscript = "i" + _EINSUM_OUT[:n_axes] + "p" + matrix_subscripts = [_EINSUM_OUT[j] + _EINSUM_IN[j] for j in range(n_axes)] + equation = ",".join(matrix_subscripts + [in_subscript]) + "->" + out_subscript + return torch.einsum(equation, *matrices, values) + + +def _component_axis_suffix(axis_name: str, prefix: str) -> tuple[bool, str]: + """Match a component-axis name and return its supported suffix.""" + suffixes = ["", "_1", "_2", "_3", "_4", "_5", "_6", "_7", "_8", "_9"] + for suffix in suffixes: + if axis_name == prefix + suffix: + return True, suffix + return False, "" + + +def _validate_component_axis_metadata( + components: list[Labels], + key: LabelsEntry, +) -> list[tuple[bool, int, int]]: + """Validate component axes and return ``(is_spherical, ell, sigma)`` metadata.""" + if len(components) > 10: + raise ValueError( + f"can not transform a tensor with {len(components)} component axes; " + "at most 10 are supported" + ) + + metadata: list[tuple[bool, int, int]] = [] + for component in components: + axis_name = component.names[0] + is_cartesian, _ = _component_axis_suffix(axis_name, "xyz") + is_spherical, suffix = _component_axis_suffix(axis_name, "o3_mu") + if is_cartesian: + expected_labels = torch.arange( + 3, + device=component.values.device, + dtype=component.values.dtype, + ) + if not torch.equal(component.values[:, 0], expected_labels): + raise ValueError( + f"Cartesian component axis '{axis_name}' must use labels " + "[0, 1, 2] in x, y, z order." + ) + metadata.append((False, 0, 1)) + elif is_spherical: + ell = _validate_nonnegative_integer( + "ell", + int(key["o3_lambda" + suffix]), + ) + sigma = int(key["o3_sigma" + suffix]) + _spherical_parity_factor(ell, sigma, is_improper=False) + + expected_labels = torch.arange( + -ell, + ell + 1, + device=component.values.device, + dtype=component.values.dtype, + ) + if not torch.equal(component.values[:, 0], expected_labels): + raise ValueError( + f"Spherical component axis '{axis_name}' for ell={ell} must use " + f"labels from {-ell} through {ell} in ascending order." + ) + metadata.append((True, ell, sigma)) + else: + raise ValueError( + f"Found a component axis '{axis_name}', which is neither a Cartesian " + "('xyz'/'xyz_1'/'xyz_2'/...) nor spherical ('o3_mu'/'o3_mu_1'/...) " + "axis; it can not be transformed." + ) + + return metadata + + +def _max_o3_lambda_in_tensor(tensor: TensorMap) -> int: + """Return the largest spherical rank in block values or attached gradients. + + A TensorMap containing only scalar or Cartesian component axes returns ``-1``. + """ + max_o3_lambda = -1 + for key, block in tensor.items(): + metadata = _validate_component_axis_metadata(block.components, key) + for is_spherical, ell, _sigma in metadata: + if is_spherical and ell > max_o3_lambda: + max_o3_lambda = ell + + for _gradient_name, gradient in block.gradients(): + gradient_metadata = _validate_component_axis_metadata( + gradient.components, + key, + ) + for is_spherical, ell, _sigma in gradient_metadata: + if is_spherical and ell > max_o3_lambda: + max_o3_lambda = ell + + return max_o3_lambda + + +def _axis_matrices_and_parity( + metadata: list[tuple[bool, int, int]], + transformation: O3Transformation, +) -> tuple[list[torch.Tensor], int]: + """Return the axis matrices and their combined spherical parity factor.""" + matrices: list[torch.Tensor] = [] + parity = 1 + for is_spherical, ell, sigma in metadata: + if is_spherical: + matrices.append(transformation._wigner_D_cache_entry(ell)) + parity *= _spherical_parity_factor( + ell, + sigma, + transformation.is_improper, + ) + else: + matrices.append(transformation._matrix) + + return matrices, parity + + +def _transform_component_values( + values: torch.Tensor, + components: list[Labels], + key: LabelsEntry, + row_indices: list[torch.Tensor], + transformations: list[O3Transformation], +) -> torch.Tensor: + """Rotate value or gradient rows with their assigned transformation.""" + metadata = _validate_component_axis_metadata(components, key) + new_values = values.clone() + for system_index, rows in enumerate(row_indices): + if len(rows) == 0: + continue + matrices, parity = _axis_matrices_and_parity( + metadata, + transformations[system_index], + ) + rotated = _contract_component_axes(values[rows], matrices) + if parity != 1: + rotated = rotated * parity + new_values[rows] = rotated + return new_values + + +def transform_block( + key: LabelsEntry, + block: TensorBlock, + systems: list[System], + transformations: list[O3Transformation], + system_ids: list[int] | torch.Tensor | None = None, +) -> TensorBlock: + """Apply per-system O(3) transformations to a block and its gradients. + + With one system, the ``"system"`` sample label is optional and ignored, as in + :py:func:`transform_tensor`. + + :param key: parent block key, supplying the O(3) labels required by spherical + component axes + :param block: block to transform + :param systems: systems corresponding positionally to ``transformations`` + :param transformations: one O(3) transformation per system, matching + ``block.values`` in dtype and device + :param system_ids: one distinct integer ``"system"`` sample label per system; + entry ``i`` is paired with ``transformations[i]``. A tensor argument must + be one-dimensional and use the same device as ``block.values``. Defaults + to ``range(len(systems))`` + :return: block with transformed values and gradients and unchanged labels; when + ``systems`` is empty, the block is unchanged + """ + system_ids = _validate_system_ids( + systems, + transformations, + system_ids, + expected_device=block.values.device, + ) + if len(systems) == 0: + return block + + _validate_transformations_dtype_device( + transformations, + expected_dtype=block.values.dtype, + expected_device=block.values.device, + ) + + return _transform_block_impl(key, block, transformations, system_ids) + + +def _transform_block_impl( + key: LabelsEntry, + block: TensorBlock, + transformations: list[O3Transformation], + system_ids: torch.Tensor, +) -> TensorBlock: + """Transform block values and gradients using validated system assignments.""" + value_sample_indices = _value_row_indices_by_system(block, system_ids) + new_block = TensorBlock( + values=_transform_component_values( + block.values, + block.components, + key, + value_sample_indices, + transformations, + ), + samples=block.samples, + components=block.components, + properties=block.properties, + ) + for gradient_name, gradient in block.gradients(): + gradient_sample_indices = _gradient_row_indices_by_system( + gradient, + block, + system_ids, + ) + new_block.add_gradient( + gradient_name, + TensorBlock( + values=_transform_component_values( + gradient.values, + gradient.components, + key, + gradient_sample_indices, + transformations, + ), + samples=gradient.samples, + components=gradient.components, + properties=gradient.properties, + ), + ) + return new_block + + +def transform_tensor( + tensor: TensorMap, + systems: list[System], + transformations: list[O3Transformation], + system_ids: list[int] | torch.Tensor | None = None, +) -> TensorMap: + """Apply per-system O(3) transformations to a TensorMap and its gradients. + + Scalar, Cartesian, and spherical data are identified by their component-axis + names, following :ref:`o3-conventions`; one :py:class:`TensorMap` may contain + all three kinds of data. At most ten component axes are supported in one + value or gradient block. + + With multiple systems, the ``"system"`` sample label assigns each value sample + to a transformation: samples labelled ``system_ids[i]`` use + ``transformations[i]``. A block may contain samples for only some of the + systems, but every ``"system"`` label present in the block must appear in + ``system_ids``. A gradient sample uses the same transformation as the parent + value sample referenced by its ``"sample"`` label. With one system, the + ``"system"`` label is optional and ignored. + + :param tensor: TensorMap to transform + :param systems: systems corresponding positionally to ``transformations`` + :param transformations: one O(3) transformation per system, matching the tensor + values in dtype and device when present + :param system_ids: one distinct integer ``"system"`` sample label per system; + entry ``i`` is paired with ``transformations[i]``. A tensor argument must + be one-dimensional and use the same device as the tensor values. Defaults + to ``range(len(systems))`` + :return: transformed TensorMap with the same keys and global information; when + ``systems`` is empty, the tensor is unchanged + """ + if len(tensor) != 0: + system_ids_device = tensor.block(0).values.device + elif len(transformations) != 0: + system_ids_device = transformations[0].device + else: + system_ids_device = None + + system_ids = _validate_system_ids( + systems, + transformations, + system_ids, + expected_device=system_ids_device, + ) + if len(systems) == 0: + return tensor + + if len(tensor) != 0: + values = tensor.block(0).values + _validate_transformations_dtype_device( + transformations, + expected_dtype=values.dtype, + expected_device=values.device, + ) + + new_blocks = [ + _transform_block_impl(key, block, transformations, system_ids) + for key, block in tensor.items() + ] + transformed = TensorMap(keys=tensor.keys, blocks=new_blocks) + for info_key, info_value in tensor.info().items(): + transformed.set_info(info_key, info_value) + + return transformed + + +def _transformation_indices( + samples: Labels, + n_transformations: int, +) -> torch.Tensor: + """Map sample rows to local transformation indices.""" + if n_transformations <= 0: + raise ValueError("n_transformations must be positive") + if n_transformations == 1: + return torch.zeros( + len(samples), + dtype=torch.long, + device=samples.device, + ) + if "system" not in samples.names: + raise ValueError("multiple transformations require a 'system' sample dimension") + + indices = samples.column("system").to(dtype=torch.long) + if bool(torch.any((indices < 0) | (indices >= n_transformations)).item()): + raise ValueError("sample system indices exceed the transformation batch") + return indices + + +def _transform_component_values_with_precomputed_matrices( + values: torch.Tensor, + components: list[Labels], + key: LabelsEntry, + transformation_indices: torch.Tensor, + matrices: torch.Tensor, + wigner_matrices: list[torch.Tensor], + is_improper: bool, +) -> torch.Tensor: + """Transform component axes with precomputed O(3) matrices.""" + metadata = _validate_component_axis_metadata(components, key) + if len(metadata) == 0: + return values.clone() + + transformed = values + parity = 1 + for component_index, (is_spherical, ell, sigma) in enumerate(metadata): + if is_spherical: + if ell >= len(wigner_matrices): + raise ValueError("spherical rank exceeds the Wigner-D storage") + axis_matrices = wigner_matrices[ell] + parity *= _spherical_parity_factor(ell, sigma, is_improper) + else: + axis_matrices = matrices + + component_axis = component_index + 1 + moved = torch.movedim(transformed, component_axis, -1) + moved_shape = moved.shape + flattened = moved.flatten(start_dim=1, end_dim=-2) + matrices_for_rows = axis_matrices.index_select( + 0, + transformation_indices, + ) + transformed = torch.bmm( + flattened, + matrices_for_rows.transpose(1, 2), + ) + transformed = transformed.reshape(moved_shape) + transformed = torch.movedim(transformed, -1, component_axis) + + if parity != 1: + transformed = transformed * parity + return transformed + + +def _transform_tensor_with_precomputed_matrices( + tensor: TensorMap, + matrices: torch.Tensor, + wigner_matrices: list[torch.Tensor], + is_improper: bool, +) -> TensorMap: + """Transform a TensorMap using precomputed matrices from one O(3) coset. + + ``matrices[i]`` is the actual Cartesian operation for local system ``i``, + while ``wigner_matrices[ell][i]`` is the Wigner-D matrix for its proper + rotational part. Every operation in the batch must be either proper or + improper, as selected by ``is_improper``. + + With multiple operations, ``"system"`` sample labels are local indices into + the matrix batch. A singleton batch does not require this sample dimension. + The caller chooses the transformation direction by supplying either the + forward matrices or their inverses. + """ + if ( + matrices.dim() != 3 + or matrices.size(0) == 0 + or matrices.size(1) != 3 + or matrices.size(2) != 3 + ): + raise ValueError("matrices must have shape (N, 3, 3) with N > 0") + if matrices.dtype != torch.float32 and matrices.dtype != torch.float64: + raise TypeError("matrices must use float32 or float64") + if len(tensor) != 0: + reference_values = tensor.block(0).values + if ( + matrices.dtype != reference_values.dtype + or matrices.device != reference_values.device + ): + raise ValueError("tensor and matrices must have the same dtype and device") + + blocks: list[TensorBlock] = [] + for key, block in tensor.items(): + value_indices = _transformation_indices( + block.samples, + matrices.size(0), + ) + new_block = TensorBlock( + values=_transform_component_values_with_precomputed_matrices( + block.values, + block.components, + key, + value_indices, + matrices, + wigner_matrices, + is_improper, + ), + samples=block.samples, + components=block.components, + properties=block.properties, + ) + + for gradient_name, gradient in block.gradients(): + parent_rows = gradient.samples.column("sample").to(dtype=torch.long) + gradient_indices = value_indices.index_select(0, parent_rows) + new_block.add_gradient( + gradient_name, + TensorBlock( + values=_transform_component_values_with_precomputed_matrices( + gradient.values, + gradient.components, + key, + gradient_indices, + matrices, + wigner_matrices, + is_improper, + ), + samples=gradient.samples, + components=gradient.components, + properties=gradient.properties, + ), + ) + blocks.append(new_block) + + transformed = TensorMap(tensor.keys, blocks) + for info_name, info_value in tensor.info().items(): + transformed.set_info(info_name, info_value) + return transformed diff --git a/python/metatomic_torch/metatomic/torch/o3/_transformations.py b/python/metatomic_torch/src/metatomic_torch/o3/_transformations.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/o3/_transformations.py rename to python/metatomic_torch/src/metatomic_torch/o3/_transformations.py diff --git a/python/metatomic_torch/metatomic/torch/o3/_utils.py b/python/metatomic_torch/src/metatomic_torch/o3/_utils.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/o3/_utils.py rename to python/metatomic_torch/src/metatomic_torch/o3/_utils.py diff --git a/python/metatomic_torch/metatomic/torch/o3/_wigner.py b/python/metatomic_torch/src/metatomic_torch/o3/_wigner.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/o3/_wigner.py rename to python/metatomic_torch/src/metatomic_torch/o3/_wigner.py diff --git a/python/metatomic_torch/metatomic/torch/serialization.py b/python/metatomic_torch/src/metatomic_torch/serialization.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/serialization.py rename to python/metatomic_torch/src/metatomic_torch/serialization.py diff --git a/python/metatomic_torch/metatomic/torch/systems_to_torch.py b/python/metatomic_torch/src/metatomic_torch/systems_to_torch.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/systems_to_torch.py rename to python/metatomic_torch/src/metatomic_torch/systems_to_torch.py diff --git a/python/metatomic_torch/metatomic/torch/utils.py b/python/metatomic_torch/src/metatomic_torch/utils.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/utils.py rename to python/metatomic_torch/src/metatomic_torch/utils.py diff --git a/python/metatomic_torch/metatomic/torch/version.py b/python/metatomic_torch/src/metatomic_torch/version.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/version.py rename to python/metatomic_torch/src/metatomic_torch/version.py diff --git a/python/metatomic_torch/tests/examples.py b/python/metatomic_torch/tests/examples.py index 6a5182ab7..c9a810421 100644 --- a/python/metatomic_torch/tests/examples.py +++ b/python/metatomic_torch/tests/examples.py @@ -14,7 +14,7 @@ EXAMPLES = os.path.abspath( - os.path.join(os.path.dirname(__file__), "..", "..", "examples") + os.path.join(os.path.dirname(__file__), "..", "..", "..", "examples") ) DOCS = os.path.abspath( @@ -24,14 +24,14 @@ def test_export_atomistic_model(tmp_path): """ - Check if the model defined in ``python/examples/1-export-atomistic-model.py`` works + Check if the model defined in ``examples/torch/1-export-atomistic-model.py`` works """ os.chdir(tmp_path) # import example from full path spec = importlib.util.spec_from_file_location( "export_atomistic_model", - os.path.join(EXAMPLES, "1-export-atomistic-model.py"), + os.path.join(EXAMPLES, "torch", "1-export-atomistic-model.py"), ) export_atomistic_model = importlib.util.module_from_spec(spec) diff --git a/python/metatomic_torchsim/pyproject.toml b/python/metatomic_torchsim/pyproject.toml index 62c55dc3b..7dd83c12e 100644 --- a/python/metatomic_torchsim/pyproject.toml +++ b/python/metatomic_torchsim/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "metatomic-torchsim" dynamic = ["version", "authors", "dependencies"] -requires-python = ">=3.10" +requires-python = ">=3.11" readme = "README.md" license = "BSD-3-Clause" diff --git a/python/metatomic_torchsim/setup.py b/python/metatomic_torchsim/setup.py index 98566e7b3..3d62963c8 100644 --- a/python/metatomic_torchsim/setup.py +++ b/python/metatomic_torchsim/setup.py @@ -1,4 +1,5 @@ import os +import pathlib import subprocess import sys @@ -7,8 +8,8 @@ from setuptools.command.sdist import sdist -ROOT = os.path.realpath(os.path.dirname(__file__)) -METATOMIC_TORCH = os.path.realpath(os.path.join(ROOT, "..", "metatomic_torch")) +ROOT = pathlib.Path(__file__).parent.resolve() +METATOMIC_TORCH = (ROOT / ".." / "metatomic_torch").resolve() METATOMIC_TORCHSIM_VERSION = "0.1.5" @@ -38,15 +39,15 @@ def git_version_info(): """ TAG_PREFIX = "metatomic-torchsim-v" - if os.path.exists("git_version_info"): + if (ROOT / "git_version_info").exists(): # we are building from a sdist, without git available, but the git # version was recorded in the `git_version_info` file - with open("git_version_info") as fd: + with open(ROOT / "git_version_info") as fd: n_commits = int(fd.readline().strip()) git_hash = fd.readline().strip() else: - script = os.path.join(ROOT, "..", "..", "scripts", "git-version-info.py") - assert os.path.exists(script) + script = (ROOT / ".." / ".." / "scripts" / "git-version-info.py").resolve() + assert script.exists() output = subprocess.run( [sys.executable, script, TAG_PREFIX], @@ -61,12 +62,15 @@ def git_version_info(): f"stdout: {output.stdout}\n" f"stderr: {output.stderr}\n" ) - elif output.stderr: + if output.stderr: print(output.stderr, file=sys.stderr) + + lines = output.stdout.splitlines() + if len(lines) < 2: + # the script gave up early (cf. `warn_and_exit`) n_commits = 0 git_hash = "" else: - lines = output.stdout.splitlines() n_commits = int(lines[0].strip()) git_hash = lines[1].strip() @@ -102,7 +106,7 @@ def create_version_number(version): if __name__ == "__main__": - with open(os.path.join(ROOT, "AUTHORS")) as fd: + with open(ROOT / "AUTHORS") as fd: authors = fd.read().splitlines() install_requires = [ @@ -113,9 +117,9 @@ def create_version_number(version): # when packaging a sdist for release, we should never use local dependencies METATOMIC_NO_LOCAL_DEPS = os.environ.get("METATOMIC_NO_LOCAL_DEPS", "0") == "1" - if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_TORCH): + if not METATOMIC_NO_LOCAL_DEPS and METATOMIC_TORCH.exists(): # we are building from a git checkout or full repo archive - install_requires.append(f"metatomic-torch @ file://{METATOMIC_TORCH}") + install_requires.append(f"metatomic-torch @ {METATOMIC_TORCH.as_uri()}") else: # we are building from a sdist/installing from a wheel install_requires.append("metatomic-torch >=0.1.12,<0.2") diff --git a/python/metatomic_torchsim/metatomic_torchsim/__init__.py b/python/metatomic_torchsim/src/metatomic_torchsim/__init__.py similarity index 100% rename from python/metatomic_torchsim/metatomic_torchsim/__init__.py rename to python/metatomic_torchsim/src/metatomic_torchsim/__init__.py diff --git a/python/metatomic_torchsim/metatomic_torchsim/_model.py b/python/metatomic_torchsim/src/metatomic_torchsim/_model.py similarity index 100% rename from python/metatomic_torchsim/metatomic_torchsim/_model.py rename to python/metatomic_torchsim/src/metatomic_torchsim/_model.py diff --git a/python/metatomic_torchsim/metatomic_torchsim/_neighbors.py b/python/metatomic_torchsim/src/metatomic_torchsim/_neighbors.py similarity index 100% rename from python/metatomic_torchsim/metatomic_torchsim/_neighbors.py rename to python/metatomic_torchsim/src/metatomic_torchsim/_neighbors.py diff --git a/python/scripts/gcc11-manylinux_2_28_aarch64/Dockerfile b/python/scripts/gcc11-manylinux_2_28_aarch64/Dockerfile deleted file mode 100644 index 3d22bae84..000000000 --- a/python/scripts/gcc11-manylinux_2_28_aarch64/Dockerfile +++ /dev/null @@ -1,9 +0,0 @@ -# Use manylinux docker image as a base -FROM quay.io/pypa/manylinux_2_28_aarch64 - -# Install an older compiler. The default compiler (gcc-14) introduces calls to -# `__cxa_call_terminate` which is not available in ubuntu 22.04 libstdc++ -ARG DEVTOOLSET_VERSION=11 -RUN yum install -y gcc-toolset-${DEVTOOLSET_VERSION}-toolchain -ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH -ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH diff --git a/python/scripts/gcc11-manylinux_2_28_x86_64/Dockerfile b/python/scripts/gcc11-manylinux_2_28_x86_64/Dockerfile deleted file mode 100644 index 5fe558f38..000000000 --- a/python/scripts/gcc11-manylinux_2_28_x86_64/Dockerfile +++ /dev/null @@ -1,9 +0,0 @@ -# Use manylinux docker image as a base -FROM quay.io/pypa/manylinux_2_28_x86_64 - -# Install an older compiler. The default compiler (gcc-14) introduces calls to -# `__cxa_call_terminate` which is not available in ubuntu 22.04 libstdc++ -ARG DEVTOOLSET_VERSION=11 -RUN yum install -y gcc-toolset-${DEVTOOLSET_VERSION}-toolchain -ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH -ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH diff --git a/python/scripts/rustc-manylinux_2_28_aarch64/Dockerfile b/python/scripts/rustc-manylinux_2_28_aarch64/Dockerfile new file mode 100644 index 000000000..8aca91391 --- /dev/null +++ b/python/scripts/rustc-manylinux_2_28_aarch64/Dockerfile @@ -0,0 +1,19 @@ +# Use manylinux docker image as a base +FROM quay.io/pypa/manylinux_2_28_aarch64 + +RUN yum install git -y +RUN git config --global --add safe.directory /code + +# Download rustup-init and install +ARG RUST_TOOLCHAIN_VERSION=1.96 +RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --profile minimal --default-toolchain ${RUST_TOOLCHAIN_VERSION} + +ENV PATH="/root/.cargo/bin:${PATH}" +ENV RUST_BUILD_TARGET="aarch64-unknown-linux-gnu" + +# Install an older C++ compiler. The default compiler (gcc-14) introduces calls +# to `__cxa_call_terminate` which is not available in ubuntu 22.04 libstdc++ +ARG DEVTOOLSET_VERSION=11 +RUN yum install -y gcc-toolset-${DEVTOOLSET_VERSION}-toolchain +ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH +ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH diff --git a/python/scripts/rustc-manylinux_2_28_x86_64/Dockerfile b/python/scripts/rustc-manylinux_2_28_x86_64/Dockerfile new file mode 100644 index 000000000..22f7a1448 --- /dev/null +++ b/python/scripts/rustc-manylinux_2_28_x86_64/Dockerfile @@ -0,0 +1,19 @@ +# Use manylinux docker image as a base +FROM quay.io/pypa/manylinux_2_28_x86_64 + +RUN yum install git -y +RUN git config --global --add safe.directory /code + +# Download rustup-init and install +ARG RUST_TOOLCHAIN_VERSION=1.96 +RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --profile minimal --default-toolchain ${RUST_TOOLCHAIN_VERSION} + +ENV PATH="/root/.cargo/bin:${PATH}" +ENV RUST_BUILD_TARGET="x86_64-unknown-linux-gnu" + +# Install an older C++ compiler. The default compiler (gcc-14) introduces calls +# to `__cxa_call_terminate` which is not available in ubuntu 22.04 libstdc++ +ARG DEVTOOLSET_VERSION=11 +RUN yum install -y gcc-toolset-${DEVTOOLSET_VERSION}-toolchain +ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH +ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH diff --git a/python/tests/run-python-tests.rs b/python/tests/run-python-tests.rs new file mode 100644 index 000000000..8d52a6f83 --- /dev/null +++ b/python/tests/run-python-tests.rs @@ -0,0 +1,23 @@ +use std::path::PathBuf; +use std::process::Command; + +#[test] +fn run_python_tests() { + let tox = which::which("tox").expect("could not find tox"); + + let mut root = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + root.pop(); + + let mut tox = Command::new(tox); + tox.arg("--"); + if cfg!(debug_assertions) { + // assume that debug assertions means that we are building the code + // in debug mode, even if optimizations could be enabled + tox.env("METATOMIC_BUILD_TYPE", "debug"); + } else { + tox.env("METATOMIC_BUILD_TYPE", "release"); + } + tox.current_dir(&root); + let status = tox.status().expect("failed to run tox"); + assert!(status.success()); +} diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 000000000..c7ad93baf --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1 @@ +disable_all_formatting = true diff --git a/scripts/check-c-api-docs.py b/scripts/check-c-api-docs.py new file mode 100755 index 000000000..73ee7d921 --- /dev/null +++ b/scripts/check-c-api-docs.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python +""" +A small script checking that all the C API functions are documented +""" + +import os +import sys + +from pycparser import c_ast, parse_file + + +ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..")) +C_API_DOCS = os.path.join(ROOT, "docs", "src", "core", "reference", "c") +FAKE_INCLUDES = [os.path.join(ROOT, "scripts", "include")] +METATOMIC_HEADER = os.path.relpath( + os.path.join(ROOT, "metatomic-core", "include", "metatomic.h") +) + + +ERRORS = 0 + + +def error(message): + global ERRORS + ERRORS += 1 + print(message) + + +def documented_functions(): + functions = [] + + for root, _, paths in os.walk(C_API_DOCS): + for path in paths: + with open(os.path.join(root, path), encoding="utf8") as fd: + for line in fd: + if line.startswith(".. doxygenfunction::"): + name = line.split()[2] + functions.append(name) + + return functions + + +def functions_in_outline(): + # function from the "miscellaneous" section of the docs don't require an outline + # (since they are not related to a specific struct type) + functions = [ + "mta_version", + "mta_last_error", + "mta_set_last_error", + "mta_string_create", + "mta_string_free", + "mta_string_view", + "mta_format_metadata", + "mta_unit_conversion_factor", + ] + + for root, _, paths in os.walk(C_API_DOCS): + for path in paths: + with open(os.path.join(root, path), encoding="utf8") as fd: + for line in fd: + if ":c:func:" in line: + name = line.split("`")[1] + functions.append(name) + return functions + + +def all_functions(): + cpp_args = ["-E"] + for path in FAKE_INCLUDES: + cpp_args += ["-I", path] + ast = parse_file(METATOMIC_HEADER, use_cpp=True, cpp_path="gcc", cpp_args=cpp_args) + + functions = [] + + class AstVisitor(c_ast.NodeVisitor): + def visit_Decl(self, node): + if not isinstance(node.type, c_ast.FuncDecl): + return + + if not node.name.startswith("mta_"): + return + + functions.append(node.name) + + visitor = AstVisitor() + visitor.visit(ast) + + return functions + + +if __name__ == "__main__": + docs = documented_functions() + outline = functions_in_outline() + for function in all_functions(): + if function not in docs: + error("Missing documentation for {}".format(function)) + if function not in outline: + error("Missing outline for {}".format(function)) + + if ERRORS != 0: + sys.exit(1) diff --git a/scripts/clean-python.sh b/scripts/clean-python.sh index ba6a9e9f5..33fafec32 100755 --- a/scripts/clean-python.sh +++ b/scripts/clean-python.sh @@ -11,12 +11,21 @@ cd "$ROOT_DIR" rm -rf dist rm -rf build rm -rf docs/build -rm -rf docs/src/examples +rm -rf docs/src/examples/*/ rm -rf docs/src/sg_execution_times.rst +rm -rf python/metatomic_core/dist +rm -rf python/metatomic_core/build + rm -rf python/metatomic_torch/dist rm -rf python/metatomic_torch/build +rm -rf python/metatomic_ase/dist +rm -rf python/metatomic_ase/build + +rm -rf python/metatomic_torchsim/dist +rm -rf python/metatomic_torchsim/build + find . -name "*.egg-info" -exec rm -rf "{}" + find . -name "__pycache__" -exec rm -rf "{}" + find . -name ".coverage" -exec rm -rf "{}" + diff --git a/scripts/git-version-info.py b/scripts/git-version-info.py index db7253223..87d384540 100755 --- a/scripts/git-version-info.py +++ b/scripts/git-version-info.py @@ -174,7 +174,14 @@ def git_hash_all_code(): git_env = os.environ.copy() git_env["GIT_INDEX_FILE"] = tmp.name - run_subprocess(["git", "add", "--all"], env=git_env, cwd=worktree) + # `core.safecrlf=false` since we are only staging files to compute a hash, + # and never write them back to the working tree: the end-of-line warnings + # git would emit here (on Windows) are only noise. + run_subprocess( + ["git", "-c", "core.safecrlf=false", "add", "--all"], + env=git_env, + cwd=worktree, + ) output = run_subprocess(["git", "write-tree"], env=git_env, cwd=worktree) short_hash = output.stdout[:7] diff --git a/scripts/include/README b/scripts/include/README new file mode 100644 index 000000000..d56dd0788 --- /dev/null +++ b/scripts/include/README @@ -0,0 +1,4 @@ +This directory contains fake headers used to allow pycparser to parse the code +without having to deal with all the complexity of actual stdlib implementations + +See https://eli.thegreenplace.net/2015/on-parsing-c-type-declarations-and-fake-headers for more information diff --git a/scripts/include/metatensor.h b/scripts/include/metatensor.h new file mode 100644 index 000000000..56c085abf --- /dev/null +++ b/scripts/include/metatensor.h @@ -0,0 +1,11 @@ +// empty header with minimal content, to be used to parse metatomic.h + +typedef struct mts_labels_t mts_labels_t; +typedef struct mts_block_t mts_block_t; +typedef struct mts_tensormap_t mts_tensormap_t; + +typedef void (*mts_create_array_callback_t)(void*); +typedef void (*mts_realloc_buffer_t)(void*); + + +typedef struct DLManagedTensorVersioned DLManagedTensorVersioned; diff --git a/python/metatomic_torch/metatomic/__init__.py b/scripts/include/metatomic/version.h similarity index 100% rename from python/metatomic_torch/metatomic/__init__.py rename to scripts/include/metatomic/version.h diff --git a/scripts/include/stdarg.h b/scripts/include/stdarg.h new file mode 100644 index 000000000..e69de29bb diff --git a/scripts/include/stdbool.h b/scripts/include/stdbool.h new file mode 100644 index 000000000..3bd41ef29 --- /dev/null +++ b/scripts/include/stdbool.h @@ -0,0 +1 @@ +typedef _Bool bool; \ No newline at end of file diff --git a/scripts/include/stddef.h b/scripts/include/stddef.h new file mode 100644 index 000000000..48b3db663 --- /dev/null +++ b/scripts/include/stddef.h @@ -0,0 +1,6 @@ +#ifndef FAKE_STDDEF_H +#define FAKE_STDDEF_H + +typedef void nullptr_t; + +#endif /* FAKE_STDDEF_H */ diff --git a/scripts/include/stdint.h b/scripts/include/stdint.h new file mode 100644 index 000000000..43ccc01dd --- /dev/null +++ b/scripts/include/stdint.h @@ -0,0 +1,7 @@ +typedef int uint64_t; +typedef int int64_t; +typedef int int32_t; +typedef int uint32_t; +typedef int uint16_t; +typedef int uint8_t; +typedef int uintptr_t; diff --git a/scripts/include/stdio.h b/scripts/include/stdio.h new file mode 100644 index 000000000..e69de29bb diff --git a/scripts/include/stdlib.h b/scripts/include/stdlib.h new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/scripts/include/stdlib.h @@ -0,0 +1 @@ + diff --git a/scripts/package-core.sh b/scripts/package-core.sh new file mode 100755 index 000000000..76cbc2f1d --- /dev/null +++ b/scripts/package-core.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash + +# This script creates an archive containing the sources for the metatomic-core +# Rust crate, and copy it to the path given as argument + +set -eux + +OUTPUT_DIR="$1" +mkdir -p "$OUTPUT_DIR" +OUTPUT_DIR=$(cd "$OUTPUT_DIR" 2>/dev/null && pwd) + +ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. && pwd) + +rm -rf "$ROOT_DIR/target/package" +cd "$ROOT_DIR/metatomic-core" + +# print the version of cargo we use for debugging purposes. +cargo --version + +# Package metatomic-core using cargo tools, and add a file for +# `n_commits_since_last_tag` +cargo package --allow-dirty --no-verify + +TMP_DIR=$(mktemp -d) + +cd "$TMP_DIR" +tar xf "$ROOT_DIR"/target/package/metatomic-core-*.crate +ARCHIVE_NAME=$(ls) + +# extract the version part of the package from the .crate file name +VERSION=${ARCHIVE_NAME:15} +ARCHIVE_NAME="metatomic-core-cxx-$VERSION" + +mv metatomic-core-* "$ARCHIVE_NAME" +cp "$ROOT_DIR/LICENSE" "$TMP_DIR/$ARCHIVE_NAME" +cp "$ROOT_DIR/AUTHORS" "$TMP_DIR/$ARCHIVE_NAME" +cp "$ROOT_DIR/README.md" "$TMP_DIR/$ARCHIVE_NAME" + +# Get the git version information, this is used when building the +# code to change the version for development builds +cd "$ROOT_DIR" +./scripts/git-version-info.py "metatomic-core-v" > "$TMP_DIR/$ARCHIVE_NAME/cmake/git_version_info" + +cd "$TMP_DIR" +# Compile metatomic-core as it's own Cargo workspace (otherwise we can not use +# the metatomic rust crate in a project using workspaces). +echo "[workspace]" >> "$ARCHIVE_NAME/Cargo.toml" + +cargo generate-lockfile --manifest-path "$ARCHIVE_NAME/Cargo.toml" + +tar --dereference --create --file "$ARCHIVE_NAME.tar" "$ARCHIVE_NAME" +gzip -9 "$ARCHIVE_NAME.tar" + +cp "$TMP_DIR/$ARCHIVE_NAME.tar.gz" "$OUTPUT_DIR/" diff --git a/scripts/update-declarations.py b/scripts/update-declarations.py new file mode 100755 index 000000000..5da40f166 --- /dev/null +++ b/scripts/update-declarations.py @@ -0,0 +1,363 @@ +#!/usr/bin/env python +""" +This script generate the ctypes declaration corresponding to the C API of metatomic. +""" + +import os + +from pycparser import c_ast, parse_file + + +ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..")) +FAKE_INCLUDES = [ + os.path.join(ROOT, "python", "scripts", "include"), + os.path.join(ROOT, "scripts", "include"), +] +METATOMIC_HEADER = os.path.relpath( + os.path.join(ROOT, "metatomic-core", "include", "metatomic.h") +) + + +# ============================================================================ # +# Shared AST parsing +# ============================================================================ # + + +class Function: + def __init__(self, name, restype): + self.name = name + self.restype = restype + self.args = [] + + def add_arg(self, name, type): + self.args.append((name, type)) + + +class Struct: + def __init__(self, name): + self.name = name + self.members = {} + + def add_member(self, name, type): + self.members[name] = type + + +class Enum: + def __init__(self, name): + self.name = name + self.values = {} + + def add_value(self, name, value): + self.values[name] = value + + +class AstVisitor(c_ast.NodeVisitor): + def __init__(self, *, include_dlpack=True): + self.functions = [] + self.enums = [] + self.structs = [] + self.types = {} + self.defines = {} + + def visit_Decl(self, node): + node_name = node.name + if node_name is None: + node_name = node.type.name + + if not node_name.startswith("mta_"): + return + + if isinstance(node.type, c_ast.Enum): + enum = Enum(node_name) + for enumerator in node.type.values.enumerators: + # Strip C unsigned/long suffixes (e.g. 0U, 1UL) + value = enumerator.value.value.rstrip("UuLl") + enum.add_value(enumerator.name, value) + self.enums.append(enum) + elif isinstance(node.type, c_ast.FuncDecl): + function = Function(node.name, node.type.type) + for parameter in node.type.args.params: + function.add_arg(parameter.name, parameter.type) + self.functions.append(function) + else: + raise RuntimeError(f"Unknown declaration type for {node_name}") + + def visit_Typedef(self, node): + # Extract metatomic stuff only + if not node.name.startswith("mta_"): + return + + if isinstance(node.type.type, c_ast.Enum): + enum = Enum(node.name) + for enumerator in node.type.type.values.enumerators: + # Strip C unsigned/long suffixes (e.g. 0U, 1UL) + value = enumerator.value.value.rstrip("UuLl") + enum.add_value(enumerator.name, value) + self.enums.append(enum) + + elif isinstance(node.type.type, c_ast.Struct): + if node.name.startswith("DLPackExchangeAPI"): + return + + struct = Struct(node.name) + for _, member in node.type.type.children(): + struct.add_member(member.name, member.type) + self.structs.append(struct) + + else: + # keep `node.type` (not `node.type.type`) so that pointer typedefs + # such as `typedef mta_opaque_string_t* mta_string_t` are translated + # to a ctypes pointer instead of the pointed-to type + self.types[node.name] = node.type + + +def _typedecl_name(type): + assert isinstance(type, c_ast.TypeDecl) + if isinstance(type.type, c_ast.Struct): + return type.type.name + elif isinstance(type.type, c_ast.Enum): + return type.type.name + else: + assert len(type.type.names) == 1 + return type.type.names[0] + + +def parse_header(file): + cpp_args = ["-E"] + for path in FAKE_INCLUDES: + cpp_args += ["-I", path] + ast = parse_file(file, use_cpp=True, cpp_path="gcc", cpp_args=cpp_args) + + visitor = AstVisitor() + visitor.visit(ast) + + # `#define` without a value associated + no_value_define = ["METATOMIC_H", "MTA_EXTERN_C"] + + with open(file) as fd: + for line in fd: + if "#define" in line: + split = line.split() + + name = split[1] + if name in no_value_define: + continue + value = split[2] + visitor.defines[name] = value + return visitor + + +# ==================================================================================== # +# Python backend # +# ==================================================================================== # + + +def _py_type_name(name): + if name.startswith("mta_") or name.startswith("mts_") or name.startswith("DL"): + return name + elif name == "uintptr_t": + return "c_uintptr_t" + elif name == "void": + return "None" + elif name == "int8_t": + return "ctypes.c_int8" + elif name == "uint8_t": + return "ctypes.c_uint8" + elif name == "int16_t": + return "ctypes.c_int16" + elif name == "uint16_t": + return "ctypes.c_uint16" + elif name == "int32_t": + return "ctypes.c_int32" + elif name == "uint32_t": + return "ctypes.c_uint32" + elif name == "int64_t": + return "ctypes.c_int64" + elif name == "uint64_t": + return "ctypes.c_uint64" + else: + return "ctypes.c_" + name + + +def _py_funcdecl(type): + restype = _py_type(type.type) + args = [_py_type(t.type) for t in type.args.params] + return f"CFUNCTYPE({restype}, {', '.join(args)})" + + +def _py_type(type): + if isinstance(type, c_ast.PtrDecl): + if isinstance(type.type, c_ast.PtrDecl): + if isinstance(type.type.type, c_ast.TypeDecl): + name = _typedecl_name(type.type.type) + if name == "char": + return "POINTER(ctypes.c_char_p)" + elif name == "uint8_t": + return "POINTER(ctypes.c_char_p)" + name = _py_type_name(name) + return f"POINTER(POINTER({name}))" + elif isinstance(type.type.type, c_ast.PtrDecl): + assert isinstance(type.type.type.type, c_ast.TypeDecl) + assert _typedecl_name(type.type.type.type) == "char" + return "POINTER(POINTER(ctypes.c_char_p))" + elif isinstance(type.type, c_ast.TypeDecl): + name = _typedecl_name(type.type) + if name == "void": + return "ctypes.c_void_p" + elif name == "char": + return "ctypes.c_char_p" + elif name == "uint8_t": + return "ctypes.c_char_p" + else: + return f"POINTER({_py_type_name(name)})" + elif isinstance(type.type, c_ast.FuncDecl): + return _py_funcdecl(type.type) + else: + if isinstance(type, c_ast.TypeDecl): + return _py_type_name(_typedecl_name(type)) + elif isinstance(type, c_ast.IdentifierType): + return _py_type_name(type.names[0]) + elif isinstance(type, c_ast.ArrayDecl): + if isinstance(type.dim, c_ast.Constant): + size = type.dim.value + else: + raise Exception("dynamically sized arrays are not supported") + return f"{_py_type(type.type)} * {size}" + elif isinstance(type, c_ast.FuncDecl): + return _py_funcdecl(type) + raise Exception("Unknown type") + + +def generate_python(data): + outpath = os.path.join( + ROOT, "python", "metatomic_core", "src", "metatomic", "_c_api.py" + ) + with open(outpath, "w") as f: + f.write( + """# fmt: off +# flake8: noqa +\"\"\" +This file declares the C-API corresponding to metatomic.h, in a way compatible +with the ctypes Python module. + +This file is automatically generated by `scripts/update-declarations.py`, +do not edit it manually! +\"\"\" + +import ctypes +import platform +from ctypes import CFUNCTYPE, POINTER + +from ctypes_dlpack import DLDataType, DLDevice, DLManagedTensorVersioned, DLPackVersion +from metatensor._c_api import ( + mts_labels_t, + mts_block_t, + mts_tensormap_t, + mts_realloc_buffer_t, + mts_create_array_callback_t, +) + + +class _EnumType(type(ctypes.c_int32)): + def __new__(metacls, name, bases, namespace): + if "_members_" not in namespace: + members = {} + for key, value in namespace.items(): + if not key.startswith("_"): + members[key] = value + namespace["_members_"] = members + else: + members = namespace["_members_"] + + namespace["_reverse_map_"] = {v: k for k, v in members.items()} + return type(ctypes.c_int32).__new__(metacls, name, bases, namespace) + + def __repr__(self): + return f"" + + +class _Enum(ctypes.c_int32, metaclass=_EnumType): + _members_ = {} + + def __repr__(self): + value_name = self._reverse_map_.get(self.value, str(self.value)) + return f"{self.__class__.__name__}.{value_name}" + + def __eq__(self, other): + if isinstance(other, int): + return self.value == other + if type(self) is type(other): + return self.value == other.value + return NotImplemented + + def __hash__(self): + return hash(self.value) + + +arch = platform.architecture()[0] +if arch == "32bit": + c_uintptr_t = ctypes.c_uint32 +elif arch == "64bit": + c_uintptr_t = ctypes.c_uint64 + +""" + ) + + # Enums + for enum in data.enums: + f.write(f"\n\nclass {enum.name}(_Enum):\n") + for name, value in enum.values.items(): + f.write(f" {name} = {value}\n") + + # structs declartions, without fields + for struct in data.structs: + f.write(f"\n\nclass {struct.name}(ctypes.Structure):\n") + f.write(" pass\n") + + # typedefs + f.write("\n\n") + for name, c_type in data.types.items(): + if name == "mta_status_t": + # this is already defined as an enum + continue + f.write(f"{name} = {_py_type(c_type)}\n") + + # structs fields definitions + f.write("\n") + for struct in data.structs: + if len(struct.members) == 0: + continue + f.write(f"\n{struct.name}._fields_ = [\n") + for name, type in struct.members.items(): + f.write(f' ("{name}", {_py_type(type)}),\n') + f.write("]\n") + + # Functions + f.write("\n\ndef setup_functions(lib):\n") + f.write(" from ._status import check_status\n") + for function in data.functions: + f.write(f"\n lib.{function.name}.argtypes = [") + args = [_py_type(arg[1]) for arg in function.args] + if args == ["None"]: + args = [] + for arg in args: + f.write(f"\n {arg},") + f.write("\n ]\n") + restype = _py_type(function.restype) + if restype == "mta_status_t" and function.name != "mta_last_error": + restype = "check_status" + f.write(f" lib.{function.name}.restype = {restype}\n") + + +# ==================================================================================== # +# main # +# ==================================================================================== # + + +def main(): + data = parse_header(METATOMIC_HEADER) + generate_python(data) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index ced9f7146..69699d06e 100644 --- a/setup.py +++ b/setup.py @@ -1,32 +1,45 @@ import os +import pathlib from setuptools import setup -ROOT = os.path.realpath(os.path.dirname(__file__)) -METATOMIC_TORCH = os.path.join(ROOT, "python", "metatomic_torch") -METATOMIC_TORCHSIM = os.path.join(ROOT, "python", "metatomic_torchsim") +ROOT = pathlib.Path(__file__).parent.resolve() +METATOMIC_CORE = (ROOT / "python" / "metatomic_core").resolve() +METATOMIC_TORCH = (ROOT / "python" / "metatomic_torch").resolve() +METATOMIC_ASE = (ROOT / "python" / "metatomic_ase").resolve() +METATOMIC_TORCHSIM = (ROOT / "python" / "metatomic_torchsim").resolve() if __name__ == "__main__": extras_require = {} + install_requires = [] # when packaging a sdist for release, we should never use local dependencies METATOMIC_NO_LOCAL_DEPS = os.environ.get("METATOMIC_NO_LOCAL_DEPS", "0") == "1" - if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_TORCH): + if not METATOMIC_NO_LOCAL_DEPS and METATOMIC_CORE.exists(): + assert METATOMIC_TORCH.exists() + assert METATOMIC_ASE.exists() + assert METATOMIC_TORCHSIM.exists() + # we are building from a git checkout - extras_require["torch"] = f"metatomic-torch @ file://{METATOMIC_TORCH}" + install_requires.append(f"metatomic-core @ {METATOMIC_CORE.as_uri()}") + extras_require["torch"] = f"metatomic-torch @ {METATOMIC_TORCH.as_uri()}" + extras_require["ase"] = f"metatomic-ase @ {METATOMIC_ASE.as_uri()}" + extras_require["torchsim"] = ( + f"metatomic-torchsim @ {METATOMIC_TORCHSIM.as_uri()}" + ) else: # we are building from a sdist/installing from a wheel - extras_require["torch"] = "metatomic-torch" + install_requires.append("metatomic-core") - if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_TORCHSIM): - extras_require["torchsim"] = f"metatomic-torchsim @ file://{METATOMIC_TORCHSIM}" - else: + extras_require["torch"] = "metatomic-torch" + extras_require["ase"] = "metatomic-ase" extras_require["torchsim"] = "metatomic-torchsim" setup( author=", ".join(open(os.path.join(ROOT, "AUTHORS")).read().splitlines()), + install_requires=install_requires, extras_require=extras_require, ) diff --git a/tox.ini b/tox.ini index 919350ccd..60e4fddbc 100644 --- a/tox.ini +++ b/tox.ini @@ -6,8 +6,7 @@ requires = tox >=4.39 # `tox` in the command-line without anything else envlist = lint - torch-tests-cxx - torch-install-tests-cxx + core-tests torch-tests docs-tests ase-tests @@ -38,92 +37,41 @@ packaging_deps = testing_deps = pytest pytest-cov + pytest-custom_exit_code metatomic_deps = + metatensor-core >=0.2.0,<0.3 + +metatomic_torch_deps = metatensor-torch >=0.10.0,<0.11 metatensor-operations >=0.5.0,<0.6 wigners >=0.4.0 ################################################################################ -##### C++ tests setup ##### +##### Python tests setup ##### ################################################################################ -[testenv:torch-tests-cxx] -description = Run the C++ tests for metatomic-torch +[testenv:core-tests] +description = Run the tests of the metatomic-core Python package deps = - cmake + {[testenv]testing_deps} + {[testenv]packaging_deps} {[testenv]metatomic_deps} - torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.13}.* +changedir = python/metatomic_core commands = - # configure cmake - cmake -B {env_dir}/build metatomic-torch \ - -DCMAKE_BUILD_TYPE=Debug \ - -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \ - -DCMAKE_PREFIX_PATH={env_site_packages_dir}/metatensor/;\ - {env_site_packages_dir}/torch/;\ - {env_site_packages_dir}/metatensor_torch/torch-{env:METATOMIC_TESTS_TORCH_VERSION:2.13}/ \ - -DMETATOMIC_TORCH_TESTS=ON - - # build code with cmake - cmake --build {env_dir}/build --config Debug --parallel - - # run all tests - ctest --test-dir {env_dir}/build --build-config Debug --output-on-failure - -[testenv:torch-install-tests-cxx] -description = Run the C++ tests for metatomic-torch -deps = - cmake - {[testenv]metatomic_deps} - torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.13}.* + pip install {[testenv]build_single_wheel} . -commands = - # configure, build and install metatomic-torch - cmake -B {env_dir}/build-metatomic-torch metatomic-torch \ - -DCMAKE_BUILD_TYPE=Debug \ - -DCMAKE_INSTALL_PREFIX={env_dir}/usr/ \ - -DCMAKE_PREFIX_PATH={env_site_packages_dir}/metatensor/;\ - {env_site_packages_dir}/torch/;\ - {env_site_packages_dir}/metatensor_torch/torch-{env:METATOMIC_TESTS_TORCH_VERSION:2.13}/ \ - -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ - -DCMAKE_INSTALL_RPATH_USE_LINK_PATH=ON - cmake --build {env_dir}/build-metatomic-torch --config Debug --parallel --target install - - # try to use the installed metatomic-torch from another CMake project - cmake -B {env_dir}/build-find-package metatomic-torch/tests/cmake-project \ - -DCMAKE_BUILD_TYPE=Debug \ - -DCMAKE_PREFIX_PATH={env_site_packages_dir}/metatensor/;\ - {env_site_packages_dir}/torch/;\ - {env_site_packages_dir}/metatensor_torch/torch-{env:METATOMIC_TESTS_TORCH_VERSION:2.13}/;\ - {env_dir}/usr/ \ - -DUSE_CMAKE_SUBDIRECTORY=OFF - - cmake --build {env_dir}/build-find-package --config Debug --parallel - ctest --test-dir {env_dir}/build-find-package --build-config Debug --output-on-failure - - # Same, but using metatomic-torch as a CMake subdirectory - cmake -B {env_dir}/build-subdirectory metatomic-torch/tests/cmake-project \ - -DCMAKE_BUILD_TYPE=Debug \ - -DCMAKE_PREFIX_PATH={env_site_packages_dir}/metatensor/;\ - {env_site_packages_dir}/torch/;\ - {env_site_packages_dir}/metatensor_torch/torch-{env:METATOMIC_TESTS_TORCH_VERSION:2.13}/ \ - -DUSE_CMAKE_SUBDIRECTORY=ON - - cmake --build {env_dir}/build-subdirectory --config Debug --parallel - ctest --test-dir {env_dir}/build-subdirectory --build-config Debug --output-on-failure + pytest --cov={env_site_packages_dir}/metatomic --cov-report= --import-mode=append {posargs} -################################################################################ -##### Python tests setup ##### -################################################################################ [testenv:torch-tests] description = Run the tests of the metatomic-torch Python package deps = {[testenv]testing_deps} {[testenv]packaging_deps} - {[testenv]metatomic_deps} + {[testenv]metatomic_torch_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.13}.* numpy @@ -134,6 +82,7 @@ deps = changedir = python/metatomic_torch commands = + pip install {[testenv]build_single_wheel} ../metatomic_core pip install {[testenv]build_single_wheel} . pip install {[testenv]build_single_wheel} ../metatomic_ase @@ -151,19 +100,30 @@ description = Run the doctests defined in any metatomic package deps = {[testenv]testing_deps} {[testenv]packaging_deps} - {[testenv]metatomic_deps} + {[testenv]metatomic_torch_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.13}.* numpy vesin >=0.6.0,<0.7 ase + torch-sim-atomistic + +setenv = + # ignore the fact that metatensor.torch.operations was loaded from a file + # not in `metatensor/torch/operations` + PY_IGNORE_IMPORTMISMATCH = 1 commands = + pip install {[testenv]build_single_wheel} python/metatomic_core pip install {[testenv]build_single_wheel} python/metatomic_torch pip install {[testenv]build_single_wheel} python/metatomic_ase + pip install {[testenv]build_single_wheel} python/metatomic_torchsim - pytest --doctest-modules --pyargs metatomic + pytest --suppress-no-test-exit-code --doctest-modules --pyargs metatomic + pytest --suppress-no-test-exit-code --doctest-modules --pyargs metatomic_torch + pytest --suppress-no-test-exit-code --doctest-modules --pyargs metatomic_ase + pytest --suppress-no-test-exit-code --doctest-modules --pyargs metatomic_torchsim ################################################################################ @@ -175,7 +135,7 @@ description = Run the tests of the metatomic-ase Python package deps = {[testenv]testing_deps} {[testenv]packaging_deps} - {[testenv]metatomic_deps} + {[testenv]metatomic_torch_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.13}.* numpy @@ -193,8 +153,9 @@ deps = changedir = python/metatomic_ase commands = - pip install {[testenv]build_single_wheel} . + pip install {[testenv]build_single_wheel} ../metatomic_core pip install {[testenv]build_single_wheel} ../metatomic_torch + pip install {[testenv]build_single_wheel} . # use the reference LJ implementation for tests {[testenv]install_lj_tests} @@ -211,7 +172,7 @@ description = Run the tests of the metatomic-torchsim Python package deps = {[testenv]testing_deps} {[testenv]packaging_deps} - {[testenv]metatomic_deps} + {[testenv]metatomic_torch_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.13}.* numpy @@ -224,8 +185,9 @@ deps = changedir = python/metatomic_torchsim commands = - pip install {[testenv]build_single_wheel} . + pip install {[testenv]build_single_wheel} ../metatomic_core pip install {[testenv]build_single_wheel} ../metatomic_torch + pip install {[testenv]build_single_wheel} . # use the reference LJ implementation for tests {[testenv]install_lj_tests} @@ -266,7 +228,7 @@ setenv = # build the docs against the CPU only version of torch PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu {env:PIP_EXTRA_INDEX_URL:} deps = - {[testenv]metatomic_deps} + {[testenv]metatomic_torch_deps} {[testenv]packaging_deps} {[testenv]testing_deps} @@ -294,6 +256,7 @@ deps = chemiscope commands = + pip install {[testenv]build_single_wheel} python/metatomic_core pip install {[testenv]build_single_wheel} python/metatomic_torch pip install {[testenv]build_single_wheel} python/metatomic_ase pip install {[testenv]build_single_wheel} python/metatomic_torchsim