ELECTRA: A Cartesian Network for 3D Charge Density Prediction with Floating Orbitals

We present the Electronic Tensor Reconstruction Algorithm (ELECTRA) - an equivariant model for predicting electronic charge densities using floating orbitals. Floating orbitals are a long-standing concept in the quantum chemistry community that promises more compact and accurate representations by placing orbitals freely in space, as opposed to centering all orbitals at the position of atoms. Finding the ideal placement of these orbitals requires extensive domain knowledge, though, which thus far has prevented widespread adoption. We solve this in a data-driven manner by training a Cartesian tensor network to predict the orbital positions along with orbital coefficients. This is made possible through a symmetry-breaking mechanism that is used to learn position displacements with lower symmetry than the input molecule while preserving the rotation equivariance of the charge density itself. Inspired by recent successes of Gaussian Splatting in representing densities in space, we are using Gaussian orbitals and predicting their weights and covariance matrices. Our method achieves a state-of-the-art balance between computational efficiency and predictive accuracy on established benchmarks. Furthermore, ELECTRA is able to lower the compute time required to arrive at converged DFT solutions - initializing calculations using our predicted densities yields an average 50.72 % reduction in self-consistent field (SCF) iterations on unseen molecules.

💡 Research Summary

The paper introduces ELECTRA (Electronic Tensor Reconstruction Algorithm), a novel equivariant deep‑learning framework for predicting three‑dimensional electronic charge densities directly from molecular geometry. Traditional approaches either expand the density in atom‑centered basis functions (LCAO) or predict scalar values on a dense real‑space grid. Atom‑centered bases require large, diffuse, high‑angular‑momentum functions to capture long‑range or highly anisotropic electron distributions, leading to high computational cost. Grid‑based methods, while expressive, suffer from massive memory and runtime demands because a single molecule can contain hundreds of thousands of grid points.

ELECTRA tackles these limitations by employing “floating orbitals” – Gaussian functions whose centers are not constrained to atomic positions but can be placed anywhere in space. Historically, floating orbitals have been used by quantum‑chemistry practitioners to improve basis‑set efficiency, but their optimal placement has required expert intuition. The authors replace this manual step with a data‑driven solution: a neural network predicts both the positions (means µ) and the parameters (weights w, covariance matrices Σ) of a set of Gaussian functions for each atom, thereby constructing a Gaussian mixture model of the electron density: \