UEIPNet: Equivariant GNN for Electronic Modeling

• UEIPNet is an equivariant graph neural network that predicts interatomic potentials and tight-binding Hamiltonians while maintaining physical consistency across mechanical and electronic domains.
• It employs e3nn-based equivariant convolutions and tensor operations to capture complex atomic interactions in materials like bilayer graphene and monolayer MoS₂.
• Trained on DFT and Wannier projection data, UEIPNet achieves near-DFT accuracy in modeling energies, forces, band structures, and strain-dependent phenomena.

UEIPNet is an equivariant graph neural network (GNN) formulated to simultaneously predict interatomic potentials and tight-binding (TB) Hamiltonians with explicit deformation–electronic coupling in atomic-scale systems. Trained on density functional theory (DFT) reference data and subsequent Wannier projection, UEIPNet produces energies and forces (as node-level targets), as well as Wannier-projected TB matrices (as edge-level targets). This dual-task capability enables unified and physically consistent modeling of both mechanical and electronic responses—achieving accuracy near that of DFT—for large-scale materials systems such as bilayer graphene (BLG) and monolayer MoS₂. UEIPNet's modular network architecture utilizes e3nn-based equivariant convolutions, permitting generalization and scalability across diverse atomic environments. Its successful demonstrations include capturing flat-band formation and strain-dependent electronic phenomena, establishing UEIPNet as a bridge between atomistic simulation and electronic-structure calculations.

1. Network Architecture and Design Principles

UEIPNet is constructed using an equivariant GNN framework with the e3nn library, adhering to rotational and inversion symmetries crucial for condensed-matter modeling. Atomic configurations are encoded as graphs, where each atom forms a node and edges connect atom pairs within a species-dependent cutoff, adjusted for periodicity. Initial node features comprise atom types, represented via one-hot vectors spanning the periodic table, and geometric information. Radial dependencies are expanded by Gaussian basis functions; angular dependencies utilize real spherical harmonics $Y_{\ell m}$, ensuring equivariance.

The network features four principal modules:

• Node Update: EquiConv operations exchange information among nodes, leveraging tensor products and gating based on atomic type and geometric orientation.
• Edge Update: Similar tensor and gating operations process paired node features to create edge representations, essential for electronic couplings.
• Interatomic Potential (IP) Module: Uses node features to predict site energies; differentiating these outputs yields atomic forces.
• Tight-Binding (TB) Module: Translates edge features into local TB Hamiltonian matrix elements via E3Linear and E3FullyCon layers, applying Clebsch–Gordan decompositions for orbital-resolved coupling ($N \times M$ for $N$ and $M$ orbitals per atom).

The complete workflow is rooted in training on DFT-computed total energies and atomic forces for node-level supervision, and Wannier90-projected Kohn–Sham Hamiltonians for edge-level TB element supervision. The TB matrix elements are specifically defined as $$e_{mn}(\mathbf{R}) = \langle w_{m0} | \hat{H} | w_{n\mathbf{R}} \rangle$$, with $w_{m0}$ and $w_{n\mathbf{R}}$ denoting Wannier functions localized on orbitals $m$ and $n$ at lattice vectors $0$ and $\mathbf{R}$.

2. Data Sources and Training Procedures

UEIPNet leverages a two-stage reference workflow:

1. Density Functional Theory (DFT): Atomic structures receive DFT evaluations to produce total energies and forces (nodes), and the Kohn–Sham Hamiltonian.
2. Wannier Projection: The Hamiltonian is projected to a localized Wannier basis using Wannier90, yielding sparse TB matrices encoding short- and long-range electronic couplings as required for edge targets.

During training, the network optimizes a multitask loss comprising both energy/force (node-level) and TB matrix element (edge-level) terms, yielding outputs consistent with underlying quantum mechanics. The methodology avoids sequential relaxation/electronic-structure steps and maintains physical consistency across the mechanical and electronic domains.

3. Functionality: Prediction of Potentials and Hamiltonians

UEIPNet is explicitly designed to perform multitask prediction in a single forward pass. The predicted interatomic potential provides accurate site energies and atomic forces, governing equilibrium geometry, structural relaxation, and elastic moduli. Simultaneously, edge-level TB predictions encode the quantum mechanical interactions requisite for simulating electronic band structures and local electronic phenomena.

Significance of the dual-target strategy includes:

• Capturing equilibrium and perturbed bonding geometries underlying mechanical properties.
• Accurately reproducing DFT-level band structure—including effects of atomic displacements, strain, and relaxation—without post hoc parameterization of TB models.

Notable system-level demonstrations include:

• Twisted Bilayer Graphene (TBG): UEIPNet identifies the coupling mechanisms—interlayer spacing, in-plane strain, and out-of-plane corrugation—that facilitate the formation of isolated flat bands, key to correlated electronic phenomena. The model further shows that modulating substrate interaction strengths can induce flat bands away from the canonical magic angle.
• Monolayer MoS₂: The network recreates the phonon dispersion curves, strain-dependent band-gap evolution, and local density of states (LDOS) fluctuations under both uniform and non-uniform strains, matching the complexity observed in direct DFT calculations.

4. Evaluation of Accuracy and Physical Fidelity

UEIPNet demonstrates near-DFT accuracy across multiple benchmarks:

• In BLG, total energy profiles, generalized stacking-fault energies (GSFE), and elastic constants computed by UEIPNet align closely with DFT reference calculations.
• Electronic band structures reconstructed from predicted TB Hamiltonians replicate high-fidelity DFT results, including for systems subject to random atomic displacements.
• In MoS₂, computed phonon dispersions and the evolution of electronic band gaps as a function of applied strain track direct DFT outputs across a range of tensor perturbations.

These results affirm the network's ability to simultaneously reproduce mechanical and electronic properties at DFT-like accuracy, from fundamental band structure features to responses under external perturbation.

5. Scalability and Generalization

UEIPNet’s graph-based representation and multitask objectives enable scalability to large atomistic systems. Networks trained on modestly sized configurations generalize to extensive supercells, e.g., TBG with thousands of atoms, without loss of physical fidelity. This supports deployment for studying local variations—such as LDOS modulations in non-uniformly strained materials—beyond the reach of conventional ab initio methods.

The architecture’s equivariant design also permits adaptation to broader material classes, including multilayer heterostructures and bulk configurations with arbitrary stacking and symmetry. The approach bridges classical atomistic simulation and electronic-structure theory, permitting real-time and physically consistent exploration of deformation–electronic interplay.

6. Implications and Future Directions

The ability of UEIPNet to model strain-tunable electronic states, flat-band engineering, and atomistic-to-electronic coupling in 2D materials has substantive implications in condensed matter and material sciences. Its unified framework facilitates rapid evaluation of strain effects on superconductivity, optoelectronic properties, and correlated electronic phenomena.

Potential future avenues include:

• Extending UEIPNet to multilayer and bulk material systems with complex stacking orders.
• Integrating UEIPNet within large-scale molecular dynamics, enabling dynamic paper of phonon–electron interactions, non-equilibrium processes, and temperature-dependent properties.
• Refining substrate interaction modeling and symmetry considerations to encompass additional degrees of freedom—such as magnetic or charge effects—supporting richer electronic phase explorations.

A plausible implication is the accelerated discovery and engineering of materials with tailored electronic responses through direct simulation of atomistic deformations and their quantum mechanical consequences.

7. Contextualization and Prospective Outlook

UEIPNet’s approach—combining equivariant GNNs, DFT data, and Wannier-tight-binding representations—represents a unified and scalable advancement in atomistic–electronic modeling. This synthesis overcomes the dichotomy between classical and quantum simulation regimes, supplying a predictive framework suitable for both fundamental research and technological innovation in two-dimensional and layered materials. The methodology is positioned to drive studies in strain engineering, electronic phase manipulation, and device-level material integration, bridging computational efficiency with quantum-accurate physical modeling.