StrainTensorNet: Predicting crystal structure elastic properties using
  SE(3)-equivariant graph neural networks
Accurately predicting the elastic properties of crystalline solids is vital
for computational materials science. However, traditional atomistic scale ab
initio approaches are computationally intensive, especially for studying
complex materials with a large number of atoms in a unit cell. We introduce a
novel data-driven approach to efficiently predict the elastic properties of
crystal structures using SE(3)-equivariant graph neural networks (GNNs). This
approach yields important scalar elastic moduli with the accuracy comparable to
recent data-driven studies. Importantly, our symmetry-aware GNNs model also
enables the prediction of the strain energy density (SED) and the associated
elastic constants, the fundamental tensorial quantities that are significantly
influenced by a material's crystallographic group. The model consistently
distinguishes independent elements of SED tensors, in accordance with the
symmetry of the crystal structures. Finally, our deep learning model possesses
meaningful latent features, offering an interpretable prediction of the elastic
properties.
Teerachote Pakornchote, Annop Ektarawong, Thiparat Chotibut
StrainNet: Predicting crystal structure elastic properties using SE(3)-equivariant graph neural networks

###### Abstract

Accurately predicting the elastic properties of crystalline solids is vital for computational materials science. However, traditional atomistic scale _ab initio_ approaches are computationally intensive, especially for studying complex materials with a large number of atoms in a unit cell. We introduce a novel data-driven approach to efficiently predict the elastic properties of crystal structures using SE(3)-equivariant graph neural networks (GNNs). This approach yields important scalar elastic moduli with the accuracy comparable to recent data-driven studies. Importantly, our symmetry-aware GNNs model also enables the prediction of the strain energy density (SED) and the associated elastic constants, the fundamental tensorial quantities that are significantly influenced by a material's crystallographic group. The model consistently distinguishes independent elements of SED tensors, in accordance with the symmetry of the crystal structures. Finally, our deep learning model possesses meaningful latent features, offering an interpretable prediction of the elastic properties.

**Keywords:** strain energy tensor, elastic constants, crystallographic group, equivariant neural networks, density functional theory

**Open data:** The dataset used in this work has been made available at

