Covariant Jacobi-Legendre expansion for total energy calculations within the projector augmented wave formalism

Machine-learning models can be trained to predict the converged electron charge density of a density functional theory (DFT) calculation. In general, the value of the density at a given point in space is invariant under global translations and rotations having that point as a center. Hence, one can construct locally invariant machine-learning density predictors. However, the widely used projector augmented wave (PAW) implementation of DFT requires the evaluation of the one-center augmentation contributions that are not rotationally invariant. Building on our recently proposed Jacobi-Legendre charge-density scheme, we construct a covariant Jacobi-Legendre model capable of predicting the local occupancies needed to compose the augmentation charge density. Our formalism is then applied to the prediction of the energy barrier for the 1H-to-1T phase transition of two-dimensional MoS2. With extremely modest training, the model is capable of performing a non-self-consistent nudged elastic band calculation at virtually the same accuracy as a fully DFT-converged one, thus saving thousands of self-consistent DFT steps. Furthermore, at variance with machine-learning force fields, the charge density is here available for any nudged elastic band image, so that we can trace the evolution of the electronic structure across the phase transition.