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..c91fe87d2 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -20,6 +20,7 @@ import lammps_lexer # noqa: E402 from sphinx.highlighting import lexers # noqa: E402 + lexers["LAMMPS"] = lammps_lexer.LAMMPSLexer(startinline=True) @@ -191,6 +192,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 +204,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/examples/index.rst b/docs/src/examples/index.rst new file mode 100644 index 000000000..77e50a720 --- /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 example 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 exising models in 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/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..23a648382 --- /dev/null +++ b/examples/c/1-create-system.c @@ -0,0 +1,271 @@ +// 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. +// +// 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, + int32_t ndim, + const int64_t *shape, + 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, + /*ndim=*/ 2, + /*shape=*/(int64_t[]){n_atoms, 3}, + /*dtype=*/(DLDataType){.code = kDLFloat, .bits = 64, .lanes = 1} +); + +DLManagedTensorVersioned* cell = tensor_from_data( + /*data=*/ cell_data, + /*ndim=*/ 2, + /*shape=*/(int64_t[]){3, 3}, + /*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, + /*ndim=*/ 1, + /*shape=*/(int64_t[]){n_atoms}, + /*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}; +int64_t pbc_shape[] = {3}; +DLManagedTensorVersioned *pbc = tensor_from_data( + /*data=*/ pbc_data, + /*ndim=*/ 1, + /*shape=*/(int64_t[]){3}, + /*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. + +mta_system_free(system); + +// %% + +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/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/tests/CMakeLists.txt b/metatomic-core/tests/CMakeLists.txt index efb9b1f7b..d40e52f84 100644 --- a/metatomic-core/tests/CMakeLists.txt +++ b/metatomic-core/tests/CMakeLists.txt @@ -68,7 +68,7 @@ endif() function(metatomic_add_test source target) add_executable(${target} ${source}) - target_link_libraries(${target} metatomic catch) + target_link_libraries(${target} metatomic) set_target_properties(${target} PROPERTIES # Ensure that the binaries find the right shared library. @@ -101,10 +101,24 @@ 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_}") +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_}") endforeach() 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/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/scripts/clean-python.sh b/scripts/clean-python.sh index 81e69b26a..33fafec32 100755 --- a/scripts/clean-python.sh +++ b/scripts/clean-python.sh @@ -11,7 +11,7 @@ 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