Equivariant Machine Learning for Electron Density Predictions

At the root of density functional theory, the first Hohenberg-Kohn theorem states that the ground-state properties of a material can be predicted uniquely by the electron density. We predict this fundamental quantity via an equivariant neural network called ChargE3Net and show that high rotation-order embeddings (up to l=4 spherical harmonics) can be used to accurately predict the charge density at any point in the unit cell. When trained on approximately 120,000 charge density grids from the Materials Project, we show that ChargE3Net generalizes well for bulk inorganic materials across the periodic table and accelerates self-consistent DFT calculations by 26% for non-magnetic materials when used as an initial charge density guess. Using the predicted charge densities in non-self-consistent calculations, we observe that about 40% of materials in a test set achieve <1meV/atom error for calculated potential energies. We show that ChargE3Net still accelerates DFT calculations on a subset of materials obtained from GNoME, further exemplifying its ability to generalize outside of its training set. We explore additional engineering improvements to the model and extensions to magnetic systems.