Dump and load extensions (#222)

.codecov.yml

@@ -9,6 +9,6 @@ coverage:
   ignore:
      - "tests/.*"
      - "examples/.*"
-     - "src/metatensor/models/experimental/.*"
+     - "src/metatrain/experimental/.*"
 
 comment: false

.github/PULL_REQUEST_TEMPLATE/pull_request_new-architecture.md

@@ -5,9 +5,9 @@
 # New architecture TODOs
 
 - [ ] Add your architecture to the experimental folder
-  `src/metatensor/models/experimental/<architecture_name>`
+  `src/metatrain/experimental/<architecture_name>`
 - [ ] Add default hyperparameter file to
-  `src/metatensor/models/cli/conf/architecture/experimental.<architecture_name>.yaml`
+  `src/metatrain/cli/conf/architecture/experimental.<architecture_name>.yaml`
 - [ ] Add a `.yml` file into github workflows `.github/workflow/<architecture_name>.yml`
 - [ ] Architecture dependencies entry in the `optional-dependencies` section in the
   `pyproject.toml`

.github/workflows/pr-docs-preview.yml

@@ -14,4 +14,4 @@ jobs:
     steps:
       - uses: readthedocs/actions/preview@v1
         with:
-          project-slug: metatensor-models
+          project-slug: metatrain

README.rst

@@ -1,36 +1,37 @@
-metatensor-models
------------------
+metatrain
+---------
 
 |tests| |codecov| |docs|
 
 .. warning::
-  **metatensor-models is still very early in the concept stage. You should not use it
+  **metatrain is still very early in the concept stage. You should not use it
   for anything important.**
 
-This is a repository for models using metatensor, in one shape or another. The only
-requirement is for these models to be able to take metatensor objects as inputs and
-outputs. The models do not need to live entirely in this repository: in the most extreme
-case, this repository can simply contain a wrapper to an external model.
+This is a repository for training machine learning models from various architecture for
+atomistic systems. The only requirement is for these models to be able to take
+`metatensor <https://docs.metatensor.org>`_ objects as inputs and outputs. The models do
+not need to live entirely in this repository: in the most extreme case, this repository
+can simply contain a wrapper
+to an external model.
 
 .. marker-introduction
 
-What is metatensor-models?
-##########################
-The idea behind ``metatensor-models`` is to have a general hub that provide an
-homogeneous environment and user interface to train, export and evaluate ML models and
-to connect those models with various MD engines (e.g. ``LAMMPS``, ``i-PI``, ``ASE``
-...). ``metatensor-models`` is the tool that transforms every ML architecture in an
-end-to-end model. Any custom ML architecture compatible with TorchScript can be
-integrated in ``metatensor-models``, gaining automatic access to a training and
-evaluation interface, as well as compatibility with various MD engines.
+What is metatrain?
+##################
+The idea behind ``metatrain`` is to have a general hub that provide an homogeneous
+environment and user interface to train, export and evaluate ML models and to connect
+those models with various MD engines (e.g. ``LAMMPS``, ``i-PI``, ``ASE`` ...).
+``metatrain`` is the tool that transforms every ML architecture in an end-to-end model.
+Any custom ML architecture compatible with TorchScript can be integrated in
+``metatrain``, gaining automatic access to a training and evaluation interface, as well
+as compatibility with various MD engines.
 
-Note: ``metatensor-models`` does not provide `per se` mathematical functionalities but
-relies on external models that implement the various architectures.
+Note: ``metatrain`` does not provide `per se` mathematical functionalities but relies on
+external models that implement the various architectures.
 
 Features
 ########
-- **Custom ML Architecture**: Integrate any TorchScriptable ML model to explore
-  innova
+- **Custom ML Architecture**: Integrate any TorchScriptable ML model to explore innova
 - **MD Engine Compatibility**: Supports various MD engines for diverse research and
   application needs.
 - **Streamlined Training**: Automated process leveraging MD-generated data to optimize
@@ -42,8 +43,8 @@ Features
 
 List of Implemented Architectures
 #################################
-Currently metatensor-models supports the following architectures for building an
-atomistic model.
+Currently ``metatrain`` supports the following architectures for building an atomistic
+model.
 
 .. list-table::
   :widths: 34 66
@@ -64,31 +65,32 @@ atomistic model.
 Documentation
 -------------
 For details, tutorials, and examples, please have a look at our
-`documentation <https://lab-cosmo.github.io/metatensor-models/latest/>`_.
+`documentation <https://lab-cosmo.github.io/metatrain/latest/>`_.
 
 .. marker-installation
 
 Installation
 ------------
-You can install metatensor-models with pip:
+You can install ``metatrain`` with pip:
 
 .. code-block:: bash
 
-    git clone https://github.com/lab-cosmo/metatensor-models
-    cd metatensor-models
+    git clone https://github.com/lab-cosmo/metatrain
+    cd metatrain
     pip install .
 
-
 In addition, specific models must be installed by specifying the model name. For
 example, to install the SOAP-BPNN model, you can run:
 
 .. code-block:: bash
 
     pip install .[soap-bpnn]
 
+You can then use ``mtt`` from the command line to train your models!
+
 Shell Completion
 ################
-`metatensor-models` comes with completion definitions for its commands for ``bash`` and
+``metatrain`` comes with completion definitions for its commands for ``bash`` and
 ``zsh``. Since it is difficult to automatically configure shell completions in a robust
 manner, you must manually configure your shell to enable its completion support.
 
@@ -98,35 +100,35 @@ exist, ``~/.profile``):
 
 .. code-block:: bash
 
