Learning the density matrix, a symmetry rich encoding of the electronic density.

In the most recent years, the electronic density has been getting attention as a target for machine learning (ML) models due to the amount of information it contains. In fact, density functional theory (DFT) proposes that all ground state electronic properties of an atomic system should be inferable from it. The last advances in ML interatomic potentials have shown that taking into account the equivariance of the data (e.g. forces should rotate when the system is rotated) greatly enhances the learning capacity while needing less data to train. In this context, equivariant models that predict the electronic density have quickly appeared. These models predict scalar values on a real space grid or coefficients for a density fitting expansion. DFT codes with atom-centered basis sets, however, compute the electronic density by products of orbitals. The coefficients of these products follow the equivariance of products of spherical harmonics, which is of higher order than the target values for the previous approaches. In our work, we target the density matrix, which contains these coefficients. By doing so, we force the model to learn more meaningful details about the atomic interactions. The computation of the density matrix scales linearly with system size and the representation is more compressed than that of a 3D grid. In this talk, we present the architecture of our models, as well as the results obtained in some common benchmarks, which are very similar to the state of the art grid-based predictions.