OneScience-Group/eSEN

eSEN

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

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

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

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

Install the Runtime Environment

DCU environment

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

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:

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:

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

The downloaded data has the following layout:

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:

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

Relax a structure:

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

Run NVT molecular dynamics:

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 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 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:

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.