v0.2.2

This release contains the following changes:

• 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.

v0.2.1

This release includes the following changes:

• Reducing memory requirements for Hessian predictions in batched inference drastically.
• Optimizing eSEN implementation, introducing (1) Quaternion method to compute Wigner D matrices (based on proposal by @bhcao, see #46), 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.

v0.2.0

This release includes the following changes:

• Introducing breaking API updates in multiple places, see Migration guide in documentation for a detailed overview. These updates make the mlip library more customizable and extendable, and enable a range of new features listed below.
• Extensive refactoring of the model implementations, updating to a Graph -> Graph signature for MLIPNetwork and for other blocks and layers used inside models.
• Introducing the e3j backend for accelerated equivariant operations for MACE and NequIP models.
• Adding the eSEN model architecture, and optionally, with a Mixture-of-Experts (MoE) formalism.
• Introducing Gaunt tensor products for MACE.
• Adding global charge conditioning, partial charge predictions, and long-range interactions via a PhysNet-inspired Coulomb term.
• Adding support for NPT ensemble MD simulations.
• Adding support for transition state search with the nudged elastic band method.
• Extending multi-head fine-tuning support to all model architectures and updating fine-tuning API and strategy.
• Enabling the training of models on Hessian labels.

v0.1.10

This version includes the following update:

• Enhancing performance of dataset preprocessing by removing JIT compilation as well as using numpy instead of jax.numpy in edge vector function inside the neighbor analysis code (contributed by @bhcao in #42).

v0.1.9

This version includes the changes listed below.

Furthermore, note that this release is planned to be one of the last of type v0.1.x, with a v0.2.0 release planned for release within the next two months. The v0.2.0 release will contain many new features, but also some breaking API changes, however, we'll provide a detailed migration guide along with the release.

• Migrating training pipeline from pmap to SPMD jax.jit with NamedSharding, enabling multi-host training on TPU and multi-GPU setups.
• Adding support for multi-host data-parallel training.
• Checkpointing now works across multiple hosts (only process 0 saves).
• Lifting restrictions on the compatible versions of the orbax-checkpoint dependency.
• Removing key field from TrainingState and random_key parameter from init_training_state. Old checkpoints containing a keyfield can still be restored; thekeyis skipped via Orbaxpartial_restore`.
• Adding multi-host utilities: create_device_mesh, create_replicated_sharding, create_dp_sharding, and sync_string for cross-host communication.
• Adding SYSTEM_METRICS log category for per-process runtime and throughput metrics. Important note: Some of these metrics were logged under TRAIN_METRICS before, so please update your custom loggers to keep logging these metrics.
• Disabling async checkpointing for multi-host compatibility.
• Deprecating should_parallelize parameter in TrainingLoop in favour of mesh.
• Deprecating devices parameter in GraphDatasetBuilder.get_splits() in favour of mesh.
• Updating pre-commit configuration to use ruff in place of isort, black and flake8.
• Adding early stopping when a simulation has exploded, meaning that its temperature is NaN or greater than 1e6.
• Fixing bug in ASE simulation engine: not using fixed random seed made runs irreproducible.
• Enabling the initialization of a simulation using 2D (num_atoms, 3) arrays for positions and velocities. This allows for initializing a simulation using a single frame, rather than restoring from 3D arrays containing a multi-step trajectory.

v0.1.8

This version includes the following changes:

• Fixing bug in ASE simulation engine to allow for passing Periodic Boundary Conditions via the box config value.
• Adapting setup line for zero shifts array to use np.zeros to guarantee correct array shape, see issue #36 for reference.

v0.1.7

This version includes the following changes:

• Fixing issues with Periodic Boundary Conditions (PBCs) during inference.
• Supporting PBCs passed from ase.Atoms during simulation with the ASE engine. Passing an orthorhombic box from configuration is still supported in both simulation engines, but might become discouraged in future releases.
• Fixing a few bugs related to batched simulations occurring in cases of reallocation of neighbor lists.

v0.1.6

This patch does not modify the library code. It only fixes the incorrect instructions for GPU-compatible installation present in README, docs, and tutorials: most shells require quotes around pip installations with extras.

v0.1.5

This version includes the following changes:

• 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).
• 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.

v0.1.4

This version includes the following changes:

• Removing constraints on some dependencies, such as numpy, jax, and flax. The mlip library now allows for more flexibility in dependency versions for downstream projects. This includes support for the newest jax versions 0.6.x and 0.7.x.
• Fixing simulation tutorial notebook by pinning versions of visualization helper libraries.
• Adding the option to pass the dataset_info of a trained model to GraphDatasetBuilder, which is important for downstream tasks. Failure to do so might lead to silent inconsistencies in the mapping from atomic numbers to specie indices, especially when the downstream data has fewer elements than the training set (see e.g. the fine-tuning tutorial).
• Fixing the stress predictions, with new formulas for the virial stress and 0 Kelvin pressure term. These features should still be seen as beta for now as we proceed to test them further (see docstrings for more details).