-  source $(metatensor-models --shell-completion)
+  source $(mtt --shell-completion)
 
 .. marker-issues
 
 Having problems or ideas?
 -------------------------
-Having a problem with metatensor-models? Please let us know by `submitting an issue
-<https://github.com/metatensor-models/issues>`_.
+Having a problem with metatrain? Please let us know by `submitting an issue
+<https://github.com/metatrain/issues>`_.
 
 Submit new features or bug fixes through a `pull request
-<https://github.com/metatensor-models/pulls>`_.
+<https://github.com/metatrain/pulls>`_.
 
 .. marker-contributing
 
 Contributors
 ------------
-Thanks goes to all people that make metatensor-models possible:
+Thanks goes to all people that make ``metatrain`` possible:
 
-.. image:: https://contrib.rocks/image?repo=lab-cosmo/metatensor-models
-  :target: https://github.com/lab-cosmo/metatensor-models/graphs/contributors
+.. image:: https://contrib.rocks/image?repo=lab-cosmo/metatrain
+  :target: https://github.com/lab-cosmo/metatrain/graphs/contributors
 
-.. |tests| image:: https://github.com/lab-cosmo/metatensor-models/workflows/Tests/badge.svg
+.. |tests| image:: https://github.com/lab-cosmo/metatrain/workflows/Tests/badge.svg
   :alt: Github Actions Tests Job Status
-  :target: https://github.com/lab-cosmo/metatensor-models/actions?query=branch%3Amain
+  :target: https://github.com/lab-cosmo/metatrain/actions?query=branch%3Amain
 
-.. |codecov| image:: https://codecov.io/gh/lab-cosmo/metatensor-models/branch/main/graph/badge.svg
+.. |codecov| image:: https://codecov.io/gh/lab-cosmo/metatrain/branch/main/graph/badge.svg
   :alt: Code coverage
-  :target: https://codecov.io/gh/lab-cosmo/metatensor-models
+  :target: https://codecov.io/gh/lab-cosmo/metatrain
 
 .. |docs| image:: https://img.shields.io/badge/documentation-latest-sucess
   :alt: Documentation
-  :target: https://lab-cosmo.github.io/metatensor-models/latest/
+  :target: https://lab-cosmo.github.io/metatrain/latest/

developer/clean-tempfiles.sh

@@ -14,8 +14,8 @@ rm -rf docs/build
 rm -rf docs/src/examples
 rm -rf docs/src/sg_execution_times.rst
 
-rm -rf src/metatensor/models/dist
-rm -rf src/metatensor/models/build
+rm -rf src/metatrain/dist
+rm -rf src/metatrain/build
 
 find . -name "*.egg-info" -exec rm -rf "{}" +
 find . -name "__pycache__" -exec rm -rf "{}" +

docs/src/advanced-concepts/index.rst

@@ -1,7 +1,7 @@
 Advanced concepts
 =================
 
-This section covers advanced concepts of the ``metatensor-models`` library,
+This section covers advanced concepts of the ``metatrain`` library,
 such as output naming, auxiliary outputs, and wrapper models.
 
 

