Thermal Conductivity Calculation using Homogeneous Non-equilibrium Molecular Dynamics Simulation with Allegro

The Allegro model [11] can be constructed using NequIP [39] which introduces an E(3) symmetry equivariant neural network (NN) coded using PyTorch [40]. However, unlike NequIP, Allegro questions the need for the message passing scheme that incorporates atomic interactions beyond the cutoff radius $R_{\rm c}$ and instead adopts the conventional method of restricting the interacting atoms to within $R_{\rm c}$ from the central atom. The atomic potential energy of Allegro $\varepsilon^{\rm Allegro}_{i}$ is calculated as the sum of the pairwise energies $\varepsilon^{\rm Allegro}_{ij}$ defined between two atoms within $R_{\rm c}$:

and the total potential energy is defined by the sum of $\varepsilon^{\rm Allegro}_{i}$ over all atoms in the system,

This definition of $E^{\rm Allegro}$ provides Allegro the advantage of parallel computation with space partitioning using the well-known Verlet neighbor list and the cell-index method required for large-scale MD simulations.

The motivation for developing Allegro was also to overcome the difficulties of ACE [13]. ACE can construct descriptors that incorporate BO expansion up to higher orders at a feasible computational cost. While most conventional descriptors are composed of up to three-body order of radial and angular parts, such as the Behler-Parrinello symmetry functions [12], ACE descriptors, which incorporates four-body order or more, have been employed [41]. Furthermore, ACE also efficiently generates descriptors with higher-order rotational symmetry characterized by irreducible representation $l$, which paved the way for MLIP construction with descriptors incorporating rotational equivariance [42]. However, a drawback is that the number of descriptors increases exponentially as the number of atomic species increases. In addition, the number of different descriptors, characterized by BO or $l$, to generate is determined by trial and error. ACE descriptors are often combined nonlinearly to improve accuracy [13], and the manner in which they are combined requires trial and error.

Allegro employs a GNN to distinguish atomic species by utilizing embedded vectors, which significantly reduce the computational cost resulting from an increase in atomic species. Furthermore, the descriptor generation process is combined with the layer structure of GNN to efficiently construct nonlinear descriptors. The number of convolution layers in GNN, which is controlled by hyperparameter $N_{\rm layer}$, is devised to correspond to the order of BO expansion. As the tensor computation of the spherical harmonics is performed in each convolution layer of the GNN, the BO order of the descriptor increases by one. For example, Allegro with $N_{\rm layer}$ = 2 produces a nonlinear combination of descriptors of up to four body orders. In other words, an Allegro model with $N_{\rm layer}$ has ($N_{\rm layer}$ + 2) body-order descriptors. In addition, the diversity of descriptors with irreducible representation $l$ can be increased through the composition of spherical harmonics via Clebsch–Gordan coefficients during the tensor computation. The hyperparameter $l_{\rm max}$ controls the maximum irreducible representation $l$ of the descriptors produced after composition. For $l_{\rm max}$ = 0, only rotation-invariant descriptors are produced, whereas, for $l_{\rm max}$ $\geq$ 1, rotation-equivariant descriptors are also produced, which can improve feature extraction for the atomic local environment.

Therefore, because $N_{\rm layer}$ and $l_{\rm max}$ play important roles in Allegro, this study focuses on these two hyperparameters and evaluates their dependence on the thermal conductivity. Six models are created using combinations of $N_{\rm layer}$ = {1,2} and $l_{\rm max}$ = {0,1,2}, that is, ($l_{\rm max}$,$N_{\rm layer}$) = (0,1), (0,2), (1,1), (1,2), (2,1), and (2,2) models.

Other hyperparameters related to the descriptors and architecture of GNN are the same for all Allegro models. The details are explained in Section 2 of SM.

