README.md

---
license: apache-2.0
tasks:
  - materials-simulation
  - molecular-dynamics
  - energy-prediction
  - force-prediction
frameworks:
  - pytorch
language:
  - en
tags:
  - OneScience
  - NequIP
  - machine-learning-potential
  - molecular-simulation
  - materials-computing
  - graph-neural-network
  - equivariant-neural-network
  - training
  - fine-tuning
  - inference
datasets:
  - OneScience-Group/FCC_Cu
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">NequIP</span>
  </strong>
</p>

# 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/<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](https://www.apache.org/licenses/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](https://huggingface.co/datasets/OneScience-Group/FCC_Cu). Its license and provenance are documented on the dataset card and by its upstream source.