File size: 3,139 Bytes
7ecf721 942af0b 7ecf721 942af0b 7ecf721 942af0b 7ecf721 942af0b 7ecf721 942af0b 7ecf721 942af0b 7ecf721 942af0b 7ecf721 942af0b 7ecf721 942af0b 7ecf721 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | ---
license: mit
tags:
- chemistry
- quantum-chemistry
- dft
- hamiltonian
- equivariant
- flow-matching
datasets:
- MD17
language:
- en
pipeline_tag: other
---
# QHFlow2 — MD17 Pre-trained Checkpoints
Pre-trained checkpoints for **QHFlow2** on the **MD17** dataset (DFT Hamiltonian prediction).
> **Paper:** [High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction](https://arxiv.org/abs/2602.16897)
> **Authors:** Seongsu Kim, Nayoung Kim, Dongwoo Kim, Sungsoo Ahn (KAIST)
> **Venue:** NeurIPS 2025
> **Code:** [github.com/seongsukim-ml/QHFlow2](https://github.com/seongsukim-ml/QHFlow2)
## Model Variants
| Size | hidden_size | num_gnn_layers | Checkpoint Size |
|------|-------------|----------------|-----------------|
| small_v2 | 64 | 3 | 313 MB |
| middle | 128 | 3 | 975 MB |
## Molecules
| Molecule | Atoms | Formula |
|----------|-------|---------|
| ethanol | 9 | C₂H₆O |
| malondialdehyde | 9 | C₃H₄O₂ |
| uracil | 12 | C₄H₄N₂O₂ |
## File Structure
```
{molecule}/
QHFlow_so2_v5_1_{size}_b10-{molecule}/
weights-epoch=79-val_loss=0.0000000.ckpt
```
## Quick Start
### 1. Install QHFlow2
```bash
git clone https://github.com/seongsukim-ml/QHFlow2.git
cd QHFlow2
pip install -e ".[fairchem]"
```
### 2. Download Checkpoints
```bash
# Download a single checkpoint
huggingface-cli download ksusu/QHFlow2-MD17 \
"ethanol/QHFlow_so2_v5_1_small_v2_b10-ethanol/weights-epoch=79-val_loss=0.0000000.ckpt" \
--local-dir ckpt/md17
# Download all checkpoints (~3.9 GB)
huggingface-cli download ksusu/QHFlow2-MD17 --local-dir ckpt/md17
```
Or in Python:
```python
from huggingface_hub import hf_hub_download
path = hf_hub_download(
repo_id="ksusu/QHFlow2-MD17",
filename="ethanol/QHFlow_so2_v5_1_small_v2_b10-ethanol/weights-epoch=79-val_loss=0.0000000.ckpt",
)
```
### 3. Download Dataset
Download from [Google Drive](https://drive.google.com/drive/folders/1d3HTu0H7gdg54kirWBqN24x-s1QW6OKV?usp=sharing) and place under `dataset/`:
```
QHFlow2/dataset/
ethanol/
malondialdehyde/
uracil/
```
### 4. Run Prediction
```bash
python -m qhflow2.experiment.train_md17 \
mode=predict \
dataset=ethanol \
ckpt=ckpt/md17/ethanol/QHFlow_so2_v5_1_small_v2_b10-ethanol/weights-epoch=79-val_loss=0.0000000.ckpt
```
### 5. Python API
```python
import torch
from qhflow2.models import get_model, get_default_model_args
args = get_default_model_args("md17")
args["version"] = "QHFlow_so2_v5_1"
args["hidden_size"] = 64
args["num_gnn_layers"] = 3
model = get_model(args)
ckpt = torch.load("weights-epoch=79-val_loss=0.0000000.ckpt", map_location="cpu")
state_dict = {k.replace("model.", ""): v for k, v in ckpt["state_dict"].items()}
model.load_state_dict(state_dict, strict=False)
model.eval()
```
## Citation
```bibtex
@inproceedings{kim2025high,
title={High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction},
author={Kim, Seongsu and Kim, Nayoung and Kim, Dongwoo and Ahn, Sungsoo},
booktitle={Advances in Neural Information Processing Systems},
year={2025}
}
```
## License
MIT
|