README.md

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