Extending OpenKIM with an Uncertainty Quantification Toolkit for Molecular Modeling

Molecular modeling is an important part of materials science, and interatomic potentials (IPs) are at the heart of most molecular modeling simulations[1]. Given the large number of IPs constructed for various applications, there is an acute need to standardize their computational implementation and evaluation and facilitate portability among researchers. To this end, the Open Knowledgebase of Interatomic Models (OpenKIM) was founded to make atomistic-scale simulations reliable, reproducible, and accessible[2, 3]. While the OpenKIM framework provides many tools for materials simulations, it does not provide a standardized tool for uncertainty quantification (UQ). In this paper, we describe a novel UQ toolbox within the OpenKIM framework, targeted at molecular modeling.

The general goal of molecular modeling is to predict properties of materials by simulating collections of atoms. In this setting, IPs are used to approximate the interaction energy of atoms as functions of the atomic positions and species[4]. IPs are used in conjunction with a simulation program, such as ASE[5] or LAMMPS[6], to model atomic behavior and extract material properties. Such simulations can be static, e.g., minimizing energy to obtain the equilibrium lattice parameter of a crystal, or dynamic, e.g. applying the fluctuation–dissipation theorem to compute thermal conductivity. The accuracy of the material properties predicted by an atomistic simulation depends strongly on the choice of IP, and considerable effort has gone into developing IPs that are accurate for specific applications.

UQ is an emerging field of applied mathematics that aims to quantify and reduce uncertainties in mathematical models[7]. In molecular modeling, UQ can help assess the reliability of conclusions drawn from atomistic simulations. It is widely recognized that the largest source of uncertainty in molecular modeling is the functional form of the IP. Additional uncertainty comes from the values of the corresponding parameters, which are typically fit to experimental data or first-principles (quantum mechanical) calculations [8]. Having been fit to data, these IPs are often used to calculate out-of-sample material properties or to conduct large scale simulations. These simulations generally suffer from inconsistent transferability, i.e., they may struggle to accurately predict material properties to which the IPs are not fit[9]. The UQ process is especially important for assessing the reliability of these out-of-sample predictions.

There are many UQ methods that have been used in molecular modeling, including $F$-statistics estimations [10], ANOVA-based methods [11], multi-objective optimization [12], and profile likelihoods[13]. Among these methods, the most frequently utilized is a Bayesian method known as Markov chain Monte Carlo (MCMC)[14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25]. As such, we focus on MCMC methods, though we anticipate future extensions that will make other tools available to the molecular modeling community.

A universal challenge in applying UQ methods to multi-parameter models, such as IPs, is that models are often sloppy [26, 27, 13]. Sloppiness refers to an extremely ill-conditioned inference problem when fitting parameters to data, which is ubiquitous in many scientific fields [28, 29, 30, 31, 32, 33, 34, 35, 36, 37]. For sloppy models, many parameter combinations are not well-constrained by available data and dubbed practically unidentifiable. This poses a number of challenges to clearly formulating and interpreting UQ results, as shown, for example, for molecular modeling [38, 13]. These subtleties need to be considered, as we demonstrate later in this paper, when performing UQ analyses.

Although many software or libraries for performing Bayesian sampling or other UQ methods exist, such as emcee[39], Chaospy[40], and EasyVVUQ[41] to name a few, but there are only few libraries that integrate UQ for molecular modeling, such as potfit[42, 25]. In this paper we introduce a UQ toolkit integrated within the OpenKIM framework to facilitate the application of UQ to molecular modeling. Integrating UQ capabilities directly into the OpenKIM framework allows for more uniform, reproducible results and helps to standardize the practice of reporting uncertainty as part of a molecular modeling workflow. In Sec. II, we introduce the theory behind UQ and describe in more detail the OpenKIM framework. We describe our specific UQ implementation in Sec. III and how to use it, with an example, in Sec. IV. Finally, we conclude and describe future directions for this UQ toolkit in Sec. V.

The KLIFF package provides a convenient environment for fitting IPs. In this work, we introduce a new toolkit for use with KLIFF that provides a framework that facilitates transparent, reproducible UQ analysis of IPs. As MCMC is the UQ method of choice in molecular modeling, we focus on this method first with the intent to add alternative sampling methods in the future. A typical workflow is as follows: First, a modeler selects (or defines) the model and formulates the posterior sampler by instantiating kliff.uq.MCMC, which requires the specification of a prior and sampling temperature. Next, an MCMC simulation is performed with the run_mcmc function to generate ensembles of parameters drawn from the target Bayesian posterior distribution. This process is represented in Fig. 1 and further discussed below.

