README.md

---
frameworks:
  - ""
language:
  - en
license: mit
tags:
  - OneScience
  - MatterGen
  - materials-science
  - crystal-generation
  - diffusion-model
  - graph-neural-network
  - conditional-generation
  - training
  - fine-tuning
tasks: []
---

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

# Model Introduction

MatterGen is a generative model for inorganic materials developed by Microsoft Research. It jointly generates elemental compositions, unit cells, and periodic atomic coordinates, and supports candidate crystal generation conditioned on material properties.

Paper: *MatterGen: a generative model for inorganic materials design*<br>
Reference implementation: https://github.com/microsoft/mattergen

# Model Description

MatterGen uses a diffusion model to learn distributions over atomic species, fractional coordinates, and crystal lattices. It supports unconditional crystal generation and conditional generation based on properties such as chemical system, space group, magnetic density, band gap, and bulk modulus.

This repository contains the OneScience-adapted MatterGen model, diffusion and sampling code, and entry points for training, property fine-tuning, crystal generation, and data conversion. Model code is located in `model/`; data processing, general-purpose network layers, and property embeddings are provided by OneScience MatChem.

# Use Cases

| Use case | Description |
| :---: | :--- |
| Unconditional crystal generation | Generate new candidate crystal structures with the base checkpoint |
| Property-conditioned generation | Generate structures for target properties such as magnetic density, band gap, and bulk modulus |
| Fixed-composition structure prediction | Search for possible crystal structures for a specified composition |
| Training from scratch | Train with MP-20 or data compatible with the MatterGen cache format |
| Property fine-tuning | Add a property adapter to a pretrained model and fine-tune it |
| Environment connectivity check | Verify OneScience, the checkpoint, and GPU/DCU availability by generating one sample |

# 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.
- A CPU can be used for import and configuration connectivity checks, but full generation and training will be slow.
- A DCU requires DTK to match the PyTorch build. Single-sample generation has been verified with DTK 25.04.2 and `torch 2.5.1+das.opt1.dtk25042`.

### Download the Model Package

```bash
hf download --model OneScience-Sugon/Mattergen --local-dir ./Mattergen
cd Mattergen
```

### Install the Runtime Environment

**DCU environment**

```bash
# Activate DTK and conda first
module load compiler/dtk/25.04.2
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

MatterGen training uses a converted data cache. The MP-20 dataset is available on Hugging Face and can be downloaded directly:

```bash
hf download --dataset OneScience-Sugon/mp20 --local-dir ./datasets/mp20
```

The downloaded MatterGen cache is located at:

```text
datasets/mp20/data/MP20/cache/mp_20/
├── train/
├── val/
└── test/
```

The default training YAML cannot use the downloaded data until you add its actual path under `args`. For example, update `demo/configs/train_8dcu.yaml` as follows:

```yaml
args:
  data_module: mp_20
  data_module.root_dir: ./datasets/mp20/data/MP20/cache/mp_20
```

For property fine-tuning, also add `data_module.root_dir` under `args` in `demo/configs/finetune_dft_mag_density_smoke.yaml`. If you downloaded the data elsewhere, use the corresponding absolute path.

Convert custom CSV data first:

```bash
python csv_to_dataset.py \
  --csv-folder /path/to/csv_folder \
  --dataset-name my_dataset \
  --cache-folder ./datasets/cache
```

### Trained Weights

After downloading the model package with the Hugging Face command above, the pretrained checkpoints are located in `weight/`:

```text
weight/
├── mattergen_base/
│   ├── config.yaml
│   └── checkpoints/
│       └── last.ckpt
└── dft_mag_density/
    ├── config.yaml
    └── checkpoints/
        └── last.ckpt
```

The repository also provides property-conditioned checkpoints for `chemical_system`, `dft_band_gap`, `dft_mag_density`, `ml_bulk_modulus`, and `space_group`. When using a demo YAML, set `checkpoint` in `demo/configs/generate_base.yaml` or `demo/configs/generate_dft_mag_density.yaml` to the corresponding `./weight/<model_name>` path. For fine-tuning, update `adapter.model_path` in `demo/configs/finetune_dft_mag_density_smoke.yaml`.

### Inference

Generate one crystal structure unconditionally:

```bash
python generate.py \
  --checkpoint ./weight/mattergen_base \
  --output outputs/generate/mattergen_base \
  --batch-size 1 \
  --num-batches 1
```

The generated outputs include:

```text
outputs/generate/mattergen_base/
├── generated_crystals.extxyz
└── generated_crystals_cif.zip
```

Magnetic-density-conditioned generation requires the matching property checkpoint:

```bash
python generate.py \
  --checkpoint ./weight/dft_mag_density \
  --output outputs/generate/dft_mag_density \
  --batch-size 1 \
  --num-batches 1 \
  --properties '{"dft_mag_density": 0.15}'
```

You can also use the demo YAML files:

```bash
cd demo
bash run.sh --config configs/generate_base.yaml
bash run.sh --config configs/generate_dft_mag_density.yaml
```

### Training

Submit an eight-DCU training job with the MP-20 cache:

```bash
cd demo
bash run.sh --config configs/train_8dcu.yaml --submit
```

You can also launch training directly from the repository root with Hydra arguments:

```bash
python train.py \
  data_module=mp_20 \
  data_module.root_dir=./datasets/mp20/data/MP20/cache/mp_20 \
  trainer.devices=1 \
  data_module.batch_size.train=4
```

### Fine-Tuning

The repository provides a smoke-test configuration for magnetic-density property fine-tuning. By default, it runs only one training batch and one validation batch:

```bash
cd demo
bash run.sh --config configs/finetune_dft_mag_density_smoke.yaml
```

For a full fine-tuning run, copy the YAML and update `adapter.model_path`, `data_module.properties`, the property embedding, batch size, and number of epochs. Remove `trainer.limit_train_batches` and `trainer.limit_val_batches`.

# 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 MatterGen-related code comes from the MatChem examples in the OneScience project and refers to the upstream MatterGen project (https://github.com/microsoft/mattergen). The upstream MatterGen code is released under the [MIT License](https://github.com/microsoft/mattergen/blob/main/LICENSE).
- If you use MatterGen training or generation results in research, please cite the original MatterGen paper, the relevant OneScience projects, and the datasets used.