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