Example: The Catalysis Problem Using the Laplace Approximation

Hide code cell source

import matplotlib.pyplot as plt
%matplotlib inline
import matplotlib_inline
matplotlib_inline.backend_inline.set_matplotlib_formats('svg')
import seaborn as sns

import urllib.request
import os
import jax
from jax import grad, jit, vmap, jacfwd, jacrev, value_and_grad, jvp
import jax.numpy as jnp
import jax.random as jr
import optax
from jax.scipy.linalg import expm
from scipy.optimize import minimize
import pandas as pd
from functools import partial

jax.config.update("jax_enable_x64", True)
colors = sns.color_palette()
key = jr.PRNGKey(0)

def download(
    url : str,
    local_filename : str = None
):
    """Download a file from a url.
    
    Arguments
    url            -- The url we want to download.
    local_filename -- The filemame to write on. If not
                      specified 
    """
    if local_filename is None:
        local_filename = os.path.basename(url)
    urllib.request.urlretrieve(url, local_filename)

Example: The Catalysis Problem Using the Laplace Approximation#

We return to the catalysis problem solved with the classical approach, using the data reported by Tsilifis et al. (2016). That treatment left several concerns.

  • The solution may not exist. (In the catalysis example, try calibrating a model that does not include the intermediate element \(X\). Check whether the reduced reaction network can reproduce the observed concentration histories.)

  • Multiple solutions may exist. (In the catalysis example, try adding one more fictitious reaction product. You will probably fit the data very well. But which model is the right one?)

  • No estimate of uncertainty.

A Bayesian treatment lets us combine prior information with the likelihood and quantify parameter uncertainty. We begin with a Laplace approximation for the catalysis problem.

Let’s start by assuming that \(\sigma = 5\). We use the prior scale \(\gamma=10\) in both calculations below, with

\[ p(x) = \mathcal{N}(0,\gamma^2I). \]

The prior contributes the following curvature to the log posterior: $\( \nabla^2 \log p(x) = -\gamma^{-2}I. \)$

data_path = Path('../../data/catalysis.csv')
if not data_path.exists():
    url = 'https://raw.githubusercontent.com/PredictiveScienceLab/advanced-scientific-machine-learning/refs/heads/main/book/data/catalysis.csv'
    download(url)
    data_path = Path('catalysis.csv')
catalysis_data = pd.read_csv(data_path)
t_exp = jnp.array(catalysis_data.loc[catalysis_data['Time'] > 0, 'Time'].values)
t0 = 0.0
t1 = t_exp[-1]

We use the same rate matrix and initial condition as the classical example. Since the system is linear, \(z(t)=\exp(tA(x))z(0)\) gives its solution. Evaluating this matrix exponential also lets us check second derivatives by automatic differentiation.

# Define the linear system
def A(x):
    """
    Return the matrix of the dynamical system.
    """
    k = jnp.exp(x) / 180.0
    res = jnp.zeros((6, 6))
    res = res.at[0, 0].set(-k[0])
    res = res.at[1, 0].set(k[0])
    res = res.at[1, 1].set(-(k[1] + k[3] + k[4]))
    res = res.at[2, 1].set(k[1])
    res = res.at[2, 2].set(-k[2])
    res = res.at[3, 2].set(k[2])
    res = res.at[4, 1].set(k[3])
    res = res.at[5, 1].set(k[4])
    return res

@jit
def solve_catalysis(t, x, z0):
    """Evaluate the same linear model using its matrix exponential."""
    matrix = A(x)
    return vmap(lambda ti: expm(ti * matrix, max_squarings=32) @ z0)(t)

But this time instead of just minimizing the sum of squared errors, we will minimize the negative log posterior.

