Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids
Peter Bjørn Jørgensen, Arghya Bhowmik
Abstract
Abstract Electron density $$\rho (\overrightarrow{{{{\bf{r}}}}})$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mi>ρ</mml:mi> <mml:mrow> <mml:mo>(</mml:mo> <mml:mrow> <mml:mover> <mml:mrow> <mml:mi>r</mml:mi> </mml:mrow> <mml:mo>→</mml:mo> </mml:mover> </mml:mrow> <mml:mo>)</mml:mo> </mml:mrow> </mml:mrow> </mml:math> 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}}}}})$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mi>ρ</mml:mi> <mml:mrow> <mml:mo>(</mml:mo> <mml:mrow> <mml:mover> <mml:mrow> <mml:mi>r</mml:mi> </mml:mrow> <mml:mo>→</mml:mo> </mml:mover> </mml:mrow> <mml:mo>)</mml:mo> </mml:mrow> </mml:mrow> </mml:math> 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}}}}})$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mi>ρ</mml:mi> <mml:mrow> <mml:mo>(</mml:mo> <mml:mrow> <mml:mover> <mml:mrow> <mml:mi>r</mml:mi> </mml:mrow> <mml:mo>→</mml:mo> </mml:mover> </mml:mrow> <mml:mo>)</mml:mo> </mml:mrow> </mml:mrow> </mml:math> . 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}}}}})$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mi>ρ</mml:mi> <mml:mrow> <mml:mo>(</mml:mo> <mml:mrow> <mml:mover> <mml:mrow> <mml:mi>r</mml:mi> </mml:mrow> <mml:mo>→</mml:mo> </mml:mover> </mml:mrow> <mml:mo>)</mml:mo> </mml:mrow> </mml:mrow> </mml:math> 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.