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