(360o) Electron Density Prediction with Graph Neural Networks on Large Catalyst Datasets

Large material datasets have recently been developed to build data-driven models for catalyst screening and discovery. Such datasets often contain a large swath of chemical systems. Graph neural networks (GNNs) have demonstrated their ability in accurately predicting molecular and material properties, notably energy and forces, but little work has been done to predict the electron density of these systems. Predicting the electron density is particularly challenging due to the large amount of data associated with a system in 3D space, O(100,000-1M+) grid points. In this work, we predict electron density using physically relevant descriptors which fit naturally into atom-centered GNNs. Namely, we fit a set of Gaussian multipole basis functions similar to those used in traditional computational chemistry calculations. This effectively reduces the electron density representation to a small number of node-level features in a graphical model. We also ensure this representation of electron density is equivariant to translations and rotations. These ideas are implemented as an extension to GemNet, a state of the art model for molecular and catalyst properties. We demonstrate our results on the Open Catalyst 2020 dataset (OC20), the largest catalyst dataset of its kind. Additionally, we show that including this additional physical information improves the overall performance of GemNet’s energy and force predictions for OC20.