A Recipe for Charge Density Prediction

Xiang Fu,Andrew Rosen,Kyle Bystrom,Rui Wang,Albert Musaelian,Boris Kozinsky,Tess Smidt,Tommi Jaakkola
2024-05-30
Abstract:In density functional theory, charge density is the core attribute of atomic systems from which all chemical properties can be derived. Machine learning methods are promising in significantly accelerating charge density prediction, yet existing approaches either lack accuracy or scalability. We propose a recipe that can achieve both. In particular, we identify three key ingredients: (1) representing the charge density with atomic and virtual orbitals (spherical fields centered at atom/virtual coordinates); (2) using expressive and learnable orbital basis sets (basis function for the spherical fields); and (3) using high-capacity equivariant neural network architecture. Our method achieves state-of-the-art accuracy while being more than an order of magnitude faster than existing methods. Furthermore, our method enables flexible efficiency-accuracy trade-offs by adjusting the model/basis sizes.
Computational Physics,Machine Learning
What problem does this paper attempt to address?
The paper discusses the problem of predicting charge density in density functional theory (DFT). Charge density is a fundamental property of atomic systems and can be used to derive all chemical properties. Existing machine learning methods show potential in accelerating charge density prediction, but lack accuracy or scalability. The paper proposes a new approach that simultaneously achieves accuracy and scalability through three key elements: 1) using atomic and virtual orbitals to represent charge density; 2) utilizing a learnable and expressive orbital basis set; 3) employing a high-capacity equivariant neural network architecture. This approach maintains high accuracy while being an order of magnitude faster than existing methods, and allows for flexible trade-offs between efficiency and accuracy by adjusting model and/or basis set size. The main goal of the paper is to effectively predict charge density using machine learning models, in order to reduce the self-consistent field convergence steps in DFT and accelerate the discovery of molecules and materials.