Accurate Prediction of Tensorial Spectra Using Equivariant Graph Neural Network

TL;DR

This work introduces the Tensorial Spectra Equivariant Neural Network (TSENN), an E(3) equivariant GNN that predicts the full frequency-dependent dielectric tensor from crystal structure by decomposing it into $\ell=0$ and $\ell=2$ spherical tensor channels, thereby enforcing crystalline symmetry. Trained on $1{,}432$ first-principles dielectric tensors of nonmagnetic semiconductors, TSENN achieves a full-tensor MAE of $0.127$, with strong capture of anisotropy and peak shapes; the real part is reliably recovered via Kramers–Kronig relations, enabling comprehensive optical-property predictions. The method offers substantial data efficiency and orders-of-magnitude speedups over direct ab initio calculations, while preserving symmetry constraints, and holds promise for rapid screening and design of anisotropic optoelectronic materials and extensions to other tensorial properties such as piezoelectric and elastic responses.

Abstract

Optical spectroscopies provide a powerful tool for harnessing light-matter interactions for unraveling complex electronic features such as the flat bands and nontrivial topologies of materials. These insights are crucial for the development and optimization of optoelectronic devices, including solar cells, light-emitting diodes, and photodetectors, where device performance is closely connected with the nature of the underlying electronic spectrum. Realistic modeling of tensor optical responses in materials, which are computationally quite demanding, however, remains challenging. Here we introduce the Tensorial Spectra Equivariant Neural Network (TSENN), which is a equivariant graph neural network architecture that maps crystal structures directly to their full photon-frequency-dependent optical tensors. By encoding the isotropic sequential scalar components along with the anisotropic sequential tensor components into l = 0 and l = 2 spherical tensor components, TSENN ensures symmetry-aware predictions that are consistent with the constraints of crystalline symmetries of materials. Trained on a dataset of frequency-dependent permittivity tensors of 1,432 bulk semiconductors computed using first-principles methods, our model achieves a mean absolute error (MAE) of 21.181 millifarads per meter (mF/m), demonstrating its potential for efficient modeling of other related properties such as the optical conductivities. Our framework opens new avenues for rational data-driven design of anisotropic optical responses for accelerating materials discovery for advancing optoelectronic applications.

Figures (6)

• Figure 1: Overview of the diversity of materials in our dataset. (a) Histogram showing the frequency of materials across different crystal symmetries (b) Distribution of band gaps (eV). (c) Periodic table highlighting the frequency of elements in the dataset, with the color gradient ranging from blue (low frequency) to yellow (high frequency). Lanthanide and Actinide series are filtered out and not displayed.
• Figure 2: Dataset visualization. Left panels illustrate the frequency-dependent dielectric tensors in Cartesian coordinates, with insets detailing the corresponding crystal systems, tensor formulas, and the non-vanishing components after Ref. powell2010chapter3. Right panels depict the transformed spectra into the spherical harmonic basis, serving as the ground truth for model training. (a) Triclinic system with all tensor components present. (b) Hexagonal system with non-vanishing components satisfying $xx = yy \neq zz$, where the $Y^0_0$ component arises from the $xx + yy + zz$ component and the $Y^2_0$ arises from the $-\frac{1}{2}(xx + yy) + zz$ component.
• Figure 3: Schematic architecture of TSENN. The model takes a periodic crystal graph as input and encodes atom types via one-hot embeddings xieCrystalGraphConvolutional2018a. Edge vectors $\mathbf{r}_{ij}$ are expanded using radial basis functions and spherical--harmonics. Features pass through multiple gated equivariant convolution layers, followed by output convolution and mean pooling. Dropout is incorporated after each convolution block to prevent overfitting. The prediction is split into $\ell=0$ (isotropic) and $\ell=2$ (anisotropic) channels and trained using the composite loss function: $\mathrm{MAE}(\ell=0) + \mathrm{MAE}(\ell=2)$.
• Figure 4: (a) Distribution of full-tensor MAE values across the test set (left) and representative spectra obtained by stratified sampling across quartiles (right). From the cumulative kernel-density-estimator (KDE) plot, we randomly selected three systems in each quartile ($Q_1$--$Q_4$), resulting in 12 representative examples that span the full error distribution. [Cubic lattices were excluded, as their isotropic spectra are already well studied.] For each sampled material, the target and predicted tensor components are shown side by side, with the corresponding MAE reported in each panel. This stratified sampling strategy ensures that both the high- and low-error cases are represented, providing a balanced and unbiased view of model's performance across crystal systems and anisotropy strengths. (b) Real parts of the dielectric tensor components for the monoclinic system Se4Te2. Black curves give the ab initio dielectric spectra based on the IPA using the Kubo formula, while red curves denote spectra reconstructed from the predicted imaginary part via the K-K relation. $\text{MAE}^{\alpha\beta}$ is reported in each panel, with values ranging from 0.082 to 0.232, demonstrating strong agreement across all non-vanishing components ($\varepsilon^{xx}, \varepsilon^{yy}, \varepsilon^{zz}, \varepsilon^{xz}$). Notably, both the off-diagonal element ($\varepsilon^{xz}$) and the anisotropic diagonal terms are faithfully recovered.
• Figure 5: Comparison of predicted and target dielectric tensors in the Cartesian and spherical--harmonics basis in five different crystal systems, ordered by increasing symmetry. In the triclinic system, all tensor components are non-vanishing. The monoclinic system retains four non-vanishing components: $xx$, $yy$, $zz$, and $xz$. The orthorhombic system exhibits anisotropy with $xx \neq yy \neq zz$, while the hexagonal system satisfies $xx = yy \neq zz$. Finally, the cubic system represents the highest symmetry with $xx = yy = zz$. Left panel shows the pairwise comparison under the spherical tensor representation, while the right panel shows the comparison under the Cartesian tensor. Inset in the right panel provides the chemical formula, crystal system, and MAE.