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

Abstract

Abstract Electron density $$\rho (\overrightarrow{{{{\bf{r}}}}})$$ ρ ( r → ) is the fundamental variable in the calculation of ground state energy with density functional theory (DFT). Beyond total energy, features and changes in $$\rho (\overrightarrow{{{{\bf{r}}}}})$$ ρ ( r → ) distributions are often used to capture critical physicochemical phenomena in functional materials. We present a machine learning framework for the prediction of $$\rho (\overrightarrow{{{{\bf{r}}}}})$$ ρ ( 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 $$\rho (\overrightarrow{{{{\bf{r}}}}})$$ ρ ( 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.