All-electron machine-learned orbital-free density functional for periodic CHNO systems

Marc K. Ickler; Andreas Hermann; Fred A. Hamprecht

physics.chem-phOct 13, 2026

All-electron machine-learned orbital-free density functional for periodic CHNO systems

Authors: Marc K. Ickler, Andreas Hermann, Fred A. Hamprecht

Organizations: Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University, Heidelberg, 69120, Germany. · Centre for Science at Extreme Conditions and SUPA, School of Physics and Astronomy, University of Edinburgh, Edinburgh, EH9 3FD, United Kingdom.

Abstract

We introduce STRUCTURES-P, a machine-learned all-electron orbital-free density functional for periodic CHNO systems. Electrostatics are evaluated analytically in an atom-centered Gaussian basis, while an equivariant network learns the kinetic and exchange-correlation energy from energies and gradients of non-ground-state densities. Minimizing the functional from a superposition of atomic densities recovers the Kohn-Sham ground state: on held-out crystals, the density error falls from 6.7% to 0.4% of the electron count and the energy error reaches 0.25 mHa per atom, half that of a same-architecture interatomic potential. A single functional can be trained for both isolated molecules and periodic systems, and the model generalizes to larger unit cells. It largely preserves relative-energy rankings and reproduces convex-hull vertices on the tested H-N composition slice. Together, these results demonstrate the viability of machine-learned all-electron orbital-free DFT for variational density optimization and energy evaluation in fixed-geometry solids.

Figures & tables

Figure 1: Overview of STRUCTURES-P for periodic crystals with Gaussian-type density basis functions, drawn here for a two-dimensional lattice. a The electron density of the primitive cell is represented by coefficients $\mathbf{p}$ of atom-centered Gaussians, $\rho(\mathbf{r})=\sum_{\mu}p_{\mu}\bar{\omega}_{\mu}(\mathbf{r})$ , where $\bar{\omega}_{\mu}$ includes the lattice images; contours show the resulting periodic density, the inset the radial profiles of the even-tempered shells. b The same structure defines a periodic graph, on which each atom of the cell receives messages from its neighbors within a cutoff $r_{\mathrm{c}}$ , including periodic images. An equivariant graph network (eSEN) predicts the kinetic and exchange–correlation energy $T_{\mathrm{S}}+E_{\mathrm{XC}}$ , while the Hartree, nuclear-attraction, and nucleus–nucleus energies are evaluated analytically from periodic two-center integrals; of these, only $E_{\mathrm{H}}$ and $E_{\mathrm{ext}}$ depend on $\mathbf{p}$ . c Together these terms define the total energy, whose gradient with respect to $\mathbf{p}$ (by automatic differentiation) drives a variational density optimization from an initial guess $\mathbf{p}^{0}$ to an optimized density $\mathbf{p}^{\ast}$ under a fixed-electron-number constraint; the update step $U_{t}$ is specified in Section 4.5 .

Figure 2: Selected test-set crystals illustrating the connectivity classes of the dataset: a finite bonded components (0D), b chains (1D), c sheets (2D), and d three-dimensional networks (3D). Surfaces show fitted Kohn–Sham reference densities at $0.20\,a_{0}^{-3}$ . Periodic images are faded, and the teal box outlines the primitive cell; the network in d is a single component, so nothing is faded. Atom colors: white (H), gray (C), blue (N), and red (O). Bond cylinders indicate geometric connectivity, not bond order.

Evaluation	MAE (mHa/cell)	Median (mHa/cell)	MAE (mHa/atom)	Mean $L^{2}$ density error $\bar{\epsilon_{\rho}}$	Mean $L^{1}$ density error $\bar{\eta_{\rho}}$ (%)
STRUCTURES-P@SAD	304	305	47.8	0.235	6.73
STRUCTURES-P@KS	1.61	0.653	0.255	0	0
STRUCTURES-P@optimized	1.58	0.627	0.254	0.0138	0.413
APBE@KS	619	500	101	0	0
eSEN MLIP	3.32	1.11	0.496	–	–

Table 1: Energy and density accuracy on all 6211 in-distribution test structures. The first three rows evaluate STRUCTURES-P at different densities; the notation method@density identifies the density at which the energy is evaluated. The eSEN MLIP predicts total energies directly from geometry and is trained without force supervision; it does not provide a density. APBE is a semilocal (GGA) kinetic-energy functional evaluated with PBE exchange–correlation and the same analytic electrostatics at the fitted KS density. The fitted KS density has zero density error by definition.

L^{2}

density errors are in

a_{0}^{-3/2}

;

L^{1}

density errors are per electron, in percent. MAE and median are the mean and median absolute energy errors relative to the PySCF KS total energy.

