GRAPH2MAT: UNIVERSAL GRAPH TO MATRIX CONVERSION FOR ELECTRON DENSITY PREDICTION

Electron density is a fundamental observable of an atomic system from which all ground state properties can be computed. As a prediction target for machine learning models, electron density is often represented on a dense real space grid, which is data heavy, or through density fitting approximations. In this work, we show the power of targeting the density matrix, a linear-scaling sparse SE(3) equivariant matrix that encodes the exact density. We introduce Graph2Mat, a universal function for converting molecular graphs into equivariant matrices. We demonstrate how a machine learning model that combines this Graph2Mat approach with state of the art molecular graph representation learning architectures can predict density matrix of molecular systems accurately, such as it achieves state-of-the-art performance surpassing even grid-based methods, while using datasets that are at least one order of magnitude smaller. Accurately predicted electronic density can also accelerate Density Functional Theory (DFT) calculations by reducing the number of self consistent field (SCF) iterations needed to converge. In particular, we get a 40\% reduction on the number of SCF steps on DFT calculations of QM9 molecules with SIESTA. The novel prediction model also allows two new and promising measures of uncertainty (total charge error and self-consistency error). These results open the door for many applications using hybrid ML-accelerated DFT methodologies, and uncertainty aware single iteration ab initio molecular dynamics.