# Negative log posterior and its gradient
def minus_log_post(x, z, y, t, sigma, gamma):

    res = solve_catalysis(t, x, z)
    flat_res = jnp.hstack([res[:, :2], res[:, 3:]]).flatten()
    
    # Negative log-likelihood 
    tmp = (flat_res - y) 
    likelihood = 0.5 * jnp.sum(tmp**2) / sigma** 2
    
    # Negative log-prior
    prior = 0.5 * jnp.sum(x**2) / gamma**2
    
    # Total negative log posterior
    posterior = likelihood + prior
    
    return posterior
# Initial guess for x
key = jr.PRNGKey(0)
x0 = jr.normal(key, shape=(5,))  

# Initial conditions
z0 = jnp.array([500., 0., 0., 0., 0., 0.0])

# Extract the experimental data
Y = catalysis_data.loc[catalysis_data['Time'] > 0, ['NO3', 'NO2', 'N2', 'NH3', 'N2O']].values
y = Y.flatten()

# Set up the optimizer
optimizer = optax.adam(learning_rate=1e-1)
opt_state = optimizer.init(x0)

x = x0  # Initialize x

sigma = 5.0
gamma = 10.0

# Use as many iterations as needed
num_iterations = 300
loss_evol = []

objective_and_gradient = jit(value_and_grad(minus_log_post))

for i in range(num_iterations):
    
    value, grads = objective_and_gradient(x, z0, y, t_exp, sigma, gamma)
    updates, opt_state = optimizer.update(grads, opt_state)
    x = optax.apply_updates(x, updates)
    loss_evol.append(value)
        # Print the loss every 100 iterations
    if i % 50 == 0:
        print(f"Iteration {i}, loss: {value}")

# Refine the optimizer endpoint before evaluating its local curvature.
fit = minimize(lambda p: objective_and_gradient(p, z0, y, t_exp, sigma, gamma), np.asarray(x), jac=True,
               method='BFGS', options={'gtol': 1e-7, 'maxiter': 1000})
x = jnp.asarray(fit.x)
gradient_norm = float(jnp.linalg.norm(objective_and_gradient(x, z0, y, t_exp, sigma, gamma)[1], ord=jnp.inf))
assert gradient_norm < 1e-5, gradient_norm
print('Final maximum absolute gradient:', gradient_norm)

print(f'The value of x is {x}')
Iteration 0, loss: 11849.38366357738
Iteration 50, loss: 256.91680399036767
Iteration 100, loss: 75.37249010327423
Iteration 150, loss: 74.17829683660295
Iteration 200, loss: 74.17546770609167
Iteration 250, loss: 74.1754368949555
Final maximum absolute gradient: 1.398124998430451e-10
The value of x is [ 1.35942816  1.65879955  1.34579466 -1.0485962  -0.16052466]
# Plot the loss evolution
fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
ax.plot(loss_evol, color=colors[0])
ax.set_xlabel("Iteration")
ax.set_ylabel("Loss")
ax.set_yscale('log')
finalize_axes(keep_box=False)
plt.show()
Negative log-posterior loss decreases on a logarithmic scale during optimization.

The small gradient confirms a stationary point of the negative log posterior. We approximate its Hessian with centered finite differences, filling both entries of every mixed partial derivative. We check the result against automatic differentiation and require positive curvature before inverting it to obtain the Laplace covariance.

