On-the-fly machine learned force fields for the study of warm dense matter: Application to diffusion and viscosity of CH Open Access

Warm dense matter (WDM) can be found in diverse physical settings, ranging from giant planets and stars to inertial confinement fusion (ICF) and other high energy density (HED) experiments.¹ The accurate modeling of WDM is therefore crucial to the understanding and design of HED experiments like ICF, as well as to the understanding of the formation, nature, and evolution of planetary and stellar systems. However, such extreme conditions of temperature and pressure present significant challenges, experimentally as well as theoretically, since both classical and quantum mechanical (degeneracy) effects contribute to the properties and behavior, with the relative importance of each noticeably varying with the conditions, i.e., temperature and density, of interest.

Kohn–Sham density functional theory (DFT)^2,3 is a widely used method for studying materials systems from the first principles of quantum mechanics, with no empirical or ad hoc parameters. However, Kohn–Sham calculations of WDM pose unique challenges, including the increase in the number of partially occupied states with temperature, whereby the cubic scaling bottleneck of such methods manifests itself at smaller system sizes. This bottleneck becomes particularly restrictive in ab initio molecular dynamics (AIMD), where the Kohn–Sham equations may need to be solved hundreds of thousands of times to reach the timescales relevant to phenomena of interest. This has motivated the development of formulations of Kohn–Sham DFT that are well suited to calculations at high temperature, including spectral quadrature (SQ) DFT,^4–7 stochastic DFT (SDFT),^8,9 mixed stochastic-deterministic DFT (MDFT),¹⁰ and the density kernel based SQ method (SQ3).¹¹ In particular, the SQ method scales linearly with system size, has a prefactor that decreases with temperature, and has excellent parallel scaling, whereby it has found a number of applications in the study of WDM.^12–14 

In spite of significant advances, the relatively large computational cost associated with ab initio methods has prompted the development of a number of approximations to Kohn–Sham DFT, including orbital-free molecular dynamics (OFMD),¹⁵ extended first principles molecular dynamics (ext-FPMD),^16,17 and spectral partitioned DFT (spDFT).¹⁸ However, though significantly reduced, these methods are still associated with significant computational cost, in the context of MD in particular. This limitation can be overcome by machine learned force field (MLFF) schemes,^19–21 which have recently found use in the study of WDM.^22–29 However, such schemes generally require an extensive training dataset, comprised of tens to hundreds of thousands of atomic configurations,^24,30–32 which is not only a computationally and labor intensive process, but generally needs to be repeated for different conditions. This limitation can be overcome by employing on-the-fly MLFF training during molecular dynamics (MD) simulations.^33–38 However, the efficiency and efficacy of such a scheme has not been explored in the context of WDM heretofore.

In this work, we develop a framework for on-the-fly MLFF MD simulations. In particular, we employ an MLFF scheme based on the kernel method and Bayesian linear regression, with the Gauss SQ method used for generation of the Kohn–Sham training data, which includes energies, atomic forces, and stresses. Through comparisons with recent Kohn–Sham results in the literature, we show that the framework is able to accelerate AIMD simulations of WDM by up to three orders of magnitude, while retaining ab initio accuracy. We apply this framework to calculate the diffusion coefficients and shear viscosity of C, H, and CH at a density of 1 g/cm³ and temperatures in the range of 75 000–750 000 K, where we find agreement with previous results for C and CH but disagreement for H, demonstrating the need for ab initio calculations as presented here.

The remainder of this paper is organized as follows. In Sec. II, we discuss the formulation and implementation of the framework for on-the-fly MLFF MD simulations of WDM. In Sec. III, we first verify its accuracy and performance and then apply it to the study of warm dense CH. Finally, we conclude in Sec. IV.

We have developed the framework outlined in Fig. 1 for performing on-the-fly MD simulations of WDM, referred to as SQ-MLFF. In particular, we employ the Gauss variant of the SQ method for generation of the Kohn–Sham DFT data for training, which includes energies, atomic forces, and stresses, as described in Sec. II A, and an MLFF scheme based on the kernel method and Bayesian linear regression for prediction of the quantities of interest, as described in Sec. II B. The quantities that form a part of both training and prediction are the energy (⁠ $E$⁠), atomic forces (f), and stress tensor (⁠ $σ$⁠), where the superscript, if present, denotes the method used for its calculation.

