UniPath/README.md

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

• Setup · Download · Sampling · Results · Training · Citation · Contact

---

1. Setup

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:

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 is optional and off by default. To enable it, export your own key before running any script:

export SWANLAB_API_KEY="<your-key>"

---

2. Download Checkpoints and Data

Both are hosted on HuggingFace:

• Checkpoints: Wthinker/UniPath (model repo, ~870 MB)
• Datasets: Wthinker/UniPath (dataset repo, ~22 GB)

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

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

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:

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:

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

Each script cds 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.

mpts52_PathRefine — 3 stages

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

mpts52_uns_PathRefine — 3 stages

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

mp20_PathRefine — 2 stages

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

mp20_uns_PathRefine — 2 stages

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

--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:

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:

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