def compute_post_cov(mu, y, sigma, gamma, z0, t, epsilon=1e-3):
    """Compute the posterior covariance matrix using the Hessian of the log-posterior."""    
    n = mu.size

    # Define the function f(mu) that computes the model outputs
    def f(x):
        sol = solve_catalysis(t, x, z0)
        f_mu = jnp.hstack([sol[:, :2], sol[:, 3:]]).flatten()
        return f_mu

    # Compute f(mu) at the given mu
    f_mu = f(mu)
    m = f_mu.size
    
    # Initialize gradient dfdx (Jacobian matrix)
    dfdx = jnp.zeros((m, n))

    # Compute gradient of our solver numerically using finite differences
    for i in range(n):
        e_i = jnp.zeros(n)
        e_i = e_i.at[i].set(epsilon)
        f_plus = f(mu + e_i)
        f_minus = f(mu - e_i)
        f_prime = (f_plus - f_minus) / (2 * epsilon)
        dfdx = dfdx.at[:, i].set(f_prime)

    # Initialize Hessian d2fdx2
    d2fdx2 = jnp.zeros((m, n, n))

    # Compute Hessian of our solver numerically using finite differences
    for i in range(n):
        for j in range(i, n):
            e_i = jnp.zeros(n)
            e_j = jnp.zeros(n)
            e_i = e_i.at[i].set(epsilon)
            e_j = e_j.at[j].set(epsilon)

            f_pp = f(mu + e_i + e_j)
            f_pm = f(mu + e_i - e_j)
            f_mp = f(mu - e_i + e_j)
            f_mm = f(mu - e_i - e_j)

            second_derivative = (f_pp - f_pm - f_mp + f_mm) / (4 * epsilon ** 2)

            # Assign the second derivative to the Hessian matrix with jax's at method
            d2fdx2 = d2fdx2.at[:, i, j].set(second_derivative)
            d2fdx2 = d2fdx2.at[:, j, i].set(second_derivative)

    # Compute the second derivative of the log-posterior (d2Ldx2)
    tmp = y - f_mu
    sigma2_inv = 1.0 / sigma ** 2
    gamma2_inv = 1.0 / gamma ** 2

    # First term: - (1 / sigma^2) * sum_i (y_i - f_i) * d2f_i/dx^2
    term1 = - jnp.einsum('i,ijk->jk', tmp, d2fdx2) * sigma2_inv

    # Second term: (1 / sigma^2) * (df/dx)^T * (df/dx)
    term2 = (dfdx.T @ dfdx) * sigma2_inv

    # Third term: (1 / gamma^2) * Identity matrix
    term3 = jnp.eye(n) * gamma2_inv

    d2Ldx2 = term1 + term2 + term3

    exact_hessian = jax.hessian(minus_log_post)(mu, z0, y, t, sigma, gamma)
    assert jnp.allclose(d2Ldx2, d2Ldx2.T, atol=1e-10)
    relative_error = jnp.linalg.norm(d2Ldx2 - exact_hessian) / jnp.linalg.norm(exact_hessian)
    assert relative_error < 1e-4, relative_error
    assert jnp.linalg.eigvalsh(d2Ldx2).min() > 0
    print('Fixed-noise Hessian relative error against AD:', float(relative_error))
    cov = jnp.linalg.inv(d2Ldx2)
    cov = 0.5 * (cov + cov.T)

    return cov
# Compute the posterior covariance
mu = x
Sigmas = compute_post_cov(mu, y, sigma, gamma, z0, t_exp)
Fixed-noise Hessian relative error against AD: 2.11671932575846e-07

Great, let’s plot our samples.

num_samples = 50
times = jnp.linspace(t0, t1, 1000)
key, subkey = jr.split(key)
samples = jax.random.multivariate_normal(subkey, mu, Sigmas, shape=(num_samples,))
# Compute the samples
sample_models = jnp.array([solve_catalysis(times, sample, z0) for sample in samples])

# Compute the median model
median_models = jnp.median(sample_models, axis=0)

# Plotting
fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])

# Define labels and colors
labels = ['NO3-', 'NO2-', 'N2', 'NH3', 'N2O', 'X']
data_cols = ['NO3', 'NO2', 'N2', 'NH3', 'N2O']
model_cols = [0, 1, 3, 4, 5, 2]
species_styles = ['-', '--', '-.', ':', (0, (5, 2)), (0, (3, 1, 1, 1))]
species_markers = ['o', 's', '^', 'D', 'v', 'P']

# Plot experimental data
for i, col in enumerate(data_cols):
    ax.plot(catalysis_data['Time'], catalysis_data[col], linestyle='none', color='black',
            marker=species_markers[i], markerfacecolor='white',
            label=f'Data {labels[i]}', markersize=6)

# Plot the mean models
for i, col in enumerate(model_cols):
    ax.plot(times, median_models[:, col], color='black',
            linestyle=species_styles[i], label=f'Median {labels[i]}')

# Plot the samples
for i in range(min(num_samples, 12)):
    for j, col in enumerate(model_cols):
        ax.plot(times, sample_models[i, :, col], color='0.65',
                linestyle=species_styles[j], linewidth=0.45, alpha=0.25)

# Add legend for the models and data
handles, labels = ax.get_legend_handles_labels()
unique_labels = {label: handle for label, handle in zip(labels, handles)}
ax.legend(unique_labels.values(), unique_labels.keys(), loc='upper center',
          bbox_to_anchor=(0.5, -0.20), ncol=2)
ax.set_ylim(0, 600)
ax.set_xlabel('Time (min)')
ax.set_ylabel(r'Concentration (mmol L$^{-1}$)')
finalize_axes(keep_box=False)
plt.show()
Observed concentrations with Laplace posterior median trajectories and a light ensemble for six species.

Exercises#

  • Investigate what happens as you go from a very large \(\sigma\) (say \(20\)) to a very small one (say \(1\).) Is there a sweet spot?

  • Investigate what happens as you change \(\gamma\) in the same way.

  • How else can you propagate uncertainty through the solver?

Estimating the Noise Level#

Having first fixed \(\sigma=5\), we now estimate the noise level jointly with the kinetic parameters using the Laplace approximation. We need a prior on \(\sigma\). Let’s pick:

\[ p(\sigma) \propto \frac{1}{\sigma}. \]

This is the Jeffreys prior for a scale parameter (Jeffreys, 1946).

Also, because of the nature of the parameterization, it probably makes sense to work with \(\log \sigma\) instead of \(\sigma\). So, let’s introduce the following unknown vector to be inferred from the data (not to be confused with the concentration vector \(z(t;x)\) of the classical example):

\[ z = (x, \theta), \]

where

\[ \theta = \log\sigma. \]

Change-of-variables term#

We need to be a little bit careful with \(p(\theta)\). We use the change-of-variables formula (Bishop, 2006). Define:

\[ \theta := g(\sigma) = \log \sigma. \]

The inverse is: $\( g^{-1}(\theta) = e^{\theta}. \)$

The formula is:

\[ p(\theta) = \left|\frac{d}{d\theta}\left(g^{-1}(\theta)\right)\right|\cdot p(\sigma=e^{\theta}) \propto e^\theta e^{-\theta} = 1. \]

So

\[ p(\theta) \propto 1. \]

Now, we just derive the posterior of \(z\):

\[ p(z|y) \propto p(y|z)p(z) \propto e^{-m\theta}\exp\left\{-\frac{\lVert y - f(x) \rVert_2^2}{2e^{2\theta}}\right\}p(x), \]

and we apply the Laplace approximation using:

\[ L(z) = -m\theta - \frac{\lVert y - f(x) \rVert_2^2}{2}e^{-2\theta} + \log p(x). \]

The first and second derivatives with respect to \(x\) are just like before, with \(\sigma^2=e^{2\theta}\). We need the derivatives with respect to \(\theta\):

\[ \frac{\partial L(z)}{\partial \theta} = -m + \lVert y - f(x) \rVert_2^2e^{-2\theta}, \]
\[ \frac{\partial^2 L(z)}{\partial \theta^2} = -2\lVert y - f(x) \rVert_2^2e^{-2\theta}, \]

and

\[ \frac{\partial^2 L(z)}{\partial x_j\partial\theta} = -2e^{-2\theta}\sum_{i=1}^m(y_i-f_i(x))\frac{\partial f_i(x)}{\partial x_j}. \]

So the log posterior will look largely the same as the previous case, but with another prior on \(\theta\).

# Negative log posterior and its gradient
def post_w_noise(params, z, y, t, gamma):

    m = y.shape[0]

    # Extract theta from our optimized parameters
    theta = params[-1] 
    # Exponential of theta is the noise
    sigma = jnp.exp(theta)

    # Extract the parameters to solve the ODE
    x = params[:-1]

    res = solve_catalysis(t, x, z)
    flat_res = jnp.hstack([res[:, :2], res[:, 3:]]).flatten()
    
    # Negative log-likelihood
    tmp = (flat_res - y) 
    likelihood = 0.5 * jnp.sum(tmp**2) / sigma** 2
    
    # Negative log-prior
    prior = 0.5 * jnp.sum(x**2) / gamma**2

    # Gaussian likelihood normalization; the prior on theta is flat.
    likelihood_normalization = theta * m
    
    # Total negative log posterior
    posterior = likelihood + prior + likelihood_normalization
    
    return posterior

Optimize again like we did before. notice how we have grouped the things we are optimizing together in params. This is nice because we can take the gradient of the log posterior with respect to this vector and pass it to the optimizer.

# Initial guess for x
key = jr.PRNGKey(0)
x0 = jr.normal(key, shape=(5,))  

# Initial conditions
z0 = jnp.array([500., 0., 0., 0., 0., 0.0])

# Extract the experimental data
Y = catalysis_data.loc[catalysis_data['Time'] > 0, ['NO3', 'NO2', 'N2', 'NH3', 'N2O']].values
y = Y.flatten()

# Initialize the noise parameter
sigma0 = 5.0
theta0 = jnp.log(sigma0)
gamma = 10.0

# Concatenate the parameters
params0 = jnp.hstack([x0, theta0])

# Set up the optimizer
optimizer = optax.adam(learning_rate=1e-1)
opt_state = optimizer.init(params0)

# Use as many iterations as needed
num_iterations = 300
loss_noise = []
params = params0 

objective_and_gradient = jit(value_and_grad(post_w_noise))

for i in range(num_iterations):
    
    value, grads = objective_and_gradient(params, z0, y, t_exp, gamma)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    loss_noise.append(value)
        # Print the loss every 100 iterations
    if i % 50 == 0:
        print(f"Iteration {i}, loss: {value}")

# Refine the optimizer endpoint before evaluating its local curvature.
fit = minimize(lambda p: objective_and_gradient(p, z0, y, t_exp, gamma), np.asarray(params), jac=True,
               method='BFGS', options={'gtol': 1e-7, 'maxiter': 1000})
params = jnp.asarray(fit.x)
gradient_norm = float(jnp.linalg.norm(objective_and_gradient(params, z0, y, t_exp, gamma)[1], ord=jnp.inf))
assert gradient_norm < 1e-5, gradient_norm
print('Final maximum absolute gradient:', gradient_norm)

post_params = params
mu_noise = params[:-1]
post_theta = params[-1]
mean_post_sigma = jnp.exp(post_theta)
Iteration 0, loss: 11897.666800950405
Iteration 50, loss: 115.89580649205872
Iteration 100, loss: 106.97548718803876
Iteration 150, loss: 103.11537573893686
Iteration 200, loss: 98.59514126999517
Iteration 250, loss: 96.39368207867776
Final maximum absolute gradient: 7.349792105330822e-07
# Plot the loss evolution
fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
ax.plot(loss_noise, color=colors[0])
ax.set_xlabel("Iteration")
ax.set_ylabel("Loss")
ax.set_yscale('log')
finalize_axes(keep_box=False)
plt.show()
Negative log-posterior loss with inferred observation noise decreases on a logarithmic scale.

We now form and check the joint Hessian for the rate coordinates and log-noise scale.

def compute_post_cov_w_noise(mu, y, gamma, z0, t, epsilon=1e-3):
    """Compute the posterior covariance matrix using the Hessian of the log-posterior."""

    n = mu.size - 1  # Number of model parameters x
    x = mu[:n]
    theta = mu[n]
    sigma = jnp.exp(theta)
    sigma2_inv = 1.0 / sigma ** 2

    # Define the function f(x) that computes the model outputs
    def f(x):
        sol = solve_catalysis(t, x, z0)
        f_x = jnp.hstack([sol[:, :2], sol[:, 3:]]).flatten()
        return f_x

    # Compute f(x) at the given x
    f_mu = f(x)
    m = f_mu.size

    # Compute tmp = y - f_mu
    tmp = y - f_mu

    # Compute gradient dfdx (Jacobian matrix), shape (m, n)
    # Use JAX's vmap to vectorize the computation over parameters
    def compute_dfdx_i(i):
        e_i = jnp.zeros_like(x)
        e_i = e_i.at[i].set(epsilon)
        f_plus = f(x + e_i)
        f_minus = f(x - e_i)
        f_prime = (f_plus - f_minus) / (2 * epsilon)
        return f_prime  # Shape: (m,)

    indices = jnp.arange(n)
    dfdx = jax.vmap(compute_dfdx_i)(indices)  # Shape: (n, m)
    dfdx = dfdx.T  # Transpose to shape (m, n)

    # Compute Hessian d2fdx2, shape (m, n, n)
    # We will compute only the upper triangle and exploit symmetry
    def compute_d2fdx2_ij(i, j):
        e_i = jnp.zeros_like(x)
        e_j = jnp.zeros_like(x)
        e_i = e_i.at[i].set(epsilon)
        e_j = e_j.at[j].set(epsilon)

        f_pp = f(x + e_i + e_j)
        f_pm = f(x + e_i - e_j)
        f_mp = f(x - e_i + e_j)
        f_mm = f(x - e_i - e_j)

        second_derivative = (f_pp - f_pm - f_mp + f_mm) / (4 * epsilon ** 2)
        return second_derivative  # Shape: (m,)

    # Generate all index pairs (i, j) with i <= j
    index_pairs = [(i, j) for i in range(n) for j in range(i, n)]
    num_pairs = len(index_pairs)

    # Vectorize over index pairs
    i_indices = jnp.array([pair[0] for pair in index_pairs])
    j_indices = jnp.array([pair[1] for pair in index_pairs])

    d2fdx2_values = jax.vmap(compute_d2fdx2_ij)(i_indices, j_indices)  # Shape: (num_pairs, m)

    # Initialize Hessian tensor
    d2fdx2 = jnp.zeros((m, n, n), dtype=jnp.float64)

    # Populate d2fdx2 using immutable updates
    for idx in range(num_pairs):
        i = i_indices[idx]
        j = j_indices[idx]
        second_derivative = d2fdx2_values[idx]  # Shape: (m,)
        d2fdx2 = d2fdx2.at[:, i, j].set(second_derivative)
        if i != j:
            d2fdx2 = d2fdx2.at[:, j, i].set(second_derivative)  # Exploit symmetry

    # Compute the components of the Hessian matrix

    # Term1: -sigma2_inv * sum_i tmp_i * d2f_i/dx^2
    term1 = - sigma2_inv * jnp.einsum('i,ijk->jk', tmp, d2fdx2)  # Shape: (n, n)

    # Term2: sigma2_inv * (dfdx.T @ dfdx)
    term2 = sigma2_inv * (dfdx.T @ dfdx)  # Shape: (n, n)

    # Compute d2Ldx2
    d2Ldx2 = term1 + term2 + jnp.eye(n) / gamma ** 2  # Shape: (n, n)

    # Compute d2Ldxdtheta: 2 * sigma2_inv * (dfdx.T @ tmp)
    d2Ldxdtheta = 2 * sigma2_inv * (dfdx.T @ tmp)  # Shape: (n,)

    # Compute d2Ldtheta2: 2 * sigma2_inv * (tmp.T @ tmp)
    d2Ldtheta2 = 2 * sigma2_inv * (tmp.T @ tmp)  # Scalar

    # Assemble the Hessian matrix Lam, shape (n + 1, n + 1)
    Lam = jnp.zeros((n+1, n+1), dtype=jnp.float64)
    Lam = Lam.at[:n, :n].set(d2Ldx2)
    Lam = Lam.at[:n, n].set(d2Ldxdtheta)
    Lam = Lam.at[n, :n].set(d2Ldxdtheta)
    Lam = Lam.at[n, n].set(d2Ldtheta2)

    exact_hessian = jax.hessian(post_w_noise)(mu, z0, y, t, gamma)
    assert jnp.allclose(Lam, Lam.T, atol=1e-10)
    relative_error = jnp.linalg.norm(Lam - exact_hessian) / jnp.linalg.norm(exact_hessian)
    assert relative_error < 1e-4, relative_error
    assert jnp.linalg.eigvalsh(Lam).min() > 0
    Sigma = jnp.linalg.inv(Lam)
    Sigma = 0.5 * (Sigma + Sigma.T)

    return Sigma

