mlip 0.2.0__tar.gz → 0.2.2__tar.gz

{mlip-0.2.0 → mlip-0.2.2}/CHANGELOG

@@ -1,5 +1,43 @@
 # Changelog
 
+## Release 0.2.2
+
+- Adding support for running well-tempered metadynamics MD simulations using the
+  JAX-MD backend.
+- Adding support for NVE ensemble MD simulations.
+- Adding support for setting a random seed for simulations via the config, and
+  introducing per-trajectory random seeds in batched JAX-MD simulations via the
+  new `independent_seeds_batched` config flag. This is now `True` by
+  default, meaning batched simulations use independent random keys per trajectory unless
+  explicitly disabled.
+- Adding support for logging of potential energy and predicted partial charges during MD
+  simulations for the JAX-MD backend, and for both MD simulations and energy
+  minimizations for the ASE backend.
+- Reducing memory requirements for the Mixture-of-Experts (MoE) dense layer.
+- Fixing a bug in graph batching where `spin_multiplicity` was missed by
+  homogenization.
+- Updating batched inference to allow for efficient energy-only predictions.
+- Fixing energy, force, and Hessian discontinuities at the cutoff in ViSNet. Set
+  `use_legacy_visnet=True` in `VisnetConfig` to reproduce models trained with
+  versions v0.2.0 and v0.2.1. Note that after loading a ViSNet model trained with
+  these versions, it can still be used for inference after updating its config via
+  `force_field = force_field.replace_config(use_legacy_visnet=True)`. The pre-trained
+  ViSNet model provided on HuggingFace has been updated after retraining with the
+  corrected version of the code. See the API documentation of
+  `VisnetConfig.use_legacy_visnet` for details of the changes to the model equations.
+
+## Release 0.2.1
+
+- Reducing memory requirements for Hessian predictions in batched inference drastically.
+- Optimizing eSEN implementation, introducing (1) Quaternion method to compute Wigner
+  D matrices, and (2) `use_remat_edgewise` flag to rematerialize `Edgewise`.
+- Improving file reading efficiency in `Hdf5Reader`.
+- Fixing `CombinedGraphDataset` by catching `StopIteration` runtime error.
+- Fixing `MSEHessianLoss` to use correct axes for aggregation.
+- Updating advanced simulation and Hessian tutorial notebooks.
+- Fixing memory bottleneck in ViSNet, adding rematerialisation option with `use_remat`.
+- Upgrading e3j dependency to minimum version of 0.1.0b3.
+
 ## Release 0.2.0
 
 - Introducing breaking API updates in multiple places, see Migration guide in

mlip-0.2.2/Dockerfile