[https://github.com/trachote/predict_elastic_tensor](https://github.com/trachote/predict_elastic_tensor)

## I Introduction

The elastic properties of crystalline solids such as elastic constants, bulk modulus, and shear modulus are important macroscopic quantities that determine materials' mechanical characteristics. Computational study of which can provide theoretical guidelines to understand various phenomena in solid materials, e.g., mechanical stability of crystal structures [1; 2], pressure-driven deformation and phase transition of materials [3; 4], the response of materials to sound wave propagation [5; 6], and the hardness of materials [7; 8; 9], to name a few. From an atomistic scale description, an _ab initio_ approach based on density functional theory has been employed to investigate the macroscopic elastic properties of materials, yielding, for example, elastic constants and elastic moduli that are in agreement with experimental results [10; 11; 12; 13].

Three atomistic scale computational methods are usually adopted to calculate the elastic constants of crystal structures: an energy-based method [14], a stress-strain method [15; 10], and a calculation from group velocity [16; 17; 18]. The energy-based and the stress-strain methods are more standard, and derive the elastic constants using respectively an energy and a stress tensor acquired from an _ab initio_ calculation. Although both methods yield a comparable prediction of elastic constants, the energy-based method is computationally less efficient, involving a larger number of plane-wave basis set and \(k\)-point meshes so that the predicted elastic constants converge to reliable values [19]. Indeed, the more efficient stress-strain method is commonly employed to determine elastic constants, for example, by Materials Project database, and by established software such as VASP and CASTEP [20; 21; 22].

Despite forming a basis for computational materials science, atomistic scale simulation can be computationally prohibitive, especially when the number of atoms in a unit cell grows large. Such constraint limits an _ab initio_ method's capability to investigate more complex crystalline solids. On the other hand, advances in machine learning bring about alternative approaches to computational materials science. This data-driven paradigm can reasonably predict the elastic properties of crystal structures, provided sufficient training data from an _ab initio_ method [23; 24]. Even when the number of atoms in a unit cell is large, such approach can efficiently predict the bulk and shear moduli of complex materials such as alloys [25; 26]. Applying machine learning together with evolutionary algorithms and _ab initio_ calculations is also potentially useful in searching for novel superhard materials [27; 28].

Machine learning models based on graph neural networks (GNNs) have received increasing attention in studying solid materials and molecules. With GNNs, it is natural to encode atomistic scaled descriptions of solids into a computational framework; atomic locations and the associated atomic attributes can be directly embedded into node attributes of a graph, whereas pairwise interactions among atoms can be encoded into the edge attributes. Efficient GNNs training procedures have also been proposed [29; 30; 31], enabling GNNs to learn the representation of a complex relationship between the input and its associated prediction [32]. These neural networks can also be endowed with translation-rotation equivariant property, so that input atomic locations of molecules
(or point clouds) in \(\mathbb{R}^{3}\) that differ only in their orientation or their centroid can be identified. Enforcing an SE(3) equivariant property helps the networks to extract a more compact (translation and rotation independent) relationship between the inputs and their predictions. Due to these appealing features, variations of SE(3) equivariant GNNs have been developed to study materials science [33].

In this work, we adopt a data-driven approach using graph neural networks (GNNs) to predict the elastic properties of crystal structures. Accounting for the symmetry of crystalline solids, our SE(3) equivariant GNNs take as an input atomistic descriptions of strained materials, such as atomic locations and their associated atomic attributes, and predict the strain energy density (SED) of the materials. The prediction of the SED, which is the energy stored in a crystal structure when distorted by either tensile or compressive strains in particular directions, can be obtained efficiently and relatively accurately, given sufficient training data. The model thus provides an alternative approach to the standard _ab initio_ energy-based prediction method.

After the SED is computed, we can then calculate the elastic constants and construct the elastic tensor with a simple analytical expression, see Sec. II.1. Other macroscopic (scalar) elastic properties including bulk modulus, shear modulus, Young's modulus, and Poisson's ratio immediately follow from the elastic constants. Sec. IV reports our model prediction results of these elastic properties. The prediction of the scalar elastic properties are comparable to those of the recent data-driven work [23; 24]. Importantly, however, the data-driven prediction of the SED and the associated elastic constants, which are fundamental tensorial quantities that depend on the crystallographic group, are first reported here. The trained model consistently reveal the independent components of strain energy tensors dictated by the symmetry of the crystal structures (see Fig. 3). Our symmetry-aware GNNs architecture as well as the dataset used to train the model are provided in Sec. III. Also, as opposed to a black-box deep learning prediction, we show in Sec. IV.3 that the latent features (data representations) that the model learned are physically meaningful, rendering our GNNs explainable. We conclude this work in Sec. V, and provide supplementary results and necessary derivations in the Supplementary Materials.

## II Linear elasticity background

### The strain energy tensor and the elastic tensor

Assuming linear elastic deformation and no external stress, the strain energy density (SED) of isotropic materials can be expressed as [34]

\[U(\epsilon_{1},\epsilon_{2},...,\epsilon_{6})=\frac{1}{2}\sum_{i,j=1}^{6}C_{ij }\epsilon_{i}\epsilon_{j}, \tag{1}\]

where \(C_{ij}\) is an _elastic constant_ and \(\epsilon_{i}\) is a strain component \(i\in\{1,2,\ldots,6\}\) in Voigt notation. If the crystal structure is strained by at most two strain components (two Voigt indices), one can consider the SED tensor of rank 2 which we denote \(U_{ij}\equiv U(\epsilon_{i},\epsilon_{j})\) for \(i\neq j\) and \(U_{ii}\equiv U(\epsilon_{i})\). In this study, we consider \(U_{ij}\) in units of eV per atom. The SED is the energy stored in distorted crystal structure when the lattice is strained by \(\epsilon_{i}\) and \(\epsilon_{j}\). Since the elastic constant tensor is symmetric, the SED tensor is also symmetric and can be expressed as

\[U_{ij} =C_{ij}\epsilon_{i}\epsilon_{j}+\frac{1}{2}\left(C_{ii}\epsilon_{ i}^{2}+C_{jj}\epsilon_{j}^{2}\right)\text{ for }i\neq j\text{ and,} \tag{2}\] \[U_{ii} =\frac{1}{2}C_{ii}\epsilon_{i}^{2}.\]

The elastic constants can then be analytically obtained as a function of the SED:

\[C_{ij}=\frac{1}{\epsilon_{i}\epsilon_{j}}\left[\left(1+\delta_{ij}\right)U_{ ij}-\left(1-\delta_{ij}\right)\left(U_{ii}+U_{jj}\right)\right], \tag{3}\]

where \(\delta_{ij}\) is the Kronecker delta.

Because the SED tensor and the elastic tensor are symmetric tensors of rank 2 with \(6\times 6\) components, there are \((6+1)6/2=21\) independent components. In the matrix form, we denote the SED tensor as

\[\mathbf{U}=\begin{bmatrix}U_{11}&U_{12}&U_{13}&U_{14}&U_{15}&U_{16}\\ &U_{22}&U_{23}&U_{24}&U_{25}&U_{26}\\ &&U_{33}&U_{34}&U_{35}&U_{36}\\ &&&U_{44}&U_{45}&U_{46}\\ &&&U_{55}&U_{56}\\ &&&U_{66}\end{bmatrix}, \tag{4}\]

where only the upper triangular elements are shown (the lower diagonal elements are left blank for brevity). For crystalline solids, however, the upper triangular elements are not completely independent, depending on the symmetry of the crystal structure. For instance, the SED tensor of a cubic lattice is

\[\mathbf{U}_{cubic}=\begin{bmatrix}U_{11}&U_{12}&U_{12}&U_{14}&U_{14}&U_{14}\\ &U_{11}&U_{12}&U_{14}&U_{14}&U_{14}\\ &&U_{11}&U_{14}&U_{14}&U_{14}\\ &&&U_{44}&U_{45}&U_{45}\\ &&&U_{44}&U_{45}\\ &&&&U_{44}\end{bmatrix}. \tag{5}\]

Note that some components of the elastic constants \(C_{ij}\) can be zero in many crystal structures, e.g. a cubic lattice has 9 zeros out of 21 independent components.
However, due to the property of SED from Eq. (2), \(U_{ij}\) is never 0. For the purpose of machine learning regression, working with SED helps avoid zero inflation problems in the training dataset. Fig. S2 shows the distributions of \(C_{ij}\) and \(U_{ij}\), which cover different ranges depending on the indices. Importantly, the distribution of elastic constants can concentrate around zero, but this is not the case for the SED. As an illustrative example of how the SED can avoid a zero inflation problem, consider the elastic constants of a diamond cubic crystal structure. For diamond, \(C_{11}\), \(C_{12}\), \(C_{14}\), \(C_{44}\), and \(C_{45}\) are 1054, 126, 0, 562, and 0 GPa, respectively. Then, Eq. (2) gives \(U_{11}\), \(U_{12}\), \(U_{14}\), \(U_{44}\), and \(U_{45}\) subjected to 2% strain to be 7.506, 16.807, 11.509, 4.002, and 8.005 meV/atom, respectively. The magnitude of \(U_{12}\) is the largest as it is the sum of \(C_{12}\) and \(C_{11}\), while \(U_{14}\) and \(U_{45}\) are smaller (but non-zero) because \(C_{14}\) and \(C_{45}\) are zero.

### Strain operations

The unit cell of a crystal can be compactly parametrized by a lattice matrix \(\mathbf{L}=[\mathbf{a}\ \mathbf{b}\ \mathbf{c}]\), where \(\mathbf{a}\), \(\mathbf{b}\), and \(\mathbf{c}\) are vectors of lattice parameters in the Cartesian coordinate. If the system is applied by strain, the lattice matrix will be deformed by

\[\mathbf{L}^{\prime}(\epsilon_{1},\epsilon_{2},...,\epsilon_{6})=\boldsymbol{ \varepsilon}_{I}\mathbf{L}, \tag{6}\]

where \(\boldsymbol{\varepsilon}_{I}\) is the strain matrix

\[\boldsymbol{\varepsilon}_{I}=\begin{bmatrix}\epsilon_{1}&\frac{\epsilon_{6}} {2}&\frac{\epsilon_{5}}{2}\\ \frac{\epsilon_{6}}{2}&\epsilon_{2}&\frac{\epsilon_{2}}{2}\\ \frac{\epsilon_{6}}{2}&\frac{\epsilon_{6}}{2}&\epsilon_{3}\end{bmatrix}+ \mathbf{I}. \tag{7}\]

In this work, we assume the crystal structures will be deformed by at most two strain components, so that \(\mathbf{L}^{\prime}=\mathbf{L}^{\prime}(\epsilon_{i},\epsilon_{j})\).

Due to an applied strain, the atomic coordinates \(\mathbf{r}\) must also be transformed accordingly as

\[\mathbf{r}^{\prime}=\boldsymbol{\varepsilon}_{I}\mathbf{L}\mathbf{r}_{f}, \tag{8}\]

where \(\mathbf{r}_{f}\) is a fractional coordinate of an atom in an unstrained lattice. Noting that \(\mathbf{r}_{f}\) typically shifts from its equilibrium value when the strain is applied (say, in density function theory or in experiments); however, the atomic relaxation is neglected for simplicity.

## III Machine learning with SE(3)-equivariant graph neural networks

### Crystal graphs

A crystal structure can be represented as a multigraph whose nodes and edges, respectively, encode atoms and their pairwise connections. A pair of nodes describing an atom of type \(m\) and an atom of type \(n\) can be connected by multiple edges, encapsulating the interactions between the atom of type \(m\) in a unit cell and atoms of type \(n\) in the unit cell as well as in other periodic cells of consideration, see Fig 1. An atom of type \(m\) can interact with an atom of the same type in periodic cells if the interaction range of consideration is large enough, represented by self-loops in the node \(m\) of the crystal graph. Each atom's location and its atomic features (e.g., atomic mass, electronegativity, polarizability, atomic radius, and etc.) in a unit cell is accounted for by the attributes of a single node.

More formally, a crystal graph \(\mathcal{G}=(\mathcal{V},\mathcal{E})\) consists of a set of nodes (vertices) \(\mathcal{V}\) and a set of edges \(\mathcal{E}\) defined as

\[\mathcal{V}=\{(\mathbf{f}_{n},\mathbf{r}_{n})\mid\mathbf{f}_{n} \in\mathbb{R}^{M},\ \mathbf{r}_{n}=\mathbf{L}\mathbf{r}_{f_{n}}\in\mathbb{R}^{3}\},\] \[\mathcal{E}=\{\Delta\mathbf{r}_{mn}^{(\mathbf{T})}\mid\Delta \mathbf{r}_{mn}^{(\mathbf{T})}=\mathbf{r}_{m}-\mathbf{r}_{n}+\mathbf{T};\ \mathbf{r}_{m},\mathbf{r}_{n}\in\mathbb{R}^{3}\},\]

where \(m\) and \(n\) are indices of the two atoms in the unit cell, \(\mathbf{f}_{n}\) is a vectorized atomic feature of an atom of type \(n\), \(M\) is the number of atomic features, \(\mathbf{r}_{f_{n}}\) is the fractional coordinate of an atom of type \(n\), and \(\mathbf{T}\) is a translation vector within the space of interest. Each node embeds two kinds of information: atomic attributes \(\mathbf{f}\) and atomic positions \(\mathbf{r}\) in the unit cell. From this definition, the unit cell and its periodic images is compactly represented as a multigraph; only a finite number of nodes describing the atoms in the unit cell is required, and the number of edges joining the two nodes grows linearly with the number of periodic images of the atoms in the structure of interest (see Fig. 1(a)).

A crystal graph describing the whole crystal structure will require infinitely many edges information. In practice, one typically embeds only partial information of the structure using the description of the unit cell and a finite number of edges encoding interactions with neighbor atoms [35]. One method to construct the edges is to connect a target atom to other atoms within a finite radius; however, still, this method generates an excessive number of edges, which is computationally infeasible for GNNs machine learning. Alternatively, some algorithms determine the connection between atoms using, e.g., the Voronoi tessellation, generating a moderate number of edges [36; 37]. However, for layered materials, using Voronoi tessellation causes the edges connecting atoms between layers to be absent. These interlayer connections are crucial in differentiating the structure strained in out-of-plane direction from the structures strained in other directions. In this work, we define an edge between a pair of atoms \(\Delta\mathbf{r}\) (dropping the translation vector superscript and the atom indices subscript for brevity) to be non-zero only if each component \(\Delta r_{\alpha}\) where \(\alpha\in\{x,y,z\}\) satisfies the following spatial cutoff criterion: \(|\Delta r_{\alpha}|\leq\min(|a_{\alpha}|+|b_{\alpha}|+|c_{\alpha}|,s)+\delta\) where \(a_{\alpha}\), \(b_{\alpha}\), and \(c_{\alpha}\) are the \(\alpha\) components of the corresponding lattice parameters constituting the lattice matrix \(\mathbf{L}\), \(s\) is a cutoff distance, and \(\delta\) is a small cutoff extension of our choice. This can reduce the number of edges from
the hard radial cutoff criterion and ensures the connections of atoms between layers for layered materials by an appropriate choice of \(s\) and \(\delta\).

### Rotational and translational equivariance

If two crystal graph inputs constructed from an identical crystal graph strained by two different sets of strained components \((\epsilon_{i},\epsilon_{j})\) and \((\epsilon_{k},\epsilon_{l})\) yielding the new vertices \(\mathcal{V}^{\prime}=\{(\mathbf{f}_{n},\mathbf{r}^{\prime}_{n})\}\) and \(\mathcal{V}^{\prime\prime}=\{(\mathbf{f}_{n},\mathbf{r}^{\prime\prime}_{n})\}\), respectively, are equivalent up to a 3D rotation and a translation, then we expect our machine learning model to predict the same SED. This symmetry-aware property can be implemented with geometric deep learning models [38]. We will equip a latent representation of our model with this symmetry-aware (equivariant) property, which is defined as follow.

A latent representation \(\phi:\mathcal{V}\rightarrow\mathcal{Y}\) is _equivariant_ if for each transformation \(T_{g}:\mathcal{V}\rightarrow\mathcal{V}\) with \(g\) being an element of an abstract group \(G\), there exists a transformation \(S_{g}:\mathcal{Y}\rightarrow\mathcal{Y}\) such that the following condition holds:

\[S_{g}[\phi(v)]=\phi\left(T_{g}[v]\right),\]

for all \(g\in G,v\in\mathcal{V}.\) In our case, the group \(G\) of interest is SE(3), providing the latent representation with a 3D rotation and translation equivariant property. This latent feature \(\phi\) can be achieved with SE(3)-equivariant GNNs, known as Tensor Field Network (TFN) [39], and with a TFN with an appropriate attention mechanism, known as SE(3)-transformers [40], see more detailed recipes of these two models in Sec. S1.

The key idea that enables SE(3)-equivariant property of these two GNNs is that its message passing kernel is constructed from translational invariant spherical harmonic bases, with the structure-preserving transformation \(S_{g}\) given by the Wigner-D matrices (see Sec. S1). Under a 3D rotation, each spherical harmonic \(J\) basis function of the kernel transforms according to

\[Y_{J}\left(\mathbf{R}_{g}^{-1}\frac{\Delta\mathbf{r}}{||\Delta\mathbf{r}||} \right)=\mathbf{D}_{J}^{*}(g)Y_{J}\left(\frac{\Delta\mathbf{r}}{||\Delta \mathbf{r}||}\right),\]

where \(\mathbf{D}_{J}(g)\) is the \(J^{th}\) Wigner-D matrix, and \(\mathbf{R}_{g}\) is a rotation matrix associated with \(g\in\) SO(3), making the learned latent representation \(\phi\) equivariant [39]. The multi-head attention mechanism in SE(3)-transformers also induces an equivariant property on the product rule between each key and value pairs, rendering the whole message passing algorithm equivariant under SE(3) [40]. In this work, we use SE(3)-transformers to first build a compressed representation of the orientations of a strained crystal structure, then such learned latent representation will be passed to other networks with a high expressive power to perform the prediction (regression) task of the SED. The following section contains the entire framework for the SED prediction.

### Model architecture and training procedure

We now describe our GNNs-based model that predicts the SED tensor. The architecture of our model, which we term _StrainNet_, is illustrated in Fig. 2. The model is designed to receive two types of input: a crystal graph representing a _strained_ crystal structure, and a one-hot vector of dimension \(6(6+1)/2=21\)[41] indicating the _strained components_, with a value \(1\) in the \(ij\) component that the strain operation is applied on and with a value \(0\) otherwise.

The crystal graph of a strained crystal structure is an input feature of the SE(3)-equivariant GNNs building block, which generates _two_ latent features. The first latent feature (\(\mathbf{f}_{\text{class},n}^{0}\) in Fig. 2) is fed into a _classification_ neural network that predicts a vector of dimension \(21\), identifying all the upper-triangular elements of the SED tensor whose values are exactly identical (by symmetry) to the component indicated by the input one-hot vector. This inherent classification informs the model to discern the degenerate structure of the SED tensor that depends on the symmetry of an input crystal graph (see Fig.3). For example, for a cubic lattice strained in the direction \(11\), the input one-hot vector \((1,0,\dots,0)^{T}\) together with the latent feature \(\mathbf{f}_{\text{class},n}^{0}\) shall give a prediction of a vector that has the value close to \(1\) in the indices \(ij=11,22,33\), and the value close to \(0\) in the other \(18\) indices. Finally, the predicted _degeneracy class vector_ (DCV) together with the second latent feature (\(\mathbf{f}_{\text{reg},n}^{0}\) in Fig. 2) will be fed into the final neural network that predicts (regresses on) the SED.

To generate an SE(3)-equivariant latent representation of an input crystal graph, as alluded to in the previous section, we first used the SE(3)-Transformer that receives the strained crystal graph input. The material struc

Figure 1: A 2-dimensional lattice and its _strained_ crystal graph embedding. Grey dashed lines are periodic boundaries. The whole structure can be spanned by the 3 atoms in 3 different bases in the unit cell, represented by 3 different nodes. Black lines are (undirected) edges connecting two neighbor nodes. The strained lattice is represented as a multigraph on the right, which is the crystal graph input into our SE(3)-equivariant GNN building block. Note that to describe the crystal graph for the unit cell and the atoms on its boundary, only three nodes are required. The multi-edges of each pair of nodes are distinguished by different crystallographic directions.
ture from the database is first transformed according to Eq. (6) by a strain of a fixed magnitude, which is then converted into the strained crystal graph input. Noting that in the _ab initio_ calculation, the atomic coordinates in the strained lattice are optimized to be in the equilibrium positions under an applied strain, but in this work, the fractional coordinates are kept as in pre-strain condition. The strained crystal graph input is given by

\[\mathcal{V}^{\prime}=\{(\mathbf{f}_{n},\mathbf{r}_{n}^{\prime}) \mid\mathbf{f}_{n}\in\mathbb{R}^{M},\;\mathbf{r}_{n}^{\prime}=\boldsymbol{ \epsilon}_{I}\mathbf{L}\mathbf{r}_{f_{n}}\in\mathbb{R}^{3}\},\] \[\mathcal{E}^{\prime}=\{\Delta\mathbf{r}_{mn}^{\prime}\mathbf{(T }^{\prime})\mid\Delta\mathbf{r}_{mn}^{\prime}\mathbf{(T}^{\prime})=\mathbf{r }_{m}^{\prime}-\mathbf{r}_{n}^{\prime}+\mathbf{T}^{\prime};\;\mathbf{r}_{m}^{ \prime},\mathbf{r}_{n}^{\prime}\in\mathbb{R}^{3}\},\]

where the _prime_ notation indicates that the atomic positions and the translation vector are of the strained lattice.

The training data of the \(\{U_{ij}\}\) is computed, via Eq. (2), from \(\{C_{ij}\}\) extracted from Materials Project database [20] (see Sec. S5). For molecules and materials, if the number of training data is large enough, only atomic numbers are sufficient for the node (atomic) attributes [42]. Nevertheless, in this work, the features such as atomic mass, electronegativity, polarizability, atomic radius, atomic number, and \(ij\) indices, are used as node attributes. Polarizability data is obtained from [43], whereas other attributes are obtained from Materials Project database. The atomic number is vectorized through an embedding layer of dimension 512, which is then concatenated with the other four node attributes (see Fig. 2). It is well known that the elastic properties are influenced by various factors, including the characteristics of acoustic phonons, which exhibit an inverse relationship with atomic mass. Additionally, the resistance to changes in bond length is correlated with both electronegativity and polarizability. These three atomic properties are thus incorporated into node attributes, and they significantly improve the model prediction accuracy.

Since the atomic features are scalar attributes, the input feature vector will be fed into the \(l=0\) channel of the transformer network since such channel transforms like a scalar (see Sec. S1). We'll denote each node's input feature vector as \(\mathbf{f}_{\text{in},n}^{0}=\mathbf{f}_{n}\). For computational feasibility, we will restrict the output feature vector of the SE(3)-transformers to consist of 4 channels \(\mathbf{f}_{\text{hid},n}^{l}\) with \(l\in\{0,1,2,3\}\). These latent features are fed into two different SE(3)-equivariant GNNs _without_ attention (TFNs), outputing a node-wise classification feature vector \(\mathbf{f}_{\text{class},n}^{0}\), and the node-wise regression feature vector \(\mathbf{f}_{\text{reg},n}^{0}\), see Fig. 2. Note that the attention mechanism is not deployed in the last GNN layer as it yields a better prediction accuracy.

To train the model to discern different orientations of a strained crystal graph and classify its appropriate degeneracy class vector, in each epoch, we draw a new realization of a crystal graph in a different orientation and of a

Figure 2: With a strained crystal graph and the strained component one-hot vector as the input, StrainNet employs a self-supervised approach combining both classification and regression networks to predict both the degeneracy class vector and the \(U_{ij}\). The \(\mathbf{f}_{\text{in},n}^{0}\equiv\mathbf{f}_{n}\) denotes the input feature of the \(n^{\text{th}}\) node, where a superscript, \(0\), indicates the input into the \(l=0\) channel of the spherical harmonics in SE(3)-Transformer, which is a channel whose feature is invariant under rotation (transforms like a scalar quantity). \(\Delta\mathbf{r}_{mn}^{\prime}\) is the edge connecting a target node \(m\) and a neighbor node \(n\). The degeneracy class vector from the classification network will be concatenated with the latent representation of the strained crystal graph (\(\mathbf{f}_{\text{reg},n}^{0}\)) for the final regression network to predict the SED.
new one-hot vector in the same class of vector associated with the original strained component. Specifically, a new crystal graph input is sampled from a uniformly random orientation (uniformly random Euler angles) around the center of mass in the unit cell of the original strained crystal graph. A new one-hot vector of the strained components is chosen such that the element 1 is drawn uniformly from one of the non-zero components of the DCV. For example, if the original strained component is 11 for a cubic material, a new one-hot vector in each epoch is drawn from the situations where only one of the strained components 11, 22, or 33 is 1, while the other components are 0.

Our self-supervised approach combines both the regression errors and the classification errors in the global loss function:

\[\mathcal{L}=\frac{1}{N}\sum_{n=1}^{N}|U_{ij}^{(n)}-U_{ij}^{(n),\text{pred}}|+ \lambda\mathcal{L}_{\text{class}}, \tag{9}\]

where the SED prediction \(U_{ij}^{(n),\text{pred}}\) is the output from the regression network, \(\mathcal{L}_{\text{class}}\) is the binary cross-entropy loss function for multi-label classification, \(\lambda\) specifies the relative significance between the classification and the regression errors, and \(N\) is the total number of training samples. We stop the training when the gradient of the loss is relatively small over multiple epochs. Training was performed on NVIDIA A100 GPUs.

### Choosing a subset of strained crystal structures for training data

By the virtue of rotational equivariance, we expect StrainNet to efficiently predict the 21 components of \(\mathbf{U}\) using training data consisting of only non-degenerate components. We have selected the number of non-degenerate components for training that depends on crystal systems, i.e. cubic, tetragonal, hexagonal, trigonal, monoclinic, and triclinic. For cubic materials, only 6 components are used for training (the number of non-degenerate components for other crystal systems can be found in Table S1 in the Supplemental Materials). Recall that, as discussed earlier, the cubic lattice has 5 distinct \(U_{ij}\) as shown in Eq. (5); however, the strain tensor transform the cubic structures into 6 distinct structures (according to the symmetry in Laue class \(m\bar{3}m\)) as follows

\[\begin{bmatrix}\mathcal{T}_{1}&\mathcal{T}_{2}&\mathcal{T}_{2}&\mathcal{O}_{ 2}&\mathcal{M}_{2}&\mathcal{M}_{2}\\ &\mathcal{T}_{1}&\mathcal{T}_{2}&\mathcal{M}_{2}&\mathcal{O}_{2}&\mathcal{M}_ {2}\\ &&\mathcal{T}_{1}&\mathcal{M}_{2}&\mathcal{M}_{2}&\mathcal{O}_{2}\\ &&\mathcal{O}_{1}&\mathcal{M}_{1}&\mathcal{M}_{1}\\ &&&&\mathcal{O}_{1}&\mathcal{M}_{1}\\ &&&&\mathcal{O}_{1}\end{bmatrix},\]

where each element in the matrix stands for a crystal system of \(\mathbf{L}^{\prime}(\epsilon_{i},\epsilon_{j})\), \(\mathcal{T}\), \(\mathcal{O}\), and \(\mathcal{M}\) are tetragonal, orthorhombic, and monoclinic lattices, respectively, and a different number labeling the subscript indicates a different structure. Despite the fact that \(U_{14}\) (\(C_{14}\)) is equal to \(U_{15}\) (\(C_{15}\)) by the cubic symmetry, \(\mathbf{L}^{\prime}(\epsilon_{1},\epsilon_{4})\) and \(\mathbf{L}^{\prime}(\epsilon_{1},\epsilon_{5})\) possess different lattice symmetries which are orthorhombic and monoclinic lattices, respectively. This is because \(\epsilon_{1}\) strains the structure in the \(x\) direction, while \(\epsilon_{4}\) and \(\epsilon_{5}\) strain the structure in both the \(y\) and the \(z\) directions, and both the \(x\) and the \(y\) directions, respectively. Since \(\mathbf{L}^{\prime}(\epsilon_{1},\epsilon_{4})\) and \(\mathbf{L}^{\prime}(\epsilon_{1},\epsilon_{5})\) are not equivalent up to a rotation, the SE(3) kernel regards the two inputs as different inputs. By exploiting SE(3) kernel to identify inputs that are equivalent up to a rotation, for cubic materials, it suffices to use the training data consisting of the distinct input structures, strained only in 6 directions, i.e., 11, 12, 14, 15, 44, and 45.

### The distance between SED components

To evaluate our model capability to predict the symmetry-dependent degeneracy pattern of the SED tensor, we use the Canberra distance between two components of the SED tensor as a metric:

\[d(U_{ij},U_{kl})\equiv\frac{1}{N}\sum_{n=1}^{N}\frac{|U_{ij}^{(n)}-U_{kl}^{(n )}|}{|U_{ij}^{(n)}|+|U_{kl}^{(n)}|}, \tag{10}\]

where \(N\) represents the number of samples, and \(ij\) and \(kl\) are Voigt indices. With this metric, the distance between \(U_{ij}\) and \(U_{kl}\) is zero if they belong to the same degeneracy class (their values are identical.) The top row of Fig. 3 shows the ground-truth distance pattern for cubic, hexagonal, and monoclinic materials. Dark blue and bright green colors indicate the \(d(U_{ij},U_{kl})\) values that are closest to zero and 0.6, respectively. It is important to note that the degeneracy pattern of \(U_{ij}\) depends on the crystal symmetry, which results in a unique Canberra distance pattern for each crystal system. For instance, in cubic materials, \(U_{11}\), \(U_{22}\), and \(U_{33}\) are identical (in the same degeneracy class), so the distance pattern of cubic materials displays a dark blue box corresponding to indices 11, 22, and 33. The distance patterns in Figs. 3 and S9 are averaged over the samples with the same crystal system in the dataset. Hence, the distance pattern of low-symmetry crystals, such as monoclinic and triclinic, will strongly depend on the dataset.

## IV Discussion and results

### Rationale for the StrainNet model architecture

The underlying principle for using an equivariant network as a core component is to construct a compact representation of the SED tensor, such that the crystal structures strained by different sets of \(\epsilon_{i}\) and \(\epsilon_{j}\) that are equivalent up to rotation (and translation) possess the same SED. Such symmetry-dependent SED tensor of
specific crystal structures is demonstrated in the ground-truth distance patterns \(d(U_{ij},U_{kl})\) of Fig. 3 (top), where the StrainNet can well approximate such symmetry-dependent patterns, see Fig. 3 (bottom).

Fig. S4 shows our earlier model development with the goal to compactly represent the SED tensor, beginning from the simplest but less expressive to StrainNet which can compactly represent the SED tensor rather accurately. All these trial GNN-based models were trained on crystal structures with fewer than 500 edges. The simplest model (Fig. S4(a)) takes as an input only the crystal graph with node attributes, i.e., atomic mass, atomic radius, electronegativity, polarizability, and embedded atomic number. This minimal equivariant model does not yield a sufficiently accurate representation of \(U_{ij}\); we found that the model is more biased towards predicting the SED tensor whose distance matrix pattern is more akin to that of the cubic crystals. This could be because the model is not sufficiently expressive, and thus attempts to fit the majority of the training dataset comprising more higher-symmetry structures (see Table S3 for the crystal system statistics of our curated dataset).

To improve the minimal model's expressiveness, we introduce the degeneracy class vector (DCV) as an additional input, to inform the model about the symmetry-dependent class of the SED tensor. The DCV is first embedded by fully-connected layers and then concatenated with global max-pooled features or node attributes, as shown in Fig. S4(b) and (c), respectively. These symmetry-informed models yield lower MAE and RMSE of the SED tensor compared to that of the minimal model, while the model in Fig. S4(b) performs slightly better than the model in Fig. S4(c) (see Table S4). Moreover, the predicted distance patterns are less biased towards those of the cubic lattice. Specifically, for crystal structures with lower symmetry such as hexagonal, tetragonal, trigonal, and orthorhombic, the predicted distance patterns are relatively similar to their ground-truth distance patterns. As expected, informing the model about the DCV of the input helps force the distance between \(U_{ij}\) and \(U_{kl}\) belonging to the same degeneracy class to be closer to zero.

While these models with a DCV included as an input give good predictions of the SED tensor, they require the knowledge of the pretrained crystal symmetry _a priori_ (to properly assign the value of the DCV). To overcome this limitation, we use a self-supervised method schematically shown Fig. 2 to modify the model to predict, rather than to require, the DCV. The classification neural network that predicts the DCV takes as an input a one-hot vector of the strained \(ij\) components, representing the \(\epsilon_{i}\) and \(\epsilon_{j}\) components in (7). This self-supervised technique enables the StrainNet to predict \(U_{ij}\) without any prior knowledge of the pretrained crystal symmetry. Importantly, StrainNet can express the degeneracy class of the SED tensor relatively well, see Figs. 3 and S9.

### StrainNet's prediction results

Table 1 summarizes StrainNet's prediction errors on the test set. To the best of our knowledge, StrainNet provides the first data-driven prediction of \(U_{ij}\) with a reasonable accuracy. Using these predicted SED tensors together with Eq. (3), elastic tensors of materials can also be calculated. Our model's prediction accuracy significantly depends on the elemental composition of the compounds. Fig. S7 depicts the MAE of \(U_{ij}\) and \(C_{ij}\), averaged over 21 \(ij\) components and categorized by elements. It is interesting to note that the compounds containing period 2 elements (Be, B, C, N, O, and F), some period 6 transition metals (Ta, W, Re, and Os), and some actinides (Th, Pa, and Pu), which have high bulk and shear moduli on average, exhibit a higher MAE for \(U_{ij}\) compared to the overall MAE of the entire test set. As a result, their MAE for \(C_{ij}\) are also significantly higher than the overall MAE of the test set. In contrast, the compounds containing lanthanides, which exhibit medium bulk and shear moduli on average, except for Gd, demonstrate a considerably lower MAE than the overall MAE of the test set, for both \(U_{ij}\) and \(C_{ij}\).

To further investigate the prediction results of the elastic constants, we plot the comparisons between the prediction (obtained from the predicted SED tensors together with Eq. (3)) and the ground truth of the elastic constants in Fig. 4. Fig. 4 (a) shows the prediction results of every \(C_{ij}\) components of all crystal structures, whereas Figs. 4(b)-(f) show the prediction results of \(C_{22}\), \(C_{23}\), \(C_{26}\), \(C_{55}\), and \(C_{56}\) of cubic materials. For cubic materials, the model is trained only on a non-degenerate subset of 21 \(U_{ij}\) components, specifically components 11, 12, 14, 15, 44, and 45. Notably, the model is able to predict other unseen components reasonably well. However, the prediction of the elastic constant components that are exactly zero is challenging. For example, the ground-truth values of \(C_{26}\) and \(C_{56}\) of cubic materials are exactly 0 GPa, while the means and standard deviations of these components from our model predictions are 0.69 and 4.45 GPa, and -0.43 and 4.85 GPa respectively (see Fig. S6 for the histograms of the prediction results). Although the standard deviations are relatively large, the means are less than 1 GPa. It is worth noting that the first-principles calculations also make some errors in these values, but these are usually not computed since they are known to be exactly 0 GPa.

To see how well StrainNet discerns the geometric relationship between different strained crystal graph inputs, we plot the Canberra distance metric of the predicted and ground-truth SED tensors for comparison in Fig. 3. The predicted distance metric of cubic materials, as shown in the left column of Fig. 3, closely resembles the ground-truth distance metric, although there are minor discrepancies in the values that belong to the same degeneracy class. For instance, the distances between the predicted \(U_{11}\), \(U_{22}\), and \(U_{33}\) are \(4.8-6.2\)\(\mu\)eV/atom, resulting in the percentage errors of \(1.83-2.43\%\) (com
puted from \(\frac{1}{N}\sum_{n=1}^{N}|U_{ij}^{(n)}-U_{kl}^{(n)}|/U_{avg}^{(n)}\times 100\%\) where \(U_{avg}^{(n)}=(|U_{ij}^{(n)}|+|U_{kl}^{(n)}|)/2\).) For hexagonal materials, \(U_{11}\) are identical to \(U_{22}\), but not to \(U_{33}\). The predicted distance metric of hexagonal materials accordingly forms a dark blue box in the 11 and 22 components. The percentage error between the predicted \(U_{11}\) and \(U_{22}\) is 2.22%. For monoclinic materials, the values of each \(U_{ij}\) are not necessarily identical to one another, and the distance metric is averaged over monoclinic materials data. The resulting distance pattern may differ significantly depending on the dataset. Notably, these agreements between the predicted and the ground-truth distance metric in various crystal structures reveal that StrainNet gives an excellent prediction of the degenerate structure of SED tensors of strained crystalline materials.

Lastly, the model can also predict other elastic properties, including Voigt's bulk modulus (\(B_{V}\)), Voigt's shear modulus (\(G_{V}\)), Young's modulus (\(Y\)), and Poisson's ratio (\(\nu\)), using the predicted \(C_{ij}\) together with Eqs. (S9)\(-\)(S12). Table 1 summarizes the MAE and RMSE of our predicted elastic properties compared to previous works. The MAE and RMSE of \(B_{V}\) are 12.59 and 20.83 GPa, respectively, which are comparable to those of Mazhnik et al. [23]. The MAE and RMSE of \(G_{V}\) are 10.49 and 18.95 GPa, respectively, with the RMSE of \(G_{V}\) being comparable to that reported by Zhao et al. [24]. Our work also yields smaller errors for \(\nu\) than Mazhnik et al. [23] Note that our dataset (see Sec. S5) differs from the dataset in the relevant work of [23, 24]. Fig. S8 presents the MAE of \(B_{V}\), \(G_{V}\), \(Y\) and \(\nu\), categorized by elements. The model produces high MAE for compounds containing period 2 elements, mainly due to their larger MAE in \(U_{ij}\) and \(C_{ij}\) (see Fig. S7). On the other hand, the model can produce smaller MAE for lanthanide compounds, as their MAE in \(U_{ij}\) and \(C_{ij}\) (see Fig. S7) are smaller.

### Interpretability of StrainNet's latent features

To investigate how latent features facilitate successful prediction of the SED tensor, we employed the diffusion map to faithfully visualize the three-dimensional data representation of the high-dimensional latent features of the entire test set [44, 45]. Figs. 5 and S10 reveal the dimensionality reduction of the latent features in the diffusion coordinates \((\varphi_{1},\varphi_{2},\varphi_{3})\). Two different latent features were considered: the global max-pooling layer from the GNN representation of the input crystal graph (denoted as \(\mathbf{A}\) in Fig. 2), and the concatenation between \(\mathbf{A}\) and the embedded DCV (denoted as \(\mathbf{B}\) in Fig. 2). These low-dimensional representations in Figs. 5 and S10 are colored by the strained component \(ij\) (middle and right columns) and by the energy scale (left column).

What does the max-pooled feature of the crystal graph input \(\mathbf{A}\) represent? Fig. S10(c) shows that the crystal graph latent representations of the _same material_ that are strained in different \(ij\) components are _almost identical_, in fact sharing greater resemblance than the representations of different materials that belong to the same symmetry class. This is consistent with the result of Bronstein et al. which shows that the latent representation in SE(3)-equivariant GNNs of the original input is very similar to those of the mildly spatially distorted inputs, and, interestingly, is less similar to those of inputs spatially translated further away [38]. Hence, the latent feature \(\mathbf{A}\) seems to uniquely encode the information of the input material, rather than the symmetry of the input. Thus \(\mathbf{A}\) alone does not suffice to meaningfully represent the symmetry-dependent SED tensor (Fig. S10(a) and (b)).

On the other hand, with the embedded DCV in the latent feature \(\mathbf{B}\), the network can differentiate different strained components within the same material class. When strained in 21 different components, the latent features \(\mathbf{A}\) of three example materials (trigonal LuTITe\({}_{2}\) phase, tetragonal YMnSi phase, and cubic Ta\({}_{3}\)Ru phase) that were not differentiable within the same material class (Fig. S10(c)) are now clearly segregated in the \(\varphi_{2}\) variable through the help of the embedded DCV, see Fig. 5(c). In fact, the latent feature \(\mathbf{B}\) which combines both the input crystal graph representation and the embedded DCV offers an interpretable latent representation of the SED (Fig. 5), which we now discuss.

The latent feature \(\mathbf{B}\) is organized such that the \((\varphi_{1},\varphi_{2})\) coordinate encodes the strain energy density that varies from a smaller value in the \((+,+)\) quadrant to a larger value in the \((-,-)\) quadrant (see Fig. 5(a)). Since \(U_{ij}\) stored by the tensile strain (\(i,j\in\{1,2,3\}\)) is typically higher than \(U_{ij}\) stored by the shear strain (\(i,j\in\{4,5,6\}\)), Figs. 5(a) and (b) consistently reveal that the latent features of the materials strained by the 11, 22, or 33 component have a negative \(\varphi_{2}\) value, whereas those of the materials strained by the 44, 55, or 66 components have a positive \(\varphi_{2}\) value. In addition, since \(U_{ij}\) for \(i\neq j\) is computed from the sum of \(C_{ii}\), \(C_{jj}\), and \(C_{ij}\), it is typically larger than \(U_{ii}\) and \(U_{jj}\). Figs. 5(a) and (b) also consistently reveal that the latent features of the materials strained by 44, 45, 11, and 12 components are respectively organized in the \(\varphi_{2}\) coordinate from a more positive to a more negative value. Additionally, Figs 5(a) and (c) show that materials with higher (lower) average SED will be represented by a larger negative (larger positive) \(\varphi_{1}\) variable. In summary, the max-pooled feature from the graph neural networks \(\mathbf{A}\) encodes the material information, and, together with the information of the DCV, the concatenated latent feature \(\mathbf{B}\) effectively encapsulates both the material information and its strained component-dependent SED.

## V Conclusions

We have demonstrated a novel data-driven framework for predicting the elastic properties of crystal structures using SE(3)-equivariant graph neural networks
(GNNs). By leveraging SE(3)-equivariant GNNs as building blocks, our self-supervised deep learning model accurately predicts Voigt's bulk modulus, Voigt's shear modulus, Young's modulus, and Poisson's ratio that are comparable to other non-equivariant deep learning studies [23, 24].

A key contribution is the prediction of the strain energy density (SED) and its associated elastic constants, which are tensorial quantities that depend on a material's crystallographic group. The similarity between the distance metrics of the SED components between the ground truths and the predictions demonstrates our model's capability to identify different symmetry groups of strained crystal structures. Requiring only a strained crystal graph and the strained component as the input, our approach offers an efficient alternative to the standard _ab initio_ method for designing new materials with tailored elastic properties.

The interpretability of the model is also a notable feature. The learned latent representations taking into account the degeneracy class vector are organized in a physically meaningful structure for predicting the SED tensors. This interpretability aspect enhances the transparency of model prediction, enabling the justification of whether the prediction is physically relevant.

The combination of interpretability and the consideration of crystallographic groups sets our model apart from recent data-driven methods for predicting elastic properties of materials. We hope this work is a stepping stone toward an efficient data-driven approach for materials discovery and design, and opens up avenues for approaching more challenging tensor prediction tasks, e.g.,

\begin{table}
\begin{tabular}{c c c c c c c c} Properties & \multicolumn{3}{c}{This work} & \multicolumn{3}{c}{Mazhnik et al. [23]} & Zhao et al. [24] \\  & Average & MAE & RMSE & Average & MAE & RMSE & RMSE \\ \hline \(U_{ij}\) (meV/atom) & 2.652 & **0.655** & **1.288** & & & & \\ \(C_{ij}\) (GPa) & 42.92 & **10.37** & **17.69** & & & & \\ \(B_{V}\) (GPa) & 107.07 & 1.148 & 19.73 & 111.83 & 11.11 & 19.54 & 16.530 \\ \(G_{V}\) (GPa) & 50.98 & 9.61 & 17.10 & 54.81 & 8.24 & 11.43 & 15.780 \\ \(Y\) & 129.62 & 22.29 & 38.63 & 138.95 & 19.15 & 26.23 & \\ \(\nu\) & 0.401 & **0.037** & 0.133 & 0.286 & 0.041 & 0.105 & \\ \end{tabular}
\end{table}
Table 1: Statistics of the dataset and the model prediction

Figure 3: The distance metrics \(d(U_{ij},U_{kl})\) computed from the test set of the ground-truth SED (top row) and the predicted SED (bottom row). The columns organize the crystal structure, appearing from left to right is for cubic, hexagonal, and monoclinic materials respectively. The model prediction shows an excellent agreement with the ground-truth for crystals with high symmetry (cubic), and a good agreement with the ground truths for crystals with lower symmetry (hexagonal and monoclinic).
predicting dielectric and piezoelectric tensors, which are second-rank and third-rank tensors, respectively.

###### Acknowledgements.

 This research is supported by the Program Management Unit for Human Resources and Institutional Development, Research and Innovation (Grant No. B05F650024) and by Thailand Science research and Innovation Fund Chulalongkorn University (IND66230005). The authors acknowledge high performance computing resources including NVIDIA A100 GPUs from Chula Intelligent and Complex Systems Lab, Faculty of Science, and from the Center for AI in Medicine (CU-AIM), Faculty of Medicine, Chulalongkorn University, Thailand. The authors also acknowledge the National Science and Technology Development Agency, National e-Science Infrastructure Consortium, Chulalongkorn University and the Chulalongkorn Academic Advancement into Its 2nd Century Project (Thailand) for providing computing infrastructure.
