NeuralUQ: A comprehensive library for uncertainty quantification in neural differential equations and operators

Consider a nonlinear ODE/PDE describing the dynamics of a physical system as follows: \cref@addtoresetequationparentequation

where $x$ is the $D_x$-dimensional space-time coordinate, $\Omega$ is a bounded domain with boundary $\Gamma$, $\xi$ is a random event in a probability space $\Xi$, while $f(x;\xi )$ and $u(x;\xi )$ represent the $D_u$-dimensional source term and sought solution evaluated at $(x,\xi )$, respectively. Further, $\text{pazocal}{F}_{\lambda ,\xi }$ is a general differential operator; $\text{pazocal}{B}_{\lambda ,\xi }$ is a boundary/initial condition operator acting on the domain boundary $\Gamma$; and $b$ is a boundary/initial condition term. Furthermore, $\lambda (x;\xi )$, with $x\in {\Omega }_{\lambda }\subset \Omega$, denotes the problem parameters. We refer to problems where the operators $\text{pazocal}{F}_{\lambda ,\xi }$ and $\text{pazocal}{B}_{\lambda ,\xi }$ are known, i.e., the differential operators that define them are specified, as neural ODEs/PDEs, because we use neural networks (NNs) throughout the present work to solve them. We refer to problems involving unknown operators $\text{pazocal}{F}_{\lambda ,\xi }$ and $\text{pazocal}{B}_{\lambda ,\xi }$ as neural operators, because we use NNs to learn them.

For neural ODEs/PDEs, $\text{pazocal}{F}_{\lambda ,\xi }$, $\text{pazocal}{B}_{\lambda ,\xi }$ are known and deterministic, i.e. $\xi$ is fixed, while $\lambda \left(x\right)$, and $u\left(x\right)$ are partially unknown. Given noisy data of $u$ and $\lambda$ sampled at random finite locations in $\Omega$ and ${\Omega }_{\lambda }$, respectively, as well as noisy data of $f$ and $b$, the objective of this problem is to obtain the solution $u\left(x\right)$ at every $x\in \Omega$ and the problem parameters $\lambda \left(x\right)$ at every $x\in {\Omega }_{\lambda }$. The dataset in this case is expressed as $\text{pazocal}D=\{\text{pazocal}{D}_f,\text{pazocal}{D}_b,\text{pazocal}{D}_u,\text{pazocal}{D}_{\lambda }\}$, where $\text{pazocal}{D}_u=\{x_i,u_i{\}}_{i=1}^{N_u}$, $\text{pazocal}{D}_{\lambda }=\{x_i,{\lambda }_i{\}}_{i=1}^{N_{\lambda }}$, $u_i=u\left(x_i\right)$, ${\lambda }_i=\lambda \left(x_i\right)$, and $\text{pazocal}{D}_f,\text{pazocal}{D}_b$ are given as above. The PINN method, developed in [raissi2019physicsinformed], addresses the PDE problems deterministically, by constructing NN approximators, parameterized by $\theta$, for modeling the solution $u\left(x\right)$. The approximator $u_{\theta }\left(x\right)$ is substituted into Eq. (1) via automatic differentiation [tensorflow2015-whitepaper] for producing $f_{\theta }=\text{pazocal}{F}_{\lambda ,\xi }\left[u_{\theta }\right]$ and $b_{\theta }=\text{pazocal}{B}_{\lambda ,\xi }\left[u_{\theta }\right]$. For solving the mixed deterministic PDE problem, an additional NN can be constructed for modeling $\lambda \left(x\right)$. Finally, the optimization problem for obtaining $\theta$ in the mixed problem scenario with the dataset $\text{pazocal}D=\{\text{pazocal}{D}_f,\text{pazocal}{D}_b,\text{pazocal}{D}_u,\text{pazocal}{D}_{\lambda }\}$ is cast as

In this equation, $\{w_f,w_b,w_u,w_{\lambda }\}$ are objective function weights for balancing the various terms in Eq. (2). Following determination of $\hat{\theta }$, $u_{\hat{\theta }}$ and ${\lambda }_{\hat{\theta }}$ can be evaluated at any $x$ in $\Omega$, ${\Omega }_{\lambda }$, respectively.

