metadata
license: mit
tags:
- chemistry
- quantum-chemistry
- dft
- hamiltonian
- equivariant
- flow-matching
datasets:
- rMD17
language:
- en
pipeline_tag: other
QHFlow2 — rMD17 Pre-trained Checkpoints
Pre-trained checkpoints for QHFlow2 on the revised MD17 (rMD17) 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 |
|---|---|---|
| aspirin | 21 | C₉H₈O₄ |
| naphthalene | 18 | C₁₀H₈ |
| salicylic_acid | 16 | C₇H₆O₃ |
File Structure
rmd-{molecule}/
QHFlow_so2_v5_1_{size}_b10-rmd-{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-rMD17 \
"rmd-aspirin/QHFlow_so2_v5_1_small_v2_b10-rmd-aspirin/weights-epoch=79-val_loss=0.0000000.ckpt" \
--local-dir ckpt/rmd17
# Download all checkpoints (~3.9 GB)
huggingface-cli download ksusu/QHFlow2-rMD17 --local-dir ckpt/rmd17
Or in Python:
from huggingface_hub import hf_hub_download
path = hf_hub_download(
repo_id="ksusu/QHFlow2-rMD17",
filename="rmd-aspirin/QHFlow_so2_v5_1_small_v2_b10-rmd-aspirin/weights-epoch=79-val_loss=0.0000000.ckpt",
)
3. Download Dataset
Download from Google Drive and place under dataset/:
QHFlow2/dataset/
rmd-aspirin/
rmd-naphthalene/
rmd-salicylic_acid/
4. Run Prediction
python -m qhflow2.experiment.train_md17 \
mode=predict \
dataset=rmd-aspirin \
ckpt=ckpt/rmd17/rmd-aspirin/QHFlow_so2_v5_1_small_v2_b10-rmd-aspirin/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