Let’s plot our results.

# Compute the posterior covariance with noise
Sigmas_noise = compute_post_cov_w_noise(post_params, y, gamma, z0, t_exp)

num_samples_noise = 50
times_noise = jnp.linspace(t0, t1, 1000)
key, subkey = jr.split(key)
samples_noise = jax.random.multivariate_normal(subkey, mu_noise, Sigmas_noise[:-1, :-1], shape=(num_samples_noise,))

# Compute the samples with noise
sample_models_noise = jnp.array([solve_catalysis(times_noise, sample, z0) for sample in samples_noise])

# Compute the median model with noise
median_models_noise = jnp.median(sample_models_noise, axis=0)

# Plotting with noise
fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])

# Define labels and colors
labels = ['NO3-', 'NO2-', 'N2', 'NH3', 'N2O', 'X']
data_cols = ['NO3', 'NO2', 'N2', 'NH3', 'N2O']
model_cols = [0, 1, 3, 4, 5, 2]
species_styles = ['-', '--', '-.', ':', (0, (5, 2)), (0, (3, 1, 1, 1))]
species_markers = ['o', 's', '^', 'D', 'v', 'P']

# Plot experimental data
for i, col in enumerate(data_cols):
    ax.plot(catalysis_data['Time'], catalysis_data[col], linestyle='none', color='black',
            marker=species_markers[i], markerfacecolor='white',
            label=f'Data {labels[i]}')

# Plot the mean models with noise
for i, col in enumerate(model_cols):
    ax.plot(times_noise, median_models_noise[:, col], color='black',
            linestyle=species_styles[i], label=f'Median {labels[i]}')

for i in range(min(num_samples_noise, 12)):
    for j, col in enumerate(model_cols):
        ax.plot(times_noise, sample_models_noise[i, :, col], color='0.65',
                linestyle=species_styles[j], linewidth=0.45, alpha=0.25)
ax.set_ylim(0, 600)
ax.set_xlabel('Time (min)')
ax.set_ylabel(r'Concentration (mmol L$^{-1}$)')
ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.20), ncol=2)
finalize_axes(keep_box=False)
plt.show()
Observed concentrations with posterior median and ensemble trajectories when observation noise is inferred.
# The Laplace marginal for theta is Gaussian, so sigma = exp(theta) is lognormal.
theta_variance = Sigmas_noise[-1, -1]
sigma_median = jnp.exp(post_theta)
sigma_mean = jnp.exp(post_theta + 0.5 * theta_variance)
sigma_std = jnp.sqrt((jnp.exp(theta_variance) - 1) * jnp.exp(2 * post_theta + theta_variance))
print(f'Noise standard deviation: median {sigma_median:.3f}, mean {sigma_mean:.3f}, posterior SD {sigma_std:.3f} mmol/L')
Noise standard deviation: median 11.116, mean 11.209, posterior SD 1.453 mmol/L

Exercises#

  • Investigate what happens as you change \(\gamma\) from smaller to bigger values.

  • Is the uncertainty we visualized above epistemic or aleatory?

  • The uncertainty visualized concerns only the model. What if you wanted to include the measurement noise in this visualization?

  • The assumption of Gaussian noise is not very good. Generate two different assumptions that you could try for the noise.