@@ -0,0 +1,19 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y git wget
+
+RUN pip install "mlip[cuda13]" 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 \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/model_addition_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/advanced_simulation_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/moe_training_and_inference_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/hessian_model_training_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/metadynamics_tutorial.ipynb
+
+EXPOSE 8888
+
+CMD ["jupyter", "notebook", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root"]

{mlip-0.2.0 → mlip-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlip
-Version: 0.2.0
+Version: 0.2.2
 Summary: Machine Learning Interatomic Potentials in JAX
 Project-URL: Homepage, https://github.com/instadeepai/mlip/
 Author-email: InstaDeep <placeholder@instadeep.com>
@@ -18,7 +18,7 @@ Requires-Python: <3.14,>=3.11
 Requires-Dist: ase>=3.22.1
 Requires-Dist: dm-haiku>=0.0.14
 Requires-Dist: dm-tree>=0.1.8
-Requires-Dist: e3j<0.1.1,>=0.1.0b0
+Requires-Dist: e3j<0.1.1,>=0.1.0b3
 Requires-Dist: e3nn-jax>=0.20.7
 Requires-Dist: flax>=0.9.0
 Requires-Dist: h5py>=3.12.1
@@ -32,20 +32,24 @@ Requires-Dist: pydantic>=2.10.2
 Requires-Dist: roundmantissa>=0.2.1
 Requires-Dist: tqdm-loggable>=0.2
 Requires-Dist: tqdm>=4.0
+Provides-Extra: cuda
+Requires-Dist: e3j-ops<0.1.1,>=0.1.0b3; extra == 'cuda'
+Requires-Dist: jax[cuda13]>=0.8.0; extra == 'cuda'
+Requires-Dist: jaxlib>=0.8.0; extra == 'cuda'
 Provides-Extra: cuda12
-Requires-Dist: e3j-ops<0.1.1,>=0.1.0b0; extra == 'cuda12'
+Requires-Dist: e3j-ops<0.1.1,>=0.1.0b3; extra == 'cuda12'
 Requires-Dist: jax[cuda12]>=0.4.32; extra == 'cuda12'
 Requires-Dist: jaxlib>=0.4.32; extra == 'cuda12'
 Provides-Extra: cuda12-local
-Requires-Dist: e3j-ops<0.1.1,>=0.1.0b0; extra == 'cuda12-local'
+Requires-Dist: e3j-ops<0.1.1,>=0.1.0b3; extra == 'cuda12-local'
 Requires-Dist: jax[cuda12-local]>=0.4.32; extra == 'cuda12-local'
 Requires-Dist: jaxlib>=0.4.32; extra == 'cuda12-local'
 Provides-Extra: cuda13
-Requires-Dist: e3j-ops<0.1.1,>=0.1.0b0; extra == 'cuda13'
+Requires-Dist: e3j-ops<0.1.1,>=0.1.0b3; extra == 'cuda13'
 Requires-Dist: jax[cuda13]>=0.8.0; extra == 'cuda13'
 Requires-Dist: jaxlib>=0.8.0; extra == 'cuda13'
 Provides-Extra: cuda13-local
-Requires-Dist: e3j-ops<0.1.1,>=0.1.0b0; extra == 'cuda13-local'
+Requires-Dist: e3j-ops<0.1.1,>=0.1.0b3; extra == 'cuda13-local'
 Requires-Dist: jax[cuda13-local]>=0.8.0; extra == 'cuda13-local'
 Requires-Dist: jaxlib>=0.8.0; extra == 'cuda13-local'
 Provides-Extra: tpu
@@ -89,10 +93,11 @@ modular API for development of new architectures
 - ⚡ Batched inference with trained MLIP models
 - 🧪 MD simulations with MLIP models using multiple simulation backends
   (currently: JAX-MD and ASE)
-- 🌡️ Support for both NVT and NPT ensembles in MD
+- 🌡️ Support for NVT, NPT and NVE ensembles in MD
 - ⛰️ Energy minimizations using the same simulation backends as for MD
 - 🚀 Batched MD simulations and energy minimizations with the JAX-MD backend
 - 🔎 Transition state search with the nudged elastic band (NEB) method
+- 🌋 Enhanced sampling via metadynamics for free energy landscape exploration
 - 🌐 Global charge conditioning, partial charge predictions, and
 support for long-range interactions
 - 📈 Training on Hessian labels
@@ -144,6 +149,9 @@ To install the **CUDA 13 version of JAX and e3j** binaries alongside *mlip*, run
 pip install "mlip[cuda13]"
 ```
 
+The alias `cuda` can be used instead of `cuda13` for backwards
+compatibility with *mlip* v1.
+
 To install the CUDA 12 version instead, run:
 
 ```bash
@@ -175,6 +183,7 @@ simple templates to build your own MLIP pipelines:
 - [Advanced simulation tutorial](https://github.com/instadeepai/mlip/blob/main/tutorials/advanced_simulation_tutorial.ipynb)
 - [MoE training tutorial](https://github.com/instadeepai/mlip/blob/main/tutorials/moe_training_and_inference_tutorial.ipynb)
 - [Hessian model training tutorial](https://github.com/instadeepai/mlip/blob/main/tutorials/hessian_model_training_tutorial.ipynb)
+- [Metadynamics tutorial](https://github.com/instadeepai/mlip/blob/main/tutorials/metadynamics_tutorial.ipynb)
 
 To run the tutorials, just install Jupyter notebooks via pip and launch it from
 a directory that contains the notebooks:
@@ -242,8 +251,8 @@ Further details can be found in our white paper (see [below](#-citing-our-work))
 **ViSNet (1,172,676 parameters):**
 | Systems   | JAX-MD       | ASE          |
 | --------- |-------------:|-------------:|
-| 1UAO      | 1.9 ms/step  | 7.1 ms/step  |
-| 1ABT      | 13.7 ms/step | 30.2 ms/step |
+| 1UAO      | 1.6 ms/step  | 7.0 ms/step  |
+| 1ABT      | 9.0 ms/step  | 24.0 ms/step |
 
 **NequIP (1,327,792 parameters):**
 | Systems   | JAX-MD       | ASE          |

{mlip-0.2.0 → mlip-0.2.2}/README.md

@@ -33,10 +33,11 @@ modular API for development of new architectures
 - ⚡ Batched inference with trained MLIP models
 - 🧪 MD simulations with MLIP models using multiple simulation backends
   (currently: JAX-MD and ASE)
-- 🌡️ Support for both NVT and NPT ensembles in MD
+- 🌡️ Support for NVT, NPT and NVE ensembles in MD
 - ⛰️ Energy minimizations using the same simulation backends as for MD
 - 🚀 Batched MD simulations and energy minimizations with the JAX-MD backend
 - 🔎 Transition state search with the nudged elastic band (NEB) method
+- 🌋 Enhanced sampling via metadynamics for free energy landscape exploration
 - 🌐 Global charge conditioning, partial charge predictions, and
 support for long-range interactions
 - 📈 Training on Hessian labels
@@ -88,6 +89,9 @@ To install the **CUDA 13 version of JAX and e3j** binaries alongside *mlip*, run
 pip install "mlip[cuda13]"
 ```
 
+The alias `cuda` can be used instead of `cuda13` for backwards
+compatibility with *mlip* v1.
+
 To install the CUDA 12 version instead, run:
 
 ```bash
@@ -119,6 +123,7 @@ simple templates to build your own MLIP pipelines:
 - [Advanced simulation tutorial](https://github.com/instadeepai/mlip/blob/main/tutorials/advanced_simulation_tutorial.ipynb)
 - [MoE training tutorial](https://github.com/instadeepai/mlip/blob/main/tutorials/moe_training_and_inference_tutorial.ipynb)
 - [Hessian model training tutorial](https://github.com/instadeepai/mlip/blob/main/tutorials/hessian_model_training_tutorial.ipynb)
+- [Metadynamics tutorial](https://github.com/instadeepai/mlip/blob/main/tutorials/metadynamics_tutorial.ipynb)
 
 To run the tutorials, just install Jupyter notebooks via pip and launch it from
 a directory that contains the notebooks:
@@ -186,8 +191,8 @@ Further details can be found in our white paper (see [below](#-citing-our-work))
 **ViSNet (1,172,676 parameters):**
 | Systems   | JAX-MD       | ASE          |
 | --------- |-------------:|-------------:|
-| 1UAO      | 1.9 ms/step  | 7.1 ms/step  |
-| 1ABT      | 13.7 ms/step | 30.2 ms/step |
+| 1UAO      | 1.6 ms/step  | 7.0 ms/step  |
+| 1ABT      | 9.0 ms/step  | 24.0 ms/step |
 
 **NequIP (1,327,792 parameters):**
 | Systems   | JAX-MD       | ASE          |

{mlip-0.2.0 → mlip-0.2.2}/docs/source/api_reference/data/batching_helpers.rst

@@ -2,7 +2,7 @@
 
 .. module:: mlip.graph.batching_helpers
 
-Batching Helpers
+Batching helpers
 ================
 
 .. autofunction:: batch_graphs

mlip-0.2.2/docs/source/api_reference/data/dataset_builder_config.rst

@@ -0,0 +1,8 @@
+.. _dataset_configs:
+
+.. module:: mlip.data.configs
+
+Dataset builder config
+======================
+
+.. autoclass:: GraphDatasetBuilderConfig

{mlip-0.2.0 → mlip-0.2.2}/docs/source/api_reference/data/index.rst

@@ -16,7 +16,7 @@ This module contains all code related to dataset reading and processing.
     chemical_systems_readers/index
     chemical_system
     dataset_info
-    dataset_configs
+    dataset_builder_config
     graph_dataset
     combined_graph_dataset
     data_prefetching

{mlip-0.2.0 → mlip-0.2.2}/docs/source/api_reference/models/esen.rst

@@ -1,7 +1,7 @@
 .. _esen:
 
-ESEN
-======
+eSEN
+====
 
 .. module:: mlip.models.esen.network
 

{mlip-0.2.0 → mlip-0.2.2}/docs/source/api_reference/simulation/index.rst

@@ -24,3 +24,6 @@ This module contains all code related to simulations (e.g., MD or energy minimiz
     batched_inference
     neb_engine
     neb_config
+    metadynamics_engine
+    metadynamics_config
+    metadynamics_potential_configs

mlip-0.2.2/docs/source/api_reference/simulation/metadynamics_config.rst

@@ -0,0 +1,10 @@
+.. _metadynamics_config:
+
+.. module:: mlip.simulation.metadynamics.config
+
+Metadynamics Simulation Config
+==============================
+
+.. autoclass:: MetadynamicsConfig
+
+.. autoclass:: JaxMDMetadynamicsSimulationConfig

mlip-0.2.2/docs/source/api_reference/simulation/metadynamics_engine.rst

@@ -0,0 +1,14 @@
+.. _jaxmd_metadynamics_engine:
+
+.. module:: mlip.simulation.metadynamics.jax_md_metad_engine
+
+Metadynamics Simulation Engine
+==============================
+
+.. autoclass:: JaxMDMetadynamicsSimulationEngine
+
+    .. automethod:: __init__
+
+    .. automethod:: run
+
+    .. automethod:: attach_logger

mlip-0.2.2/docs/source/api_reference/simulation/metadynamics_potential_configs.rst

@@ -0,0 +1,20 @@
+.. _metadynamics_potential_configs:
+
+.. module:: mlip.simulation.metadynamics.potential_terms
+
+Metadynamics Potential Configs
+==============================
+
+.. autoclass:: DistanceCVConfig
+
+.. autoclass:: AngleCVConfig
+
+.. autoclass:: DihedralCVConfig
+
+.. autoclass:: CoordinationNumberCVConfig
+
+.. autoclass:: DistanceWallConfig
+
+.. autoclass:: AngleWallConfig
+
+.. autoclass:: PositionalRestraintConfig

{mlip-0.2.0 → mlip-0.2.2}/docs/source/api_reference/simulation/simulation_state.rst

@@ -12,3 +12,7 @@ Simulation State
 .. autoclass:: NEBSimulationState
 
     .. automethod:: __init__
+
+.. autoclass:: mlip.simulation.metadynamics.states.MetadynamicsSimulationState
+
+    .. automethod:: __init__

{mlip-0.2.0 → mlip-0.2.2}/docs/source/conf.py

@@ -18,6 +18,9 @@ extensions = ["sphinx.ext.autodoc", "sphinx.ext.napoleon", "sphinx_design"]
 exclude_patterns = []
 default_role = "code"
 
+# Suppress cross-reference warnings for attribute names shared across multiple classes.
+suppress_warnings = ["ref.python"]
+
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 

{mlip-0.2.0 → mlip-0.2.2}/docs/source/index.rst

@@ -40,8 +40,8 @@ It provides tools for:
       :link: user_guide/simulations
       :link-type: doc
 
-      Molecular dynamics, energy minimization, and transition state search with
-      multiple backends
+      Molecular dynamics, energy minimization, transition state search, and
+      metadynamics with multiple backends
 
    .. grid-item-card:: ⚡ State-of-the-art speed
 

{mlip-0.2.0 → mlip-0.2.2}/docs/source/installation/index.rst

@@ -19,6 +19,9 @@ Use this command instead to install the GPU-compatible version:
 
 **This command installs the CUDA 13 version of JAX.**
 
+Note that the alias `cuda` can be used instead of `cuda13` for backwards compatibility
+with *mlip* v1.
+
 To install the CUDA 12 version instead, run:
 
 .. code-block:: bash

{mlip-0.2.0 → mlip-0.2.2}/docs/source/user_guide/index.rst

@@ -78,6 +78,7 @@ deep-dive tutorials :ref:`below <tutorials>` to help you with getting onboarded
 * `Addition of new models <https://github.com/instadeepai/mlip/blob/main/tutorials/model_addition_tutorial.ipynb>`_
 * `Training on Hessian labels <https://github.com/instadeepai/mlip/blob/main/tutorials/hessian_model_training_tutorial.ipynb>`_
 * `Advanced simulation <https://github.com/instadeepai/mlip/blob/main/tutorials/advanced_simulation_tutorial.ipynb>`_
+* `Metadynamics <https://github.com/instadeepai/mlip/blob/main/tutorials/metadynamics_tutorial.ipynb>`_
 
 To run the tutorials, install Jupyter notebooks via pip and launch it from
 a directory that contains the notebooks: