Higher Order Equivariant Graph Neural Networks for Charge Density Prediction

The calculation of electron density distribution in materials and molecules is central to the study of their quantum and macro-scale properties, yet accurate and efficient calculation remains a long-standing challenge in the field of material science. This work introduces ChargE3Net, an E(3)-equivariant graph neural network for predicting electron density in atomic systems. Unlike existing methods, ChargE3Net achieves equivariance through the use of higher-order tensor representations, and directly predicts the charge density at a set of desired locations. We demonstrate the effectiveness of ChargE3Net on large and diverse sets of molecules and materials, where it achieves state-of-the-art performance over existing methods, and scales to larger systems than what is feasible to compute with density functional theory. Through additional experimentation, we demonstrate the effect of introducing higher-order equivariant representations, and why they yield performance improvements in the charge density prediction setting.