OneScience-Group/NequIP

NequIP

Model Introduction

NequIP is a machine-learning interatomic potential (MLIP) for molecular and materials systems. Built on an E(3)-equivariant graph neural network, it predicts the energies and forces of atomic structures.

Reference implementation: https://github.com/mir-group/nequip

Model Description

This repository provides the OneScience-integrated NequIP model code, OAM-L model weights, and runnable examples for training, fine-tuning, and inference. The model/ directory corresponds only to src/onescience/models/nequip/ in the main OneScience repository; training utilities, data-processing tools, and other shared modules are provided by the installed OneScience package.

The included OAM-L weights are:

File	Purpose
weight/NequIP-OAM-L-0.1.nequip.pth	Compiled model for ASE single-point energy, atomic force, and stress inference
weight/NequIP-OAM-L-0.1.nequip.zip	NequIP package for OAM-L fine-tuning and checkpoint inference

Use Cases

Use case	Description
Interatomic-potential training	Train a NequIP model using the example configurations and ASE extxyz data
Pretrained-model fine-tuning	Fine-tune the OAM-L package using data labeled with energy and forces
Single-point energy and force inference	Predict the energy, atomic forces, and stress of a structure with a compiled model or fine-tuned checkpoint
Structure relaxation	Optimize atomic positions with ASE
Energy-volume curve	Scan the volume of a periodic crystal and calculate the corresponding energy
Slurm/DCU training	Submit single-device or multi-device jobs using the included configurations and launch scripts

Usage

1. Using OneCode

Try intelligent, one-click AI4S programming in the OneCode online environment:

Try intelligent, one-click AI4S programming

2. Manual Installation and Usage

Hardware requirements

• A GPU or DCU is recommended for training.
• A CPU can be used for import checks and small-configuration connectivity tests, but full training will be slow.
• DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.

Download the Model Package

hf download --model OneScience-Group/NequIP --local-dir ./NequIP
cd NequIP

Install the Runtime Environment

DCU environment

# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is also supported
pip install onescience[matchem-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

GPU environment

# Activate Conda first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# uv installation is also supported
pip install onescience[matchem-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Training Data

Training data is not bundled with this repository. Using the introductory FCC Cu dataset as an example, download it from Hugging Face to data/ in the repository root:

hf download --dataset OneScience-Group/FCC_Cu --local-dir ./data

After downloading, the raw data is located at data/data/FCC_Cu/raw/fcu.xyz. The dataset contains 6,855 structures, each with 52 atoms of C, H, O, and Cu. It uses an ASE-readable extxyz format and includes periodic cells together with energy and force labels. For production training or fine-tuning, use data consistent with the target system, label definitions, and units.

The training scripts read models and data from shared directories. Set the paths for your cluster before training or fine-tuning:

export ONESCIENCE_MODELS_DIR=/path/to/onescience-models
export ONESCIENCE_DATASETS_DIR=/path/to/onescience-datasets

To use the OAM-L weights included in this repository, copy them into the shared model directory:

mkdir -p "$ONESCIENCE_MODELS_DIR/NequIP"
cp weight/NequIP-OAM-L-0.1.nequip.pth "$ONESCIENCE_MODELS_DIR/NequIP/"
cp weight/NequIP-OAM-L-0.1.nequip.zip "$ONESCIENCE_MODELS_DIR/NequIP/"

Training

Generate minimal smoke-test data and run local training:

python demo/prepare_smoke_data.py
bash demo/run.sh --config configs/tutorial_smoke.yaml

Download the official FCU tutorial data and submit a training job:

python demo/download_tutorial_data.py
bash demo/run.sh --config configs/tutorial_fcu.yaml --submit

The eight-DCU configurations run locally or submit to Slurm automatically, depending on the currently available resources:

bash demo/run.sh --config configs/tutorial_smoke_8dcu.yaml
bash demo/run.sh --config configs/tutorial_fcu_8dcu.yaml

Outputs are written to outputs/ by default. The actual wait time for a training job depends on the cluster queue and available resources.

Model Weights

This repository includes the OAM-L trained weights:

e83a1d656f8b19b55d2f05708c83e054612f713e9a1b06266aa010db58e56517  weight/NequIP-OAM-L-0.1.nequip.pth
5d01a4fab228abb3cdb6ace0033f93993729956bca6a42234a2a8816825b9a0f  weight/NequIP-OAM-L-0.1.nequip.zip

Fine-Tuning

Validate the OAM-L fine-tuning workflow with generated smoke-test data:

python demo/prepare_smoke_data.py
bash demo/run.sh --config configs/oam_l_finetune_smoke.yaml --submit

Use the production fine-tuning configuration:

bash demo/run.sh --config configs/oam_l_finetune.yaml --submit

Provide production fine-tuning data through ONESCIENCE_DATASETS_DIR in an ASE-readable extxyz format. Every frame must contain at least energy and forces; element types, units, and label definitions must be consistent with the OAM-L package and configuration.

Inference

Use the compiled model for single-point energy, atomic force, and stress prediction:

python single_point.py --compiled-model weight/NequIP-OAM-L-0.1.nequip.pth
python single_point.py \
  --compiled-model weight/NequIP-OAM-L-0.1.nequip.pth \
  --input structure.cif \
  --output outputs/single_point.json

Calculate an energy-volume curve and perform structure relaxation:

python energy_volume.py
python structure_relaxation.py --fmax 0.05 --steps 100 --output-dir outputs/oam_l_relax

Run inference with a checkpoint produced by fine-tuning:

python single_point.py \
  --checkpoint outputs/<run>/checkpoints/best.ckpt \
  --package weight/NequIP-OAM-L-0.1.nequip.zip \
  --output outputs/<run>/single_point.json

Official OneScience Resources

Platform	OneScience Main Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

• The NequIP-related code comes from the OneScience MatChem integration and refers to the upstream NequIP project (https://github.com/mir-group/nequip). The OneScience integration code follows the Apache License 2.0 used by the main repository.
• If you use NequIP or OAM-L training results in research, please cite the original NequIP paper, the relevant OneScience projects, and the datasets used.
• Redistribution rights for the OAM-L model weights are governed by the original OneScience/OAM-L release terms. Confirm the applicable rights and restrictions before use.
• The FCC Cu dataset is published separately at OneScience-Group/FCC_Cu. Its license and provenance are documented on the dataset card and by its upstream source.