Bayesian Inference: Biochemical Oxygen Demand (BOD) Example

This example demonstrates Bayesian parameter estimation for a biochemical oxygen demand model using transport maps. The problem comes from environmental engineering and was originally presented in [13] and later used as a benchmark in transport map applications [1].

The model describes the evolution of biochemical oxygen demand (BOD) in a river system using an exponential growth model with two uncertain parameters controlling growth and decay rates.

using TransportMaps
using Plots
using Distributions

The Forward Model

The BOD model is given by:

$$\mathcal{B}(t)=A(1-\exp (-Bt))+\epsilon$$

where the parameters $A$ and $B$ unknown material parameters and $\epsilon \sim \mathcal{N}(0,{10}^{-3})$ represents measurement noise.

The parameters $A$ and $B$ follow the prior distributions

$$A\sim \mathcal{U}(0.4,1.2),B\sim \mathcal{U}(0.01,0.31)$$

In order to describe the input parameters in an unbounded space, we transform them into a space with standard normal prior distributions $p({\theta }_i)\sim \mathcal{N}(0,1)$. We write $A$ and $B$ as functions of the new variables ${\theta }_1$ and ${\theta }_2$ with the help of a density transformation with the standard normal CDF $\mathrm{\Phi }({\theta }_i)$:

$$\begin{array}{rl}A& =0.4+(1.2-0.4)\cdot \mathrm{\Phi }({\theta }_1)\\ B& =0.01+(0.31-0.01)\cdot \mathrm{\Phi }({\theta }_2)\end{array}$$

We define the forward model as a function of $\theta$ and $t$ in Julia:

function forward_model(t, θ)
    A = 0.4 + (1.2 - 0.4) * cdf(Normal(), θ[1])
    B = 0.01 + (0.31 - 0.01) * cdf(Normal(), θ[2])
    return A * (1 - exp(-B * t))
end

Experimental Data

We have BOD measurements at five time points:

t = [1, 2, 3, 4, 5]
D = [0.18, 0.32, 0.42, 0.49, 0.54]
σ = sqrt(1e-3)

Let's visualize the data along with model predictions for different parameter values:

s = scatter(t, D, label="Data", xlabel="Time (t)", ylabel="Biochemical Oxygen Demand (D)",
    size=(600, 400), legend=:topleft)
# Plot model output for some parameter values
t_values = range(0, 5, length=100)
for θ₁ in [-0.5, 0, 0.5]
    for θ₂ in [-0.5, 0, 0.5]
        plot!(t_values, [forward_model(ti, [θ₁, θ₂]) for ti in t_values],
            label="(θ₁ = $θ₁, θ₂ = $θ₂)", linestyle=:dash)
    end
end

Bayesian Inference Setup

We define the posterior distribution using a standard normal prior on both parameters and a Gaussian likelihood for the observations:

$$\pi (\mathbf{y}|\mathit{\theta })=\prod _{t=1}^5\frac{1}{\sqrt{2\pi {\sigma }^2}}\exp (-\frac{1}{2{\sigma }^2}(y_t-\mathcal{B}(t){)}^2)$$

function logposterior(θ)
    # Calculate the likelihood
    likelihood = sum([logpdf(Normal(forward_model(t[k], θ), σ), D[k]) for k in 1:5])
    # Calculate the prior
    prior = logpdf(Normal(), θ[1]) + logpdf(Normal(), θ[2])
    return prior + likelihood
end

target = MapTargetDensity(x -> logposterior(x))

MapTargetDensity(backend=ADTypes.AutoForwardDiff())

Creating and Optimizing the Transport Map

We use a 2-dimensional polynomial transport map with degree 3 and Softplus rectifier:

M = PolynomialMap(2, 3, :normal, Softplus(), LinearizedHermiteBasis())

PolynomialMap:
  Dimensions: 2
  Total coefficients: 14
  Reference density: Distributions.Normal{Float64}(μ=0.0, σ=1.0)
  Maximum degree: 3
  Basis: LinearizedHermiteBasis
  Rectifier: Softplus
  Components:
    Component 1: 4 basis functions
    Component 2: 10 basis functions
  Coefficients: min=0.0, max=0.0, mean=0.0

Set up Gauss-Hermite quadrature for optimization:

quadrature = GaussHermiteWeights(10, 2)

GaussHermiteWeights:
  Number of points: 100
  Dimensions: 2
  Quadrature type: Tensor product Gauss-Hermite
  Reference measure: Standard Gaussian
  Weight range: [1.858172610271843e-11, 0.11877833902739371]

Optimize the map coefficients:

res = optimize!(M, target, quadrature)
println("Optimization result: ", res)

