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](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/