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

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.

💡 Research Summary

The paper introduces ELECTRAFI (Electronic Tensor Reconstruction Algorithm with Fourier Inversion), a novel machine‑learning framework for predicting periodic charge densities in crystalline materials with unprecedented speed and competitive accuracy. Traditional density‑prediction approaches either evaluate the charge density on a dense real‑space grid or expand it in atom‑centered orbital bases (often involving spherical harmonics). Both strategies struggle with the intrinsically non‑local nature of periodic systems: real‑space grid methods require costly evaluation of many grid points and explicit summation over periodic images, while orbital‑based methods need large, high‑angular‑momentum expansions to capture long‑range features, leading to high computational overhead.

ELECTRAFI circumvents these issues by representing the valence charge density as a mixture of anisotropic Gaussian functions whose parameters (weights, centers, and covariance matrices) are predicted by a graph neural network (GNN) conditioned on the atomic graph of the crystal. Crucially, Gaussian functions are self‑reciprocal under Fourier transformation: the Fourier transform of a Gaussian is another Gaussian with analytically known parameters. By exploiting this closed‑form property, the model computes the reciprocal‑space coefficients ρ̂(G) directly from the predicted Gaussian parameters:

ρ̂(G) = Σ_j w_j exp(−½ GᵀΣ_j G) e^{−i G·μ_j}.

Instead of summing over an infinite set of real‑space lattice translations to enforce periodicity, ELECTRAFI applies Poisson’s summation formula, which states that the sum over real‑space images equals a sum over reciprocal‑lattice vectors. The truncated reciprocal lattice acts as a low‑pass filter, which aligns with the fact that DFT charge densities are already band‑limited. After obtaining the finite set of ρ̂(G), a single inverse fast Fourier transform (IFFT) reconstructs the smooth, periodic real‑space density ρ(r). This pipeline eliminates any explicit real‑space grid probing, avoids costly image summations, and requires only O(N log N) operations (where N is the number of reciprocal‑space points).

The GNN encoder respects rotational and reflection symmetries, ensuring data‑efficiency while remaining computationally lightweight. By learning full covariance matrices, the model captures anisotropic charge distributions around each atom, making it suitable for metals, semiconductors, and insulators alike.

Benchmarking against state‑of‑the‑art grid‑based density networks and orbital‑based Gaussian models on a diverse set of periodic crystals shows that ELECTRAFI matches or slightly improves mean absolute error (MAE) and root‑mean‑square error (RMSE) while delivering inference times on the order of 10⁻³ seconds for a 64³ grid. This corresponds to speed‑ups of 200–633× relative to the strongest competing method.

Beyond raw accuracy, the authors evaluate the practical impact on density‑functional theory (DFT) calculations. Using the predicted density as the initial guess for the self‑consistent field (SCF) cycle reduces the number of SCF iterations by an average of 15 % and cuts total DFT wall‑time by up to 20 % in realistic workflows. The key insight is that both high prediction fidelity and ultra‑fast inference are required for end‑to‑end speed‑ups; slower but more accurate models do not yield net savings because their inference overhead dominates.

The paper also discusses broader implications. The analytical Fourier‑transform trick can be extended to other quantities that are naturally expressed in reciprocal space (e.g., electrostatic potentials, exchange‑correlation kernels). Moreover, the approach bridges the gap between machine‑learning surrogates and traditional electronic‑structure methods, positioning ELECTRAFI as a universal accelerator for a wide range of quantum‑chemical simulations, from high‑throughput materials screening to long‑time molecular dynamics.

In summary, ELECTRAFI introduces a mathematically elegant and computationally efficient scheme: (1) predict anisotropic Gaussian parameters with a symmetry‑aware GNN, (2) analytically transform them to reciprocal space, (3) enforce periodicity via Poisson’s summation, and (4) recover the real‑space density with a single IFFT. This combination yields state‑of‑the‑art accuracy, sub‑millisecond inference, and demonstrable reductions in full DFT runtimes, marking a significant step toward integrating ML‑based density predictions into routine electronic‑structure pipelines.