UniPath/README.md

<div align="center">

# One Path to Model Them All: Learnable-Time Flow Matching for Crystal Structure and Energy Prediction

[![Paper](https://img.shields.io/badge/Paper-ACM%20DL-B31B1B.svg)](https://dl.acm.org/doi/pdf/10.1145/3770855.3818902)
[![Checkpoints](https://img.shields.io/badge/%F0%9F%A4%97%20Checkpoints-Wthinker%2FUniPath-FFD21E.svg)](https://huggingface.co/Wthinker/UniPath)
[![Datasets](https://img.shields.io/badge/%F0%9F%A4%97%20Datasets-Wthinker%2FUniPath-FFD21E.svg)](https://huggingface.co/datasets/Wthinker/UniPath)

<a href="https://huggingface.co/BAAI/AIDD/resolve/main/UniPath/assets/UniPath_Poster.png"><img src="https://huggingface.co/BAAI/AIDD/resolve/main/UniPath/assets/UniPath_Poster.png" alt="UniPath KDD 2026 poster" width="900"></a>

</div>

- [Setup](#1-setup) · [Download](#2-download-checkpoints-and-data) · [Sampling](#3-run-sampling) · [Results](#4-results) · [Training](#5-training-from-scratch) · [Citation](#citation) · [Contact](#contact)

---

## 1. Setup

```bash
git clone <repo_url>
cd crystal-uniug

conda env create -f environment.yml
conda activate crystal_uniug

pip install -e remote/cdvae
pip install -e remote/DiffCSP-official
pip install -e remote/riemannian-fm
pip install -e .
```

Then set four environment variables. The easiest way is to copy them into a `.env` file at the project root and `source .env`:

```bash
export PROJECT_ROOT="/path/to/crystal-uniug"
export HYDRA_JOBS="/path/to/crystal-uniug"
export WANDB_DIR="/path/to/crystal-uniug"
export DATA_DIR="/path/to/data_download/dataset"   # set in step 2
```

`PROJECT_ROOT` is required: the training and sampling scripts `cd` to it and resolve
checkpoint paths against it.

Experiment logging via [SwanLab](https://swanlab.cn) is optional and off by default. To enable it,
export your own key before running any script:

```bash
export SWANLAB_API_KEY="<your-key>"
```

---

## 2. Download Checkpoints and Data

Both are hosted on HuggingFace:

- Checkpoints: [`Wthinker/UniPath`](https://huggingface.co/Wthinker/UniPath) (model repo, ~870 MB)
- Datasets: [`Wthinker/UniPath`](https://huggingface.co/datasets/Wthinker/UniPath) (dataset repo, ~22 GB)

If `huggingface.co` is unreachable, set `export HF_ENDPOINT=https://hf-mirror.com` first.

```bash
# Checkpoints -> ckpt_download/ckpt/
hf download Wthinker/UniPath --repo-type model --local-dir ./ckpt_download

# Datasets -> data_download/dataset/
hf download Wthinker/UniPath --repo-type dataset --local-dir ./data_download

export DATA_DIR="$PWD/data_download/dataset"
```

### Checkpoint paths

| Name | Path under `ckpt_download/ckpt/` |
|---|---|
| `mpts52_PathRefine` | `mpts52_PathRefine/every_n_epochs/epoch=149-step=80976.ckpt` |
| `mpts52_uns_PathRefine` | `mpts52_uns_PathRefine/every_n_epochs/epoch=149-step=36225.ckpt` |
| `mp20_PathRefine` | `mp20_PathRefine/every_n_epochs/epoch=499-step=190587.ckpt` |
| `mp20_uns_PathRefine` | `mp20_uns_PathRefine/every_n_epochs/epoch=299-step=87163.ckpt` |

> ⚠️ **Do not move a `.ckpt` out of its directory.** Model loading resolves the architecture by
> walking up the directory tree to find the sibling `.hydra/config.yaml`. Moving the file alone
> will fail with `Failed to find hydra config!`.

### Expected data layout

Each experiment needs one **CSP** dataset plus one **MD** dataset:

```
$DATA_DIR/
├── mpts_52/                 # CSP — both mpts52 checkpoints
├── mp_20/                   # CSP — both mp20 checkpoints
├── mptsubmpts52/            # MD  — mpts52_PathRefine
├── mptsubmpts52_uns_V0/     # MD  — mpts52_uns_PathRefine
├── mptsubmp20_v0/           # MD  — mp20_PathRefine
└── mptsubmp20_uns_V0/       # MD  — mp20_uns_PathRefine
```

CSP dirs contain `{train,val,test}.csv` + `{train,val,test}_ori.pt`; MD dirs contain `{train,val}.csv` + `{train,val}_ori.pt`.

> Sampling only predicts on the test split, but the data module still loads the **train** split to
> compute normalization scalers, and the MD splits because these are joint CSP+MD models. All of the
> above must be present.

---

## 3. Run Sampling

Sampling runs four steps in sequence. Set `CKPT` to any checkpoint from the table above:

```bash
CKPT=./ckpt_download/ckpt/mp20_PathRefine/every_n_epochs/epoch=499-step=190587.ckpt
SUBDIR=reproduce_1

python ./scripts_model/evaluate.py reconstruct ${CKPT} \
    --subdir ${SUBDIR} --inference_anneal_slope 10.0 \
    --stage test --num_evals 1 --batch_size 1024 --sample_way flow --single_gpu && \
python ./scripts_model/evaluate.py consolidate ${CKPT} --subdir ${SUBDIR} && \
python ./scripts_model/evaluate.py old_eval_metrics ${CKPT} --subdir ${SUBDIR} --stage test --do_not_log_wandb && \
python ./scripts_model/evaluate.py lattice_metrics ${CKPT} --subdir ${SUBDIR} --stage test
```

To do all four checkpoints, loop over them:

```bash
for name in mpts52_PathRefine mpts52_uns_PathRefine mp20_PathRefine mp20_uns_PathRefine; do
    CKPT=$(ls ./ckpt_download/ckpt/${name}/every_n_epochs/*.ckpt)
    echo "=== ${name} ==="
    python ./scripts_model/evaluate.py reconstruct   "${CKPT}" --subdir reproduce_1 \
        --inference_anneal_slope 10.0 --stage test --num_evals 1 \
        --batch_size 1024 --sample_way flow --single_gpu       && \
    python ./scripts_model/evaluate.py consolidate      "${CKPT}" --subdir reproduce_1 && \
    python ./scripts_model/evaluate.py old_eval_metrics "${CKPT}" --subdir reproduce_1 --stage test --do_not_log_wandb && \
    python ./scripts_model/evaluate.py lattice_metrics  "${CKPT}" --subdir reproduce_1 --stage test
done
```

A ready-made template is also available at `UniPath/script/sample.sh` — edit `ckpt` and `target_dir_name`, then `bash UniPath/script/sample.sh --run_id 1`.

### Key flags

| Flag | Meaning |
|---|---|
| `--inference_anneal_slope` | annealing slope for flow inference (`10.0` reproduces the paper) |
| `--batch_size` | sampling batch size (`1024`) |
| `--sample_way flow` | flow-based sampler |
| `--num_evals` | number of sampling repeats per structure (`1`) |
| `--single_gpu` | run on one GPU; **omit** for multi-GPU DDP |

> **Runtime:** `reconstruct` dominates — roughly 1.5 h (mp_20) to 3.5 h (mpts_52) per checkpoint on a
> single A100. The other three steps take a few minutes each.

---

## 4. Results

Each run writes to `<ckpt_dir>/<SUBDIR>/`:

| File | Contents |
|---|---|
| `old_eval_metrics_reconstruct_multi.json` | `match_rate`, `rms_dist` (multi-eval) |
| `old_eval_metrics_reconstruct_single.json` | `match_rate`, `rms_dist` (single-eval) |
| `consolidated_reconstruct.pt` | consolidated raw predictions |
| `lengths.png` / `angles.png` | lattice parameter distributions |

Compare the reported `match_rate` and `rms_dist` against the values in the paper. With
`--num_evals 1` the multi- and single-eval JSONs are identical.

---

## 5. Training from Scratch

Only needed if you want to retrain rather than use the released checkpoints. Budget several GPU-days
per checkpoint.

### Script layout

```
UniPath/script/uni_train/main/
├── mpts52/       -> mpts52_PathRefine       (3 stages)
├── mpts52_uns/   -> mpts52_uns_PathRefine   (3 stages)
├── mp20/         -> mp20_PathRefine         (2 stages)
└── mp20_uns/     -> mp20_uns_PathRefine     (2 stages)
```

Outputs are written to `runs/trash/main/<experiment_name>/`.

### Launching

Scripts launch `UniPath/run_uni.py` through `torchrun` and read the standard distributed variables
(`WORLD_SIZE` = nodes, `NPROC_PER_NODE` = GPUs per node, `RANK`, `MASTER_ADDR`, `MASTER_PORT`).
Each script declares a **total** batch size and divides it by the world size, so the per-device batch
size adapts to your GPU count — but the total must be divisible by it.

Single-node, 4 GPUs:

```bash
export WORLD_SIZE=1 NPROC_PER_NODE=4 RANK=0
export MASTER_ADDR=127.0.0.1 MASTER_PORT=29500
```

Each script `cd`s to the project root and sources `UniPath/running_setup.sh`, so it can be invoked from anywhere.

### Training order

Each checkpoint is produced by a **multi-stage pipeline**: every stage after the first warm-starts
from the previous stage's checkpoint (`partial_ckpt_load=True`, a partial weight load rather than a
Lightning full-state resume). Stages must therefore run **in order**.

<details>
<summary><b>mpts52_PathRefine</b> — 3 stages</summary>

```bash
T=UniPath/script/uni_train/main/mpts52

# 1. Joint CSP+MD pretraining, from scratch (150 epochs)
bash $T/uni_E_tclearn_freeze.sh

# 2. Unsupervised-flow path refinement, warm-started from stage 1 @ epoch 109
bash $T/mpts52_mptsub_path_refine_wunstune.sh --w_uns 0.1

# 3. + consistency loss, warm-started from stage 2 @ epoch 129   -> RELEASED
bash $T/mpts52_mptsub_path_refine_consist_wunstune.sh --w_uns 0.1 --w_consist 0.05
```
</details>

<details>
<summary><b>mpts52_uns_PathRefine</b> — 3 stages</summary>

```bash
T=UniPath/script/uni_train/main/mpts52_uns

bash $T/uni_E_tclearn_freeze.sh
bash $T/mpts52_uns_path_refine_wunstune.sh --w_uns 0.02
bash $T/mpts52_uns_path_refine_consist_wunstune.sh --w_uns 0.02 --w_consist 0.02   # -> RELEASED
```
</details>

<details>
<summary><b>mp20_PathRefine</b> — 2 stages</summary>

```bash
T=UniPath/script/uni_train/main/mp20

bash $T/bz256_tclearn_freeze_epoch500.sh
bash $T/mp20_wtune.sh --w_uns 0.001 --w_consist 0.01   # -> RELEASED
```
</details>

<details>
<summary><b>mp20_uns_PathRefine</b> — 2 stages</summary>

```bash
T=UniPath/script/uni_train/main/mp20_uns

bash $T/uni_E_tclearn_freeze.sh
bash $T/mp20_uns_wunte.sh --w_uns 0.02 --w_consist 0.5   # -> RELEASED
```
</details>

`--w_uns` sets the unsupervised-flow loss weight, `--w_consist` the consistency loss weight. The
values above are the ones that produced the released checkpoints.

### Before you train

- **`partial_ckpt_path` is built from `${PROJECT_ROOT}`** in each warm-starting script, so exporting
  `PROJECT_ROOT` is enough — no need to edit the paths by hand.
- **Parent checkpoints are referenced by exact filename** (e.g. `epoch=109-step=424160.ckpt`). The
  step count depends on your world size and batch size, so a retrained parent produces a *different*
  filename — update the referencing script accordingly.
- **`runs/` is gitignored**, so training outputs are never committed.

### Sampling your own checkpoint

Use the Section 3 commands with `CKPT` pointing into your run directory, keeping the checkpoint beside its `.hydra/config.yaml`:

```bash
CKPT=runs/trash/main/mpts52_PathRefine_wuns0.1_wc0.05_e109_e129/every_n_epochs/epoch=149-step=80976.ckpt
```

---

## Citation

If you find this work useful, please cite:

```bibtex
@inproceedings{li2026one,
  title={One Path to Model Them All: Learnable-Time Flow Matching for Crystal Structure and Energy Prediction},
  author={Li, Songyou and Li, Mingze and Liu, Qianpu and Cen, Jiacheng and Wu, Liming and Xu, Tingyang and Rong, Yu and Zhao, Deli and Huang, Wenbing},
  booktitle={Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. 2},
  pages={11351--11362},
  year={2026}
}
```

---

## Contact

Questions and issues are welcome — open a GitHub issue or reach out to
[songyou_li@126.com](mailto:songyou_li@126.com).