The first step in this workflow is to create the model. The model comprises the parameterized IP as well as atomic configurations, reference data to calibrate the IP parameters, and a simulation that calculates the corresponding material properties. KLIFF is used to construct a family of parameterized models based on an OpenKIM IP and read the configuration files that contain the atomic configurations and reference data to calibrate the IP. By default, KLIFF uses a simulation that computes the energy, forces, and stress for each configuration, although extending KLIFF to use other simulation and the corresponding material properties is also possible (see Ref. [60]). Then, KLIFF uses this information to construct a weighted least squares cost function for the model (see Eq. (1)). An example of this process is shown in the Python script in Fig. 2.

In previous versions (version 0.3.3 or earlier), KLIFF allowed for a single weight for each type of material property in the cost function, e.g., a single weight for all energies, a single weight for all forces, and so on. To enable more flexible UQ sampling, we expand KLIFF to enable specification of custom weights for each data point. Specific support is added for the weight calculation as defined in Eq. (2) that takes arguments $c_{1}$ and $c_{2}$. The default values are $c_{1}=1.0$ and $c_{2}=0.0$, which corresponds to a uniform weight for each data point. (See MagnitudeInverseweight in Fig. 2 for setting $c_{1}$ and $c_{2}$ to values other than the defaults.) Alternatively, users can define their own method to compute the weights. This update is included in KLIFF version 0.4.0.

The UQ framework is implemented as module uq in KLIFF. The posterior sampler is constructed by creating a kliff.uq.MCMC instance. Internally, this action computes $T_{0}$, generates a temperature ladder, and defines the untempered log-likelihood and log-prior functions. As a default, this class inherits from the sampler in the ptemcee Python package to perform PTMCMC. The algorithm simulates multiple chains (several at each sampling temperature) in parallel and mixes chains from different sampling temperatures to allow the MCMC walkers to explore a wider range of parameters. The number of chains or walkers can be specified through an optional argument nwalkers; the default value is twice the number of parameters in the model. This process is illustrated in the listing in Fig. 3.

The arguments to instantiate kliff.uq.MCMC are (1) a kliff.loss.Loss instance, which defines the untempered ($T=1$) likelihood function in Eq. (4), (2) the prior, and (3) the sampling temperatures. Since uniform priors are a common choice, the constructor implements this by default. In this case, the boundaries of the prior support need to be specified by the user via the logprior_args argument. However, the user can also pass a custom prior as the logprior_fn argument.

There are two options to specify the temperature ladder. First, the user may specify the number of temperatures (ntemps) and the ratio between the maximum temperature and $T_{0}$ in Eq. (10) (Tmax_ratio). In this case, an internal function generates a logarithmically spaced temperature between $T=1.0$ and $T=T_{\text{max\_ratio}}\times T_{0}$, inclusive. Alternatively, the user may specify the complete list of temperature values (Tladder). When Tladder is specified, then this list overrides the values passed for ntemps and Tmax_ratio.

The UQ implementation in KLIFF supports parallelization over the configurations for the cost function evaluation and over the walkers for the MCMC sampling. However, the current implementation only supports OpenMP-style parallelization for the cost function evaluation and both OpenMP and MPI for the MCMC sampling, with a future work to allow MPI in the earlier. In this paper we only show an example of parallelization in the Bayesian sampling, which is done by declaring a multiprocessing pool after instantiating kliff.uq.MCMC. The optimal number of parallel processes in this case is the product of the numbers of the sampling temperatures and the walkers.

After constructing the posterior sampler, MCMC sampling is run by calling the run_mcmc method of the kliff.uq.MCMC class (see Fig. 4). This function requires the initial position of each walker as an $L\times J\times N$ array, where $L,J,$ and $N$ are the number of temperatures, walkers, and parameters, respectively. The user also needs to specify the number of iterations to run the MCMC simulation.

The convergence of the parameter chains is assessed by calculating $\hat{R}^{p}$. In our implementation, this is realized by the function kliff.uq.rhat, demonstrated by the listing in Fig. 5. This function takes an array containing the MCMC samples for one sampling temperature. The $\hat{R}^{p}$ is calculated for each sampling temperature separately. Note that in practice, some sampling temperatures may converge much sooner than others. If the resulting values are larger than some threshold (typically 1.1), then the MCMC algorithm should continue to iterate. If $\hat{R}^{p}$ is less than the target threshold, the samples are assumed to have converged to the posterior and the calculation is terminated.

This UQ framework is integrated in KLIFF and we provide some examples in an online repository[64]. In the next section, we describe the use and interpretation of a UQ calculation for a Stillinger–Weber potential.

Having described the basic interface to this UQ toolkit, we now consider the concrete example given in the combined listings in Figs. 2–5. In this section we discuss the specific model and MCMC setup in these scripts. We then present and discuss the sampling results and some of the subtleties associated with their interpretation.