For neural operators, $\text{pazocal}{F}_{\lambda ,\xi }$, $\text{pazocal}{B}_{\lambda ,\xi }$ are unknown and stochastic, while $\lambda (x;\xi ),u(x;\xi )$ are partially unknown. Note that the stochasticity of $\text{pazocal}{F}_{\lambda ,\xi }$ and $\text{pazocal}{B}_{\lambda ,\xi }$ arises from the stochasticity in $\lambda (x;\xi )$ with $\xi \in \Xi$. The available data consists of paired realizations of $f,b,u$, and $\lambda$, each pair associated with a different $\xi \in \Xi$. The realizations are considered noisy in general for $u$, and clean for $f,b$, and $\lambda$. Further, the objective of this problem is to learn the operator mapping defined by Eq. (1). The DeepONet method, developed in [lu2021learning] and designed for general nonlinear operator learning, addresses this problem by constructing an approximator using NNs, denoted by $G_{\theta }$ and parameterized by $\theta$, to model the mapping from the input function domain to the output function domain. In the following and without loss of generality, we consider both $f$ and $b$ in Eq. (1) as known and the problem becomes learning the mapping from the parameter $\lambda$ to the solution $u$. Further, each value of $\xi \in \Xi$ in Eq. (1) corresponds to a different $\lambda (x;\xi )$ function, and thus to a different solution $u(x;\xi )$. For any $\xi \in \Xi$, the approximator $G_{\theta }$ takes as input a representation of $\lambda (x;\xi ),x\in {\Omega }_{\lambda }$, and outputs a function, $G_{\theta }\left(\lambda \right)(x;\xi ),x\in \Omega$. The DeepONet consists of two sub-networks, one that takes the coordinates $x\in \Omega$ as input, called trunk net, and one that takes the discretized function $\lambda$ as input, called branch net. Considering a one-dimensional $u$, i.e., $D_u=1$, each of the two networks produces a $w$-dimensional output. The two outputs are merged by computing their inner product for producing the final prediction $G_{\theta }\left(\lambda \right)\left(x\right),x\in \Omega$. For learning the mapping from $\lambda$ to $u$ using the vanilla DeepONet [lu2021comprehensive], the minimization problem with the mean squared error (MSE) loss

is solved, using the training dataset $\text{pazocal}{D}^{\prime }=\{x_j,{\xi }_i,{\lambda }^{\left(i\right)},u_j^{\left(i\right)}{\}}_{i,j=1,1}^{N,N_u}$, where $\{x_j{\}}_{j=1}^{N_u}$ are the locations, $\{{\xi }_i{\}}_{i=1}^N$ are the sampled events and in this case represent the unique pairing between $u$ and $\lambda$, $u_j^{\left(i\right)}:=u(x_j,{\xi }_i)$ is the measurement of $u$ at the location $x_j$ for event ${\xi }_i$, and ${\lambda }^{\left(i\right)}$ is the representation of $\lambda (x;{\xi }_i)$. Note that, to be consistent with neural ODEs/PDEs, we use $u_{\theta }$ to denote the approximator for operators in the rest of this paper. In this case, $u_{\theta }=u_{\theta }(x;\lambda (\cdot ;\xi \left)\right)$. DeepONets, once trained, could act as PDE-agnostic surrogates encoding the hidden physics for further learning and inferences [meng2022learning]. Finally, note that in addition to supervised learning using data as described above, DeepONet can alternatively be trained by employing only the PINN loss during training. As shown in [goswami2021physicsinformed], the most effective training is the hybrid one, where both physics and data are used in the training loss; see also [wang2021learning].

In contrast to the point estimates obtained via the optimization problems of Eqs. (2) and (3), the goal of UQ is to estimate the predictive distribution

which is often referred to as Bayesian model average (BMA). In Eq. (4), $p\left(\theta |\text{pazocal}D\right)$ denotes the posterior density, which can be obtained using Bayes’ rule via

In Eq. (5), $p\left(\text{pazocal}D|\theta \right)$ is the likelihood of the data, i.e., $p\left(\text{pazocal}D|\theta \right)={\prod }_{i=1}^{N}p(u_i|x_i,\theta )$ for independent and identically distributed (i.i.d.) data; $p\left(\theta \right)$ is the probability density function for the prior distribution of the parameters $\theta$; and $p\left(\text{pazocal}D\right)$ is called marginal likelihood or evidence, because it represents the probability that we observe $\text{pazocal}D$. The evidence is given as

