🔬 materials science

MANDALA: An E(3)-Equivariant Graph Neural Network Framework for Learning Electronic-Structure Operators with Observable Guidance

Mandala is a modular, E(3)-equivariant graph neural network framework that bridges the gap between machine learning and electronic structure calculations by learning block-sparse quantum operators to directly predict electronic observables like band structures and densities of states, thereby complementing traditional machine learning interatomic potentials.

Original authors: Bartosz Brzoza, Wiktoria Szopa, Zakaria Elabid, Vincent Martinetto, Varadarajan Rengaraj, Mani Lokamani, Thomas D. Kühne, Attila Cangi

Published 2026-07-29

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Bartosz Brzoza, Wiktoria Szopa, Zakaria Elabid, Vincent Martinetto, Varadarajan Rengaraj, Mani Lokamani, Thomas D. Kühne, Attila Cangi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand how a massive, complex machine works, like a futuristic city or a giant biological cell. To do this, scientists use a powerful tool called "Density Functional Theory" (DFT). Think of DFT as a high-resolution camera that takes pictures of how electrons (the tiny, negatively charged particles that hold atoms together) move and arrange themselves. This camera is so good that it helps chemists and material scientists design new medicines, stronger batteries, and better solar panels. However, there's a catch: taking these pictures is incredibly slow and expensive. It's like trying to film a movie in 8K resolution on a smartphone that only has enough battery for ten seconds of recording. Because of this, scientists can only simulate very small systems or very short moments in time.

To speed things up, many researchers have started using "Machine Learning" to guess the answers. Usually, these AI models act like a weather forecaster: they look at the atoms and predict the "weather" (how much energy the system has and how hard the atoms are pushing on each other). This is fast and useful, but it's like knowing the temperature without knowing why it's hot or cold. You miss the underlying physics, like the wind patterns or the humidity. If you want to understand the "band structure" (which tells you if a material is a metal or an insulator) or how electricity flows through it, you need the full picture, not just the temperature. This is where the new paper comes in, offering a way to teach AI to see the whole movie, not just the weather report.

Meet Mandala: The AI That Learns the "Blueprints" of Matter

In this paper, a team of scientists introduces Mandala, a new software framework that acts like a master architect for the quantum world. While most AI models in materials science are content to guess the final score of a game (the energy), Mandala is designed to learn the actual rules and playbook (the mathematical operators) that govern how electrons behave.

