Efficient E(3)-equivariant framework for universal charge density prediction

Abstract

Electronic structure is ubiquitously obtained via density functional theory (DFT), where the charge density plays a central role. This work presents EdenGNN (Equivariant Density Graph Neural Network), a machine learning (ML) charge density model for electronic structure. Current universal ML charge density models are hampered by prohibitive computational costs. Furthermore, despite being trained on projector augmented-wave (PAW) based DFT datasets, they predict only the pseudo charge density, which is insufficient to reconstruct the electronic structure. In contrast, EdenGNN overcomes these limitations. It additionally predicts the augmentation occupancies, enabling electronic structure calculations with PAW accuracy. Critically, by employing a basis-expansion formulation with fully trainable radial basis functions and a $\Delta$-learning strategy to capture charge transfer, it is over an order of magnitude faster. Trained on the Materials Project database, our universal model, EdenGNN-Uni, accurately predicts the band structures for the majority of materials across a vast chemical space. These findings establish the ML charge density model as a scalable \textit{ab initio} method for large-scale electronic structure calculations and high-throughput screening.