This example is based on the Stillinger–Weber potential in OpenKIM[65, 66]. This is a cluster potential originally introduced to model silicon [67]. For a system with $n$ atoms, the potential energy of atom $i$, $\mathcal{V}_{i}$, is given by

where $\phi_{m}$ denotes the $m$-body interactions. The Stillinger–Weber potential includes both two-body and three-body interactions. The two-body term only depends on the distance between atom $i$ and $j$, denoted by $r_{ij}$,

The three-body term additionally depends on $\beta_{ijk}$, which is the bond angle between the $i-j$ and $i-k$ bonds,

This IP includes nine parameters: $A$, $B$, $\sigma$, $p$, $q$, $r^{\text{cut}}$, $\lambda$, $\beta^{0}$ and $\gamma$. The total energy of the system $\mathcal{V}$ is

The atomic forces are calculated by taking the negative gradient of Eq. (13) with respect to the atomic positions.

In this example, we choose the tunable parameters to be $A,B,\sigma,\lambda,$ and $\gamma$. The other parameters are fixed to the default values given in OpenKIM[66]. Physically, the tunable parameters $A$ and $\lambda$ set energy scales in the potential, and $\sigma$ and $\gamma$ set length scales. To be physically relevant, these parameters are constrained to be positive. Parameter $B$ controls the relative scale of the repulsive part of the interaction and, thus, is also constrained to be positive. To enforce this constraint, we use a log-transform, i.e., the tunable parameters are $\theta=(\log(A),\log(B),\log(\sigma),\log(\lambda),\log(\gamma))$.

The training set consists of the energies and forces of silicon in the diamond cubic crystal structure. We use 400 atomic configurations, including stretched and compressed cells with random perturbations. The reference energy and force data were generated using the environment-dependent interatomic potential (EDIP) [68, 69, 70].²²2Since our objective is to explore UQ rather than develop an accurate model for silicon, we take the EDIP potential to be the “exact” ground truth. This greatly reduces the computational cost of generating the training set relative to DFT. Next, we compute the weights using Eq. (2), with $c_{1}=10^{-2}$ eV/Å for the forces and zero for the energies, and $c_{2}=10^{-1}$ for both.

The training produces the following best fit parameters for the SW potential:

The natural temperature for this model is $T_{0}=1.324$. Notice that the reference data were generated using another IP. Thus, while there is some model inadequacy, it is relatively small compared to some other real-world examples, which explains the low $T_{0}$ value; in practice $T_{0}$ is typically orders of magnitude larger[13]. In our analysis, we extend the temperature ladder to a much higher temperature to mimic typical effects users may encounter in real-world applications.

We use a uniform prior where $\pi(\theta)$ is constant if $-8<\theta<8$ and zero otherwise. The support of this prior is chosen to be sufficiently wide to ensure that sampling is not artificially restricted to regions near the best fit.

Next, PTMCMC sampling is performed for 150,000 iterations. From each walker, we discard the first 10,000 steps as the burn-in time. We thin the remaining chains by keeping every 200-th step to ensure uncorrelated samples. From the remaining samples, the maximum value of $\hat{R}^{p}$ is 1.046.

The posterior distribution from which samples are drawn is a joint distribution for the five parameters. Because we cannot directly visualize such a high dimensional space, it is common to instead plot marginal distributions. The marginal distribution of a high-dimensional, joint probability distribution is the projection of the probability onto a lower dimensional subspace.³³3The marginal distribution of a parameter is calculated by summing the conditional distributions of that parameter over all possible values of the other parameters. We project the five-dimensional posterior probability distribution onto each of the one-dimensional parameter axes in Fig. 6. The marginal distributions at different sampling temperatures are shown in different colors and superimposed to enable direct comparison. Each column corresponds to a different parameter in the potential. In the second row, a log scale is used for the vertical axis to bring out the details of the sparser distributions that result at higher sampling temperatures.

Examining Fig. 6, first note the general trend that distributions becomes wider as the temperature increases. This is expected, since higher temperatures correspond to smaller weights, i.e., larger effective error bars, in Eq. (1). However, the effect is parameter-dependent. For example, consider the distributions of $\log(\lambda)$ and $\log(\gamma)$ at $T=10^{2}$. At this temperature, the distributions are relatively localized around their best fit values. However, at the next highest temperature, $T=10^{3}$, the distributions extend to the boundary of their respective priors. This phenomenon, in which the marginal posterior distribution of some parameters abruptly transitions away from being localized at a specific sampling temperature, is called parameter evaporation. Inspecting Fig. 6 reveals that $\log(\lambda)$ and $\log(\gamma)$ evaporate around $T=10^{3}$ while parameters $A$ and $B$ evaporate around $T=10^{4}$. At a sufficiently high sampling temperature, all parameters would evaporate, regardless of the prior, as show in Ref. [13] for SW potential for a molybdenum disulfide system.

