An Equivariant Model for Quantum Chemical Tensor Learning

We present a graph-based, SE(3)-equivariant model designed for general tensor prediction tasks in both DFT and correlated wavefunction methods. The model maps molecular graphs containing elemental and geometric information to quantum mechanical tensors in a fully equivariant manner. When available, low-cost quantum mechanical inputs can be incorporated to reduce model size and enhance accuracy. We demonstrate the model's high precision and data efficiency across diverse applications, including the prediction of (i) mean-field one-electron density matrices (1RDMs) at the KS-DFT and HF levels, (ii) correlated 1RDMs at the MP2 level, and (iii) wavefunction amplitudes from linear-response TDDFT. We illustrate the use of these learned tensors in downstream tasks such as accelerating SCF convergence, enabling reduced-cost coupled cluster calculations via frozen natural orbitals, and predicting neutral excitation energies.