release: create release-0.1.5 branch

Author: github-actions[bot]

.github/workflows/deploy_docs.yaml

@@ -14,18 +14,17 @@ jobs:
        image: python:3.12-slim-bullseye
     steps:
       - uses: actions/checkout@v4
-      - name: Install dependencies
+
+      - name: Setup
         run: |
           apt-get update && apt-get install -y coreutils git
-          POETRY_VERSION=1.8.4
-          pip install -U pip setuptools
-          pip install poetry==${POETRY_VERSION}
-          poetry install
-          poetry run pip install git+https://github.com/jax-md/jax-md.git
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
 
       - name: Sphinx build
         run: |
-          poetry run sphinx-build -b html docs/source _build
+          uv run sphinx-build -b html docs/source _build
 
       - name: Deploy to GitHub Pages
         uses: peaceiris/actions-gh-pages@v3

.github/workflows/publish.yaml

@@ -17,14 +17,13 @@ jobs:
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: install poetry
-        run: |
-          POETRY_VERSION=1.8.4
-          pip install -U pip setuptools
-          pip install poetry==${POETRY_VERSION}
+      - name: setup
+        run: apt-get update && apt-get install -y coreutils
+
+      - name: install uv
+        uses: astral-sh/setup-uv@v6
 
       - name: build and publish
         run: |
-          export POETRY_PYPI_TOKEN_PYPI=${{secrets.POETRY_PYPI_TOKEN}}
-          poetry build
-          poetry publish
+          uv build
+          uv publish --token ${{secrets.POETRY_PYPI_TOKEN}}

.github/workflows/tests_and_linters.yaml

@@ -22,16 +22,16 @@ jobs:
       - name: checkout code 📦
         uses: actions/checkout@v4
 
-      - name: install poetry
-        run: |
-          POETRY_VERSION=1.8.4
-          pip install -U pip setuptools
-          pip install poetry==${POETRY_VERSION}
-          poetry install
+      - name: install uv
+        uses: astral-sh/setup-uv@v6
+
+      - name: setup env
+        run: uv sync --only-dev
 
       - name: run linters 🖌️
         run: |
-          poetry run pre-commit run --all-files --verbose
+          git init
+          uv run --no-sync pre-commit run --all-files --verbose
 
   tests:
     runs-on: ubuntu-latest
@@ -41,21 +41,25 @@ jobs:
       - name: checkout code 📂
         uses: actions/checkout@v4
 
-      - name: install poetry
+      - name: setup
         run: |
-          POETRY_VERSION=1.8.4
-          pip install -U pip setuptools
-          pip install poetry==${POETRY_VERSION}
-          poetry install
-          poetry add git+https://github.com/jax-md/jax-md.git
+          apt-get update && apt-get install -y \
+          coreutils \
+          git
+
+      - name: install uv
+        uses: astral-sh/setup-uv@v6
+
+      - name: setup env
+        run: uv sync
 
       - name: run tests 🧪
         run: |
-          poetry run pytest --ignore tests/experiments --verbose \
-                                            --cov-report xml:coverage.xml \
-                                            --cov-report term-missing \
-                                            --junitxml=pytest.xml \
-                                            --cov=mlip tests/
+          uv run --no-sync pytest --ignore tests/experiments --verbose \
+                                  --cov-report xml:coverage.xml \
+                                  --cov-report term-missing \
+                                  --junitxml=pytest.xml \
+                                  --cov=mlip tests/
 
       - name: pytest coverage comment
         id: coverageComment

CHANGELOG

@@ -1,5 +1,24 @@
 # Changelog
 