docs/src/advanced-concepts/output-naming.rst

@@ -1,7 +1,7 @@
 Output naming
 =============
 
-The name and format of the outputs in ``metatensor-models`` are based on
+The name and format of the outputs in ``metatrain`` are based on
 those of the `<metatensor.torch.atomistic
 https://lab-cosmo.github.io/metatensor/latest/atomistic/outputs.html>`_
 package. An immediate example is given by the ``energy`` output.

docs/src/architectures/alchemical-model.rst

@@ -31,7 +31,7 @@ Default Hyperparameters
 -----------------------
 The default hyperparameters for the Alchemical Model model are:
 
-.. literalinclude:: ../../../src/metatensor/models/experimental/alchemical_model/default-hypers.yaml
+.. literalinclude:: ../../../src/metatrain/experimental/alchemical_model/default-hypers.yaml
    :language: yaml
 
 

docs/src/architectures/index.rst

@@ -3,7 +3,7 @@
 Available Architectures
 =======================
 
-This is a list of all architectures available in ``metatensor-models``.
+This is a list of all architectures available in ``metatrain``.
 
 .. toctree::
    :maxdepth: 1

docs/src/architectures/pet.rst

@@ -17,5 +17,5 @@ Default Hyperparameters
 -----------------------
 The default hyperparameters for the PET model are:
 
-.. literalinclude:: ../../../src/metatensor/models/experimental/pet/default-hypers.yaml
+.. literalinclude:: ../../../src/metatrain/experimental/pet/default-hypers.yaml
    :language: yaml

docs/src/architectures/soap-bpnn.rst

@@ -119,7 +119,7 @@ Default Hyperparameters
 -----------------------
 The default hyperparameters for the SOAP-BPNN model are:
 
-.. literalinclude:: ../../../src/metatensor/models/experimental/soap_bpnn/default-hypers.yaml
+.. literalinclude:: ../../../src/metatrain/experimental/soap_bpnn/default-hypers.yaml
    :language: yaml
 
 

docs/src/conf.py

@@ -5,7 +5,7 @@
 
 import tomli  # Replace by tomllib from std library once docs are build with Python 3.11
 
-import metatensor.models
+import metatrain
 
 
 # When importing metatensor-torch, this will change the definition of the classes
@@ -40,7 +40,7 @@
 copyright = f"{datetime.now().date().year}, {author}"
 
 # The full version, including alpha/beta/rc tags
-release = metatensor.models.__version__
+release = metatrain.__version__
 
 
 # -- General configuration ---------------------------------------------------
@@ -53,7 +53,7 @@ def generate_examples():
     # include the corresponding output later.
     del os.environ["METATENSOR_IMPORT_FOR_SPHINX"]
     script = os.path.join(ROOT, "docs", "generate_examples", "generate-examples.py")
-    subprocess.run([sys.executable, script], capture_output=False)
+    subprocess.run([sys.executable, script])
     os.environ["METATENSOR_IMPORT_FOR_SPHINX"] = "1"
 
 

docs/src/dev-docs/architecture-life-cycle.rst

@@ -3,15 +3,15 @@ Life Cycle of an Architecture
 
 .. TODO: Maybe add a flowchart later
 
-Architectures in `metatensor-models` undergo different stages based on their
+Architectures in `metatrain` undergo different stages based on their
 development/functionality level and maintenance status. We distinguish three distinct
 stages: **experimental**, **stable**, and **deprecated**. Typically, an architecture
 starts as experimental, advances to stable, and eventually becomes deprecated before
 removal if maintenance is no longer feasible.
 
 .. note::
     The development and maintenance of an architecture must be fully undertaken by the
-    architecture's authors or maintainers. The core developers of `metatensor-models`
+    architecture's authors or maintainers. The core developers of `metatrain`
     provide infrastructure and implementation support but are not responsible for the
     architecture's internal functionality or any issues that may arise therein.
 
@@ -44,7 +44,7 @@ satisfied:
 1. Provision of regression prediction tests with a small (not exported) checkpoint file.
 2. Comprehensive architecture documentation
 3. If an architecture has external dependencies, all must be publicly available on PyPI.