The Adam optimizer [43], which is a stochastic gradient descent method, was adopted to train all the Allegro models. The default parameters of Adam were set to $\beta_{1}$ = 0.9, $\beta_{2}$ = 0.999, and $\epsilon$ = 10^-8 without weight decay. The total number of epochs was 1,000. All the models were trained with float64 precision by minimizing the following cost function $C$, which comprises three loss functions associated with the total potential energy (first term), atomic force (second term), and total virial (third term):

where $B$ denotes the batch size used for the Adam optimizer and was set to $B$ = 5 in this study. As these terms differ in dimension and size, $p_{E}$, $p_{F}$, and $p_{W}$ were introduced as adjustable parameters; $p_{E}$ = $p_{F}$ = 1.0 and $p_{W}$ = 0.0005 were used. Note that the training for ${\rm W}^{\rm Allegro}_{I}$ is equivalent to that of the potential part of virial stress tensor calculated by dividing ${\rm W}^{\rm Allegro}_{I}$ by the volume $\varOmega$. The training was performed on 80% of the data mentioned in Section 2.1 with the remaining 20% of the data used for validation. Hereafter, the former and latter are referred to as the Train and Validation datasets, respectively. The training was performed using System C at the Institute for Solid State Physics, the University of Tokyo. The training wall times for the six Allegro models described in Section 2.2.1 are shown in Fig. S2(a) in SM.

There is a concern that the physical quantities calculated from the Allegro models depend on the initial weight parameters. Therefore, for statistical evaluation, five models with different initial weight parameters were constructed.

The heat flux ${\bf{J}}_{Q}$ in Eq. (1.2) requires an atomic virial tensor ${\rm W}_{i}$. However, this is not present in the original Allegro model [11]. The derivation of an atomic virial tensor that is appropriate for the heat flux requires careful consideration of the nature of Allegro as a many-body potential. The extended code [20] proposed by Yu et al. is helpful for this purpose. They introduced relative coordinates $\{{\bf{r}}_{ij}$ ($\equiv$ ${\bf{r}}_{j}$ $-$ ${\bf{r}}_{i}$)} into the computational graph of the atomic potential energy $\varepsilon^{\rm Allegro}_{i}$ in PyTorch. This allows the calculation of its derivatives $\frac{\partial\varepsilon^{\rm Allegro}_{i}}{\partial{\bf{r}}_{ij}}$ by automatic differentiation, and Yu et al. defined the following type of atomic virial tensor ${\rm W}^{\rm sym}_{i}$:

This definition is based on the objective of finding an atomic virial tensor that is symmetric as well as the total virial tensor. In practice, this can be expressed as follows:

However, ${\rm W}^{\rm sym}_{i}$ would not be suitable for thermal conductivity calculations. According to the work by Fan et al. on Tersoff potential [16], the rigorous heat flux J_Q for many-body potentials can be derived from the following equation [44].

The second term of Eq. (2.6), which coincides with the third term of Eq. (1.2), yields a rigorous formula for the atomic virial tensor. As the atomic potential energy of Allegro $\epsilon^{\rm Allegro}_{i}$ is only a function of relative coordinates ${\bf{r}}_{ij}$ within a sphere of radius $R_{\rm c}$

the second term of Eq. (2.6) can be expressed as [16]:

Thus, the following atomic virial tensor ${\rm W}^{\rm asym}_{i}$ would be appropriate for the heat flux [10],

which is an asymmetric tensor. Although the mathematical properties of the two atomic virial tensors are different, their summation over $N_{\rm atom}$ corresponds to the same total virial tensor ${\rm W}^{\rm Allegro}$. To demonstrate that the derivation of ${\rm W}^{\rm asym}_{i}$ is indispensable, the thermal conductivities obtained using these two atomic virial tensors are compared later. The code of the Allegro model, extended to the output ${\rm W}^{\rm asym}_{i}$ and ${\rm W}^{\rm sym}_{i}$, is available in Ref. [45].