Figure 3: Density optimization reduces density and energy errors on the 6211-cell CHNO test set. a Optimized density for a representative test cell with errors near the test-set medians (selection in Section 4.5 ). The molecule is centered in the cell. The isosurface encloses 75% of the electrons; the shaded plane is used in b . Atom colors: white (H), gray (C), blue (N), and red (O). b Fitted Kohn–Sham density $\rho_{\mathrm{KS}}$ (logarithmic scale; contours at 0.01, 0.1, and $1\,a_{0}^{-3}$ ), SAD and optimized-density errors, and the latter magnified by 20. All error maps use a common symmetric color scale, and RMS denotes the root-mean-square density difference over the displayed plane in $a_{0}^{-3}$ . Circles mark atoms within $1.1\,a_{0}$ of the plane. c Density and absolute energy errors during optimization: median and interquartile range over 400 randomly sampled test cells, with the representative cell overlaid. Converged runs retain their final values. d Distribution of the density-error reduction $\epsilon_{\rho}^{\mathrm{SAD}}/\epsilon_{\rho}^{\mathrm{opt}}$ over all 6211 structures; dashed and red lines mark no improvement and the median, respectively.

Test domain	$N$	MAE (mHa/cell)	MAE/atom (mHa)	$L^{2}$ density error	Converged (%)
QM9	10831	0.245	0.0146	0.0111	100.0
Periodic CHNO	6211	1.77	0.280	0.0197	99.5

Table 2: Variational test-set accuracy of one jointly trained STRUCTURES-P model on isolated QM9 molecules and periodic CHNO crystals. All model parameters are shared between the two domains. Density errors are in

a_{0}^{-3/2}

; convergence refers to the regularized-gradient criterion.

Figure 4: The mean optimized-energy error remains below 1 mHa/atom for cells larger than those seen in training. Each point is one structure, with absolute error against the Kohn–Sham reference plotted against electron count: a in-distribution and b larger-cell test. Horizontal jitter of at most 0.4 electrons separates coincident points; it does not affect the statistics. Solid and dashed lines show the median and 90th percentile in fixed 15-electron bins containing at least 12 structures, using all errors. Both panels share the same axes. Triangles on the upper boundary mark errors above 2 mHa/atom, with their counts shown above each panel.

Figure 5: STRUCTURES-P reproduces the relative enthalpies of its PySCF training reference more closely than plane-wave CASTEP does. Per-composition enthalpy rankings for CASTEP, PySCF (nominal

0.10

Å -1 k-point spacing), and STRUCTURES-P, using

H=E+pV

at the

\sim

1 GPa structure-search pressure. a For each composition (with

n

structures), the enthalpy of every structure relative to each method’s own composition minimum (meV/atom). Lines connect the same structure between CASTEP and PySCF (orange) and between PySCF and STRUCTURES-P (blue); horizontal segments indicate equal relative enthalpies for the two methods. b Notched box plots of relative-enthalpy deviations from PySCF (boxes: interquartile range; whiskers: 1.5 times that range; notches: approximate 95% confidence intervals for the median, which may extend beyond the quartiles). c Percentage of within-composition structure pairs whose order changes relative to PySCF. Candidates are selected within 150 meV/atom of the PySCF composition minimum.

Figure 6: Formation-enthalpy convex hull along the H–N edge of the CHNO composition space, using

H=E+pV

at the

\sim

1 GPa search pressure and method-specific elemental references from the evaluated candidate set. a Lowest formation enthalpy at each stoichiometry and its convex hull for CASTEP (circles, dash-dotted), PySCF (squares, solid), and STRUCTURES-P (diamonds, dashed). Faint gray points show individual PySCF candidates with formation enthalpies up to 300 meV/atom. Filled diamonds mark the three stoichiometries where STRUCTURES-P and PySCF select different structures. All three methods give the H, H 3 N, and N hull vertices. b STRUCTURES-P minus PySCF for the composition minima (diamonds) and hull (dashed line), resolved on a finer energy scale.

Setting	Value
eSEN layers; angular limits $l_{\mathrm{max}},m_{\mathrm{max}}$	4; 4, 4
Sphere, edge, and hidden channels	128 each
Radial basis functions; neighbor cutoff	128; $6\,a_{0}$
Maximum neighbors per atom	300
Optimizer	AdamW
Initial learning rate; schedule	$10^{-3}$ ; cosine decay to zero

Table 3: Training settings for the large CHNO model. The joint QM9–CHNO model uses the same architecture and optimizer, with a batch size of 128 and 25 epochs.

