| --- |
| license: lgpl-3.0 |
| tags: |
| - chemistry |
| - electron-density |
| - graph-neural-network |
| - equivariant |
| - dft |
| library_name: pytorch |
| --- |
| |
| # BOA — Basis Overlap Architecture |
|
|
| Trained checkpoints for the ICLR 2026 paper |
| [**A Function-Centric Graph Neural Network Approach For Predicting Electron Densities**](https://openreview.net/forum?id=HDdkFjFEZd). |
|
|
| BOA is an equivariant graph neural network that predicts ground-state electron |
| densities. Its message passing uses the overlap matrix of the basis functions |
| that represent the predicted density, rather than treating the basis coefficients |
| as generic node features. |
|
|
| Code: https://github.com/sciai-lab/boa |
|
|
| ## Checkpoints |
| | file | dataset | NMAE [%] *(this checkpoint)* | NMAE [%] *(mean ± standard error)* | seeds | |
| | --- | --- | --- | --- | --- | |
| | `qm9_pyscf_large.ckpt` | QM9 (PySCF) | 0.106 | 0.116 ± 0.006 | 5 | |
| | `qm9_pyscf_small.ckpt` | QM9 (PySCF) | 0.113 | 0.13 ± 0.01 | 3 | |
| | `qm9_vasp_large.ckpt` | QM9 (VASP) | 0.132 | 0.1339 ± 0.0005 | 5 | |
| | `qm9_vasp_small.ckpt` | QM9 (VASP) | 0.137 | 0.1381 ± 0.0003 | 3 | |
| | `qm9_pyscf_small_small_cutoff.ckpt` | QM9 (PySCF) | 0.121 | - | 1 | |
| | `benzene.ckpt` | MD | 0.355 | 0.361 ± 0.003 | 3 | |
| | `resorcinol.ckpt` | MD | 0.362 | 0.371 ± 0.004 | 3 | |
| | `phenol.ckpt` | MD | 0.494 | 0.56 ± 0.03 | 3 | |
| | `malonaldehyde.ckpt` | MD | 0.585 | 0.61 ± 0.01 | 3 | |
| | `ethanol.ckpt` | MD | 0.705 | 0.710 ± 0.004 | 3 | |
| | `ethane.ckpt` | MD | 0.767 | 0.772 ± 0.002 | 3 | |
|
|
| **These are the best seeds** The paper reports the mean |
| over multiple seeds for each dataset; each checkpoint here is the single best-performing seed, so its NMAE |
| is better than the published figure. |
|
|
| ### The reduced-cutoff model (extrapolation) |
|
|
| `qm9_pyscf_small_small_cutoff.ckpt` is the model from §3.2 of the paper, trained |
| for **extrapolation to molecules far larger than those seen in training**. It is |
| the `small` configuration on QM9/PySCF with two reduced radii: message passing |
| 6 Å → **3 Å** and edge features 3 Å → **2 Å** (config |
| `configs/experiment/qm9_pyscf_small_small_cutoff.yaml`). |
|
|
| ## Usage |
|
|
| Install the code from https://github.com/sciai-lab/boa, then: |
|
|
| ```python |
| from boa.model.module import ChgLightningModule |
| |
| model = ChgLightningModule.load_from_checkpoint("qm9_pyscf_large.ckpt", map_location="cpu") |
| model.eval() |
| ``` |
|
|
| **On the weights.** BOA is trained with an exponential moving average, and all |
| reported numbers were measured with the EMA weights rather than the raw ones. In |
| these released files the `state_dict` **is** the EMA weights, so a plain load |
| gives you the evaluated model — no extra step. (Training checkpoints are not |
| like this: there `state_dict` holds the raw weights and `ema.copy_to` must be |
| applied first. If you evaluate one of these files with the repository's |
| `boa/test.py`, its `ema.copy_to` call is a no-op and remains correct.) |
|
|
