README.md

---
datasets:
  - OneScience-Group/oxides
frameworks:
  - ""
language:
  - en
license: mit
tags:
  - OneScience
  - eSEN
  - materials-science
  - machine-learning-potential
  - equivariant-graph-neural-network
  - structure-relaxation
  - molecular-dynamics
  - fine-tuning
tasks: []
---

<p align="center">
  <strong>
    <span style="font-size: 30px;">eSEN</span>
  </strong>
</p>

# Model Introduction

eSEN (equivariant Smooth Energy Network) is an equivariant graph neural network interatomic potential proposed by FAIR Chemistry. It predicts the energy, atomic forces, and stress of material structures. OneScience provides an ASE calculator and fine-tuning entry points for single-point calculations, structure relaxation, molecular dynamics, and fine-tuning on custom materials data.

Upstream implementation: [FAIR-Chem/fairchem](https://github.com/FAIR-Chem/fairchem)

# Model Description

eSEN takes periodic atomic structure graphs as input, learns a smooth potential energy surface through rotationally equivariant representations, and obtains conservative forces from energy gradients. Different pretrained checkpoints correspond to different materials domains. Select weights whose target elemental systems and DFT labeling settings are close to your use case.

This repository contains the eSEN model code, inference scripts, an oxide PBE fine-tuning example, and multi-device launch configurations. Restricted pretrained eSEN checkpoints are not distributed with this repository.

# Use Cases

| Use case | Description |
| :---: | :--- |
| Single-point calculation | Predict the total energy, atomic forces, and stress of a periodic structure |
| Structure relaxation | Optimize atomic positions and, optionally, the unit cell with ASE BFGS |
| Molecular dynamics | Run NVT trajectories with ASE Langevin dynamics |
| Oxide PBE fine-tuning | Fine-tune an MPTrj checkpoint using energy and force labels |
| Distributed fine-tuning | Support single-node multi-device and multi-node Slurm DDP |

This repository does not provide a complete pretraining workflow from random initialization. Its training entry point is intended for checkpoint fine-tuning.

# 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 inference and fine-tuning.
- A CPU can be used for import and configuration checks but is not recommended for production workloads.
- A DCU requires a DTK runtime matching the PyTorch build.

### Download the Model Package

```bash
hf download --model OneScience-Group/eSEN --local-dir ./eSEN
cd eSEN
```

### Install the Runtime Environment

**DCU environment**

```bash
conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[matchem-dcu] \
  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
  --trusted-host mirrors.onescience.ai
```

**GPU environment**

```bash
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
pip install onescience[matchem-gpu] \
  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
  --trusted-host mirrors.onescience.ai
```

### Pretrained Weights

Access to pretrained eSEN checkpoints must be requested from the official FAIR Chemistry model release page. This repository neither provides nor redistributes those weights. After obtaining access, place the required checkpoint at:

```text
weight/
├── Jd.pt                        # Rotation-basis file included in this repository
├── esen_30m_mptrj.pt            # Add after obtaining access
├── esen_30m_omat.pt             # Add after obtaining access
└── esen_30m_oam.pt              # Add after obtaining access
```

Checkpoint details:

| Checkpoint | Training domain | Recommended use |
| --- | --- | --- |
| `esen_30m_mptrj.pt` | MPTrj | PBE/PBE+U inference for inorganic crystals and fine-tuning in a similar labeling domain |
| `esen_30m_omat.pt` | OMat24 | A broader range of non-equilibrium inorganic structures |
| `esen_30m_oam.pt` | OAM | General-purpose materials pretraining starting point |

Access and download resources:

- FAIR Chemistry model page: https://huggingface.co/fairchem
- Official FAIR Chemistry repository: https://github.com/FAIR-Chem/fairchem

`Jd.pt` is included in the repository under `weight/`. The inference and fine-tuning scripts use it automatically, with no additional download or UMA dependency required.

### Fine-Tuning Dataset

This repository does not include the training data. The oxide PBE example dataset is published separately on Hugging Face:

```bash
hf download --dataset OneScience-Group/oxides \
  --local-dir ./datasets/oxides
```

The downloaded data has the following layout:

```text
datasets/oxides/data/OXIDES/prepared/
├── train.db
├── val.db
└── test.db
```

The data has already been converted to ASE DB format. The default configuration uses energy and force supervision and does not train on stress. Use `prepare_oxide_dataset.py` to regenerate the data from the official oxide JSON files.

### Inference

Predict single-point energy, forces, and stress:

```bash
python single_point.py --checkpoint weight/esen_30m_mptrj.pt
```

Relax a structure:

```bash
python relax.py \
  --checkpoint weight/esen_30m_mptrj.pt \
  --fmax 0.05 --steps 100 --output relaxed.cif
```

Run NVT molecular dynamics:

```bash
python md.py \
  --checkpoint weight/esen_30m_mptrj.pt \
  --steps 100 --temperature 300 --timestep 1.0 --output md.traj
```

Use `--input` to read CIF, POSCAR, XYZ, or any other structure format supported by ASE.

### Fine-Tuning

Fine-tune on oxide PBE data with one device:

```bash
bash demo/run.sh --config configs/finetune_1dcu.yaml
```

Main configuration fields:

| YAML field | Purpose |
| --- | --- |
| `checkpoint` | Path to the initialization checkpoint |
| `train`, `val` | Paths to ASE DB or ASE-LMDB data |
| `epochs`, `batch_size`, `workers` | Number of training epochs, batch size, and data-loading workers |
| `lr` | AdamW learning rate |
| `energy_weight`, `force_weight`, `stress_weight` | Loss weight for each supervised target |
| `fit_element_references` | Whether to refit elemental energy references on the training set |
| `launch.num_nodes`, `launch.num_gpus` | Number of nodes and devices per node |
| `launch.mode` | Run directly with `local` or submit to Slurm with `submit` |
| `slurm.*` | Slurm partition, time limit, and CPU resources |

Use the same entry point for multi-device and multi-node jobs by selecting another YAML file:

```bash
bash demo/run.sh --config configs/finetune_2dcu.yaml
bash demo/run.sh --config configs/finetune_16dcu.yaml
```

When the current node has fewer visible devices than requested, or when the configuration requests multiple nodes, `demo/run.sh` automatically submits the job through `sbatch`.

Evaluate a fine-tuned checkpoint:

```bash
python evaluate.py \
  --checkpoint outputs/<run>/checkpoints/esen_oxides_1dcu_finetuned.pt \
  --data datasets/oxides/data/OXIDES/prepared/test.db
```

# File Reference

| Path | Purpose |
| --- | --- |
| `model/` | Source code for the eSEN backbone, prediction heads, and graph interface |
| `single_point.py` | Single-point energy, force, and stress inference |
| `relax.py` | Periodic structure relaxation |
| `md.py` | NVT molecular dynamics |
| `finetune.py` | Low-level checkpoint fine-tuning entry point |
| `evaluate.py` | Error evaluation on an independent dataset |
| `prepare_oxide_dataset.py` | Convert oxide JSON data to ASE DB format |
| `demo/run.sh` | YAML-driven local and Slurm fine-tuning entry point |
| `demo/configs/` | Single-device, multi-device, and multi-node fine-tuning configurations |

# 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 eSEN model code is adapted from FairChem Core and follows the upstream FairChem MIT License. The model code and the included `Jd.pt` file are used under that model license.
- eSEN checkpoints are not distributed with this repository. Follow the access conditions and licenses specified on the FAIR Chemistry model pages.
- When using OMat24, MPTrj, OAM, or oxide data, cite the datasets and corresponding model work actually used.