Improving Quantum Mechanical Calculations Using Graph Neural Networks to Predict Energies from Atomic Structure

Graph neural networks (GNN) offer a potential method of bypassing the Kohn-Sham equations in density functional theory (DFT) calculations by learning both the Hohenberg-Kohn (HK) mapping of electron density to energy, allowing for calculations of much larger atomic systems and time scales and enabling large-scale MD simulations with DFT-level accuracy. In this work, we investigate the feasibility of GNNs to learn the HK map from the external potential approximated as Gaussians to the electron density 𝑛(𝑟), and the mapping from 𝑛(𝑟) to the energy density 𝑒(𝑟) using Pytorch Geometric. We develop a graph representation for densities on radial grid points and determine that a k-nearest neighbor algorithm for determining node connections is an effective approach compared to a distance cutoff model, having an average graph size of 6.31 MB and 32.0 MB for datasets with 𝑘 = 10 and 𝑘 = 50 respectively. Furthermore, we develop two GNNs in Pytorch Geometric, and demonstrate a decrease in training losses for a 𝑛(𝑟) to 𝑒(𝑟) of 8.52 · 10^14 and 3.10 · 10^14 for 𝑘 = 10 and 𝑘 = 20 datasets respectively, suggesting the model could be further trained and optimized to learn the electron density to energy functional.