EricBoi/dcmnet-qm9-esp

DCMNet: distributed charges for the QM9 electrostatic potential

E(3)-equivariant message-passing networks that predict three point charges per atom and their positions for organic molecules (H, C, N, O, F). The charges reproduce the molecular electrostatic potential (ESP) through plain Coulomb sums, so they can be used directly as a distributed-charge electrostatics model in a force field.

Trained with mmml (DCMNetCharges) on the ESP of all 133,885 QM9 molecules predicted by the PET-OMol25-L electron-density model, prepared by Paolo Pegolo. The targets are machine-learned, not DFT.

Models

Each directory holds model.safetensors (EMA weights, 98,967 parameters) and config.json (architecture, training settings, metrics). RMSE in kcal/mol/e on every point of the held-out splits; "Lebedev" is a Merz-Kollman grid (110-point shells at 1.4, 1.6, 1.8, 2.0 x r_vdW), "ec" is a surface grid straddling the solvent-accessible surface plus far-field points.

Model	Trained on	Total charge	1.4-shell loss weight	Test RMSE, own grid	Valid RMSE, Lebedev	Valid RMSE, ec
leb-n	lebedev	0	1.0	0.630	0.628	0.787
leb-a	lebedev	0	1.0	0.636	0.635	0.792
ec-a	ec, area-weighted	0	1.0	0.766	0.650	0.767
leb-q	lebedev	-0.0113 e	1.0	0.685	0.683	0.751
leb-q-w	lebedev	-0.0113 e	0.25	0.693	0.692	0.745
ec-q-w	ec, area-weighted	-0.0113 e	0.25	0.674	0.773	0.676

Validation error by Lebedev shell, RMSE / mean error (model - teacher), kcal/mol/e:

Model	1.4 r_vdW	1.6	1.8	2.0
leb-n	0.46 / -0.18	0.56 / +0.50	0.75 / +0.72	0.78 / +0.76
leb-a	0.48 / -0.18	0.57 / +0.50	0.75 / +0.72	0.79 / +0.76
ec-a	0.55 / -0.19	0.57 / +0.48	0.74 / +0.71	0.78 / +0.76
leb-q	1.10 / -1.00	0.44 / -0.31	0.28 / -0.05	0.24 / +0.02
leb-q-w	1.13 / -1.04	0.43 / -0.33	0.24 / -0.07	0.20 / +0.01
ec-q-w	1.26 / -1.10	0.53 / -0.39	0.25 / -0.11	0.17 / -0.01

Against DFT

For 11 molecules, wB97M-V/def2-TZVPD (the OMol25 level) was computed on the same grid points with gpu4pyscf. Median RMSE over the 8 ordinary molecules, Lebedev grid, kcal/mol/e: def2-QZVPD 0.04, def2-SVPD 0.25, PET-OMol25-L teacher 1.09 (mostly the -0.89 offset from its net charge), MBIS atomic multipoles from the DFT density l <= 2 0.64, leb-n 0.81, leb-a 0.81, ec-a 0.82, leb-q-w neutralised 0.86 (as trained 1.53). The neutral models are closer to DFT than the teacher is. On three hard molecules (a strained cage, a zwitterion, a charged pair) the networks reach 1.6-6.6, so treat out-of-distribution chemistry with care. Varying the charges per atom (atom-centred, 1, 2, 3, 4; separate runs, not uploaded) shows the gains stop at 3: test RMSE against the teacher 1.21, 0.72, 0.66, 0.63, 0.64.

Usage

pip install "jax==0.10.2" "flax==0.12.8" "e3x==1.0.2" safetensors huggingface_hub

from huggingface_hub import snapshot_download
import sys, numpy as np
path = snapshot_download("EricBoi/dcmnet-qm9-esp")
sys.path.insert(0, path)
from dcmnet_esp import load, predict, esp

model, params, cfg = load(f"{path}/leb-n")
Z = np.array([8, 1, 1])                                    # water
R = np.array([[0, 0, 0.1173], [0, 0.7572, -0.4692], [0, -0.7572, -0.4692]])   # Angstrom
q, sites = predict(model, params, cfg, Z, R)               # 9 charges (e), 9 positions (A)
V = esp(q, sites, np.array([[0.0, 0.0, 3.12]]))            # kcal/mol/e

dcmnet_esp.py is self-contained (jax, flax, e3x) and reproduces mmml's DCMNetCharges with the same parameter names; checked against the training code on 9 molecules to within GPU TF32 precision (charges 1e-3 e, ESP 0.05 kcal/mol/e). Inputs are one molecule at a time, unpadded.

Architecture

Embedding	element -> 64 scalar features
Interaction blocks	3 x (message passing over all atom pairs with 32 reciprocal-Bernstein radial functions x spherical harmonics l <= 2, smooth 6 A cutoff; residual atom-wise MLP Dense-SiLU-Dense). Features: 64 channels each of l = 0, 1, 2
Readout	TensorDense to 3 channels (l <= 1). l = 0 -> Dense + per-element bias -> 3 charges per atom; l = 1 -> hard-tanh x 0.5 A -> 3 offset vectors per atom
Charge sites	atom position + offset, each Cartesian component within +-0.5 A
Total charge	exact projection: every site shifted by (sum q - Q) / n_sites
ESP	V(p) = sum q / abs(p - site), no learned parameters

Training

• Data: QM9 (133,885 molecules), one random split shared by both grids (seed 42): 110,000 train, 10,000 validation, 13,885 test. About 1.3% of molecules are zwitterions; all are neutral overall.
• Loss: mean squared ESP error over 512 random grid points per molecule per step (ec grid: weighted by quadrature surface area), plus optional down-weighting of the innermost shell.
• AdamW (weight decay 1e-5), gradient clip 1.0, warmup then cosine decay, batch 64, EMA 0.999; one NVIDIA GH200 per model (CSCS Alps, clariden), JAX 0.10.2, TF32 matmuls.
• full-v1 models: 500 epochs, peak learning rate 5e-4. full-v2 models: 300 more epochs from the full-v1 weights, peak 3e-4.

Known limitations

• Teacher net charge. The PET-OMol25-L ESP carries a near-constant net charge of -0.0113 e per molecule (spread 1e-4 e), an artefact of the density model. Neutral models cannot represent it and sit about +0.7 kcal/mol/e above the teacher on the outer shells. Models with total charge -0.0113 e reproduce it; call predict(..., neutral=True) to remove it (a uniform shift of every site) for force-field use, where molecules should be neutral.
• Charge penetration. At 1.4 x r_vdW the teacher is about 1 kcal/mol/e more positive than any point-charge model can be (the electron cloud overlaps the grid). Exact multipole fits to each molecule hit the same floor (0.59 kcal/mol/e on Lebedev with neutral atomic quadrupoles), so this is a limit of point charges, not of the network.
• Chemistry is QM9's: up to 9 heavy atoms, H/C/N/O/F, closed-shell, neutral, gas-phase equilibrium geometries. Errors are largest for zwitterions.

Acknowledgements

Teacher data: Paolo Pegolo. Model code: mmml (Eric Boittier). Training on CSCS Alps (clariden, Swiss AI Initiative project ab010).