license: mit
tags:
- chemistry
- electron-density
- flow-matching
- equivariant-neural-networks
- qm9
datasets:
- divelab/OrbFlow-data
OrbFlow checkpoints
Trained checkpoints for OrbFlow, which predicts atomic-orbital expansion coefficients (and the resulting electron density) with equivariant flow matching.
- Code: https://github.com/TODO/OrbFlow
- Data: divelab/OrbFlow-data
- Paper: TODO
Contents
Each folder is one training run with its best-validation checkpoint, config.yaml and metadata.json.
Optimizer state has been removed, so these files are for evaluation and inference; they cannot resume training.
| Folder | Model | Seeds | Checkpoint size |
|---|---|---|---|
qm9_qhf_eqv3_gt_2ph_final, …_s42, …_s123 |
QM9: EquiformerV3, 8 layers, ℓmax = 6; 75k endpoint → 450k integrated-density training | 0, 42, 123 | 1.2 GB |
md_<mol>_k3l3_slim_10k100k, …_seed42, …_seed505 |
MD: slim EquiformerV3, 3 layers, ℓmax = 3; 10k → 110k | 0, 42, 505 | 155 MB |
MD molecules: benzene, ethane, ethanol, malonaldehyde, phenol, resorcinol. The seed-505 phenol and resorcinol runs stopped at ~60k of their 110k training steps (job time limit); their best-validation checkpoints are from step ~60k.
Usage
Download into the root of the OrbFlow code repository, so the folder names match what the evaluation scripts expect
(the excludes keep this model card from overwriting the code repository's README.md):
git clone https://github.com/TODO/OrbFlow.git && cd OrbFlow
# set up the environment and .env as described in the code README, then:
hf download divelab/OrbFlow --local-dir . --exclude README.md --exclude .gitattributes
# QM9 test set (Euler K = 2 and 3; full test split and last 1600 molecules)
CKPT_PATH=qm9_qhf_eqv3_gt_2ph_final sbatch vision/test_qhf_eqv3_gt_2ph_final_full.slurm
# MD test sets (6 molecules); for other seeds use EXP_SUFFIX=md_%s_k3l3_slim_10k100k_seed42
sbatch vision/test_flow_md_paper_2phase_beta13_all6.slurm
To download a single run: hf download divelab/OrbFlow --include "qm9_qhf_eqv3_gt_2ph_final/*" --local-dir .
Results
Test nMAPE (%) with Euler K = 2, mean ± std over training seeds, computed with the evaluation scripts above.
QM9 (10,000 test molecules, seeds 0 / 42 / 123)
| Split | nMAPE (%) |
|---|---|
| Full test set | 0.1518 ± 0.0008 |
| Last 1,600 molecules | 0.1646 ± 0.0008 |
MD (seeds 0 / 42 / 505)
| Molecule | nMAPE (%) |
|---|---|
| benzene | 0.4210 ± 0.0017 |
| ethane | 0.8706 ± 0.0084 |
| ethanol | 1.0784 ± 0.0229 |
| malonaldehyde | 1.1771 ± 0.0127 |
| phenol | 0.5753 ± 0.0083 |
| resorcinol | 0.6626 ± 0.0136 |
| average (6 molecules) | 0.7975 ± 0.0067 |
Citation
TODO
OrbFlow builds on SCDP (Fu et al., NeurIPS 2024).