The MD simulation starts with a few SQ calculations that provide the initial training data. This is then followed by predictions from the MLFF model in subsequent MD steps, except when the uncertainty in forces so computed, i.e., Bayesian error, exceeds the threshold $s tol$⁠, at which point an SQ calculation is again performed, data from which is added to the training dataset. The threshold $s tol$ is not fixed, but is set to the maximum value of the Bayesian error in forces for the MLFF based MD step subsequent to the training step. To avoid the cubic scaling bottleneck in training, a two-step data selection procedure is adopted, wherein only those atoms whose Bayesian error in forces exceeds a predefined threshold are added to the training dataset, and then, CUR³⁹ is performed on the resulting dataset for downsampling.

In its complete form, the Gauss SQ method exploits the nearsightedness of matter,⁴² neglected here for simplicity of presentation. In particular, within a real-space representation, the finite-temperature density matrix has exponential decay away from its diagonal,^43–46 a consequence of the locality of electronic interactions. This can be exploited to restrict the bilinear forms to be spatially localized, i.e., use of nodal Hamiltonians^6,7,43 rather than the full Hamiltonian, which enables the Gauss SQ method to scale linearly with system size, with increasing efficiency at higher temperatures as the density matrix becomes more localized and the Fermi operator becomes smoother.^6,7,46 On increasing quadrature order and localization radius, convergence to exact cubic scaling diagonalization results is readily obtained.^6,7,40 The SQ method also provides results corresponding to the infinite-crystal without recourse to Brillouin zone integration or large supercells,^6,7,41 a technique referred to as the infinite-cell method.

The SQ-MLFF framework is amenable to the calculation of ionic transport properties, i.e., diffusion coefficients, shear viscosity, and ionic thermal conductivity. In particular, these quantities can be determined from MD simulations using the Green–Kubo relations, wherein macroscopic dynamics properties are written as time-integrals of microscopic time correlation functions.⁵⁰ Though the ionic thermal conductivity—a quantity that is incompatible with the DFT formalism—becomes accessible in SQ-MLFF, it is not computed in the current work, since the electronic contribution is expected to be more dominant, particularly for the conditions of interest.³² 

We have developed a parallel implementation the SQ-MLFF formalism in the SPARC electronic structure code.^54–57 Choosing diffusion coefficient and shear viscosity as the target properties of interest, we first verify the accuracy and performance of SQ-MLFF for warm dense carbon (C) in Sec. III A and then apply it to the study of a warm dense carbon hydrogen mixture (CH) in Sec. III B. The associated data can be found in  Appendix. In all cases, we perform isokinetic ensemble (NVK) MD simulations with the Gaussian thermostat⁵⁸ for 500 000 steps, with the initial 5000 steps used for equilibration and the remaining used for production. In so doing, the statistical errors in the self-diffusion coefficient, inter-diffusion coefficient, and viscosity are reduced to within $4 × 10 − 4$ cm²/s, 0.005 cm²/s, and 0.1 mPa s, respectively.

In the DFT calculations, we adopt the local density approximation (LDA) for the exchange-correlation,^2,59 and optimized norm-conserving Vanderbilt (ONCV) pseudopotentials⁶⁰ with all electrons in valence and cutoff radii of 0.5 and 0.63  bohr for H and C, respectively, consistent with previous work at such conditions.^12,61 All DFT parameters, including the mesh size, and the Gauss SQ parameters, namely, truncation radius and quadrature order, are chosen such that the forces and stresses are converged to within 1% and 2%, respectively. This translates to the discretization error in the diffusion coefficients and viscosity being within 0.5% and 1%, respectively. Note that the error further decreases with higher grid resolution, truncation radius, and quadrature order. In the MLFF calculations, the hyperparameters are chosen to be the same as that found optimal in previous work.³⁸ Indeed, the performance of the MLFF scheme has been found to relatively insensitive to their choice, even for the extreme conditions considered here.

We consider C at a density of 10 g/cm³ and temperatures in the range of 100 000–2 000 000 K, specifically, 100 000, 200 000, 500 000, 750 000, 1 000 000, and 2 000 000 K. In the MD simulations, we consider system sizes of 200 and 64 atoms for the lowest and highest temperatures, respectively, and interpolate linearly for temperatures in between. MD time steps ranged from 0.09 and 0.019 fs, depending on temperature.

