Global Plane Waves From Local Gaussians: Periodic Charge Densities in a Blink

TL;DR

This work addresses the high computational cost of plane-wave DFT for periodic materials by learning a fast, accurate initial charge-density guess. It introduces ELECTRAFI, which represents the density as a finite mixture of anisotropic Gaussians whose parameters are learned from a graph neural backbone, and reconstructs the periodic density in reciprocal space using analytic Fourier transforms and the Poisson summation formula, followed by a single inverse FFT. The key contributions are (i) sub-second inference with up to $633\times$ end-to-end speedups over dense-grid density models, (ii) consistent end-to-end DFT time reductions when used as SCF initialization, and (iii) an element-wise error analysis showing complementary regimes relative to prior models. The framework demonstrates that combining accuracy with inference efficiency is crucial for practical ML-accelerated DFT workflows and argues for close integration with target DFT codes to maximize end-to-end gains.

Abstract

We introduce ELECTRAFI, a fast, end-to-end differentiable model for predicting periodic charge densities in crystalline materials. ELECTRAFI constructs anisotropic Gaussians in real space and exploits their closed-form Fourier transforms to analytically evaluate plane-wave coefficients via the Poisson summation formula. This formulation delegates non-local and periodic behavior to analytic transforms, enabling reconstruction of the full periodic charge density with a single inverse FFT. By avoiding explicit real-space grid probing, periodic image summation, and spherical harmonic expansions, ELECTRAFI matches or exceeds state-of-the-art accuracy across periodic benchmarks while being up to $633 \times$ faster than the strongest competing method, reconstructing crystal charge densities in a fraction of a second. When used to initialize DFT calculations, ELECTRAFI reduces total DFT compute cost by up to ~20%, whereas slower charge density models negate savings due to high inference times. Our results show that accuracy and inference cost jointly determine end-to-end DFT speedups, and motivate our focus on efficiency.

Figures (8)

• Figure 1: The electron density is parametrized by a set of Gaussians, where each component $j$ has a signed weight $w^{(j)}$, displacement from anchor $\mathbf{d}^{(j)}$, and covariance $\boldsymbol{\Sigma}^{(j)}$.
• Figure 2: Logarithmic scaling plots showing wall-clock time versus system size for ELECTRAFI, DFT, and ChargE3Net. Top: Materials Project (MP) test set. Bottom: GNoME test set.
• Figure 3: End-to-end (TOTAL, i.e., including model inference time) vs DFT-only wall-time reduction as a function of system size. Top: Materials Project (MP) test set. Bottom: GNoME test set.
• Figure 4: Error trends across elements for ELECTRAFI and ChargE3Net on the MP-Full and GNoME test sets. The error for each element has been calculated as the average error for all structures including the element.
• Figure 5: Lowest error structure from the MP-Full test set. (a,b) Predicted densities for ELECTRAFI and ChargE3Net. (c–f) $\Delta\rho=\rho_{\mathrm{pred}}-\rho_{\mathrm{ref}}$ at two iso-levels. (g) Reference density from SCF calculations. (h) Crystal structure AlMg30NaO32.