Optimization result:  * Status: success

 * Candidate solution
    Final objective value:     -6.228606e+00

 * Found with
    Algorithm:     L-BFGS

 * Convergence measures
    |x - x'|               = 1.53e-09 ≰ 0.0e+00
    |x - x'|/|x'|          = 9.40e-10 ≰ 0.0e+00
    |f(x) - f(x')|         = 0.00e+00 ≤ 0.0e+00
    |f(x) - f(x')|/|f(x')| = 0.00e+00 ≤ 0.0e+00
    |g(x)|                 = 4.99e-09 ≤ 1.0e-08

 * Work counters
    Seconds run:   4  (vs limit Inf)
    Iterations:    78
    f(x) calls:    254
    ∇f(x) calls:   254
    ∇f(x)ᵀv calls: 0

Generating Posterior Samples

Generate samples from the standard normal distribution and map them to the posterior:

samples_z = randn(1000, 2)

1000×2 Matrix{Float64}:
  0.808288    1.29539
 -1.12207    -0.0401347
 -1.10464     0.875608
 -0.416993    0.585995
  0.287588   -0.927587
  0.229819   -0.414203
 -0.421769   -0.0856216
 -1.35559    -0.570313
  0.0694591  -0.242781
 -0.117323   -3.05756
  ⋮          
 -0.905438   -1.73076
  1.84802     0.432257
 -1.29846     0.183927
  0.596723    1.55975
  0.431956    1.60559
 -0.797097    0.723346
 -0.0625319   0.474741
 -1.25053    -0.0100866
 -0.34637    -1.01294

Map the samples through our transport map:

mapped_samples = evaluate(M, samples_z)

1000×2 Matrix{Float64}:
  0.249818    0.562582
 -0.302227    1.51078
 -0.299553    1.85239
 -0.171311    1.26676
  0.0285589   0.572644
  0.0085126   0.673603
 -0.172387    1.10523
 -0.335899    1.44635
 -0.043264    0.79005
 -0.0971228   0.491073
  ⋮          
 -0.267245    0.992314
  1.0237     -0.0146809
 -0.328008    1.68644
  0.150253    0.744218
  0.0822056   0.873555
 -0.248183    1.5726
 -0.0819886   1.00285
 -0.321222    1.58806
 -0.155038    0.880915

Quality Assessment

Compute the variance diagnostic to assess the quality of our approximation:

var_diag = variance_diagnostic(M, target, samples_z)
println("Variance Diagnostic: ", var_diag)

Variance Diagnostic: 0.018012021800394785

Visualization

Plot the mapped samples along with contours of the true posterior density:

θ₁ = range(-0.5, 1.5, length=100)
θ₂ = range(-0.5, 3, length=100)

posterior_values = [pdf(target, [θ₁, θ₂]) for θ₂ in θ₂, θ₁ in θ₁]

scatter(mapped_samples[:, 1], mapped_samples[:, 2],
    label="Mapped Samples", alpha=0.5, color=1,
    xlabel="θ₁", ylabel="θ₂", title="Posterior Density and Mapped Samples")
contour!(θ₁, θ₂, posterior_values, colormap=:viridis, label="Posterior Density")

Finally, we can compute the pullback density to visualize how well our transport map approximates the posterior:

posterior_pullback = [pullback(M, [θ₁, θ₂]) for θ₂ in θ₂, θ₁ in θ₁]

contour(θ₁, θ₂, posterior_values ./ maximum(posterior_values);
    levels=5, colormap=:viridis, colorbar=false,
    label="Target", xlabel="θ₁", ylabel="θ₂")
contour!(θ₁, θ₂, posterior_pullback ./ maximum(posterior_pullback);
    levels=5, colormap=:viridis, linestyle=:dash,
    label="Pullback")

We can also visually observe a good agreement between the true posterior and the TM approximation.

Conditional Density and Samples

We can compute conditional densities and generate conditional samples using the transport map. Therefore, we make use of the factorization of the structure of triangular maps given by the Knothe-Rosenblatt rearrangement [1], [2]:

$$\pi (\mathrm{x})=\underset{T_1^{-1}\left(z_1\right)}{\underbrace{\pi \left(x_1\right)}}\underset{T_2^{-1}(z_2;x_1)}{\underbrace{\pi (x_2\mid x_1)}}\underset{T_3^{-1}(z_3;x_1,x_2)}{\underbrace{\pi (x_3\mid x_1,x_2)}}\cdots \underset{T_K^{-1}(z_K;x_1,\dots ,x_{K-1})}{\underbrace{\pi (x_K\mid x_1,\dots ,x_{K-1})}},$$

This allows us to compute the conditional density $\pi ({\theta }_2|{\theta }_1)$ and generate samples from this conditional distribution efficiently. We only need to invert the second component of the map to obtain ${\theta }_2$ given a fixed value of ${\theta }_1$ and then push forward samples from the reference distribution.

We can use the conditional_sample function to generate samples from the conditional distribution. Therefore, we samples from the standard normal distribution for $z_2$ and push them through the conditional map. We use the previously generated samples for $z_2$ and fix ${\theta }_1$.

θ₁ = 0.
conditional_samples = conditional_sample(M, θ₁, randn(10_000))

Then, we compute the conditional density of ${\theta }_2$ given ${\theta }_1$ first analytically by integrating out ${\theta }_1$ from the joint posterior. We use numerical integration for this purpose and evaluate the conditional density on a grid.

θ_range = 0:0.01:2
int_analytical = gaussquadrature(ξ -> pdf(target, [θ₁, ξ]), 1000, -10., 10.)
posterior_conditional(θ₂) = pdf(target, [θ₁, θ₂]) / int_analytical
conditional_analytical = posterior_conditional.(θ_range)

Then we compute the conditional density using the transport map, which is given as:

$$\pi ({\theta }_2|{\theta }_1)={\rho }_2(T^2({\theta }_1,{\theta }_2{)}^{-1})\left|\frac{\partial T^2({\theta }_1,{\theta }_2{)}^{-1}}{\partial {\theta }_2}\right|$$

conditional_mapped = conditional_density(M, θ_range, θ₁)

Finally, we plot the results:

histogram(conditional_samples, bins=50, normalize=:pdf, α=0.5,
    label="Conditional Samples", xlabel="θ₂", ylabel="π(θ₂ | θ₁=$θ₁)")
plot!(θ_range, conditional_analytical, lw=2, label="Analytical Conditional PDF")
plot!(θ_range, conditional_mapped, lw=2, label="TM Conditional PDF")

---

This page was generated using Literate.jl.