The homogeneous non-equilibrium MD (HNEMD) method [22, 23, 24], which is based on the GK formula, successfully reduces the computational cost by enabling the calculation of the thermal conductivity $\kappa$ in the $x$ direction using the time average of the heat flux ${\bf{J}}_{Q}$ instead of integrating the autocorrelation function in Eq. (1.1):

where $F_{\rm ext}$ denotes the magnitude of the perturbation along the $x$ direction of the system. An appropriate $F_{\rm ext}$ is determined from the linear response regime by evaluating the thermal conductivity as a function of $F_{\rm ext}$. Based on the evaluation results (Fig. S3 in SM), we selected $F_{\rm ext}$ = 0.005 bohr^-1 in this study.

The SHC method, originally developed for the NEMD method [46], was recently extended to HNEMD by Fan et al. [26]. First, we define the cross-correlation function ${\bf{K}}(t)$ between atomic enthalpy ${\rm h}_{i}$ and atomic velocity ${\bf{v}}_{i}$:

${\rm h}_{i}$ is given by:

where I denotes the identity matrix. Note that Fan et al. constructed a correlation function using only ${\rm W}_{i}$ [26], whereas we used ${\rm h}_{i}$ for the construction. When the Fourier transform of ${\bf{K}}(t)$ is expressed as

${\bf{K}}(t)$ can also be written using the inverse Fourier transform as

where Re[$c$] is the real part of the complex number $c$. As ${\bf{K}}(t=0)$ of Eq. (2.11) corresponds to the heat flux formula in Eq. (1.2), the thermal conductivity $\kappa$ can be expressed as follows via Eqs. (2.10) and (2.14):

where $\tilde{\kappa}(\omega)\equiv\frac{2\varOmega}{F_{\rm ext}T}{\rm Re}\left[\tilde{K}_{x}(\omega)\right]$; $K_{x}(0)$ and $\tilde{K}_{x}(\omega)$ denote the $x$ direction components of ${\bf{K}}(0)$ and ${\tilde{\bf{K}}}(\omega)$, respectively. Thus, using spectral thermal conductivity $\tilde{\kappa}(\omega)$, the total thermal conductivity $\kappa$ can be decomposed in frequency space $\omega$. For convenience, the cumulative $\kappa^{\rm cum}(\omega)$ is defined as:

All HNEMD simulations were executed using the $NVT$ ensemble. The values of the cell vectors of the MD cell were averaged over those of the FPMDs with the $NPT$ ensemble that were used as the training data. For each HNEMD simulation, after relaxation for 100,000 steps, a 1,000,000-MD step simulation was performed to sample ${\bf{J}}_{Q}$ of Eq. (1.2) and calculate K(t) in Eq. (2.11). Due to its anisotropic structure, the thermal conductivity of $\beta$-Ag₂Se may depend on the direction. The thermal conductivity along the $a$-axis of the unit cell (see Fig. S1 in SM) was calculated as an example. The thermal conductivities were calculated using Eqs. (2.10), (2.15), and (2.16). HNEMD simulations were performed using System C at the Institute for Solid State Physics, the University of Tokyo. The wall times of the simulations for the six Allegro models described in section 2.2.1 are shown in Fig. S2(a) in SM.