Think of the electrons in a material as a giant, chaotic dance party. Traditional AI models just count how many people are dancing and how much energy the music is using. Mandala, however, learns the specific steps of the dance, the way partners connect, and the rhythm of the room. It does this by predicting three specific "blueprints" or matrices: the Hamiltonian (which describes the energy and rules of the dance), the Overlap (which describes how the dancers' personal space overlaps), and the Density (which shows where the dancers are crowded).

The Secret Sauce: Symmetry and Sparsity

Why is Mandala special? It uses two clever tricks to make the impossible possible.

First, it uses E(3)-equivariance. Imagine you are looking at a snowflake. If you rotate the snowflake, it looks exactly the same, just turned. A normal AI might get confused if you rotate the snowflake in the picture, thinking it's a new, different object. Mandala, however, is "symmetry-aware." It understands that the laws of physics don't change just because you turn your head. It learns the dance steps in a way that works no matter how you rotate or flip the atoms. This makes it incredibly efficient and accurate, like a dancer who knows the choreography so well they can perform it in any direction.

Second, it uses sparsity. In a giant city, you don't talk to everyone at once; you only interact with your neighbors. Similarly, electrons mostly interact with the atoms right next to them. Mandala realizes that most of the "blueprints" are empty (zero) because distant atoms don't affect each other. Instead of trying to memorize a massive, empty spreadsheet, it focuses only on the active connections. This is like a detective who only interviews the witnesses who were actually at the scene, ignoring the rest of the city, making the investigation much faster.

What Did They Find?

The authors tested Mandala on three very different types of materials to see if it could handle the chaos of the real world:

ZnCu2Sn(SeS)2: A complex material made of five different elements, like a soup with five distinct ingredients that need to be mixed perfectly.
Amorphous Silicon Oxide: Glassy, messy, and disordered, like a pile of sand that hasn't been sorted.
Crystalline Silicon: A neat, orderly grid, like a perfectly stacked brick wall.

In all these cases, Mandala successfully learned to predict the "blueprints" (the matrices) with high accuracy. But the real magic happened when they used those predictions to calculate things they didn't explicitly train for. Because Mandala learned the underlying rules, it could automatically generate:

Band Structures: The "roadmap" of energy levels that tells us if a material conducts electricity.
Density of States: A histogram showing how many electron "seats" are available at different energy levels.
Band Energies and Electron Counts: The total energy and number of electrons, derived directly from the predicted blueprints.

The paper shows that Mandala can do this without needing a separate AI for each task. It's a "one-stop shop" that learns the physics and then derives the answers.

The Limits and the Future

The authors are careful to note that Mandala isn't a magic wand that solves everything instantly. It works best when the atoms are close enough to interact within a certain distance (a "cutoff radius"), and it currently focuses on specific types of materials where the electrons are localized (stuck near atoms) rather than floating freely everywhere.

They also point out that while Mandala can predict the "band energy" (a specific part of the total energy), it does not yet calculate the entire total energy of the system. This means it cannot currently predict the forces that move atoms around in a simulation (like in a video game physics engine) with perfect accuracy. That is a goal for the future.

However, the results are promising. By learning the "operators" (the mathematical rules) instead of just the "observables" (the final numbers), Mandala bridges the gap between fast machine learning and deep quantum physics. It suggests that we can have our cake and eat it too: the speed of AI with the detailed, physics-rich insights of traditional quantum mechanics.

In short, Mandala is a new kind of toolkit that teaches computers to understand the language of electrons, not just the vocabulary. It allows scientists to simulate larger, more complex materials and see the hidden electronic details that were previously too expensive to calculate, opening the door to designing better materials for our future.

Technical Summary: Mandala – An E(3)-Equivariant Graph Neural Network Framework for Learning Electronic-Structure Operators

1. Problem Statement

Electronic-structure calculations based on Kohn-Sham density functional theory (DFT) are indispensable in computational materials science but are limited by high computational costs, restricting accessible system sizes and simulation times. While machine-learned interatomic potentials (MLIPs) have successfully accelerated simulations for energies and forces, they typically operate at the scalar or vector level. Consequently, conventional MLIPs fail to provide the operator-level information required to reconstruct quantum-mechanical observables such as band structures, densities of states (DOS), spatial charge distributions, and other electronic properties. There is a methodological gap for a scalable framework that learns the electronic-structure operators themselves (Hamiltonian, overlap, and density matrices) while retaining the ability to derive these complex observables directly from the learned operators.

2. Methodology

Mandala is a modular software framework designed to learn block-sparse electronic-structure matrices using E(3)-equivariant graph neural networks (GNNs). The methodology is built on the following pillars:

Operator Representation: Instead of treating electronic structure as dense global arrays, Mandala represents Hamiltonian ( $H$ ), overlap ( $S$ ), and density ( $D$ ) matrices as atom-pair-resolved sparse block tensors. This leverages the locality of electronic structure in localized-orbital bases.
E(3)-Equivariance: The framework utilizes E(3)-equivariant GNNs to respect the geometric transformation laws of 3D space (translations, rotations, and reflections). Hidden features are decomposed into irreducible representations (irreps) of the rotation group $O(3)$ , labeled by angular momentum $\ell$ and parity $p$ . This ensures that the learned operators transform correctly under symmetry operations without relying on handcrafted invariants.
Unified Workflow: The framework integrates several components into a single pipeline:
- Data Ingestion: Parsers for multiple backends (OpenMX, FHI-aims, PySCF) convert diverse outputs into a unified Snapshot format.
- Basis Conversion: A BlockIrrepMapper converts backend-specific orbital orderings and sign conventions into a common internal representation compatible with the e3nn library.
- Graph Construction: Periodic boundary conditions are handled by constructing edges with explicit lattice shift vectors, ensuring exact alignment between graph edges and sparse matrix blocks.
- Model Architecture: The network employs equivariant message passing, tensor products, and configurable readout heads to predict irrep vectors, which are then reconstructed into sparse operator blocks.
Observable Guidance: A key feature is the ability to use "observable-guided" training. Because the predicted operators are differentiable, quantities like band energy ( $E_b = \text{Tr}(DH)$ ) and electron count ( $N_e = \text{Tr}(DS)$ ) can be computed directly from the predicted matrices. These can be included as loss terms alongside matrix reconstruction errors. For periodic systems, spectral guidance is also supported, comparing generalized eigenvalue spectra derived from predicted operators against reference spectra on a k-point mesh.
Postprocessing: The framework includes physics-aware postprocessing steps, such as symmetrization of matrices, gauge alignment (overlap-shift correction) for Hamiltonians, and density normalization to ensure electron count consistency.

3. Key Contributions

Framework Design: Mandala provides a reusable, modular software stack that unifies sparse operator prediction, observable guidance, multi-backend data handling, and broad architecture exploration. It allows researchers to exchange data formats, physical targets, and model variants without rewriting the core code.
Operator-Level Learning: Unlike standard MLIPs that predict scalar energies, Mandala predicts the underlying quantum operators. This enables the direct reconstruction of band structures, DOS, and charge distributions from the learned model.
E(3)-Equivariant Implementation: The framework implements a rigorous E(3)-equivariant approach for block-sparse matrices, mapping atom-pair blocks to irreducible representations to ensure geometric consistency.
Flexible Training Objectives: The system supports pure matrix learning, mixed matrix-and-observable learning, and fully observable-guided training. It allows for the simultaneous prediction of Hamiltonian, overlap, and density matrices with shared latent representations.
Scalability and Extensibility: The framework is designed to support heterogeneous chemical compositions and various neural architecture variants within one workflow. It includes tools for systematic architecture search via configuration parameters.

4. Results

The paper demonstrates Mandala's capabilities on three distinct datasets:

ZnCu2Sn(SeS)2: A multicomponent chalcogenide with five elements and chemically distinct environments. The model successfully resolved heterogeneous local chemistry and coupling between environments.
Amorphous SiO2: A 150-atom amorphous glass configuration, demonstrating the framework's ability to handle disordered, non-periodic-like environments within a periodic framework.
Perturbed Silicon: Crystalline silicon with perturbations, used to validate performance on a standard periodic system.

Performance Metrics:

The models achieved low Mean Absolute Errors (MAE) for Hamiltonian, overlap, and density matrices on validation sets.
Band energies and electron counts were accurately predicted directly from the operators.
Density of States (DOS) and band structures were successfully reconstructed and matched reference DFT calculations.
Inference Scaling: Benchmarks on periodic silicon supercells (up to 4096 atoms) on a single NVIDIA B200 GPU showed that inference scales linearly with system size up to the tested limit. The primary bottlenecks identified were graph preparation and model-evaluation memory, not the neural network forward pass itself.

5. Significance and Claims

The paper positions Mandala as a complementary tool to atomistic interatomic potential workflows. Its primary significance lies in bridging the gap between efficient machine learning and the need for detailed electronic-structure analysis.

Complementarity: Mandala is intended to complement existing MLIP workflows by resolving electronic structure and operator-derived observables within a single scalable implementation, rather than replacing them.
Methodological Advancement: By learning operators directly, the framework retains a representation tied to quantum-mechanical principles, allowing for the analysis of phenomena such as charge transfer, orbital hybridization, and defect-localized states, which are inaccessible to scalar-only MLIPs.
Reproducibility and Extensibility: The framework emphasizes reproducibility through deterministic edge ordering, serialized configurations, and modular design, making it a robust platform for future methodological experimentation in scientific machine learning.

The authors note current limitations, including the restriction to real-valued, non-spin-polarized matrices, the reliance on localized orbital bases, and the fact that the band energy observable does not yet include all contributions to the total Kohn-Sham energy (limiting the direct use of derivatives for total-energy forces and stresses). However, the framework is designed to accommodate these extensions in future development.