---
license: mit
tags:
  - chemistry
  - quantum-chemistry
  - dft
  - hamiltonian
  - equivariant
  - flow-matching
datasets:
  - QH9
language:
  - en
pipeline_tag: other
---

# QHFlow2 — QH9 Pre-trained Checkpoints

Pre-trained checkpoints for **QHFlow2** on the **QH9** 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 | Params | Checkpoint Size |
|------|-------------|----------------|--------|-----------------|
| small | 64 | 3 | ~35M | 139 MB |
| middle | 128 | 3 | ~125M | 494 MB |
| large | 256 | 4 | ~530M | 2.0 GB |
| extra_large | 512 | 4 | ~2.8B | 11 GB |

## Dataset Splits

| Split | Description | Epochs |
|-------|-------------|--------|
| QH9Stable-random | Random train/test split | 78 |
| QH9Stable-size_ood | Size out-of-distribution | 78 |
| QH9Dynamic-300k-geometry | Geometry OOD (300k conformers) | 33 |
| QH9Dynamic-300k-mol | Molecular OOD (300k conformers) | 33 |

## File Structure

```
{split}/
  QHFlow_so2_v5_1_{size}-{split}/
    weights-epoch=XX-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
pip install huggingface_hub

# Download a single checkpoint
huggingface-cli download ksusu/QHFlow2-QH9 \
  "QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt" \
  --local-dir ckpt/QH9

# Download all checkpoints for a specific split
huggingface-cli download ksusu/QHFlow2-QH9 \
  --include "QH9Stable-random/*" \
  --local-dir ckpt/QH9

# Download everything (~48 GB)
huggingface-cli download ksusu/QHFlow2-QH9 --local-dir ckpt/QH9
```

Or in Python:

```python
from huggingface_hub import hf_hub_download

path = hf_hub_download(
    repo_id="ksusu/QHFlow2-QH9",
    filename="QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt",
)
```

### 3. Run Prediction

```bash
cd QHFlow2

# Predict on QH9Stable test set
python -m qhflow2.experiment.train_qh9 \
  mode=predict \
  dataset=QH9Stable \
  dataset.split=random \
  model=QHFlow_so2_v5_1_small \
  ckpt=ckpt/QH9/QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt
```

### 4. Python API

```python
import torch
from qhflow2.models import get_model, get_default_model_args

# Build model
args = get_default_model_args("qh9")
args["version"] = "QHFlow_so2_v5_1"
args["hidden_size"] = 64         # small
args["num_gnn_layers"] = 3
model = get_model(args)

# Load checkpoint
ckpt = torch.load("weights-epoch=78-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()
```

### 5. Load via Lightning (full pipeline)

```python
from qhflow2.pl_module import _get_model_by_pl_type
from omegaconf import OmegaConf

conf = OmegaConf.load("configs/qh9/config_flow_v2_simple.yaml")
LitModel = _get_model_by_pl_type("flow")
lit_model = LitModel.load_from_checkpoint(
    "weights-epoch=78-val_loss=0.0000000.ckpt",
    conf=conf,
    strict=False,
)
lit_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