metadata
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 Authors: Seongsu Kim, Nayoung Kim, Dongwoo Kim, Sungsoo Ahn (KAIST) Venue: NeurIPS 2025 Code: 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
git clone https://github.com/seongsukim-ml/QHFlow2.git
cd QHFlow2
pip install -e ".[fairchem]"
2. Download Checkpoints
# 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:
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 and place under dataset/:
QHFlow2/dataset/
ethanol/
malondialdehyde/
uracil/
4. Run Prediction
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
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
@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