NequIP

0.17.1 · active · verified Fri Apr 17

NequIP is an open-source Python library for building E(3)-equivariant interatomic potentials, primarily used in molecular dynamics and materials science simulations. It leverages PyTorch and e3nn for deep learning models that respect physical symmetries, enabling accurate and robust simulations. The current stable version is 0.17.1, with releases typically tied to significant feature additions or dependency updates.

Common errors

```
AttributeError: module 'e3nn' has no attribute 'o3' (or similar e3nn related errors like 'TensorProduct' object has no attribute 'irrep_out')
```
cause This typically indicates an incompatibility between your installed `e3nn` and `nequip` versions. Older `e3nn` versions (e.g., v0.4.x) lack symbols or have different API signatures required by newer NequIP.

fix Upgrade `e3nn` to the version required by your NequIP installation. For NequIP v0.16.0+, use `pip install e3nn>=0.5.0`.
```
KeyError: 'some_config_key' or 'yaml.scanner.ScannerError: while scanning a block scalar'
```
cause Your configuration file (YAML) is either malformed or uses a schema/keys that are no longer valid for your current NequIP version. This is common when upgrading NequIP and reusing old config files.

fix Review the NequIP documentation for your installed version and compare your YAML config file to the latest examples. Update keys, structure, or re-create the config based on the current schema.
```
RuntimeError: Error(s) in loading state_dict for NequIP: Missing key(s) in state_dict: "_c_rescalers.0.x_in"...
```
cause This error occurs when trying to load a pre-trained model checkpoint that was saved with an older NequIP version (pre-0.17.0) into a newer NequIP installation. The model's internal parameter structure has changed.

fix The model checkpoint is fundamentally incompatible. You must re-train the model with NequIP v0.17.0+ or revert your NequIP installation to the version the model was trained with.

Warnings

breaking The internal structure of the `NequIP` model's parameters and buffers significantly changed in v0.17.0. Saved models trained with NequIP versions <0.17.0 are incompatible and will fail to load or produce incorrect results. They require re-training or conversion.
Fix: Re-train your models using NequIP v0.17.0+, or refer to the NequIP documentation for any experimental conversion scripts if available (not officially supported for all versions).
gotcha NequIP has a strict dependency on `e3nn` versions. Specifically, NequIP >=0.16.0 requires `e3nn>=0.5.0`. Mismatched `e3nn` versions can lead to `AttributeError` or unexpected behavior during model definition or computation.
Fix: Ensure your `e3nn` installation is compatible with your `nequip` version. For NequIP v0.16.0+, run `pip install e3nn>=0.5.0` to guarantee the correct version.
gotcha Configuration file formats (YAML) have changed across major NequIP versions. Using an older configuration file with a newer NequIP version can lead to `KeyError`, schema validation errors, or incorrect model behavior.
Fix: Always refer to the documentation for your specific NequIP version for the correct configuration file schema. Adapt older YAML files to the new format, or regenerate them using updated examples.

Install

pip install nequip Install NequIP

Imports

NequIP
```
from nequip.model import NequIP
```
Config
```
from nequip.utils import Config
```
train_main
```
from nequip.scripts.train import main as train_main
```
Commonly used for programmatically invoking the training script.

Quickstart

Initializes a basic NequIP model from a configuration dictionary. In practice, configurations are often loaded from YAML files and models are either trained or loaded from pre-existing checkpoints. This example demonstrates the fundamental model instantiation.

import torch
from nequip.utils import Config
from nequip.model import NequIP

# Example minimal configuration for a NequIP model
# In real applications, this would typically be loaded from a YAML file.
config_dict = {
    'r_max': 5.0,
    'max_ell': 1,
    'num_layers': 3,
    'chemical_species_set': ['H', 'O'], # Example species
    'num_species': 2,
    'l_max_hidden': 1,
    'nonlinearity_type': 'silu',
    'resnet': True,
    'env_embed_multi_head': True,
    'mlp_output_irreps': '64x0e+64x1o+64x1e',
    'weight_init': 'xavier_uniform',
    'irreps_out_per_structure': '1x0e',
    'cutoff_type': 'polynomial', 
    'activation': 'silu'
}

# Create a Config object
config = Config(config_dict)

# Initialize the NequIP model
# This model is a placeholder; it would need to be trained or loaded from a checkpoint.
model = NequIP(config)

print(f"NequIP model initialized with {sum(p.numel() for p in model.parameters())} parameters.")
# Example of model structure for a single atom
# x = torch.zeros(1, 3)
# z = torch.tensor([1]) # Atomic number for H
# model(x, z) # Requires a full batch of positions and atomic numbers

view raw JSON →