Figure 4 shows the results of the thermal conductivities obtained from the HNEMD simulations using two types of atomic virial tensors ${\rm W}^{\rm sym}_{i}$ and ${\rm W}^{\rm asym}_{i}$ defined in Section 2.3.1. As obtaining the reference value from FPMD is quite difficult because of the high computational cost, the thermal conductivities of the models were compared with the experimental value. The lattice thermal conductivity of $\beta$-Ag₂Se was experimentally estimated to be 0.3 $\pm$ 0.1 Wm^-1K^-1 [50], which is quite low despite the crystalline phase. In terms of the results calculated using ${\rm W}^{\rm sym}_{i}$ in Eq. (2.5), the four (0,1), (0,2), (1,1), and (1,2) models had almost the same thermal conductivities $\sim$0.45 Wm^-1K^-1, which were slightly larger than the experimental value. However, the thermal conductivities of the (2,1) and (2,2) models further increased by $\sim$0.54 and $\sim$0.64 Wm^-1K^-1, respectively, clearly deviating from the experimental value. By contrast, the thermal conductivities calculated using ${\rm W}^{\rm asym}_{i}$ in Eq. (2.9) are $\sim$0.35 Wm^-1K^-1 for the two models with $l_{\rm max}$ = 0, and $\sim$0.40 Wm^-1K^-1 for the four models with $l_{\rm max}$ = 1 and 2. These values fall within the range of experimental error. The fact that the former had a smaller thermal conductivity by $\sim$0.05 Wm^-1K^-1 would be attributed to the deviation of the structure from the FPMD simulation, however, the latter appeared to be converging, unlike those of ${\rm W}^{\rm sym}_{i}$. From these observations, we conclude that our derivation of ${\rm W}^{\rm asym}_{i}$ would be essential for the thermal conductivity calculations using Allegro.

The increase in thermal conductivity by $\sim$0.05 Wm^-1K^-1 between the models with $l_{\rm max}$ = 0 and $l_{\rm max}$ = 1 and 2 (filled black circles in Fig. 4) would reflect the differences in the atomic structures, as shown in previous Section 3.1. Using the HNEMD-SHC method [26], we can delve deeper into the cause of such small differences in thermal conductivity. The differences in the atomic structures are reflected in the phonon distribution. The vibrational density of states (VDOSs) has often been employed in many previous studies to elucidate the affected frequency regions [27, 51, 52, 53, 54, 55]. The partial VDOSs of Ag and Se (VDOS_Ag and VDOS_Se) were calculated using the velocity correlation function [27], as shown in Figs. 5(a) and 5(b). The profiles of VDOSs for the (0,1) and (0,2) models were similar. However, the VDOSs for the (1,1), (1,2), (2,1), and (2,2) models were indistinguishable from one another. Interestingly, a significant influence was observed not only on VDOS_Se but also on VDOS_Ag, although the contributions of $\Delta G_{\rm AgAg}$ and $\Delta G_{\rm AgSe}$ to $\Delta G$ were quite small, as shown in Fig. 3(b). The peak near 8 meV and shoulder near 4 meV in VDOS_Ag of the (0,1) model shift inward in the (1,1) model. In addition, the peak near 12 meV in the (0,1) model had disappeared. The peaks at approximately 5 and 18 meV for VDOS_Se in the (0,1) model shift slightly outward in the (1,1) model. This may have caused the difference in total thermal conductivity; however, the contribution of each is unknown.

To this end, the HNEMD-SHC method is useful because the thermal conductivity can be decomposed in frequency space. Figures 5(c) and 5(d) show the spectral thermal conductivity $\tilde{\kappa}(\omega)$ in Eq. (2.15) and its cumulative $\kappa^{\rm cum}(\omega)$ in Eq. (2.16), respectively, for the (0,1) and (1,1) models. We found that $\tilde{\kappa}(\omega)$ from 15 to 20 meV in the (1,1) model was larger than that in the (0,1) model, which mainly contributes to the deviation in $\kappa$ $\sim$0.05 Wm^-1K^-1 in Fig. 4. This can be understood more easily by $\kappa^{\rm cum}(\omega)$. Both models showed similar profiles from 0 to 15 meV. After 15 meV, a difference emerged, and the values converged with a difference of $\sim$ 0.05 Wm^-1K^-1 after 20 meV. The region from 15 to 20 meV corresponds to the peak shift of VDOS_Se in Fig. 5(b).

Considering these results, Allegro, which provides a high level description of atomic structures, is combined with the HNEMD-SHC method, which is expected to be even more useful for the thermal conductivity analysis of more complex systems, such as those with structural defects.