neural_paw_dft — Neural Electronic Initialization for PAW-DFT

Trained weights for the paper Complete Neural Electronic Initialization Accelerates Materials DFT (arXiv:2609.21759).

neural_paw_dft predicts a complete starting point for a VASP calculation from the structure alone: the total electron density and the spin-difference density on the FFT grid (ELECTRAFI, EScAIP backbone) and the PAW augmentation occupancies (AugNet, MACE backbone). Together they fill every block of a CHGCAR, so VASP starts from a near-converged density instead of a superposition of atomic charges.

Code: https://github.com/aerte/neural_paw_dft
ELECTRAFI architecture: https://github.com/Jotels/ELECTRAFI

Checkpoints

The file stem is the model's name in the neural_paw_dft registry (neural_paw_dft.models.REGISTRY); select models by that name.

ELECTRAFI — total density

Grid models built on the ELECTRAFI architecture (EScAIP backbone), predicting the total valence electron density on the VASP FFT grid, normalised to the exact electron count.

file	description
electrafi_total.safetensors	Default total-density model.
electrafi_total_v2.safetensors	Same architecture and target, from an independent second training run.

Spin-ELECTRAFI — total and spin-difference density

ELECTRAFI extended with a spin head: predicts the total density and the spin-difference density ρ↑ − ρ↓ on the grid.

file	description
electrafi_spin_constrained.safetensors	Net moment rescaled to a prescribed value at inference; the pipeline supplies the sum of CHGNet's site moments.
electrafi_spin_unconstrained.safetensors	No moment input; the net moment is whatever the network predicts.

AugNet — total augmentation occupancies

MACE-backbone models predicting the PAW augmentation occupancies for the total density, i.e. the per-atom coefficient blocks in the augmentation occupancies section of a CHGCAR. Each file has a *.config.json sidecar with the architecture needed to rebuild it.

file	training structures	description
augnet_total_full.safetensors	full	Default AugNet.
augnet_total_50k.safetensors	50k	Data-scaling ablation.
augnet_total_10k.safetensors	10k	Data-scaling ablation.
augnet_total_1k.safetensors	1k	Data-scaling ablation.

Spin-AugNet — spin augmentation occupancies

Same backbone, trained on the magnetisation channel: PAW augmentation occupancies for the spin-difference density. Pairs with the Spin-ELECTRAFI models.

file	training structures	description
augnet_spin_full.safetensors	full	Default spin AugNet.

Usage

Install the code from https://github.com/aerte/neural_paw_dft, then:

import torch
from ase.build import bulk
from pymatgen.io.ase import AseAtomsAdaptor
from neural_paw_dft.pipeline.config import ElectrafiConfig
from neural_paw_dft.pipeline.electrafi import load_electrafi, predict_density

structure = AseAtomsAdaptor.get_structure(bulk("Fe", "bcc", a=2.87))
cfg = ElectrafiConfig(checkpoint="electrafi_total", spin=False)
model = load_electrafi(cfg, torch.device("cpu"), n_atoms=len(structure))
rho, _ = predict_density(model, structure, (32, 32, 32), n_elec=8.0)

The spin models return the spin-difference grid as well. electrafi_spin_constrained pins its net moment to m_total and raises without it; the pipeline uses the sum of CHGNet's unsigned site moments (the paper's convention), which you can reproduce directly:

from chgnet.model import CHGNet

m_total = float(CHGNet.load(use_device="cpu").predict_structure(structure)["m"].sum())
cfg = ElectrafiConfig(checkpoint="electrafi_spin_constrained", spin=True)
model = load_electrafi(cfg, torch.device("cpu"), n_atoms=len(structure))
rho, rho_spin = predict_density(model, structure, (32, 32, 32), n_elec=8.0, m_total=m_total)

electrafi_spin_unconstrained needs no m_total (any value passed is ignored). Densities are in e/ų on the grid; n_elec is the valence electron count of the POTCARs (8 for the one Fe_pv atom in this primitive cell).

or run the whole pipeline (CHGNet moments → ELECTRAFI grids → AugNet occupancies → VASP-ready directory with CHGCAR, INCAR, POSCAR, POTCAR, KPOINTS):

ndi config-template > ndi.yaml
ndi build POSCAR --config ndi.yaml --out seed/

Always select a model by registry name, not by file path. The training config (spin arm, spin_renorm) is resolved from the name; a bare path cannot be resolved and raises.

On the weights

These are inference-only safetensors exports. The Lightning training checkpoints (with optimizer state) are not published.

ELECTRAFI weights are bitwise identical to the training checkpoints. safetensors is a lossless container, so predicted densities are bit-for-bit those of the original runs.

AugNet weights were converted out of the cuEquivariance layout and reproduce the original runs to ~1e-7 relative. The published files carry no cuEquivariance dependency, so AugNet inference runs on torch + mace-torch + e3nn alone. The conversion un-fuses cuEquivariance's packed symmetric-contraction tensor into per-contraction e3nn weights (mace.cli.convert_cueq_e3nn), then changes basis from the reduced Clebsch–Gordan basis (which MACE only builds when cuequivariance is importable) to the full _wigner_nj basis. The reduced basis lies exactly inside the span of the full one, so the change of basis is exact (least-squares residual ~1e-16); the remaining ~1e-7 is float32 round-off from running different but mathematically equivalent kernels. Verified on NaCl, Si and FeO for all five AugNet models with cuequivariance uninstallable. If you need bitwise reproduction of the original AugNet runs, ask for the cuEquivariance checkpoints.

The AugNet sidecars record "use_reduced_cg": false. MACE otherwise picks its Clebsch–Gordan basis from whether cuequivariance is importable, which changes the symmetric-contraction weight shapes; neural_paw_dft pins the value so the same weights load the same way everywhere. U_matrix_* buffers are omitted from the files because they are constants recomputed at construction.

License

The weights are released under CC BY-NC 4.0: free for noncommercial use with attribution. If you use them in a publication, please cite the paper above.

Citation

If you use this work, please cite:

Ærtebjerg, Felix, Jonas Elsborg, and Arghya Bhowmik. "Complete Neural Electronic Initialization Accelerates Materials DFT." arXiv preprint arXiv:2609.21759 (2026).

@article{aertebjerg2026complete,
  title   = {Complete Neural Electronic Initialization Accelerates Materials DFT},
  author  = {{\AE}rtebjerg, Felix and Elsborg, Jonas and Bhowmik, Arghya},
  journal = {arXiv preprint arXiv:2609.21759},
  year    = {2026}
}

---

Copyright (c) 2026 Felix Ærtebjerg