diff --git a/src/metatensor/README.md b/src/metatensor/README.md index 4459644d2f..6672f16250 100644 --- a/src/metatensor/README.md +++ b/src/metatensor/README.md @@ -3,82 +3,8 @@ ## Building the code -1. You'll need to fist install libtorch, either by installing PyTorch itself - with Python, or by downloading the prebuilt C++ library from - https://pytorch.org/get-started/locally/. - -```bash -# point this to the path where you extracted the C++ libtorch -TORCH_PREFIX=../../.. -# if you used Python to install torch, you can do this: -TORCH_CMAKE_PREFIX=$(python -c "import torch; print(torch.utils.cmake_prefix_path)") -TORCH_PREFIX=$(cd "$TORCH_CMAKE_PREFIX/../.." && pwd) - -TORCH_INCLUDES="-I$TORCH_PREFIX/include -I$TORCH_PREFIX/include/torch/csrc/api/include" -``` - -2. a) build and install metatensor-torch from source. You'll need a rust - compiler on your system, the easiest way is by using https://rustup.rs/ - -```bash -# patch a bug from torch's MKL detection -cd -./src/metatensor/patch-torch.sh "$TORCH_PREFIX" - -cd - -# define a location where metatensor should be installed -METATENSOR_PREFIX=<...> - -METATENSOR_TORCH_PREFIX="$METATENSOR_PREFIX" - -git clone https://github.com/lab-cosmo/metatensor --branch=metatensor-torch-v0.4.0 -cd metatensor - -mkdir build && cd build -cmake -DBUILD_SHARED_LIBS=ON \ - -DCMAKE_INSTALL_PREFIX="$METATENSOR_PREFIX" \ - -DCMAKE_PREFIX_PATH="$TORCH_PREFIX" \ - -DBUILD_METATENSOR_TORCH=ON \ - -DMETATENSOR_INSTALL_BOTH_STATIC_SHARED=OFF \ - .. - -cmake --build . --target install --parallel -``` - -2. b) alternatively, use metatensor-torch from Python (`pip install metatensor[torch]`) - -```bash -METATENSOR_CMAKE_PREFIX=$(python -c "import metatensor; print(metatensor.utils.cmake_prefix_path)") -METATENSOR_PREFIX=$(cd "$METATENSOR_CMAKE_PREFIX/../.." && pwd) - -METATENSOR_TORCH_CMAKE_PREFIX=$(python -c "import metatensor.torch; print(metatensor.torch.utils.cmake_prefix_path)") -METATENSOR_TORCH_PREFIX=$(cd "$METATENSOR_TORCH_CMAKE_PREFIX/../.." && pwd) -``` - -3. build Plumed itself - -```bash -cd - -# set the rpath to make sure plumed executable will be able to find the right libraries -RPATH="-Wl,-rpath,$TORCH_PREFIX/lib -Wl,-rpath,$METATENSOR_PREFIX/lib -Wl,-rpath,$METATENSOR_TORCH_PREFIX/lib" - -# configure PLUMED with metatensor -./configure --enable-libtorch --enable-metatensor --enable-modules=+metatensor \ - LDFLAGS="-L$TORCH_PREFIX/lib -L$METATENSOR_PREFIX/lib -L$METATENSOR_TORCH_PREFIX/lib $RPATH" \ - CPPFLAGS="$TORCH_INCLUDES -I$METATENSOR_PREFIX/include -I$METATENSOR_TORCH_PREFIX/include" - -# If you are on Linux and use a pip-installed version of libtorch, or the -# pre-cxx11-ABI build of libtorch, you'll need to add "-D_GLIBCXX_USE_CXX11_ABI=0" -# to the compilation flags: -./configure --enable-libtorch --enable-metatensor --enable-modules=+metatensor \ - LDFLAGS="-L$TORCH_PREFIX/lib -L$METATENSOR_PREFIX/lib -L$METATENSOR_TORCH_PREFIX/lib $RPATH" \ - CPPFLAGS="$TORCH_INCLUDES -I$METATENSOR_PREFIX/include -I$METATENSOR_TORCH_PREFIX/include" \ - CXXFLAGS="-D_GLIBCXX_USE_CXX11_ABI=0" - -make -j && make install -``` +See [the main documentation](../../user-doc/METATENSORMOD.md) for more +information on how to compile and use this module. diff --git a/src/metatensor/metatensor.cpp b/src/metatensor/metatensor.cpp index a36f35f6fa..981cbfcbd7 100644 --- a/src/metatensor/metatensor.cpp +++ b/src/metatensor/metatensor.cpp @@ -12,6 +12,58 @@ #include "core/ActionRegister.h" #include "core/PlumedMain.h" +//+PLUMEDOC METATENSORMOD_COLVAR METATENSOR +/* +Use arbitrary machine learning models as collective variables. + +Note that this action requires the metatensor-torch library. Check the +instructions in the \ref METATENSORMOD page to enable this module. + +This action enables the use of fully custom machine learning models — based on +the [metatensor atomistic models][mts_models] interface — as collective +variables in PLUMED. Such machine learning model are typically written and +customized using Python code, and then exported to run within PLUMED as +[TorchScript], which is a subset of Python that can be executed by the C++ torch +library. + +Metatensor offers a way to define such models and pass data from PLUMED (or any +other simulation engine) to the model and back. For more information on how to +define such model, have a look at the [corresponding tutorials][mts_tutorials], +or at the code in `regtest/metatensor/`. Each of the Python scripts in this +directory defines a custom machine learning CV that can be used with PLUMED. + +\par Examples + +TODO + +\par Collective variables and metatensor models + +Collective variables are not yet part of the [known outputs][mts_outputs] for +metatensor models. Until the output format is standardized, this action expects +the following: + +- the output name should be `"plumed::cv"`; +- the output should contain a single [block][mts_block]; +- the output samples should be named `["system", "atom"]` for per-atom outputs; + or `["system"]` for global outputs. The `"system"` index should always be 0, + and the `"atom"` index should be the index of the atom (between 0 and the + total number of atoms); +- the output should not have any components; +- the output can have arbitrary properties; +- the output should not have any explicit gradients, all gradient calculations + are done using autograd. + +*/ /* + +[TorchScript]: https://pytorch.org/docs/stable/jit.html +[mts_models]: https://lab-cosmo.github.io/metatensor/latest/atomistic/index.html +[mts_tutorials]: https://lab-cosmo.github.io/metatensor/latest/examples/atomistic/index.html +[mts_outputs]: https://lab-cosmo.github.io/metatensor/latest/atomistic/outputs.html +[mts_block]: https://lab-cosmo.github.io/metatensor/latest/torch/reference/block.html + +*/ +//+ENDPLUMEDOC + #if !defined(__PLUMED_HAS_LIBTORCH) || !defined(__PLUMED_HAS_METATENSOR) diff --git a/user-doc/METATENSORMOD.md b/user-doc/METATENSORMOD.md new file mode 100644 index 0000000000..2df18f5d0d --- /dev/null +++ b/user-doc/METATENSORMOD.md @@ -0,0 +1,157 @@ +\page METATENSORMOD Metatensor + + + +# Overview + +This module implements the interface between PLUMED and [metatensor], allowing +to use arbitrary machine learning models as collective variables. These machine +learning models are defined using custom Python code — following the [metatensor +atomistic models][mts_models] interface — and then exported to [TorchScript]. +The exported model is then loaded inside PLUMED and executed during the +simulation. + + +# Installation + +This module requires two main dependencies: the C++ torch library (i.e. +`libtorch`); and the C++ metatensor_torch library. There are multiple ways of +installing both libraries, which are discussed below. + +## Installing the libraries through Python's package manager (`pip`) + +The easiest way to get all dependencies on your system is to download the +pre-built Python wheels with `pip`. This is the same set of wheels you will need +to define custom models. + +```bash +pip install "metatensor-torch ==0.4.0" # change this version to get newer releases + +# optional: get the other metatensor tools to define models (these are only usable from Python). +pip install metatensor-operations metatensor-learn + +# export the location to all the libraries: +TORCH_CMAKE_PREFIX=$(python -c "import torch; print(torch.utils.cmake_prefix_path)") +TORCH_PREFIX=$(cd "$TORCH_CMAKE_PREFIX/../.." && pwd) + +METATENSOR_CMAKE_PREFIX=$(python -c "import metatensor; print(metatensor.utils.cmake_prefix_path)") +METATENSOR_PREFIX=$(cd "$METATENSOR_CMAKE_PREFIX/../.." && pwd) + +METATENSOR_TORCH_CMAKE_PREFIX=$(python -c "import metatensor.torch; print(metatensor.torch.utils.cmake_prefix_path)") +METATENSOR_TORCH_PREFIX=$(cd "$METATENSOR_TORCH_CMAKE_PREFIX/../.." && pwd) + +# The torch library installed by pip uses a pre-cxx11 ABI +TORCH_CPPFLAGS="-D_GLIBCXX_USE_CXX11_ABI=0" +``` + +That's it, you can now jump to [the last part](#building-plumed-with-metatensor) +of the installation instructions. + +## Using pre-built libraries + +If you only want to use existing models, you can download pre-built versions of +the libraries and build PLUMED against these. First, you'll need to download +libtorch (see also \ref installation-libtorch for other instructions on +installing a pre-built libtorch): + +```bash +# Download torch 2.2.2 for x86_64 (Intel) Linux. +# +# Variations of this link for other operating systems (macOS, Windows), CPU +# architecture (Apple Silicon, arm64), CUDA versions, and newer versions of +# libtorch can be found at https://pytorch.org/get-started/locally/ + +wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.2.2%2Bcpu.zip +unzip libtorch-cxx11-abi-shared-with-deps-2.2.2+cpu.zip + +# alternatively if you have a CUDA-enabled GPU, you can use the corresponding +# pre-built library (here for CUDA 12.1): +wget https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.2%2Bcu121.zip +unzip libtorch-cxx11-abi-shared-with-deps-2.2.2+cu121.zip + +# Make the location of libtorch visible +TORCH_PREFIX=$(pwd)/libtorch + +# if you are using a library with pre-cxx11 ABI, you need an extra flag: +TORCH_CPPFLAGS="-D_GLIBCXX_USE_CXX11_ABI=0" +``` + +Once you acquire libtorch, it is time to build metatensor and metatensor_torch +from sources. There is currently no standalone pre-built library for these +(although you can use the pre-built version that comes with `pip`). For this, +you'll need a rust compiler on your system, which you can get with +[rustup](https://rustup.rs/) or any other method at your convenience. + +```bash +# patch a bug from torch's MKL detection in CMake +cd +./src/metatensor/patch-torch.sh "$TORCH_PREFIX" + +cd + +# define a location where metatensor should be installed +METATENSOR_PREFIX=<...> + +METATENSOR_TORCH_PREFIX="$METATENSOR_PREFIX" + +git clone https://github.com/lab-cosmo/metatensor +# or a more recent release of metatensor-torch +git checkout metatensor-torch-v0.4.0 +cd metatensor + +mkdir build && cd build +cmake -DBUILD_SHARED_LIBS=ON \ + -DCMAKE_INSTALL_PREFIX="$METATENSOR_PREFIX" \ + -DCMAKE_PREFIX_PATH="$TORCH_PREFIX" \ + -DBUILD_METATENSOR_TORCH=ON \ + -DMETATENSOR_INSTALL_BOTH_STATIC_SHARED=OFF \ + .. + +cmake --build . --target install --parallel +``` + +## Building plumed with metatensor + +Once you installed all dependencies with one of the methods above, you can now +configure PLUMED: + +```bash +# set include search path for the compilers +TORCH_INCLUDES="-I$TORCH_PREFIX/include -I$TORCH_PREFIX/include/torch/csrc/api/include" +CPPFLAGS="$TORCH_INCLUDES $TORCH_CPPFLAGS -I$METATENSOR_PREFIX/include -I$METATENSOR_TORCH_PREFIX/include $CPPFLAGS" + +# set library search path for the linker +LDFLAGS="-L$TORCH_PREFIX/lib -L$METATENSOR_PREFIX/lib -L$METATENSOR_TORCH_PREFIX/lib $LDFLAGS" + +# set the rpath to make sure plumed executable will be able to find the right libraries +LDFLAGS="$LDFLAGS -Wl,-rpath,$TORCH_PREFIX/lib" +LDFLAGS="$LDFLAGS -Wl,-rpath,$METATENSOR_PREFIX/lib -Wl,-rpath,$METATENSOR_TORCH_PREFIX/lib" + +# configure PLUMED +./configure --enable-libtorch --enable-metatensor --enable-modules=+metatensor \ + LDFLAGS="$LDFLAGS" CPPFLAGS="$CPPFLAGS" +``` + +Pay close attention to the output, it should contain **both** a line about +`checking libtorch` and a line about `checking metatensor`, both ending with +`...yes`. If this is not the case, you'll get a warning about `cannot enable +__PLUMED_HAS_LIBTORCH` or `cannot enable __PLUMED_HAS_METATENSOR`. If you get +any of these warnings, you should check `config.log` to know more about what's +going on and why these libraries can't be found. + +# Module Contents + +This module defines the following actions: + +@METATENSORMOD_COLVAR@ + + + + +[TorchScript]: https://pytorch.org/docs/stable/jit.html +[metatensor]: https://lab-cosmo.github.io/metatensor/latest/index.html +[mts_models]: https://lab-cosmo.github.io/metatensor/latest/atomistic/index.html