Figure 7: Computational cost over the evaluated cell-size range. a Amortized end-to-end wall time per structure for STRUCTURES-P with full-pair and local-pair regularization, including startup and compilation, and mean wall time per matched single-thread PySCF calculation. b Median time per structure for density optimization, construction of the local-pair regularizer, and the remaining static setup; open squares show full-pair regularizer construction.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

	MAE (mHa/cell)		Mean $L^{2}$		Conv. (%)
Training setup / test	Small	Large	Small	Large	Small	Large
STRUCTURES-P@optimized, CHNO / CHNO	1.981	1.583	0.0174	0.0138	99.871	99.920
STRUCTURES-P@optimized, joint / QM9	0.354	0.245	0.0128	0.0111	100.000	100.000
STRUCTURES-P@optimized, joint / CHNO	2.299	1.771	0.0240	0.0197	98.905	99.549
eSEN MLIP, CHNO / CHNO	3.361	3.316	–	–	–	–

Appendix

Supplementary Table 1: Model-size ablation for STRUCTURES-P on the complete test sets. The small and large models have approximately 2.73 and 22.75 million parameters, respectively. The variational density optimizations use float64 arithmetic, soft regularization

\epsilon=0.1

, and the same convergence threshold and cycle limit. Mean

L^{2}

is the mean cell-level density error

\epsilon_{\rho}

of Eq. (4) of the main text, in

a_{0}^{-3/2}

; Conv. (%) is the percentage of structures meeting the regularized-gradient convergence criterion within the cycle limit (main-text Methods, Density optimization ). The eSEN MLIP predicts total energies directly from geometry and is trained without force supervision. Density error and convergence do not apply to this baseline.

Supplementary Figure 1: CH 2 rankings before and after k-point refinement. Lines connect the same structures between a) the two PySCF settings, b) original PySCF and STRUCTURES-P, and c) refined PySCF and STRUCTURES-P. Each method is referenced to its own minimum over the same 20 candidates. The highlighted structures are A, the original PySCF and model minimum (CHNO-140055-4417-4), and B, the refined PySCF minimum (CHNO-59219-1291-63). Spacings are in Å -1 .

Supplementary Figure 2: Distributions of reference changes and model errors. Absolute per-atom energy differences for the same 516 screened cells: a) PySCF at nominal spacings of

0.10

and

0.03

Å -1 ; b) STRUCTURES-P versus the original

0.10

Å -1 PySCF reference. Both histograms use the same 30 logarithmically spaced bins and axis scales. Solid black and dashed red lines mark the mean and median, respectively. At fixed geometry and

pV

, reference energy changes are also enthalpy changes.

Density Functional Theory

See all 25 papers

A fully differentiable framework for training proxy Exchange Correlation Functionals for periodic systems
Feb 17, 2026Rakshit Kumar Singh, Aryan Amit Barsainyan, Bharath RamsundarComputational Quantum Chemistry Density Functional Theory
Deep Forest Sciences
Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory
Apr 22, 2026Roman Remme, Fred A. HamprechtSurrogate Modeling Quantum Chemistry
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University
Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory
Aug 24, 2026Danish Khan, Maurice D. Hanisch, Nikolai Argatoff +3Density Functional Theory Computational Quantum Chemistry
Department of Computing and Mathematical Sciences, California Institute of Technology · Department of Mathematics, ETH Zürich · Division of Chemistry and Chemical Engineering, California Institute of Technology +1

Feb 17, 2026J/K move · Enter open · S save

A fully differentiable framework for training proxy Exchange Correlation Functionals for periodic systems

Rakshit Kumar Singh, Aryan Amit Barsainyan, Bharath Ramsundar

Deep Forest Sciences

Computational Quantum Chemistry Density Functional Theory Quantum Chemistry Computational Materials Science Differentiable Programming

Open paper AskPDF arXiv

Density Functional Theory (DFT) is widely used for first-principles simulations in chemistry and materials science, but its computational cost remains a key limitation for large systems. Motivated by recent advances in ML-based exchange-correlation (XC) functionals, this paper introduces a differentiable framework that integrates machine learning models into density functional theory (DFT) for solids and other periodic systems. The framework defines a clean API for neural network models that can act as drop in replacements for conventional exchange-correlation (XC) functionals and enables gradients to flow through the full self-consistent DFT workflow. The framework is implemented in Python using a PyTorch backend, making it fully differentiable and easy to use with standard deep learning tools. We integrate the implementation with the DeepChem library to promote the reuse of established models and to lower the barrier for experimentation. In initial benchmarks against established electronic structure packages (GPAW and PySCF), our models achieve relative errors on the order of 5-10%.

Abstract

Density Functional Theory

Explore similar work

A fully differentiable framework for training proxy Exchange Correlation Functionals for periodic systems

Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory

Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory