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 [9] 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
using SpecialFunctions

The Forward Model

The BOD model is given by:

\[\mathcal{B}(t) = A(1-\exp(-Bt))+\mathcal{E}\]

where the parameters A and B are functions of the uncertain inputs θ₁ and θ₂:

\[\begin{aligned} A &= 0.4 + 0.4\left(1 + \text{erf}\left(\frac{\theta_1}{\sqrt{2}} \right)\right) \\ B &= 0.01 + 0.15\left(1 + \text{erf}\left(\frac{\theta_2}{\sqrt{2}} \right)\right) \end{aligned}\]

and ℰ ~ 𝒩(0, 10⁻³) represents measurement noise.

function forward_model(t, x)
    A = 0.4 + 0.4 * (1 + erf(x[1] / sqrt(2)))
    B = 0.01 + 0.15 * (1 + erf(x[2] / sqrt(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 x₁ in [-0.5, 0, 0.5]
    for x₂ in [-0.5, 0, 0.5]
        plot!(t_values, [forward_model(ti, [x₁, x₂]) for ti in t_values],
              label="(x₁ = $x₁, x₂ = $x₂)", linestyle=:dash)
    end
end

BOD Realizations

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}|\boldsymbol{\theta}) = \prod_{t=1}^{5} \frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{1}{2\sigma^2}(y_t - \mathcal{B}(t))^2\right)\]

function posterior(x)
    # Calculate the likelihood
    likelihood = prod([pdf(Normal(forward_model(t[k], x), σ), D[k]) for k in 1:5])
    # Calculate the prior
    prior = pdf(Normal(0,1), x[1]) * pdf(Normal(0,1), x[2])
    return prior * likelihood
end

target = MapTargetDensity(posterior, :auto_diff)

MapTargetDensity(density=posterior, gradient_type=auto_diff)

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())

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

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]
  Weight sum: 0.9999999999999994
  Points (first 5):
    [-4.859462828332312, -4.859462828332312] → weight: 1.858172610271843e-11
    [-3.581823483551927, -4.859462828332312] → weight: 3.2677804672640156e-9
    [-2.484325841638953, -4.859462828332312] → weight: 8.238338476283024e-8
    [-1.4659890943911627, -4.859462828332312] → weight: 5.840231806713457e-7
    [-0.48493570751550125, -4.859462828332312] → weight: 1.4856333877315967e-6
    ... and 95 more

Optimize the map coefficients:

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

 12.506762 seconds (168.44 M allocations: 9.810 GiB, 5.30% gc time, 3.59% compilation time)
Optimization result:  * Status: success

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

 * Found with
    Algorithm:     L-BFGS

 * Convergence measures
    |x - x'|               = 2.18e-09 ≰ 0.0e+00
    |x - x'|/|x'|          = 1.35e-09 ≰ 0.0e+00
    |f(x) - f(x')|         = 1.12e-13 ≰ 0.0e+00
    |f(x) - f(x')|/|f(x')| = 1.80e-14 ≰ 0.0e+00
    |g(x)|                 = 6.68e-09 ≤ 1.0e-08

 * Work counters
    Seconds run:   12  (vs limit Inf)
    Iterations:    123
    f(x) calls:    368
    ∇f(x) calls:   368

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}:
 -1.81761     1.539
 -1.2094      0.412999
  0.0280115   0.0477759
 -0.0261609  -0.945836
  1.16377    -0.454578
  1.14597    -0.0249904
 -0.364864    0.633504
  0.583589   -0.97581
 -0.46293     0.589189
 -1.72181    -0.162177
  ⋮          
 -1.01781     0.768779
 -0.21267     0.762467
  0.177849   -0.783171
 -0.200454   -0.31037
 -0.290224    0.310542
  0.762085   -0.329262
 -0.64664    -1.45032
  0.0235214   0.49285
 -1.74864    -0.918771

Map the samples through our transport map:

mapped_samples = evaluate(M, samples_z)

1000×2 Matrix{Float64}:
 -0.392813   2.84286
 -0.313125   1.71808
 -0.0600533  0.869273
 -0.0752472  0.734279
  0.425555   0.218161
  0.414593   0.263077
 -0.15969    1.24027
  0.130869   0.43516
 -0.181176   1.29395
 -0.381391   1.74869
  ⋮          
 -0.283681   1.73501
 -0.123849   1.1698
 -0.0152137  0.657153
 -0.120831   0.93175
 -0.142511   1.11309
  0.209077   0.419551
 -0.218421   0.932792
 -0.0613323  0.955441
 -0.384626   1.46972

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.011423391411532185

Visualization

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

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

posterior_values = [posterior([x₁, x₂]) for x₂ in x₂, x₁ in x₁]

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

BOD Samples

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

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

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

BOD Pullback Density

Model Parameter Interpretation

The posterior samples provide uncertainty quantification for the BOD model parameters:

x₁: Controls the maximum BOD level (parameter A)
x₂: Controls the rate of BOD development (parameter B)

The correlation structure in the posterior reflects the interdependence between these parameters in explaining the observed data.