In Fig. 2, we present the diffusion coefficient and viscosity computed using SQ-MLFF and compare it with recent Kohn–Sham DFT results in the literature,¹² obtained using the SQ method. We observe that there is very good agreement between the SQ-MLFF and SQ-DFT results, with a maximum difference of 0.002 cm²/s in the diffusion coefficient. The viscosity, which was only computed at a density of 16 g/cm³ and 2 000 000 K in the cited reference, is also in very good agreement when the same conditions are simulated here, with a difference of only 0.4 mPa s, i.e., 43.2 mPa s from SQ-MLFF vs 42.8 mPa s from SQ-DFT. We also observe that both the diffusion coefficient and viscosity demonstrate a power law behavior with temperature, in agreement with previous results for dense plasmas.^62–64 In particular, we obtain a exponent of 0.96 for the diffusion coefficient, which is in very good agreement with the exponent of 0.95 obtained from SQ-DFT.¹² These results demonstrate the accuracy of SQ-MLFF for the calculation of ionic transport properties of WDM. Though not the focus of the work, the accuracy of SQ-MLFF extends to equations of state calculations, with the pressure computed in the above simulations being in agreement with SQ-DFT results to within 0.6% ( Appendix).

In Table I, we present the performance of SQ-MLFF for the MD simulations described above. We observe that the number of DFT steps, which constitutes of only ∼0.06% and ∼0.02% of the total number of MD steps for the lowest and highest temperatures, respectively, decreases with temperature. This can likely be attributed to the decrease in quantum mechanical effects with temperature, making it more amenable to MLFF development. The occurrence of these DFT steps during the MD simulation can be found in Fig. 3. We observe that most of the DFT calculations occur toward the beginning of the MD, with decreasing frequency as the simulation progresses. This is similar to the observations made for ambient conditions,^33,34,38 with one noticeable difference being that in the current simulations, a significantly larger number of DFT steps are performed even after a few thousand MD steps, likely due to the increased movement of atoms causing new configurations to be encountered later in the MD simulation. Even though there are so few DFT steps, they still constitute ∼90%–95% of the total CPU time. In terms of efficiency, it is estimated that there is speedup of three orders of magnitude and more by SQ-MLFF relative to SQ-DFT in both CPU and wall time. Though the SQ formalism scales to many tens of thousands of processors and beyond,^4,65 the current MLFF implementation scales only up to the number of processors equal to the number of atoms. Therefore, the percentage of total CPU time taken by DFT and the speedup in wall time achieved by MLFF will reduce as the number of processors increases beyond the number of atoms in the system. Indeed, the MLFF code can be parallelized to scale to an order of magnitude larger number of processors, ensuring that three orders of magnitude speedup is also achieved in the wall time for such simulations. The speedup of SQ-MLFF increases with the number of MD steps; therefore, the numbers for simulations targeting the diffusion coefficient alone will be significantly smaller, e.g., the speedup to achieve a statistical error of 1%, which requires around ∼25 000 steps, is estimated to be around two orders of magnitude.

We consider CH at a density of 1 g/cm³ and temperatures in the range of 75 000–750 000 K, specifically, 75 000, 100 000, 200 000, 500 000, and 750 000 K. In the MD simulations, we consider system sizes of 216 and 108 atoms for the lowest and highest temperatures, respectively, and again interpolate linearly for temperatures in between. MD time steps ranged from 0.04 and 0.004 fs, depending on temperature. The Coulomb coupling parameter (Γ) and electron degeneracy parameter (Θ) so calculated from the DFT steps for the conditions considered here range from 0.92 to 1.60 and 0.90 to 4.90, respectively ( Appendix), which falls within the WDM regime.⁶⁶ The computed pressure can also be found in  Appendix.

In Fig. 4, we present the variation of the diffusion coefficients for CH—self-diffusion coefficients for C and H, and inter-diffusion coefficient for CH—with temperature. We also compare these results with those obtained for C and H at the same conditions, again computed using SQ-MLFF. We observe that the inter-diffusion coefficient has a power law behavior with temperature of $D CH ∼ T 1.09$⁠, with values that are closer to the self-diffusion coefficient of H than C in CH. Such a power law behavior is also observed for C and H, consistent with observations for other materials in the literature.^{62–64,67,68} In addition, we find that at each temperature, the value of the inter-diffusion coefficient is relatively close to the molar fraction weighted average of the self-diffusion coefficients of the individual species, i.e., $D CH ∼ 1 2 ( D C + D H )$⁠, with a maximum difference of 0.06 cm²/s, suggesting a relatively weak coupling between C and H in CH. This is substantiated by the self-diffusion coefficients of C and H being relatively close to their corresponding values in CH. We also observe that at 232 120 K, the value of $D CH$ as obtained from the power law fit is 0.21 cm²/s, which is in relatively good agreement with recent work,⁵¹ where the mean value between different methods, including the average atom method and orbital-free DFT, is 0.24 cm²/s, with a relatively large standard deviation of 0.12 cm²/s. Though this is also true for C, there is a significant difference in the predictions for H, suggesting that higher fidelity is required for simulations of H in the WDM regime.