Parameter evaporation has been observed in Ref. [13] for IPs and has important implications for UQ analysis. To see this, first notice that while the distribution for $A$ and $B$ remain localized at $T=10^{3}$, they have shifted relative to their distributions at $T=10^{2}$, i.e., they are localized around different values. We can explain this shift in terms of the evaporation of the other parameters. At $T=10^{3}$, the range of values sampled for the parameters $\lambda$ and $\gamma$ are very different from those at $T=10^{2}$. Consequently, the sampled values for the parameters $A$ and $B$ shift to minimize the cost with this more diffuse distribution of $\lambda$ and $\gamma.$ That is to say, the values for parameters $A$ and $B$ are biased as a consequence of having the sub-optimal values of $\lambda$ and $\gamma$. While this shift is not inherently problematic, notice that the distributions for $A$ and $B$ at $T=10^{3}$ have very little overlap with their counterparts at $T=10^{2}$. Because the lower temperature distribution is dominated by samples near the best fit parameters, we infer that the higher temperature distribution has very few samples near the best fit. This is potentially problematic since it implies that parameter values that best fit the data are not represented in the sample.

Figure 7 shows the distribution of the untempered costs at each sampling temperature. As expected, the average value of the cost increases at higher temperatures. However, this increase is not the result of stretching the distribution, as is the case for linear regression model. Rather, the entire distribution shifts to the right at each rung in the temperature ladder. As the temperature increases and parameters evaporate, the posterior is dominated by regions of the parameter space that are poor fits to the data. These are regions of parameter space in which the prior places considerable weight, that is, they are high-entropy regions. This is a nonlinear effect due to the subtle interplay between the sampling temperature, the prior, and the degenerate modes of sloppy models. The resulting posterior distribution can depend very strongly on details of the problem, such as the choice of prior, the error bars for the data, and the sampling temperature.

In general, we recommend that practitioners check their results for robustness for several priors over a range of temperatures. In this work, we have used uniform priors in log-transformed parameters as a pedagogical example. This is a reasonable choice in general; however, we could have also used uniform priors in the original, untransformed parameterization as well as Gaussian priors in both log-transformed and untransformed parameters. We caution against using Jeffreys prior, as its interaction with the degenerate modes can lead to strong biases[71]. In this example, we have sampled up to a relatively high temperature relative to $T_{0}$ as defined in Eq. (10). This was also a pedagogical choice to illustrate important effects and mimic more realistic examples (for example, previous work has seen values of $T_{0}>10^{6}$[13]). In practice, we suggest including sampling temperatures up to a few times larger than $T_{0}$. The highest sampling temperatures are not for the purposes of UQ, but including them in the sampling algorithm improves convergence rates. For UQ analysis, we suggest considering ensembles generated by temperatures from $50\%$ below to $50\%$ above $T_{0}$.

IP developers should also use uncertainty quantification tools throughout the model development cycle. For example, they could extend the training data to better constrain the sloppy parameters. Alternatively, they could use model reduction methods[72] to remove the degenerate modes from the model. For the latter, care must be taken to not remove unidentifiable parameters that would be important for downstream applications.

In this work, we describe a UQ toolkit that extends the KLIFF package as part of the OpenKIM environment. This toolkit provides a framework that facilitates transparent, reproducible UQ analysis for both the development and application of IPs in molecular modeling. We focus on a Bayesian, MCMC approach for quantifying parametric uncertainty and model inadequacy. We use a parallel-tempered MCMC method, which improves convergence and allows us to mimic the effect of model inadequacy in our uncertainty estimates.

In our implementation we focused on ease of use in order to lower the barrier to entry for molecular modeling practitioners. However, we caution users not to treat these methods as off-the-shelf black boxes. UQ is an emerging field with many open questions, especially surrounding model inadequacy which is often the dominant problem in atomistic simulations. We encourage IP practitioners and developers alike to familiarize themselves with statistical subtleties related to sloppiness and parameter unidentifiability[47, 13] and check their conclusions for robustness over a range of sampling temperatures for multiple priors.

In future work, we plan to integrate other UQ methods within KLIFF, such as the frequentist profile likelihood[73]. To address the problems associated with UQ of sloppy IPs[13], we plan to integrate a model reduction scheme motivated by information geometry[72].

This work is supported by the National Science Foundation under awards DMR-1834332 and DMR-1834251. We would like to acknowledge the computational facilities provided by the Brigham Young University Office of Research Computing.