File size: 2,984 Bytes
a552579 4a293ff a552579 0896165 a552579 0896165 a552579 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | ---
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/divelab/AIRS/tree/main/OpenDFT/OrbFlow
- Data: [divelab/OrbFlow-data](https://huggingface.co/datasets/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`):
```bash
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](https://github.com/kyonofx/scdp) (Fu et al., NeurIPS 2024).
|