The uncertainty of parameters $\theta$ is typically referred to as epistemic uncertainty.

Generally, inferring $\theta$ given $\text{pazocal}D$, i.e., computing the posterior exactly via Eq. (5) is computationally and analytically intractable. This is addressed by approximate posterior inference, which approximates the posterior by another distribution and/or obtains representative samples from the posterior. Specifically, the Bayesian and the deterministic frameworks are the two main families of UQ methods for SciML. The former employs either a sampling method or a variational inference (VI) method to obtain samples from the posterior distribution, whereas the latter can be interpreted as an extension of standard training of deterministic NNs. Various UQ methods have been developed for SciML within these two frameworks. We adopt the unified framework proposed in [psaros2022uncertainty] to include all the UQ models employed in this study in Fig. 2, and present a short overview for them in the rest of this section.

For neural ODEs/PDEs, the solution $u$ and/or the problem parameters $\lambda$ are modeled with surrogates parameterized by $\theta$ and denoted by $\text{pazocal}{M}_1$ in Fig. 2. The Bayesian PINNs (B-PINNs) [yang2020bayesian] and the functional prior (FP) based on the physics-informed generative adversarial networks (PI-GANs) [meng2022learning, yang2020physicsinformed] are two recently proposed physics-informed learning approaches with UQ within the Bayesian framework. The Bayesian neural networks (BNNs) and the generators from trained PI-GANs are employed as the surrogate for $u/\lambda$ in these two methods, respectively. Further, either sampling or variational inference (VI) approaches can be utilized for obtaining the posterior samples for $\theta$. In addition to the aforementioned methods, the Monte Carlo dropout [gal2016dropout] (MCD) as well as the deep ensemble [lakshminarayanan2017simple] (DEns) approaches, which employ standard, i.e., deterministic, NNs as surrogates for $u$ and/or $\lambda$, can be combined with PINNs to quantify uncertainties in problems with known physics [zhang2019quantifying, psaros2022uncertainty].

For neural operators, a DeepONet parameterized by $\theta$ is employed as the surrogate for the solution $u$ and denoted by $\text{pazocal}{M}_2$ in Fig. 2. In the Bayesian framework, [lin2021accelerated] employed BNNs in DeepONet and a sampling method is then utilized to obtain the posterior samples for the parameters in BNNs. As for deterministic methods, extensions of standard training of DeepONets, such as DEns, are used for quantifying uncertainties in predictions [psaros2022uncertainty, yang2022scalable, pickering2022discovering]. These models correspond to Fig. 2 with $\text{pazocal}{M}_2$ being a DeepONet and the dashed arrows are not used. Further, the inference input in these models, i.e. the input to the branch net, should be represented in the same way as used in the training, i.e., the number of the discrete values as well as their locations should be the same. For a more general case in which the inference input and/or output data are noisy and measured at a certain number of random locations, [meng2022learning] proposed to combine a pre-trained DeepONet with a GAN-based FP to quantify the uncertainties in predictions. In [psaros2022uncertainty], a BNN is utilized as an alternative to the generator of a GAN in [meng2022learning]. These two models correspond to Fig. 2 with $\text{pazocal}{M}_1$ being either a BNN or the generator from a pre-trained GAN and $\text{pazocal}{M}_2$ a pre-trained DeepONet.

In all cases, different $\theta$ values correspond to different functions $u_{\theta }$. We refer to the surrogate combined with a distribution $p\left(\theta \right)$ for $\theta$ as a process. Following approximate posterior inference, the posterior samples for $\theta$ are used for obtaining the statistics of the function samples. In this regard, the value of $u$ at a location $x$ given the data $\text{pazocal}D$ is a random variable denoted by $(u|x,\text{pazocal}D)$. The mean of $(u|x,\text{pazocal}D)$ is denoted by $\hat{u}\left(x\right)=E[u|x,\text{pazocal}D]$ and approximated by $\overline{\mu }\left(x\right)$ as

where $\left\{u_{{\hat{\theta }}_j}\right(x){\}}_{j=1}^M$ is the set of NN predictions corresponding to the samples $\{{\hat{\theta }}_j{\}}_{j=1}^M$. Further, the diagonal part of the covariance matrix of $(u|x,\text{pazocal}D)$ is given by

