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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) ## 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 ```bash hf download --model OneScience-Group/NequIP --local-dir ./NequIP cd NequIP ``` ### Install the Runtime Environment **DCU environment** ```bash # 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** ```bash # 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: ```bash 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: ```bash 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: ```bash 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: ```bash 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: ```bash 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 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: ```text 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: ```bash 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 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: ```bash 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: ```bash 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: ```bash python single_point.py \ --checkpoint outputs/