+## Release 0.1.5
+
+- Adding batched simulations feature for MD simulations and energy minimizations
+  with the JAX-MD backend.
+- Removing now useless `stress_virial` prediction.
+- Fixing correctness of `stress` and 0K `pressure` predictions. In 0.1.4,
+  the stress computation actually involved a derivative with respect to
+  cell but with fixed positions. Now, the strain also acts on positions within
+  the unit cell, thus deforming the material homogeneously. This rigorously
+  translation-invariant stress exempts from any Virial term correction of
+  cell boundary effects. See for instance
+  [Thompson, Plimpton and Mattson 2009, eq (2)](https://doi.org/10.1063/1.3245303).
+- Migrating from poetry to uv for dependency and package management.
+- Improving inefficient logging strategy in ASE simulation backend.
+- Clarifying in the documentation that we recommend a smaller value for the timestep
+  when running energy minimizations with the JAX-MD simulation backend.
+- Removing need for separate install command for JAX-MD dependency.
+- Adding easier install method for GPU-compatible JAX.
+
 ## Release 0.1.4
 
 - Removing constraints on some dependencies, such as numpy, jax, and flax. The mlip

Dockerfile

@@ -1,10 +1,10 @@
-FROM python:3.10.12-slim-bullseye
+FROM python:3.12-slim-bullseye
 
 WORKDIR /app
 
 RUN apt-get update && apt-get install -y git wget
 
-RUN pip install mlip "jax[cuda12]==0.6.2" huggingface_hub git+https://github.com/jax-md/jax-md.git notebook
+RUN pip install mlip[cuda] huggingface_hub 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 \

README.md

@@ -1,5 +1,9 @@
 # 🪩 MLIP: Machine Learning Interatomic Potentials 🚀
 
+[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+[![Python 3.11](https://img.shields.io/badge/python-3.11%20%7C%203.12%20%7C%203.13-blue)](https://www.python.org/downloads/release/python-3110/)
+[![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit)](https://github.com/pre-commit/pre-commit)
+[![Tests and Linters 🧪](https://github.com/instadeepai/mlip/actions/workflows/tests_and_linters.yaml/badge.svg?branch=main)](https://github.com/instadeepai/mlip/actions/workflows/tests_and_linters.yaml)
 ![badge](https://img.shields.io/endpoint?url=https://gist.githubusercontent.com/mlipbot/b6e4bf384215e60775699a83c3c00aef/raw/pytest-coverage-comment.json)
 
 ## 👀 Overview
@@ -12,6 +16,7 @@ the following functionality:
 - Training and fine-tuning MLIP models
 - Batched inference with trained MLIP models
 - MD simulations with MLIP models using multiple simulation backends (for now: JAX-MD and ASE)
+- Batched MD simulations and energy minimizations with the JAX-MD simulation backend.
 - Energy minimizations with MLIP models using the same simulation backends as for MD.
 
 The purpose of the library is to provide users with a toolbox
@@ -49,29 +54,18 @@ pip install mlip
 
 However, this command **only installs the regular CPU version** of JAX.
 We recommend that the library is run on GPU.
-This requires also installing the necessary versions
-of [jaxlib](https://pypi.org/project/jaxlib/) which can also be installed via pip. See
-the [installation guide of JAX](https://docs.jax.dev/en/latest/installation.html) for
-more information.
-At time of release, the following install command is supported:
+Use this command instead to install the GPU-compatible version:
 
 ```bash
-pip install -U "jax[cuda12]"
+pip install mlip[cuda]
 ```
 
-Note that using the TPU version of *jaxlib* is, in principle, also supported by
-this library. However, it has not been thoroughly tested and should therefore be
-considered an experimental feature.
+**This command installs the CUDA 12 version of JAX.** For different versions, please
+install *mlip* without the `cuda` flag and install the desired JAX version via pip.
 
-Also, some tasks in *mlip* will
-require [JAX-MD](https://github.com/jax-md/jax-md>) as a dependency. As the newest
-version of JAX-MD is not available on PyPI yet, this dependency will not
-be shipped with *mlip* automatically and instead must be installed
-directly from the GitHub repository, like this:
-
-```bash
-pip install git+https://github.com/jax-md/jax-md.git
-```
+Note that using the TPU version of JAX is, in principle, also supported by
+this library. You need to install it separately via pip. However, it has not been
+thoroughly tested and should therefore be considered an experimental feature.
 
 ## ⚡ Examples
 

docs/source/index.rst

@@ -15,6 +15,7 @@ in JAX. It contains the following features:
 * Batched inference with trained MLIP models
 * MD simulations with MLIP models using multiple simulation backends
 * Energy minimizations with MLIP models using multiple simulation backends
+* Batched MD simulations and energy minimizations with the JAX-MD simulation backend
 * Fine-tuning of pre-trained MLIP models
 
 As a first step, we recommend that you check out our page on :ref:`installation`

docs/source/installation/index.rst

@@ -11,26 +11,15 @@ The *mlip* library can be installed via pip:
 
 However, this command **only installs the regular CPU version** of JAX.
 We recommend that the library is run on GPU.
-This requires also installing the necessary versions
-of `jaxlib <https://pypi.org/project/jaxlib/>`_ which can also be installed via pip. See
-the `installation guide of JAX <https://docs.jax.dev/en/latest/installation.html>`_ for
-more information.
-At time of release, the following install command is supported:
+Use this command instead to install the GPU-compatible version:
 
 .. code-block:: bash
 
-    pip install -U "jax[cuda12]"
+    pip install mlip[cuda]
 
-Note that using the TPU version of *jaxlib* is, in principle, also supported by
-this library. However, it has not been thoroughly tested and should therefore be
-considered an experimental feature.
+**This command installs the CUDA 12 version of JAX.** For different versions, please
+install *mlip* without the `cuda` flag and install the desired JAX version via pip.
 
-Also, some tasks in *mlip* will
-require `JAX-MD <https://github.com/jax-md/jax-md>`_ as a dependency. As the newest
-version of JAX-MD is not available on PyPI yet, this dependency will not
-be shipped with *mlip* automatically and instead must be installed
-directly from the GitHub repository, like this:
-
-.. code-block:: bash
-
-    pip install git+https://github.com/jax-md/jax-md.git
+Note that using the TPU version of JAX is, in principle, also supported by
+this library. You need to install it separately via pip. However, it has not been
+thoroughly tested and should therefore be considered an experimental feature.

docs/source/user_guide/simulations.rst

@@ -82,12 +82,21 @@ class for more details. Most importantly, the `simulation_type` needs to be set
 `SimulationType.MINIMIZATION` (see
 :py:class:`SimulationType <mlip.simulation.enums.SimulationType>`).
 
+.. note::
+
+    The default timestep of 1.0 fs that is common for MD simulations may not be optimal
+    for energy minimizations. We recommend to set this value to 0.1 fs when using the
+    `SimulationType.MINIMIZATION` mode with the JAX-MD backend.
+
 **Algorithms**: For MD, the NVT-Langevin algorithm is used
 (see `here <https://jax-md.readthedocs.io/en/main/jax_md.simulate.html#jax_md.simulate.nvt_langevin>`_).
 For energy minimization, the FIRE algorithm is used
 (see `here <https://jax-md.readthedocs.io/en/main/jax_md.minimize.html#jax_md.minimize.fire_descent>`_).
 We plan to provide more options in future versions of the library.
 
+Furthermore, for MD simulations, we support running them in a **batched manner**.
+See :ref:`this <batched_simulations>` section below for more information.
+
 .. note::
 
    A special feature of the JAX-MD backend is that a simulation is divided into
@@ -203,6 +212,45 @@ The logger must be attached before starting the simulation.
 In ASE, this logging function will be called depending on the logging interval set,
 and in JAX-MD, it will be called after every episode.
 
+.. _batched_simulations:
+
+Batched simulations with JAX-MD
+-------------------------------
+
+For MD simulations or energy minimizations with JAX-MD, we support running them in a
+batched manner for multiple systems. The API for this is straightforward,
+instead of passing a single `ase.Atoms` object to the engine, we pass a list of them.
+After the simulation, the simulation state will contain lists of properties,
+for example, a list of position arrays (i.e., the trajectories) instead of a single
+position array. Note that it is also supported that the input molecules have
+varying sizes. See example code below:
+
+.. code-block:: python
+
+    from ase.io import read as ase_read
+    from mlip.simulation.jax_md import JaxMDSimulationEngine
+
+    systems = []
+    for path in ["/path/to/mol_1", "/path/to/mol_2", "/path/to/mol_3"]:
+        atoms = ase_read(path)
+        systems.append(atoms)
+
+    force_field, md_config = _get_from_somewhere()  # placeholder
+    md_engine = JaxMDSimulationEngine(systems, force_field, md_config)
+    md_engine.run()
+
+    # Fetch results:
+    # Get trajectory and temperatures for "/path/to/mol_2" (indexing starts at 0)
+    md_state = md_engine.state
+    print(md_state.positions[1])
+    print(md_state.temperature[1])
+
+    # Compute time, for example, is not a list
+    print(md_state.compute_time_seconds)
+
+The example above works for both energy minimizations and MD simulations in the same
+way.
+
 .. _batched_inference:
 
 Batched inference