A MCMC implementation for a non-standard normal¶

This algorithm samples from the posterior distribution of the mean ( $\mu$ ) and precision ( $\tau = 1/\sigma^2$ ) for a normal model, given a set of observed data and conjugate priors.

Model Setup¶

Likelihood:
For data $x_1, \dots, x_n$ assumed iid normal:
$x_i \sim \mathcal{N}(\mu, \tau^{-1})$
(1)
The likelihood (joint density) is:
$p(\mathbf{x} | \mu, \tau) = \prod_{i=1}^n \left[ \sqrt{\frac{\tau}{2\pi}} \exp\left(-\frac{\tau}{2}(x_i - \mu)^2\right) \right]$
(2)
Priors:
- For $\mu$ :
  $\mu \sim \mathcal{N}(\mu_0, \tau_0^{-1})$
  (3)
  $p(\mu) = \sqrt{\frac{\tau_0}{2\pi}} \exp\left(-\frac{\tau_0}{2}(\mu - \mu_0)^2\right)$
  (4)
- For $\tau$ (precision):
  $\tau \sim \text{Gamma}(a, b)$
  (5)
  $p(\tau) = \frac{b^a}{\Gamma(a)} \tau^{a-1} \exp(-b\tau)$
  (6)
Posterior:
$p(\mu, \tau | \mathbf{x}) \propto p(\mathbf{x} | \mu, \tau)\, p(\mu)\, p(\tau)$
(7)

Algorithm Steps and Equations¶

Update $\mu$ (Metropolis-Hastings step):
- Proposal: Sample $\mu^* \sim \mathcal{N}(\mu, \delta^2)$
- Acceptance ratio:
  Compute the log-acceptance ratio:
  $\log R = \left[ \log p(\mathbf{x}|\mu^*,\tau) + \log p(\mu^*) \right] - \left[ \log p(\mathbf{x}|\mu,\tau) + \log p(\mu) \right]$
  (8)
  Where
  $\log p(\mathbf{x}|\mu,\tau) = \sum_{i=1}^n \left( -\frac{1}{2}\log(2\pi/\tau) - \frac{\tau}{2}(x_i-\mu)^2 \right)$
  (9)
  $\log p(\mu) = -\frac{1}{2}\log(2\pi/\tau_0) - \frac{\tau_0}{2}(\mu - \mu_0)^2$
  (10)
  The code implements these terms using exponentials and logs for numerical stability.
- Decision:
  Accept $\mu^*$ if $\log U \leq \log R$ , where $U \sim \text{Uniform}(0,1)$ .
Update $\tau$ (Gibbs step):
- The conditional posterior for $\tau$ is:
  $\tau | \mu, \mathbf{x} \sim \text{Gamma} \left(a + \frac{n}{2},\ b + \frac{1}{2} \sum_{i=1}^n (x_i - \mu)^2 \right)$
  (11)
- The code samples $\tau$ from this gamma distribution at each iteration.

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import scipy.stats as stats

def MHnormal(sample, priors, nIter, delta):
    # Monitors
    mon_mu = np.zeros(nIter)
    mon_tau = np.zeros(nIter)

    # Initial values
    mu = np.mean(sample)
    tau = 1/np.var(sample)

    # Priors
    mu0 = priors['mu0']
    tau0 = priors['tau0']
    a = priors['a']
    b = priors['b']
    n = len(sample)

    for i in range(nIter):
        # Proposal for mu
        mu_hat = np.random.normal(mu, delta)
        # Log acceptance ratio
        logR = (
            np.sum(
                np.log(
                    np.exp(-0.5 * tau * (sample - mu_hat)**2) / np.sqrt(2 * np.pi / tau)
                )
            )
            - np.sum(
                np.log(
                    np.exp(-0.5 * tau * (sample - mu)**2) / np.sqrt(2 * np.pi / tau)
                )
            )
            + np.log(
                np.exp(-0.5 * tau0 * (mu_hat - mu0)**2) / np.sqrt(2 * np.pi / tau0)
            )
            - np.log(
                np.exp(-0.5 * tau0 * (mu - mu0)**2) / np.sqrt(2 * np.pi / tau0)
            )
        )
        logU = np.log(np.random.uniform(0, 1))
        if logU <= logR:
            mu = mu_hat

        # Gibbs step for tau (shape, rate)
        tau = np.random.gamma(a + n / 2, 1 / (b + 0.5 * np.sum((sample - mu) ** 2)))

        # Store
        mon_mu[i] = mu
        mon_tau[i] = tau

    output = pd.DataFrame({'mu': mon_mu, 'tau': mon_tau})

    # Plotting
    fig, axs = plt.subplots(2, 2, figsize=(10, 8))
    # Traceplots
    axs[0, 0].plot(output['mu'])
    axs[0, 0].set_title(r'$\mu$ trace')
    axs[0, 0].set_ylabel('value')
    
    axs[0, 1].plot(1 / np.sqrt(output['tau']))
    axs[0, 1].set_title(r'$\sigma$ trace')
    axs[0, 1].set_ylabel('value')
    
    # Density plots
    mu_density = stats.gaussian_kde(output['mu'])
    mu_vals = np.linspace(min(output['mu']), max(output['mu']), 200)
    axs[1, 0].plot(mu_vals, mu_density(mu_vals), label='Posterior Density')
    axs[1, 0].set_title(r'$\mu$ density')
    axs[1, 0].set_xlabel('Value')
    axs[1, 0].set_ylabel('Density')
    
    sigma_samples = 1 / np.sqrt(output['tau'])
    sigma_density = stats.gaussian_kde(sigma_samples)
    sigma_vals = np.linspace(min(sigma_samples), max(sigma_samples), 200)
    axs[1, 1].plot(sigma_vals, sigma_density(sigma_vals), label='Posterior Density')
    axs[1, 1].set_title(r'$\sigma$ density')
    axs[1, 1].set_xlabel('Value')
    axs[1, 1].set_ylabel('Density')
    
    plt.tight_layout()
    plt.show()

    return output

# Example usage:
np.random.seed(42)
sample = np.random.normal(40, 13, 100)
priors = {
    'mu0': np.mean(sample),
    'tau0': float(np.std(sample)),
    'a': 0.01,
    'b': 0.01
}
#print(sample)

output = MHnormal(sample = sample, priors = priors, nIter = 500000, delta = 1)