In this equation, ${\overline{\sigma }}^2\left(x\right)$ represents the approximate total uncertainty of $(u|x,\text{pazocal}D)$, while ${\overline{\sigma }}_a^2\left(x\right)$ and ${\overline{\sigma }}_e^2\left(x\right)$ denote the aleatoric and the epistemic parts of the total uncertainty, respectively. In practice, a Gaussian (or, e.g., a Student-t) distribution is commonly fitted to the posterior samples. For example, we denote as $\text{pazocal}N\left(\overline{\mu }\right(x),{\overline{\sigma }}^2(x\left)\right)$ the Gaussian approximation to the samples $\left\{u_{{\hat{\theta }}_j}\right(x){\}}_{j=1}^M$ at an arbitrary location $x$, where $\overline{\mu }\left(x\right)$ and ${\overline{\sigma }}^2\left(x\right)$ are obtained via Eqs. (7) and (8), respectively.

Performance evaluation is necessary for model selection, e.g., prior/architecture optimization, comparison between UQ methods, and overall quality evaluation of the UQ design procedure [pickering2022structure]. As reviewed in [psaros2022uncertainty], performance evaluation metrics include, but are not limited to, the relative ${\ell }_2$ error for evaluating accuracy, the mean predictive likelihood for evaluating predictive capacity, and the root mean squared calibration error for evaluating statistical consistency; i.e., how closely the predicted model matches the data-generating distribution.

Further, calibration approaches can be adopted to improve statistical consistency in UQ for SciML [psaros2022uncertainty]. Various approaches have been developed for calibration, e.g., [levi2019evaluating] re-weights the predicted variance optimally to calibrate the output uncertainty; [kuleshov2018accurate] approximates the mis-calibration function and applies it to all predictive distributions, and thus modifies both the output mean and variance; and [zelikman2020crude] fits an empirical distribution to the scaled residuals, and also calibrates both the output mean and variance. To use these approaches, an additional dataset is required after training. This can be a left-out calibration dataset, similar to validation and test sets.

As mentioned in the previous section, the value of $u$ at a location $x$ given the data $\text{pazocal}D$ is a random variable denoted by $(u|x,\text{pazocal}D)$. The aim of UQ in the library is to obtain, calibrate, compute the statistics, and evaluate the quality of samples drawn from the posterior distribution $p\left(u|\text{pazocal}D\right)$. To that end, in the design phase a model is constructed that includes surrogates representing the solution $u$ and potentially $f$ and/or $\lambda$, as well as likelihood functions representing the data distribution. Surrogates are parameterized by $\theta$ and when combined with probability distributions of $\theta$ are represented as stochastic processes; i.e., a proposed family of solutions. To obtain the function samples for the solutions, a posterior inference method is combined with the model. Finally, the obtained function samples can be calibrated using additional data. In the revision phase, metrics are used for evaluating the performance and re-designing the solution procedure. A schematic workflow of the library is provided in Fig. 3 and Procedure 1.

As discussed in Section 2.2, there are two general UQ frameworks for SciML, namely the Bayesian and the deterministic frameworks. The former relies on Bayes’ theorem, and includes sampling methods, such as Monte Carlo Markov Chain (MCMC), and VI methods, such as mean-field variational inference (MFVI). The latter relies on standard NN training, i.e., minimizing/maximizing an objective function, typically using the MSE or the maximum a posteriori (MAP) metric, between measurements and NN predictions.

The procedural differences of the aforementioned methods necessitate the development of three different UQ method families in our library, i.e., the Samplable, Variational, and Trainable to include all the sampling-, VI- and standard-NN-based methods, respectively. Specifically, the sampling methods target directly the unnormalized posterior density function, without employing and optimizing any trainable variables. The VI-based methods employ pre-defined distributions parameterized by trainable variables, which are optimized to approximate the true posterior distributions. Finally, the standard-NN-based UQ methods obtain point estimates for the parameters in the surrogates using an optimizer, rather than directly sampling from or approximating the posterior distributions. A summary of the currently supported UQ methods in NeuralUQ is provided in Table 1.

In this section, we describe the main components, built-in functionalities and how to use NeuralUQ.

