Allegro Equivariant Deep Learning Interatomic Potentials

0.8.2 · active · verified Fri Apr 17

Allegro is an open-source Python library for building highly scalable and accurate equivariant deep learning interatomic potentials, often used in materials science and molecular dynamics. It is currently at version 0.8.2 and sees active development with frequent minor releases and occasional breaking changes in major versions.

Common errors

```
AttributeError: module 'allegro.model' has no attribute 'AllegroEnergyModel'
```
cause The `AllegroEnergyModel` class was removed in Allegro v0.8.0.

fix Use `from allegro.model import AllegroModel` and set `do_derivatives=False` in its constructor if you need an energy-only model: `AllegroModel(..., do_derivatives=False)`.
```
RuntimeError: The installed version of nequip (X.Y.Z) is incompatible with allegro (requires >=0.17.0)
```
cause The installed `nequip` package version does not meet the minimum requirement for the installed `nequip-allegro` version.

fix Uninstall existing `nequip` and `nequip-allegro` and reinstall `nequip` with the required version, then `nequip-allegro`. For example, for Allegro 0.8.2, run `pip uninstall nequip nequip-allegro` then `pip install nequip>=0.17.0` followed by `pip install nequip-allegro`.
```
ImportError: cannot import name 'AllegroModel' from 'allegro.model' (.../allegro/model/__init__.py)
```
cause The `nequip-allegro` package might not be correctly installed, or there's a typo in the import statement.

fix Ensure `nequip-allegro` is installed by running `pip install nequip-allegro`. Verify the import statement `from allegro.model import AllegroModel` is correct and matches the official documentation.

Warnings

breaking The `AllegroEnergyModel` class was removed in version 0.8.0. For energy-only models, use `AllegroModel` with `do_derivatives=False`.
Fix: Replace `AllegroEnergyModel(...)` with `AllegroModel(..., do_derivatives=False)`.
breaking Allegro has strict compatibility requirements with `nequip` versions. Installing an incompatible `nequip` version will lead to runtime errors or unexpected behavior.
Fix: Always check the `nequip-allegro` release notes (e.g., `minimum nequip version is v0.17.0` for Allegro 0.8.2) and install the specified `nequip` version using `pip install nequip==X.Y.Z` before or alongside `nequip-allegro`.
breaking The mechanism for using custom Triton TP kernels was changed in version 0.6.0, potentially breaking existing custom implementations.
Fix: Refer to the official Allegro documentation for the updated method of integrating custom Triton TP kernels, or revert to versions <0.6.0 if custom kernels are critical and unmigrated.
gotcha Integrating CuEquivariance for acceleration requires specific configuration modifications within your `nequip` config file (e.g., using `nequip.model.modify` with the `training_module`). It is not activated by direct Python imports alone.
Fix: Consult the 'CuEquivariance Acceleration' section in the official Allegro documentation for detailed instructions on configuring your NequIP training pipeline to leverage CuEquivariance.

Install

pip install nequip-allegro Install latest version
pip install nequip-allegro[cu-equiv] Install with CuEquivariance acceleration (requires NVIDIA GPU)

Imports

AllegroModel
```
from allegro.model import AllegroModel
```

EMALightningModule

from allegro.train import EMALightningModule

Quickstart

This quickstart demonstrates how to programmatically define and instantiate a basic `AllegroModel` instance. It sets up a minimal configuration dictionary and passes it to the `AllegroModel` constructor. Note that this only builds the model architecture; training or inference requires additional steps using the NequIP framework, including data loading and graph creation.

import torch
from allegro.model import AllegroModel

# Minimal configuration for an Allegro model
# In a real-world scenario, this would often be loaded from a YAML config file.
config = {
    "r_max": 3.0,
    "l_max": 1,
    "parity_l_max": 1,
    "chemical_species": ["H", "O"], # Define the species the model will handle
    "output_keys": {"energy": "total_energy"},
    "avg_num_neighbors": 10.0,
    "num_hidden_layers": 1,
    "layer_l_maxs": [1, 1],
    "feature_irreps_hidden": "16x0e + 16x1o",
    "nonlinearity_type": "gate",
    "cutoff_type": "wigner_d",
    "radial_basis": "gaussian",
    "radial_basis_kwargs": {"max_radial": 10, "num_basis": 8},
    "num_basis_f": 8,
    "embedding_units": 16,
    "num_layers": 1
}

# Create an instance of the Allegro model
# This initializes the model structure, but it is not yet trained.
model = AllegroModel(config)

print("AllegroModel created successfully.")
print(f"Model output irreps: {model.num_irreps_out}")

# Example of forward pass (requires a nequip.data.Graph object)
# For a full example, refer to nequip and allegro documentation
# on creating datasets and training graphs.

view raw JSON →