README.md

---
license: mit
tasks:
  - materials-simulation
  - catalysis
  - molecular-dynamics
  - energy-prediction
  - force-prediction
frameworks:
  - pytorch
language:
  - en
tags:
  - OneScience
  - UMA
  - materials-science
  - catalysis
  - molecular-simulation
  - machine-learning-potential
  - graph-neural-network
  - equivariant-neural-network
  - training
  - inference
datasets:
  - OneScience-Sugon/oc20
---

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

# Model Introduction

UMA (Universal Materials Interaction Model) is a general-purpose machine-learning interatomic potential for materials and catalytic systems. Built on an equivariant graph neural network, it predicts the energies and forces of atomic structures.

# Model Description

UMA uses an equivariant graph neural network architecture trained on multiple materials and catalysis datasets, including OC20, OC22, OC25, OMat, OMOL, ODAC, and OMC. It supports energy and force prediction and structure optimization for catalytic adsorption, inorganic materials, molecular systems, and MOFs.

# Use Cases

| Use case | Description |
| :---: | :--- |
| OC20 energy/force fine-tuning | Read OC20 fine-tuning data with the standard configuration and train a UMA model |
| OC22 oxide-catalysis fine-tuning | Fine-tune on OC22 with the standard configuration (1P2 only) |
| OC25 (electro)catalysis fine-tuning | Fine-tune on OC25 with the standard configuration (1P2 only) |
| OMat inorganic-material fine-tuning | Fine-tune on OMat with the standard configuration |
| OMOL molecule/polymer fine-tuning | Fine-tune on OMOL with the standard configuration |
| ODAC MOF fine-tuning | Fine-tune on ODAC with the standard configuration |
| OMC molecular-crystal fine-tuning | Fine-tune on OMC with the standard configuration |
| Training preflight | Check configurations, data paths, run scripts, and checkpoint locations |
| Catalytic adsorption modeling | Use OC20/OC22/OC25 tasks as references for training and inference on adsorption and catalytic-surface systems |
| Inference-script reference | Adapt upstream examples for crystal relaxation, adsorption-system relaxation, or molecular MD |
| Custom data migration | Convert ASE-readable structures into UMA fine-tuning data, then replace the training and validation paths |

# 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 fine-tuning.
- A CPU can be used to check configurations and data paths, but is not recommended for production UMA training.
- 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-Sugon/UMA --local-dir ./UMA
cd UMA
```

### 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 Datasets

Training data is not bundled with this repository. The following example uses **OC20 fine-tuning**. OC22, OC25, OMat, OMOL, ODAC, and OMC follow the same workflow; only `--uma-task` and the data paths need to be changed.

```bash
hf download --dataset OneScience-Sugon/oc20 --local-dir ./data
```

**Data conversion**

The downloaded raw data usually consists of `.extxyz` files. Convert them to ASE-LMDB format with `scripts/create_uma_finetune_dataset.py`, which also computes `elem_refs` and `normalizer_rmsd`. The script supports these tasks:

| Task | Description |
| --- | --- |
| `oc20` | Catalysis (example) |
| `oc22` | Oxide catalysis (1P2 only) |
| `oc25` | (Electro)catalysis (1P2 only) |
| `omat` | Inorganic materials |
| `omol` | Molecules and polymers |
| `odac` | MOFs |
| `omc` | Molecular crystals |

```bash
python scripts/create_uma_finetune_dataset.py \
    --train-dir data/oc20/s2ef_200k_uncompressed \
    --val-dir data/oc20/s2ef_val_id_uncompressed \
    --uma-task oc20 \
    --regression-tasks ef \
    --output-dir data/oc20_finetune \
    --num-workers 8
```

The conversion produces:

```text
data/oc20_finetune/
├── train/                 # ASE-LMDB training data
├── val/                   # ASE-LMDB validation data
└── data/                  # Generated data configuration YAML
    └── uma_conserving_data_task_energy_force.yaml
```

Next, use `scripts/update_demo_config.py` to copy the generated `elem_refs`, `normalizer_rmsd`, and data paths into the demo configuration:

```bash
python scripts/update_demo_config.py --demo-config demo/configs/oc20_ef_4dcu.yaml
```

`demo/run.sh` automatically sets the repository root as `ONESCIENCE_DATASETS_DIR`, so relative paths in the configuration resolve automatically.

### Trained Weights

This repository includes the rotation-basis file `weight/Jd.pt`. Download a pretrained UMA checkpoint, such as `uma-s-1p1_converted.pt`, from the official fairchem repository, convert it to UMA format, and place it at the following location (upload pending):

```text
weight/uma-s-1p1_converted.pt
```

- Official fairchem repository: https://github.com/facebookresearch/fairchem

`demo/run.sh` and the example scripts under `inference/` automatically detect `weight/Jd.pt` and set `ONESCIENCE_UMA_JD_PATH`.

### Fine-Tuning

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

### Inference

```bash
python inference/run_molecular_md.py
```

Specify the path to the pretrained weights before running inference.

## 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 UMA example code comes from the MatChem examples in the OneScience project and refers to the upstream fairchem project (https://github.com/facebookresearch/fairchem). The upstream fairchem software is released under the [MIT License](https://fair-chem.github.io/core/install.html#license). Individual fairchem model checkpoints and datasets may have separate licenses; follow the applicable terms.
- If you use UMA fine-tuning or inference results in research, please cite UMA or the relevant universal materials interaction model method, the fairchem and OneScience projects, and the datasets used.