The components of NeuralUQ, including surrogates, variables, likelihoods, and inferences, are highly flexible and customizable. In this section, we describe how to configure the library for extending current capabilities and addressing problem scenarios beyond the ones considered in this study.

The Kraichnan-Orszag model consisting of three coupled nonlinear ODEs describes the temporal evolution of a system with several interacting inviscid shear waves, whose initial conditions are drawn from a Gaussian distribution [wan2006multi]. Here, we use NeuralUQ to solve the inverse Kraichnan-Orszag problem, i.e., identify the unknowns in the governing equations given sparse and noisy measurements on the solutions without the knowledge of the initial conditions. The governing equations are:

where $a=1$ and $b=1$ are constants.

Particularly for this problem, we assume that $a$ and $b$ are unknown, whereas 11 noisy measurements of $x_1$ and $x_3$ are available, and 7 of $x_2$; see also Fig. 5. The noise for all measurements is assumed to be Gaussian with zero mean and a known standard deviation set to 0.05. The objective here is to identify $a$ and $b$, as well as to reconstruct the complete solutions $x_1$, $x_2$, and $x_3$ for $t\in [0,10]$, with uncertainties.

We present the predictions for $x_1,x_2,x_3$ and $a,b$ using different approaches in Fig. 5 and Table 3, respectively. As shown in Fig. 5, for both HMC and DEns, the predicted uncertainty increases at locations with no available data. Further, the error between the predicted mean and the reference solution for $x_1,x_2,x_3$ is bounded by the total uncertainty. Furthermore, as shown in Table 3, the predicted values for $a$ and $b$, using both approaches, are close to the reference values. In the table, the results using MFVI and MCD are also provided for comparisons. In Table 4, we evaluate the predictions for $x_1,x_2,x_3$ using the RL2E and the NLL metrics. As shown, HMC and DEns provide the most reasonable predictions for the mean as well as uncertainties according to the computed metrics.

Scientific models are critical tools for understanding and forecasting the spread of the COVID-19. However, constructing models that are capable of accurately describing the COVID-19 still remains a challenge. Here, we demonstrate the capability of NeuralUQ for learning the corresponding model with uncertainties given measurements. In particular, we consider the SIRD compartmental model with time-dependent parameters/coefficients/rates for the COVID-19 [chen2020time, calafiore2020time, sen2021use], expressed as:

where $\beta$, the transmission rate, $\gamma$, the recovery rate, and $\nu$, the death rate, are unknown and assumed to be time-dependent functions, and $n$ is the total population, which is assumed to be constant and equal to the summation of $S,I,R,D$ throughout the time span. $S\left(t\right),I\left(t\right),R\left(t\right)$ and $D\left(t\right)$ are the number of individuals susceptible to the disease, the number of individuals infected, the (cumulative) number of individuals recovered, and the (cumulative) number of individuals deceased from the disease, respectively, at time $t$. In this example, we use the COVID-19 dataset from Italy [dong2020interactive], which is the same as in [kharazmi2021identifiability].

In this problem, we utilize the data on $I$, $R$, and $D$ in the first 160 days for training. The objective is to predict $I$, $R$, and $D$ for $t\in [161,200]$, and identify the unknown functions $\beta$, $\gamma$, and $\nu$ for $t\in [0,200]$. The results from HMC and DEns are presented in Fig. 6. We observe that the error between the predicted mean and the reference solution of $I/R/D$ from both approaches is bounded by the total uncertainties. Further, HMC and DEns yield different predictions for $\beta$ and $\nu$ for times $t<100$, while they are quite similar for $t>100$. Also, both methods predict a decrease of $\beta$ and $\nu$ with time. In addition, the predictions for $\gamma$ from both methods are similar, i.e., they are close to zero for $t\in [0,200]$. Note that reference solutions for the unknown functions $\beta$, $\gamma$, and $\nu$, are not available and thus corresponding comparisons are not provided here.

The Korteweg-de Vries (KdV) equation, derived more than 150 years ago, can be used to model surface waves in a shallow canal. It has been applied in various disciplines for modeling surface gravity waves, internal solitons in the ocean, magma flows, and conduit waves, for example. In this section, we solve an inverse KdV problem using NeuralUQ. Specifically, we discover the KdV equation given data for a system with interactions between two solitary waves. The KdV equation is expressed in the general form:

