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.
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 p of atom-centered Gaussians, ρ(r)=∑μpμωˉμ(r) , where ωˉμ 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 rc , including periodic images. An equivariant graph network (eSEN) predicts the kinetic and exchange–correlation energy TS+EXC , while the Hartree, nuclear-attraction, and nucleus–nucleus energies are evaluated analytically from periodic two-center integrals; of these, only EH and Eext depend on p . c Together these terms define the total energy, whose gradient with respect to p (by automatic differentiation) drives a variational density optimization from an initial guess p0 to an optimized density p∗ under a fixed-electron-number constraint; the update step Ut 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.20a0−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 L2 density error ϵρˉ
Mean L1 density error ηρˉ (%)
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. L2 density errors are in a0−3/2 ; L1 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 ρKS (logarithmic scale; contours at 0.01, 0.1, and 1a0−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 a0−3 . Circles mark atoms within 1.1a0 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 ϵρSAD/ϵρ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)
L2 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 a0−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 ∼ 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 ∼ 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 lmax,mmax
4; 4, 4
Sphere, edge, and hidden channels
128 each
Radial basis functions; neighbor cutoff
128; 6a0
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 L2
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 ϵ=0.1 , and the same convergence threshold and cycle limit. Mean L2 is the mean cell-level density error ϵρ of Eq. (4) of the main text, in a0−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 (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%.
Rakshit Kumar Singh, Aryan Amit Barsainyan, Bharath Ramsundar
We introduce surrogate functionals: machine-learned energy functionals for orbital-free density functional theory (OF-DFT) which are defined not by universal fidelity to a physical reference, but merely by the requirement that density optimization with a fixed procedure yields the true ground-state density. Helpfully, training surrogate functionals requires only ground-state densities, no energies or gradients away from the ground state. We here propose a gradient-descent-improvement loss that guarantees exponential convergence of the density to the ground state, and combine it with an adaptive sampling scheme that concentrates learning around the optimization trajectories actually visited during inference. On the QM9 and QMugs benchmarks, surrogate functionals achieve density errors competitive with or improving upon the state of the art for fully supervised machine-learned OF-DFT, while eliminating the need for the O(N3) orthononormalization step required by prior work, yielding improved runtime scaling for larger systems.
Roman Remme, Fred A. Hamprecht
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University
Kohn--Sham density functional theory (DFT) underpins electronic-structure simulations, but repeated orbital diagonalizations lead to cubic scaling, restricting quantum calculations to modest scales only. Eliminating these auxiliary orbitals while retaining Kohn--Sham accuracy is the central goal of orbital-free DFT, but both analytical and machine-learning methods have so far fallen short. Prior learning approaches either try to learn the variational kinetic-energy functionals, which are ill-conditioned, or directly predict the ground state, which extrapolate poorly to larger systems. Instead, we identify the Kohn--Sham map as the right learning target for orbital-free DFT. It maps a Kohn--Sham potential directly to the corresponding density and noninteracting kinetic energy, quantities otherwise obtained through an orbital diagonalization. Focusing on the density component in this work, a domain-invariant SE(3)-equivariant Fourier neural operator learns to predict it from the potential as input on real-space grids, enabling stable quasi-linear scaling SCFs. Trained jointly on 8,504 molecules and solids, a single model generalizes to out-of-distribution organic molecules, insulators, and metals. For the first time, the same method converges SCFs across these systems without explicitly constructing Kohn--Sham orbitals, while reproducing densities, electronic spectra, and structural observables at Kohn--Sham DFT accuracy. Linear-scaling SCFs additionally allow converging magnesium dislocation densities containing up to 82,500 valence electrons on a single GPU.
Danish Khan, Maurice D. Hanisch, Nikolai Argatoff +3
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