release: create release-0.0.1a5 branch

Author: github-actions[bot]

.github/workflows/tests_and_linters.yaml

@@ -1,4 +1,4 @@
-name: Open-source Tests and Linters 🧪
+name: Tests and Linters 🧪
 
 on:
   push:

CHANGELOG

@@ -1,5 +1,5 @@
 # Changelog
 
-## Release 1.0.0
+## Release 0.1.0
 
 To be added upon release...

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.10.12-slim-bullseye
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y git wget
+
+RUN pip install mlip "jax[cuda12]==0.4.33" huggingface_hub git+https://github.com/jax-md/jax-md.git notebook
+
+RUN wget https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/simulation_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/model_training_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/model_addition_tutorial.ipynb
+
+EXPOSE 8888
+
+CMD ["jupyter", "notebook", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root"]

Dockerfile.release

@@ -96,9 +96,22 @@ pip install notebook && jupyter notebook
 The installation of *mlip* itself is included within the notebooks. We recommend to
 run these notebooks with GPU acceleration enabled.
 
-## 🤗 Foundation models (via HuggingFace)
+Alternatively, we provide a `Dockerfile` in this repository that you can use to
+run the tutorial notebooks. This can be achieved by executing the following lines
+from any directory that contains the downloaded `Dockerfile`:
 
-We have prepared foundation models pre-trained on a subset of the
+```bash
+docker build . -t mlip_tutorials
+docker run -p 8888:8888 --gpus all mlip_tutorials
+```
+
+Note that this will only work on machines with NVIDIA GPUs.
+Once running, you can access the Jupyter notebook server by clicking on the URL
+displayed in the console of the form "http[]()://127.0.0.1:8888/tree?token=abcdef...".
+
+## 🤗 Pre-trained models (via HuggingFace)
+
+We have prepared pre-trained models trained on a subset of the
 [SPICE2 dataset](https://zenodo.org/records/10975225) for each of the models included in
 this repo. They can be accessed directly on [InstaDeep's MLIP collection](https://huggingface.co/collections/InstaDeepAI/ml-interatomic-potentials-68134208c01a954ede6dae42),
 along with our curated dataset or directly through
@@ -112,13 +125,40 @@ hf_hub_download(repo_id="InstaDeepAI/visnet-organics", filename="visnet_organics
 hf_hub_download(repo_id="InstaDeepAI/nequip-organics", filename="nequip_organics_01.zip", local_dir="")
 hf_hub_download(repo_id="InstaDeepAI/SPICE2-curated", filename="SPICE2_curated.zip", local_dir="")
 ```
-Note that the foundation models are released on a different license than this library,
+Note that the pre-trained models are released on a different license than this library,
 please refer to the model cards of the relevant HuggingFace repos.
 
+## 🚀 Inference time benchmarks
+
+In order to showcase the runtime efficiency, we conducted benchmarks across all three models
+on two different systems: 1UAO (138 atoms) and 1ABT (1205 atoms), both run for 1ns on a H100
+NVidia GPU. All model implementations are our own, including the Torch + ASE benchmarks, and
+should not be considered representative of the performance of the code developed by the
+original authors of the methods. Further details can be found in our whitepaper (see below).
+
+**MACE (2,139,152 parameters):**
+| Systems   | JAX + JAX MD | JAX + ASE    | Torch + ASE  |
+| --------- |-------------:|-------------:|-------------:|
+| 1UAO      | 6.3 ms/step  | 11.6 ms/step | TBC ms/step  |
+| 1ABT      | TBC ms/step  | TBC ms/step  | TBC ms/step  |
+
+**ViSNet (1,137,922 parameters):**
+| Systems   | JAX + JAX MD | JAX + ASE    | Torch + ASE  |
+| --------- |-------------:|-------------:|-------------:|
+| 1UAO      | 2.9 ms/step  | 6.2 ms/step  | 33.8 ms/step |
+| 1ABT      | 25.4 ms/step | TBC ms/step  | TBC ms/step  |
+
+**NequIP (1,327,792 parameters):**
+| Systems   | JAX + JAX MD | JAX + ASE    | Torch + ASE  |
+| --------- |-------------:|-------------:|-------------:|
+| 1UAO      | 3.8 ms/step  | 8.5 ms/step  | 38.7 ms/step |
+| 1ABT      | TBC ms/step  | TBC ms/step  | TBC ms/step  |
+
 ## 🙏 Acknowledgments
 
-We would like to acknowledge beta testers for this library: Leon Wehrhan,
-Sebastien Boyer, Massimo Bortone, Tom Barrett, and Alex Laterre.
+We would like to acknowledge beta testers for this library: Isabel Wilkinson,
+Nick Venanzi, Hassan Sirelkhatim, Leon Wehrhan, Sebastien Boyer, Massimo Bortone,
+Tom Barrett, and Alex Laterre.
 
 ## 📚 Citing our work
 
@@ -128,4 +168,4 @@ C. Brunken, O. Peltre, H. Chomet, L. Walewski, M. McAuliffe, V. Heyraud,
 S. Attias, M. Maarand, Y. Khanfir, E. Toledo, F. Falcioni, M. Bluntzer,
 S. Acosta-Gutiérrez and J. Tilly, *Machine Learning Interatomic Potentials:
 library for efficient training, model development and simulation of molecular systems*,
-uploaded to arXiv soon.
+available on the arXiv.

README.md

@@ -7,28 +7,28 @@ Model fine-tuning
 
    Currently, fine-tuning is only available for MACE models.
 
-A common use case is fine-tuning a pre-trained MLIP foundation model
+A common use case is fine-tuning a pre-trained MLIP model
 on additional data to improve its accuracy for specific types of chemical systems.
 
 In the following, we describe how to fine-tune an MLIP model with this library. We
 recall that an MLIP model can be trained using multiple read-out heads. Note that
 currently, this is just implemented for the MACE architecture. The number of read-out
 heads can be set via ``num_readout_heads`` in
 :py:func:`MaceConfig <mlip.models.mace.config.MaceConfig>`.
-By default, one trains a foundation model with only one read-out head. However, it does
+By default, one trains a model with only one read-out head. However, it does
 not matter for this fine-tuning step whether a model already has *N* read-out heads,
 it can be fine-tuned by adding more heads and optimizing their associated weights only.
 Note that the final energy prediction of a model is obtained by summing the outputs
 of the *N* read-out heads.
 
 To fine-tune a given model, set up the new model with at least one more read-out head
-than the foundation model you already have pre-trained.
+than the pre-trained model.
 
 .. code-block:: python
 
     from mlip.models import Mace, ForceField
 
-    foundation_model_params = _get_params_for_pretrained_model()  # placeholder
+    pretrained_model_params = _get_params_for_pretrained_model()  # placeholder
 
     # Make sure the new model you create has at least one more read-out head
     mace = Mace(Mace.Config(num_readout_heads=2), dataset_info)
@@ -43,7 +43,7 @@ the function
     from mlip.models.params_transfer import transfer_params
 
     transferred_params, finetuning_blocks = transfer_params(
-        foundation_model_params,
+        pretrained_model_params,
         initial_force_field.params,
         scale_factor=0.1,
     )
@@ -64,7 +64,7 @@ to experiment with this hyperparameter themselves.
 
 The resulting ``transferred_params`` have the shape of your new model, but the new
 heads are not yet optimized. The other parameters are taken from the pre-trained
-foundation model. The second output of the function ``finetuning_blocks`` holds a list
+model. The second output of the function ``finetuning_blocks`` holds a list
 of module names inside the parameters that correspond to the blocks of untrained
 parameters. This list will be needed for the subsequent step.
 
@@ -99,15 +99,15 @@ the transferred parameters works like this:
 **To summarize, there are only three additional steps that are**
 **required for fine-tuning in contrast to a regular model training:**
 
-* Loading the original foundation model parameters *and* setting up a new model that
+* Loading the original pre-trained model parameters *and* setting up a new model that
   has the same configuration but with one or more additional read-out heads.
 * Transfer the parameters using the function
   :py:func:`transfer_params() <mlip.models.params_transfer.transfer_params>`.
 * Mask the optimizer using the function
   :py:func:`mask_optimizer_for_finetuning() <mlip.training.finetuning_utils.mask_optimizer_for_finetuning>`.
 
 **Additional note:** When fine-tuning on datasets that are quite different to the
-original dataset which the foundation model was trained on, we recommend to add a subset
+original dataset which the pre-trained model was trained on, we recommend to add a subset
 of the original dataset to the dataset the fine-tuning is performed on. The proportion
 to which the original dataset should extend the new data points (e.g., 50:50 or
 90:10 ratio) is a hyperparameter to experiment with and the optimal choice

docs/source/user_guide/finetuning.rst

@@ -9,7 +9,7 @@ Create a model and force field
 --------------------------------
 
 This section discusses how to initialize an MLIP model for subsequent training.
-If you are just interested in loading a pretrained model for application in simulations,
+If you are just interested in loading a pre-trained model for application in simulations,
 please see the dedicated section :ref:`below <load_zip_model>`.
 
 Our MLIP models exist in two abstraction levels:
@@ -109,7 +109,7 @@ Load a model from a zip archive
 -------------------------------
 
 To load a model (e.g., MACE) from our lightweight zip format that we ship our
-foundation models with, you can use the function
+pre-trained models with, you can use the function
 :py:func:`load_model_from_zip <mlip.models.model_io.load_model_from_zip>`:
 
 .. code-block:: python

docs/source/user_guide/models.rst

@@ -58,8 +58,8 @@ can be accessed after the run like this:
 
 However, the final parameters are not always the ones with the best
 performance on the validation set, and hence,
-you can also access these with ``training_loop.best_params`` or directly
-use `training_loop.best_model` to get the
+you can also access these with ``training_loop.best_model.params``.
+Therefore, use `training_loop.best_model` to get the
 :py:class:`ForceField <mlip.models.force_field.ForceField>` instance that holds
 the best parameters. If you want to save a
 trained force field not only via the checkpointing API described further below,

docs/source/user_guide/training.rst

@@ -12,16 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from typing import Optional, Union
 
 import jax.numpy as jnp
 
 from mlip.data.dataset_info import DatasetInfo
 from mlip.data.helpers.atomic_number_table import AtomicNumberTable
 
-logger = logging.getLogger("mlip")
-
 
 def get_atomic_energies(
     dataset_info: DatasetInfo,
@@ -69,19 +66,13 @@ def get_atomic_energies(
             z_table.z_to_index(z): energy
             for z, energy in dataset_info.atomic_energies_map.items()
         }
-        logger.debug(
-            f"Computed average atomic energies using least "
-            f"squares, taken from dataset info: {atomic_energies_dict}"
-        )
         atomic_energies = jnp.array(
             [atomic_energies_dict[i] for i in range(len(z_table.zs))]
         )
     elif atomic_energies_input == "zero":
-        logger.debug("Not using atomic energies, setting them to zero.")
         atomic_energies = jnp.zeros(num_species)
     elif isinstance(atomic_energies_input, dict):
         atomic_energies_dict = atomic_energies_input
-        logger.debug(f"Use Atomic Energies that are provided: {atomic_energies_dict}")
         atomic_energies = jnp.array(
             [atomic_energies_dict.get(z, 0.0) for z in range(num_species)]
         )

mlip/models/atomic_energies.py

@@ -1,3 +1,7 @@
+# MIT License
+# Copyright (c) 2022 mace-jax
+# See https://github.com/ACEsuit/mace-jax/blob/main/MIT.md
+#
 # Copyright 2025 InstaDeep Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");