In Fig. 5, we present the variation of the viscosity with temperature for CH. We also compare these results with the viscosity of C and H at these conditions, again computed using SQ-MLFF. We observe that the viscosity has a power law behavior with temperature of $η ∼ T 0.94$⁠, with values that are closer to the viscosity of C. Such a power law behavior has been observed for other materials in the literature.^62,64 We also observe that at 232 120 K, the value of the viscosity as obtained from the power law fit is 4.14 mPa s, which is in relatively good agreement with recent work,⁵¹ where the mean value between different methods, including the average atom method and orbital-free DFT, is 5.33 mPa s, with a relatively large standard deviation of 3.03 mPa s. Though this is also true for C, there is again a significant difference for H, further substantiating the need for higher fidelity in the simulations of H in the WDM regime.

In this work, we have developed a framework for on-the-fly MLFF MD simulations of WDM, wherein we have employed an MLFF scheme based on the kernel method and Bayesian linear regression, with the required training data generated from Kohn–Sham DFT using the Gauss SQ method, within which we calculate energies, atomic forces, and stresses. Choosing warm dense carbon as a representative example, we verified the accuracy of the formalism through comparisons with recent Kohn–Sham DFT results in the literature. In so doing, we have demonstrated that AIMD simulations of WDM can be accelerated up to three orders of magnitude while retaining ab initio accuracy. We have applied this framework to calculate the diffusion coefficients and shear viscosity of CH at a density of 1 g/cm³ and temperatures in the range of 75 000–750 000 K. We found that the self- and inter-diffusion coefficients and the viscosity obey a power law with temperature, and that the diffusion coefficients suggest that the coupling between C and H is relatively weak in CH. In addition, we found agreement within standard deviation with previous results for C and CH but disagreement for H, demonstrating the need for ab initio calculations as presented here. Overall, the proposed formalism promises to enable simulations of WDM with ab initio accuracy at a small fraction of the original cost.

The inclusion of internal energy within the machine learned model will help accelerate equation of state calculations such as the Hugoniot, making it a worthy subject for future research. Extending the current MLFF implementation to enable efficient scaling on large scale supercomputers, with GPU acceleration of the key computational kernels,⁶⁹ will further reduce the wall time for large MD simulations, making it another worthy subject for future research. Finally, extending this framework to enable Δ-machine learning with respect to orbital-free DFT,³⁸ the accuracy of which increases with temperature, is likely to provide substantial gains provided that sufficiently accurate pseudopotentials are developed for orbital-free DFT, making it another worthy subject for future research.

The authors gratefully acknowledge support from Grant No. DE-NA0004128 funded by the U.S. Department of Energy (DOE), National Nuclear Security Administration (NNSA). J.E.P gratefully acknowledges support from the U.S. DOE, NNSA: Advanced Simulation and Computing (ASC) Program at Lawrence Livermore National Laboratory (LLNL). This work was performed in part under the auspices of the U.S. DOE by LLNL under Contract No. DE-AC52-07NA27344. This research was also supported by the supercomputing infrastructure provided by Partnership for an Advanced Computing Environment (PACE) through its Hive (U.S. National Science Foundation through Grant No. MRI-1828187) and Phoenix clusters at Georgia Institute of Technology, Atlanta, Georgia.

The authors have no conflicts to disclose.

Shashikant Kumar: Conceptualization (equal); Investigation (lead); Methodology (equal); Software (lead); Validation (lead); Writing – original draft (lead); Writing – review & editing (equal). Xin Jing: Investigation (supporting); Methodology (equal); Software (equal); Validation (supporting); Writing – review & editing (supporting). John E. Pask: Conceptualization (equal); Funding acquisition (supporting); Investigation (supporting); Methodology (equal); Writing – original draft (supporting); Writing – review & editing (equal). Phanish Suryanarayana: Conceptualization (lead); Funding acquisition (lead); Investigation (equal); Methodology (equal); Project administration (lead); Supervision (lead); Validation (equal); Writing – original draft (supporting); Writing – review & editing (lead).

The data that support the findings of this study are available within the article and from the corresponding author upon reasonable request.