Publish AIDD open-source resources

4947683 verified about 1 month ago

11.2 kB

<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).

	<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).