where ${\lambda }_1$ and ${\lambda }_2$ are constants, and $f=0$ is considered. The exact solution can be expressed as [benes2006decompositions]

where

and

where ${\lambda }_1=1.5$ and ${\lambda }_2=0.25$ are considered.

For the inverse problem, the objective is to identify ${\lambda }_1$ and ${\lambda }_2$, which are considered unknown, and to predict $u$ in the entire spatial-temporal domain, given partial noisy measurements on $u$ and $f$. Specifically, 200 and 100 random noisy measurements on $u$ and $f$, respectively, are considered available; see also Fig. 7. The noise for the measurements on both $u$ and $f$ is assumed to be Gaussian with zero mean and known standard deviation set to 0.1.

Shown in Figs. 7-7 are the errors between the HMC and DEns predicted means, for $u$ at $t=-1.5$ and 0.4, and the reference solutions bounded by the predicted total uncertainties. Further, the predicted values for ${\lambda }_1$ and ${\lambda }_2$, which are provided in Table 5, are close to the reference values for both methods.

Flow through porous media is studied in various disciplines, such as chemical engineering, e.g., for flows in packed-bed catalytic reactors, oil recovery, e.g., for the displacement of oil with water, and in geophysics e.g., for the transport of pollutant in soil. A steady flow through porous media in two dimensions is considered here, which is described by the Darcy’s law as follows:

where $\lambda$ is the hydraulic conductivity field, $u$ denotes the hydraulic head, and $f=50$ is considered. The boundary conditions are: \cref@addtoresetequationparentequation

where $n$ denotes the unit normal vector of the boundary.

In general, $\lambda (x,y;\xi )$ is determined by the pore structure. In this study, we employ a stochastic model for $\lambda (x,y;\xi )$ to consider porous media with different micro-structures. In particular, $\lambda (x,y;\xi )=\exp \left(\tilde{\lambda }\right(x,y;\xi \left)\right)$, where $\tilde{\lambda }(x,y;\xi )$ is a truncated Karhunen-Loève expansion of a Gaussian process with zero mean and the following kernel: \cref@addtoresetequationparentequation

and we only keep the first 100 leading terms of the expansion.

We consider two different scenarios here, i.e., Case A, in which we employ a DeepONet mapping $\tilde{\lambda }$ to $u$ pre-trained with 9900 different paired data ($\tilde{\lambda }$, $u$), and the test data are partial noisy measurements on $\lambda$ and $u$; and Case B, in which we employ the DEns method to train a DeepONet using the same training data as in Case A, which also maps $\tilde{\lambda }$ to $u$. Complete $\lambda$ measurements are used as test data, i.e., a measurement is available at each required location of the DeepONet input. Particularly, we refer to the DeepONet trained with DEns as uncertain DeepONet (U-DeepONet) following [psaros2022uncertainty].

In Table 6, we provide a glossary table of frequently used terms in this paper, and in ML/SciML more generally.

We demonstrate the surrogate customization, using a one-dimensional regression problem. The target function is:

We assume that we have 3 noisy measurements of $u$, which are equidistantly distributed in $[-0.7,-0.3]$. The noise for the measurements is assumed to be Gaussian with zero mean and known standard deviation set to 0.05.

We employ three different surrogates for different scenarios: (1) We have neither historical data nor prior knowledge on the formulation of the target function. We then employ a BNN with one hidden layer which has 50 neurons and the hyperbolic tangent activation function as the surrogate. In addition, the standard normal distribution is utilized as the prior distribution for each weight and bias in the BNN; (2) the expression describing the target function is not considered as known, but we have historical data for it. We then train a GAN to learn the functional prior using historical data; and (3) we have prior knowledge regarding the target function expressing that it is a sine function with unknown amplitude and frequency, i.e., $A\sin \left(\omega x\right)$, where $A$ and $\omega$ are parameters to be inferred given measurements. The details regarding the pre-trained GAN in (2) can be seen in the GitHub repository, and the posteriors for the parameters in each surrogate are computed using HMC.

