Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids

Electron density ρ(r→) is the fundamental variable in the calculation of ground state energy with density functional theory (DFT). Beyond total energy, features and changes in ρ(r→) distributions are often used to capture critical physicochemical phenomena in functional materials. We present a machine learning framework for the prediction of ρ(r→). The model is based on equivariant graph neural networks and the electron density is predicted at special query point vertices that are part of the message-passing graph, but only receive messages. The model is tested across multiple datasets of molecules (QM9), liquid ethylene carbonate electrolyte (EC) and LixNiyMnzCo(1-y-z)O2 lithium ion battery cathodes (NMC). For QM9 molecules, the accuracy of the proposed model exceeds typical variability in ρ(r→) obtained from DFT done with different exchange-correlation functionals. The accuracy on all three datasets is beyond state of the art and the computation time is orders of magnitude faster than DFT.