Here we build a Bayesian problem in which the forward model is a partial differential equation (PDE) model, the 1D heat problem in particular.
Try to at least run through part 1 to 4 before working on the optional exercises.
Table of contents¶
1.1. Learning objectives
1.2. Loading the PDE test problem
1.3. Building and solving the Bayesian inverse problem
1.4. Parametrizing the unknown parameters via step function expansion
1.5. ★ Observation on a part of the domain
1.6. ★ Parametrizing the unknown parameters via KL expansion
1.1. Learning objectives ¶
Solve PDE-based Bayesian problem using CUQIpy.
Use different parametrizations of the Bayesian parameters (e.g. KL expansion, non-linear maps).
★ Indicates optional section.
This notebook was run on a local machine and not using github actions for this book due to its long execution time.
1.2. Loading the PDE test problem ¶
We first import the required python standard packages that we need:
import numpy as np
import matplotlib.pyplot as plt
from math import floor
import cuqi
# Disable progress bar dynamic update for cleaner output for the book. You can
# enable it again by setting it to True to monitor sampler progress.
cuqi.config.PROGRESS_BAR_DYNAMIC_UPDATE = FalseFrom CUQIpy we import the classes that we use in this exercise:
from cuqi.testproblem import Heat1D
from cuqi.distribution import Gaussian, JointDistribution
from cuqi.sampler import PCN, MH, CWMHWe load the test problem Heat1D which provides a one dimensional (1D) time-dependent heat model with zero boundary conditions. The model is discretized using finite differences.
The PDE is given by:
where is the temperature and is the thermal diffusivity (assumed to be 1 here). We assume the source term is zero. The unknown parameters (random variable) for this test problem is the initial temperature profile .
The data is a random variable containing the temperature measurements everywhere in the domain at the final time corrupted by Gaussian noise:
where is the forward model that maps the initial condition to the final time solution via solving the 1D time-dependent heat problem. is the measurement noise.
Given observed data the task is to infer the initial temperature profile .
Before we load the Heat1D problem, let us set the parameters: final time , number of finite difference nodes , and the length of the domain
N = 30 # Number of finite difference nodes
L = 1 # Length of the domain
tau_max = 0.02 # Final timeWe assume that the exact initial condition that we want to infer is a step function with three pieces. Here we use commonly used Python and NumPy functions to define the initial condition.
n_steps = 3
n_steps_values = [0,1,2]
myExactSolution = np.zeros(N)
start_idx=0
for i in range(n_steps):
end_idx = floor((i+1)*N/n_steps)
myExactSolution[start_idx:end_idx] = n_steps_values[i]
start_idx = end_idxWe plot the exact solution for each node :
plt.plot(myExactSolution)
plt.xlabel("i")
While it is possible to create a PDE model from scratch, for this notebook we use the Heat1D test problem that is already implemented in CUQIpy.
We may extract components of a Heat1D instance by calling the get_components method.
model, data, problemInfo = Heat1D(
dim=N,
endpoint=L,
max_time=tau_max,
exactSolution=myExactSolution
).get_components()Let us take a look at what we obtain from the test problem. We view the model:
modelCUQI PDEModel: Continuous1D[30] -> Continuous1D[30].
Forward parameters: ['x'].
PDE: TimeDependentLinearPDE.Note that the forward model parameter, named x here, is the unknown parameter which we want to infer. It represents the initial condition above, or a a parameterization of it.
We can look at the returned data:
dataCUQIarray: NumPy array wrapped with geometry.
---------------------------------------------
Geometry:
Continuous1D[30]
Parameters:
True
Array:
CUQIarray([0.02997649, 0.10236499, 0.14642851, 0.19658966, 0.19123706,
0.22876182, 0.30480228, 0.38611469, 0.41670538, 0.49612066,
0.61614546, 0.69627186, 0.78527292, 0.82681347, 0.9109078 ,
1.03047837, 1.10974723, 1.13783192, 1.22819756, 1.29139258,
1.27662701, 1.30455305, 1.24459808, 1.14410371, 1.1007695 ,
0.95443519, 0.82253866, 0.6349439 , 0.42422807, 0.22622538])And the problemInfo:
problemInfoProblemInfo with the following set attributes:
['exactSolution', 'exactData', 'infoString']
infoString: Noise type: Additive i.i.d. noise with mean zero and signal to noise ratio: 200Now let us plot the exact solution (exact initial condition) of this inverse problem and the exact and noisy data (the final time solution before and after adding observation noise):
problemInfo.exactSolution.plot()
problemInfo.exactData.plot()
data.plot()
plt.legend(['exact solution', 'exact data', 'noisy data']);
Note that the values of the initial temperature profile and the data at 0 and are not included in this plot.
# Your code here
In case the parameters were changed in the exercise, we reset them here
N = 30 # Number of finite difference nodes
L = 1 # Length of the domain
tau_max = 0.02 # Final time
model, data, problemInfo = Heat1D(
dim=N,
endpoint=L,
max_time=tau_max,
exactSolution=myExactSolution
).get_components()1.3. Building and solving the Bayesian inverse problem ¶
The joint distribution of the data and the parameter (where represents the unknown values of the initial condition at the grid nodes in this section) is given by
Where is the prior pdf, is the data distribution pdf. We start by defining the prior distribution :
mean = 0
std = 1.2
x = Gaussian(mean, cov=std**2, geometry=model.domain_geometry) # The prior distribution# Your code here
Assuming an i.i.d. additive Gaussian noise and given the fact that we know the exact data in this case, we can compute the noise level as follows:
sigma_noise = np.std(problemInfo.exactData - data)*np.ones(model.range_dim) # noise levelAnd then define the data distribution :
y = Gaussian(mean=model(x), cov=sigma_noise**2, geometry=model.range_geometry)Now that we have all the components we need, we can create the joint distribution , from which the posterior distribution can be created by setting (we named the observed data as data in our code):
First, we define the joint distribution :
joint = JointDistribution(y, x)
print(joint)JointDistribution(
Equation:
p(y,x) = p(y|x)p(x)
Densities:
y ~ CUQI Gaussian. Conditioning variables ['x'].
x ~ CUQI Gaussian.
)
The posterior distribution pdf is given by the Bayes rule:
By setting in the joint distribution we obtain the posterior distribution:
posterior = joint(y=data)
print(posterior)Posterior(
Equation:
p(x|y) ∝ L(x|y)p(x)
Densities:
y ~ CUQI Gaussian Likelihood function. Parameters ['x'].
x ~ CUQI Gaussian.
)
We can now sample the posterior. Let’s try the preconditioned Crank-Nicolson (pCN) sampler (~40 seconds):
MySampler = PCN(posterior, scale=0.01)
MySampler.warmup(10000, tune_freq=0.01)
MySampler.sample(20000)
posterior_samples = MySampler.get_samples()Warmup: 100%|██████████| 10000/10000 [00:17<00:00, 584.18it/s, acc rate: 49.03%]
Sample: 100%|██████████| 20000/20000 [00:33<00:00, 604.52it/s, acc rate: 43.48%]
Note here in the sampler warmup phase, we set the tuning frequency to 0.01, which means that the scale of the proposal distribution is adapted every 1% of the samples.
Let’s look at the credible interval:
posterior_samples.plot_ci(95, exact=problemInfo.exactSolution)
We can see that the mean reconstruction of the initial temperature profile matches the general trend of the exact solution to some extent but it does not capture the piece-wise constant nature of the exact solution.
Also we note that since the heat problem has zero boundary conditions, the reconstruction tends to go to zero at the right boundary.
1.4. Parametrizing the unknown parameters via step function expansion ¶
One way to improve the solution of this Bayesian problem is to use better prior information. Here we assume the prior is a step function with three pieces (which is exactly how we created the exact solution). This also makes the Bayesian inverse problem simpler because now we only have three unknown parameters to infer.
To test this case we pass field_type='Step' to the constructor of Heat1D, which creates a StepExpansion domain geometry for the model during initializing the Heat1D test problem. It is also possible to change the domain geometry manually after initializing the test problem, but we will not do that here.
The parameter field_params is a dictionary that is used to pass keyword arguments that the underlying domain field geometry accepts. For example StepExpansion has a keyword argument n_steps and thus we can set field_params={'n_steps': 3}.
n_steps = 3 # number of steps in the step expansion domain geometry
N = 30
model, data, problemInfo = Heat1D(
dim=N,
endpoint=L,
max_time=tau_max,
field_type="Step",
field_params={"n_steps": n_steps},
exactSolution=myExactSolution,
).get_components()
Let’s look at the model in this case:
modelCUQI PDEModel: StepExpansion[3: 30] -> Continuous1D[30].
Forward parameters: ['x'].
PDE: TimeDependentLinearPDE.We then continue to create the Bayesian inverse problem (prior, data distribution and then posterior) with a prior of dimension = n_steps.
# Prior
x = Gaussian(mean, std**2, geometry=model.domain_geometry)
# Data distribution
sigma_noise = np.std(problemInfo.exactData - data)*np.ones(model.range_dim) # noise level
y = Gaussian(mean=model(x), cov=sigma_noise**2, geometry=model.range_geometry)And the posterior:
joint = JointDistribution(y, x)
posterior = joint(y=data)We then sample the posterior using Metropolis Hastings MH sampler (~40 seconds)
MySampler = MH(posterior, scale=0.01)
MySampler.warmup(10000, tune_freq=0.01)
MySampler.sample(20000)
posterior_samples = MySampler.get_samples()Warmup: 0%| | 0/10000 [00:00<?, ?it/s]Warmup: 100%|██████████| 10000/10000 [00:12<00:00, 774.24it/s, acc rate: 25.42%]
Sample: 100%|██████████| 20000/20000 [00:30<00:00, 651.59it/s, acc rate: 23.06%]
Let’s take a look at the posterior:
posterior_samples.plot_ci(95, exact=problemInfo.exactSolution)
posterior_samples.shape(3, 30000)
We show the trace plot: a plot of the kernel density estimator (left) and chains (right) of the n_steps variables:
posterior_samples.plot_trace()array([[<Axes: title={'center': 'x0'}>, <Axes: title={'center': 'x0'}>],
[<Axes: title={'center': 'x1'}>, <Axes: title={'center': 'x1'}>],
[<Axes: title={'center': 'x2'}>, <Axes: title={'center': 'x2'}>]],
dtype=object)
We show pair plot of 2D marginal posterior distributions:
posterior_samples.plot_pair()array([[<Axes: ylabel='x1'>, <Axes: >],
[<Axes: xlabel='x0', ylabel='x2'>, <Axes: xlabel='x1'>]],
dtype=object)
We notice that there seems to be some burn-in, which is the initial part of the chain and appears as a trail leading to the high-density region. We show the pair plot after removing 1000 burn-in:
posterior_samples.burnthin(1000).plot_pair()array([[<Axes: ylabel='x1'>, <Axes: >],
[<Axes: xlabel='x0', ylabel='x2'>, <Axes: xlabel='x1'>]],
dtype=object)
We can see that the burn-in, the trail part, is indeed removed. Another observation here is the correlation (or inverse correlation) between each pair of the variables.
We can also see the effect of removing the burn-in if we look at the credible interval before and after removing the burn-in. Let us look at the credible interval before:
posterior_samples.plot_ci(100, exact=problemInfo.exactSolution)
And after removing the burn-in:
posterior_samples.burnthin(1000).plot_ci(100, exact=problemInfo.exactSolution)
We compute the effective sample size (ESS) which approximately gives the number of independent samples in the chain, Roy (2020):
posterior_samples.compute_ess()array([392.0780591 , 422.49001125, 452.82444651])# Your code here
1.5. ★ Observation on a part of the domain ¶
Here we solve the same problem as in section 3 but with observing the data only on the right half of the domain.
We chose the number of steps to be 4:
N = 30
n_steps = 4 # Number of steps in the StepExpansion geometry. Then we write the observation_nodes map which can be passed to the Heat1D.
It is a lambda function that takes the forward model range grid (range_grid) as an input and generates a sub grid of the nodes where we have observations (data).
# observe in the right half of the domain
observation_nodes = lambda range_grid: range_grid[np.where(range_grid>L/2)] We load the Heat1D problem. Note in this case we do not pass an exactSolution. If no exactSolution is passed, the Heat1D test problem will create an exact solution.
model, data, problemInfo = Heat1D(
dim=N,
endpoint=L,
max_time=tau_max,
field_type="Step",
field_params={"n_steps": n_steps},
observation_grid_map=observation_nodes,
).get_components()
Now let us plot the exact solution of this inverse problem and the exact and noisy data:
problemInfo.exactSolution.plot()
problemInfo.exactData.plot()
data.plot()
plt.legend(['exact solution', 'exact data', 'noisy data']);
We then continue to create the Bayesian inverse problem (prior, data distribution and then posterior) with a prior of dimension = 4.
# Prior
x = Gaussian(1, std**2, geometry=model.domain_geometry)
# Data distribution
sigma_noise = np.std(problemInfo.exactData - data)*np.ones(model.range_dim) # noise level
y = Gaussian(mean=model(x), cov=sigma_noise**2, geometry=model.range_geometry)And the posterior:
joint = JointDistribution(y, x)
posterior = joint(y=data)We then sample the posterior using the Metropolis Hastings sampler (~50 seconds)
MySampler = MH(posterior, initial_point=np.ones(posterior.dim), scale=0.01)
MySampler.warmup(10000, tune_freq=0.01)
MySampler.sample(20000)
posterior_samples = MySampler.get_samples()Warmup: 0%| | 0/10000 [00:00<?, ?it/s]Warmup: 100%|██████████| 10000/10000 [00:12<00:00, 815.29it/s, acc rate: 28.21%]
Sample: 100%|██████████| 20000/20000 [00:30<00:00, 660.48it/s, acc rate: 23.59%]
Let’s take a look at the posterior:
posterior_samples.burnthin(1000).plot_ci(95, exact=problemInfo.exactSolution)
We see that the credible interval is wider on the side of the domain where data is not available (the left side) and narrower as we get to the right side of the domain.
1.6. ★ Parametrizing the unknown parameters via KL expansion ¶
Here we explore the Bayesian inversion for a more general initial temperature profile. We parametrize the unknown initial temperature profile using Karhunen–Loève (KL) expansion, which we introduced in section 3. More forward model examples. We denote by now the expansion coefficients of the KL expansion of the initial temperature profile we are trying to infer.
Let’s load the Heat1D test case and pass field_type = 'KL', which behind the scenes will set the domain geometry of the model to be a KL expansion geometry (See KLExpansion documentation for details):
N = 35
model, data, problemInfo = Heat1D(
dim=N,
endpoint=L,
max_time=tau_max,
field_type="KL"
).get_components()
Now we inspect the model:
modelCUQI PDEModel: KLExpansion[35] -> Continuous1D[35].
Forward parameters: ['x'].
PDE: TimeDependentLinearPDE.And the exact solution and the data:
problemInfo.exactSolution.plot()
problemInfo.exactData.plot()
data.plot()
plt.legend(['exact solution', 'exact data', 'noisy data']);
Note that the exact solution here is a general signal that is not constructed from the basis functions. We define the prior as an i.i.d. multivariate Gaussian with variance 32. We determined the choice of this variance value through trial and error approach to get a reasonable Bayesian reconstruction of the exact solution:
sigma_prior = 3*np.ones(model.domain_dim)
x = Gaussian(mean, sigma_prior**2, geometry=model.domain_geometry)We define the data distribution:
sigma_noise = np.std(problemInfo.exactData - data)*np.ones(model.range_dim) # noise level
y = Gaussian(mean=model(x), cov=sigma_noise**2, geometry=model.range_geometry)And the posterior distribution:
joint = JointDistribution(y, x)
posterior = joint(y=data)We sample the posterior, here we use Component-wise Metropolis Hastings (CWMH) sampler, see section 1.2.2. Story 2 - Not all unknowns are created equal: Component-wise Metropolis Hastings CWMH for more details about CWMH. The sampling takes ~70 seconds:
MySampler = CWMH(posterior, initial_point=np.ones(N))
MySampler.warmup(500, tune_freq=0.01)
MySampler.sample(1000)
posterior_samples = MySampler.get_samples()Warmup: 100%|██████████| 500/500 [00:15<00:00, 31.67it/s, acc rate: 84.18%]
Sample: 100%|██████████| 1000/1000 [00:31<00:00, 32.25it/s, acc rate: 84.03%]
Then we plot the credible interval (you can try plotting different credible intervals, e.g. )
posterior_samples.plot_ci(95, exact=problemInfo.exactSolution)
Using plot_ci with default options results in a credible interval with vanishing width at some locations where the upper and lower bounds of the credible interval seem to cross. To look into what actually happens here, let us first plot the upper and the lower bounds of the credible interval:
upper, lower = posterior_samples.compute_ci(95)
model.domain_geometry.plot(upper, color='r', label='Upper limit of 95% CI')
model.domain_geometry.plot(lower, color='b', label='Lower limit of 95% CI')
plt.legend()
One might expect the upper bound to be always above the lower bound, but here we see that they intersect and switch order at some locations. This is because the credible interval is computed from the samples, which are KL coefficients samples, then converted to the original function values, the initial temperature profile over . The resulting upper and lower bounds are not guaranteed to follow the expected order everywhere in . Let us look at the upper and lower bounds of the credible interval in the parameter space (the KL coefficients) which are guaranteed to follow the expected order:
model.domain_geometry.plot(upper, plot_par=True, color='r', label='Upper limit of 95% CI')
model.domain_geometry.plot(lower, plot_par=True, color='b', label='Lower limit of 95% CI')
plt.legend()
Note that, indeed, each coefficient in the lower bound is smaller than the corresponding coefficient in the upper bound. Let us use a different approach to plotting the credible interval. We can first convert the samples to function values first, i.e., the temperature profiles over , then compute the credible interval from the resulting samples. To convert the samples to function values we run:
funvals_samples = posterior_samples.funvalsThen we plot the credible interval computed from the function values:
funvals_samples.plot_ci(95, exact=problemInfo.exactSolution)
We can see that the width of the resulting credible interval is larger than zero everywhere over . We can also plot the samples, converted to function values, to ensure that this credible interval is consistent with how the samples look like.
posterior_samples_burnthin = posterior_samples.burnthin(0,10)
for i, s in enumerate(posterior_samples_burnthin):
model.domain_geometry.plot(s)
The samples indeed convey consistent information with the latter credible interval plot.
Let’s look at the effective sample size (ESS):
posterior_samples.compute_ess()array([112.32077156, 121.77632586, 360.34784175, 55.02953875,
24.00783406, 55.96779181, 29.94162333, 38.26841749,
43.07419138, 58.91602579, 32.50637055, 36.29476571,
20.96031105, 47.50622953, 56.40834103, 41.23495218,
22.82706052, 36.87458645, 22.4457361 , 52.40607866,
33.87163859, 21.14304201, 53.25813134, 29.01124206,
48.04381648, 34.38726782, 34.24121519, 34.32485159,
42.2986886 , 4.14403911, 52.22690903, 23.08267447,
40.34121 , 14.3707869 , 30.69438323])We note that the ESS varies considerably among the variables. We can view the trace plot for, let’s say, the first and the second variables:
posterior_samples.plot_trace([0,1])array([[<Axes: title={'center': 'x0'}>, <Axes: title={'center': 'x0'}>],
[<Axes: title={'center': 'x1'}>, <Axes: title={'center': 'x1'}>]],
dtype=object)
We note that (as expected from the values of ESS) the chain quality of , which corresponds to the second coefficient in the KL expansion, is much better than that of the first variable . Low quality chain (where chain samples are highly correlated) indicates difficulty in exploring the corresponding parameter and possibly high sensitivity of the model to that particular parameter. Sampling methods that incorporate gradient information (which we do not explore in this notebook) are expected to work better in this situation.
A third way of looking at the credible intervals, is to look at the expansion coefficients credible intervals. We plot the credible intervals for these coefficients from both prior and posterior samples by passing the flag plot_par=True to plot_ci function:
The prior:
plt.figure()
x.sample(1000).plot_ci(95, plot_par=True)
plt.xticks(np.arange(x.dim)[::5]);
The posterior:
posterior_samples.plot_ci(95, plot_par=True)
plt.xticks(np.arange(x.dim)[::5]);
By comparing the two plots above, we see that the first few coefficients are inferred with higher certainty than the remaining coefficients. In parts, this is due to the nature of the Heat problem where high oscillatory initial temperature features (corresponding to the higher modes in the expansion) will be smoothed out (lost) faster and thus are harder to retrieve based on measurements from the final solution.
- Roy, V. (2020). Convergence diagnostics for markov chain monte carlo. Annual Review of Statistics and Its Application, 7(1), 387–412.