Gaussian Plane-Wave Neural Operator for Electron Density Estimation

TL;DR

GPWNO addresses efficient electron density estimation for periodic molecular systems by decomposing the density into plane-wave and Gaussian-type orbital components. It combines PW-based reciprocal-space convolutions with atom-centered GTO-based equivariant message passing, ensuring SE(3) equivariance and adherence to periodic boundary conditions. The approach demonstrates strong, multi-dataset improvements over ten baselines and through comprehensive ablations validates the complementary roles of PW, GTO, and high-frequency masking. The method scales as $O(d_{atom} N + d_{probe} M + M \log M)$, enabling scalable density prediction for larger systems with controllable cost.

Abstract

This work studies machine learning for electron density prediction, which is fundamental for understanding chemical systems and density functional theory (DFT) simulations. To this end, we introduce the Gaussian plane-wave neural operator (GPWNO), which operates in the infinite-dimensional functional space using the plane-wave and Gaussian-type orbital bases, widely recognized in the context of DFT. In particular, both high- and low-frequency components of the density can be effectively represented due to the complementary nature of the two bases. Extensive experiments on QM9, MD, and material project datasets demonstrate GPWNO's superior performance over ten baselines.

Figures (7)

• Figure 1: Illustration of Gaussian plane wave neural operator (GPWNO).(left) Our GPWNO decomposes the electron density estimation into two regions: high-frequency and fast-decaying regions for GTO basis and low-frequency and slow-decaying regions for the PW basis. (right) Our GPWNO makes two predictions based on the GTO and the PW bases. The PW-based prediction is from the lattice-based discretization via the probe nodes, while the GTO-based prediction is from the atom-wise message-passing scheme. Final output is evaluated at an arbitrary query point.
• Figure 2: Neural architecture for the PWNO layer. Starting from a molecule ${\mathcal{M}}$, our framework makes a PW-based prediction $\rho_{\text{PW}}$ for given query points ${\mathcal{Q}}$. The PWNO layer processes the discretized signal using convolution, which is expressed as applying FT, frequency-wise linear transformation $\mathcal{R}$ and inverse FT in sequence. Point-wise activation function $\sigma$ is introduced for the non-linear transformation of the discretized signals.
• Figure 3: Visualization of ground truth (GT), density prediction, and error in NMAE (%). For electron density prediction (Pred.) on QM9 and MP datasets, red and blue colors indicate higher electron density, respectively. For the error (Error), lighter colors indicates lower error.
• Figure 4: Comparison between the baselines on QM9 (lower left is better).
• Figure 5: Increasing number of probe nodes $M$ for PW-only prediction and comparing with GPWNO.