Graph2Mat: Universal graph to matrix conversion for electron density prediction

The 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 representations can accurately predict the density matrix of molecular systems. The models achieve state-of-the-art performance on electron density prediction by matching the accuracy of grid-based methods, while using datasets that are at least one order of magnitude smaller. Accurately predicted electron densities can also accelerate Density Functional Theory (DFT) calculations by reducing the number of self-consistent field (SCF) iterations needed to converge. In this work, we get an average 40% reduction on the number of SCF steps in DFT calculations of QM9 molecules with SIESTA. The novel prediction model also allows for two new and promising measures of uncertainty (total charge error and self-consistency error) that will facilitate its practical usage, e.g. within active learning workflows. These results open the door for many applications using hybrid ML-accelerated DFT methodologies, and uncertainty aware single iteration ab initio molecular dynamics.