-4. Adherence to the standard output infrastructure of `metatensor-models`, including
+4. Adherence to the standard output infrastructure of `metatrain`, including
    logging and model save locations.
 
 Deprecated Architectures

docs/src/dev-docs/cli/eval.rst

@@ -1,7 +1,7 @@
 Eval
 ####
 
-.. automodule:: metatensor.models.cli.eval
+.. automodule:: metatrain.cli.eval
     :members:
     :undoc-members:
     :show-inheritance:

docs/src/dev-docs/cli/export.rst

@@ -1,7 +1,7 @@
 Export
 ######
 
-.. automodule:: metatensor.models.cli.export
+.. automodule:: metatrain.cli.export
     :members:
     :undoc-members:
     :show-inheritance:

docs/src/dev-docs/cli/formatter.rst

@@ -1,7 +1,7 @@
-formatter
+Formatter
 #########
 
-.. automodule:: metatensor.models.cli.formatter
+.. automodule:: metatrain.cli.formatter
     :members:
     :undoc-members:
     :show-inheritance:

docs/src/dev-docs/cli/index.rst

@@ -2,7 +2,7 @@ CLI API
 =======
 
 This is the API for the command line interface ``cli`` functions for the ``train``,
-the ``eval`` and the ``export`` functions of ``metatensor-models``.
+the ``eval`` and the ``export`` functions of ``metatrain``.
 
 .. toctree::
    :maxdepth: 1

docs/src/dev-docs/cli/train.rst

@@ -1,7 +1,7 @@
 Train
 #####
 
-.. automodule:: metatensor.models.cli.train
+.. automodule:: metatrain.cli.train
     :members:
     :undoc-members:
     :show-inheritance:

docs/src/dev-docs/index.rst

@@ -1,7 +1,7 @@
 Developer documentation
 =======================
 
-This is a collection of documentation for developers of the metatensor-models package.
+This is a collection of documentation for developers of the `metatrain` package.
 It includes documentation on how to add a new model, as well as the API of the utils
 module.
 

docs/src/dev-docs/new-architecture.rst

@@ -3,8 +3,8 @@
 Adding a new architecture
 =========================
 
-To work with` metatensor-models` any architecture has to follow the same public API to
-be called correctly within the :py:func:`metatensor.models.cli.train` function to
+To work with` metatrain` any architecture has to follow the same public API to
+be called correctly within the :py:func:`metatrain.cli.train` function to
 process the user's options. In brief, the core of the ``train`` function looks similar
 to these lines
 
@@ -86,12 +86,12 @@ Note that the ``ModelInterface`` does not necessary inherit from
 :py:class:`torch.nn.Module` since training can be performed in any way.
 ``__supported_devices__`` and ``__supported_dtypes__`` can be defined to set the
 capabilities of the model. These two lists should be sorted in order of preference since
-`metatensor-models` will use these to determine, based on the user request and
+`metatrain` will use these to determine, based on the user request and
 machines's availability, the optimal `dtype` and `device` for training.
 
 The ``export()`` method is required to transform a trained model into a standalone file
 to be used in combination with molecular dynamic engines to run simulations. We provide
-a helper function :py:func:`metatensor.models.utils.export.export` to export a torch
+a helper function :py:func:`metatrain.utils.export.export` to export a torch
 model to an :py:class:`MetatensorAtomisticModel
 <metatensor.torch.atomistic.MetatensorAtomisticModel>`.
 
@@ -111,7 +111,7 @@ methods for ``train()``.
             train_datasets: List[Union[Dataset, torch.utils.data.Subset]],
             validation_datasets: List[Union[Dataset, torch.utils.data.Subset]],
             checkpoint_dir: str,
-        ): ...
+        ) -> None: ...
 
 The names of the ``ModelInterface`` and the ``TrainerInterface`` are free to choose but
 should be linked to constants in the ``__init__.py`` of each architecture. On top of
@@ -135,7 +135,7 @@ these two constants the ``__init__.py`` must contain constants for the original
 
 
 :param __model__: Mapping of the custom ``ModelInterface`` to a general one to be loaded
-    by metatensor-models
+    by metatrain
 :param __trainer__: Same as ``__MODEL_CLASS__`` but the Trainer class.
 :param __authors__: Tuple denoting the original authors with email address and Github
     handle of an architecture. These do not necessary be in charge of maintaining the