As shown in Fig. 11, the predictions from the GAN-FP and $A\sin \left(\omega x\right)$ surrogates are quite similar for both the predicted means and uncertainties, which are more accurate than the results from the BNN surrogate in terms of the means. Also, the errors between the predicted means and the reference solution are mostly bounded by the predicted uncertainties from GAN-FP and $A\sin \left(\omega x\right)$, whereas the errors are only partly bounded by the predicted uncertainties in BNN. Overall, we demonstrate the flexibility and customizability of NeuralUQ by customizing problem-dependent surrogates to achieve better predicted accuracy.

We demonstrate the customization for prior in Samplable and posterior estimation in Variational, using a one-dimensional differential equation problem. Specifically, we consider a nonlinear diffusion-reaction system expressed as follows:

where $u$ represents the concentration of a certain solute, $D$ denotes the diffusion coefficient, $k_r$ is the reaction rate, and $f$ is the source term.

We first demonstrate the customization of prior using the example of an inverse problem described by Eq. (20). In this problem, $D=0.01$ is a known constant, and we have a few noisy measurements on $u$, and $f$. The objective is to infer $u$ and $f$ for $x\in [-1,1]$, and the reaction rate $k_r$ which is an unknown constant. Specifically, we have noisy measurements of $u$ at 5 random locations, and 17 noisy equidistant measurements of $f$. We then employ a BNN as well as a Samplable variable as the surrogate for $u$ and $k_r$, respectively. The equation can then be encoded using the automatic differentiation as in PINNs. The employed BNN has 3 hidden layers with 50 neurons per layers, and the hyperbolic tangent activation function is used as the activation function. We assume that each weight/bias has the same prior distribution in this BNN, which is independent Gaussian here, i.e., $N(0,1^2)$. Further, we employ three different prior distributions for $k_r$, i.e., $k_r\sim N(0,1^2)$, $k_r\sim HalfN(0.5,{0.5}^2)$ and $\log k_r\sim N(0,1^2)$, where $HalfN$ refers to half-normal distribution. The posterior samples for the parameters in BNN as well as $k_r$ are obtained using HMC. We only present the predictions for $k_r$ in Table 7, for simplicity. As shown, the predictions for $k_r$ are strongly affected by its prior distribution. Overall, NeuralUQ enables the users to utilize different prior distributions to achieve better accuracy.

We proceed to show the customization of posterior estimate in Variational using the example of a forward problem governed by Eq. (20). In particular, $D=0.01$, $k_r=0.2$, and the boundary condition is $u(-1)=u\left(1\right)=0$. In addition, we assume that we have 10 noisy measurements of $f$ which are randomly distributed at $x\in [-1,1]$. The noise for the measurements, which is $N(0,{0.01}^2)$, is assumed to be known. The objective is to infer $u$ given data of $f$. We employ a PI-GAN, which is the same as in [meng2022learning], to solve this problem. The same historical data in [meng2022learning] is utilized here to train the PI-GAN to obtain the functional priors for $u$ and $f$. Note that HMC is employed in [meng2022learning] for the posterior estimation of the parameters in PI-GAN. Here we utilize VI with two different parameterized posterior distributions to demonstrate the flexibility of NeuralUQ, i.e., independent Gaussian with unknown means and standard deviations (mean-field Gaussian), and independent Bernoulli distributions with fixed $p=0.5$ as the probability of success/failure, multiplied by unknown variables (mean-field Bernoulli). The predictions for $u$ and $f$ are displayed in Fig. 12, it is observed that (1) the predicted means for $u$ and $f$ from the two parameterized posteriors are similar, and (2) the errors between the predicted means for $u$ and $f$ are mostly bounded by the predicted uncertainties in both methods.

Finally, note that we employ these specific prior/posterior distributions for demonstration purposes. Interested users of NeuralUQ can also customize other problem-dependent prior/posterior distributions by following the examples presented above.

Regarding data generation, in Section 4.1, the nonstiff differential equations solver ode45 in Matlab is utilized for solving Eq. (9) to generate the training data as well as the reference solutions; in Section 4.3, the training data are generated from the exact solution of $u$ from [benes2006decompositions]; and in Section 4.4, the finite-element-based Partial Differential Equation Toolbox in Matlab is utilized for solving Eq. (16) to generate the training data as well as the reference solutions. Note that all data as well as the employed NNs related to this study can be accessed in the GitHub repository of NeuralUQ.