optimiz-rs/docs/mcmc.md

MCMC Sampling API

Overview

The MCMC (Markov Chain Monte Carlo) module implements the Metropolis-Hastings algorithm for sampling from arbitrary probability distributions. This is particularly useful for Bayesian parameter estimation, posterior inference, and uncertainty quantification.

Function: mcmc_sample

from optimizr import mcmc_sample

Signature

mcmc_sample(
    log_likelihood_fn: Callable[[List[float], List[float]], float],
    data: np.ndarray,
    initial_params: np.ndarray,
    param_bounds: List[Tuple[float, float]],
    n_samples: int = 10000,
    burn_in: int = 1000,
    proposal_std: float = 0.1,
) -> np.ndarray

Parameters

• log_likelihood_fn (callable): Function that computes log P(data | params).

  • Signature: log_likelihood_fn(params: list, data: list) -> float
  • Should return the natural logarithm of the likelihood.
  • Higher values indicate better fit.

• data (np.ndarray): Observed data passed to the log-likelihood function.

• initial_params (np.ndarray): Starting parameter values. Should be a 1D array.

• param_bounds (List[Tuple[float, float]]): List of (min, max) bounds for each parameter. Must have the same length as initial_params.

• n_samples (int, optional): Number of samples to generate after burn-in. Default is 10,000.

• burn_in (int, optional): Number of initial samples to discard. Default is 1,000.

• proposal_std (float, optional): Standard deviation of Gaussian random walk proposals. Default is 0.1.

Returns

• samples (np.ndarray): Array of shape (n_samples, n_params) containing parameter samples from the posterior distribution.

Raises

• ValueError: If initial_params and param_bounds have different lengths.

Basic Example

import numpy as np
from optimizr import mcmc_sample

# Define log-likelihood for Gaussian model
def log_likelihood(params, data):
    mu, sigma = params
    if sigma <= 0:
        return -np.inf
    residuals = (data - mu) / sigma
    return -0.5 * np.sum(residuals**2) - len(data) * np.log(sigma)

# Generate synthetic data
np.random.seed(42)
true_mu, true_sigma = 2.5, 1.2
data = np.random.normal(true_mu, true_sigma, 100)

# Sample from posterior
samples = mcmc_sample(
    log_likelihood_fn=log_likelihood,
    data=data,
    initial_params=np.array([0.0, 1.0]),
    param_bounds=[(-10, 10), (0.1, 10)],
    n_samples=10000,
    burn_in=1000,
    proposal_std=0.1
)

# Analyze results
print(f"True mean: {true_mu:.2f}, Estimated: {np.mean(samples[:, 0]):.2f}")
print(f"True std: {true_sigma:.2f}, Estimated: {np.mean(samples[:, 1]):.2f}")

# Posterior credible intervals
print(f"Mean 95% CI: {np.percentile(samples[:, 0], [2.5, 97.5])}")
print(f"Std 95% CI: {np.percentile(samples[:, 1], [2.5, 97.5])}")

Advanced Examples

1. Linear Regression

import numpy as np
from optimizr import mcmc_sample

# Log-likelihood for linear regression
def log_likelihood(params, data):
    x, y = data
    slope, intercept, sigma = params
    
    if sigma <= 0:
        return -np.inf
    
    predictions = slope * x + intercept
    residuals = (y - predictions) / sigma
    
    return -0.5 * np.sum(residuals**2) - len(y) * np.log(sigma)

# Generate data
np.random.seed(42)
x = np.linspace(0, 10, 100)
y = 2.5 * x + 1.0 + np.random.normal(0, 0.5, 100)

# Sample from posterior
samples = mcmc_sample(
    log_likelihood_fn=log_likelihood,
    data=[x, y],
    initial_params=np.array([1.0, 0.0, 1.0]),
    param_bounds=[(-10, 10), (-10, 10), (0.01, 10)],
    n_samples=20000,
    burn_in=2000,
    proposal_std=0.05
)

print(f"Slope: {np.mean(samples[:, 0]):.3f} ± {np.std(samples[:, 0]):.3f}")
print(f"Intercept: {np.mean(samples[:, 1]):.3f} ± {np.std(samples[:, 1]):.3f}")
print(f"Sigma: {np.mean(samples[:, 2]):.3f} ± {np.std(samples[:, 2]):.3f}")

2. Mixture Model

import numpy as np
from optimizr import mcmc_sample
from scipy.stats import norm

def log_likelihood(params, data):
    """Two-component Gaussian mixture"""
    mu1, sigma1, mu2, sigma2, weight1 = params
    
    # Ensure valid parameters
    if sigma1 <= 0 or sigma2 <= 0:
        return -np.inf
    if not (0 <= weight1 <= 1):
        return -np.inf
    
    weight2 = 1 - weight1
    
    # Mixture likelihood
    likelihood = (weight1 * norm.pdf(data, mu1, sigma1) + 
                  weight2 * norm.pdf(data, mu2, sigma2))
    
    return np.sum(np.log(likelihood + 1e-10))

# Generate mixture data
np.random.seed(42)
data = np.concatenate([
    np.random.normal(0, 1, 300),
    np.random.normal(5, 1.5, 200)
])

samples = mcmc_sample(
    log_likelihood_fn=log_likelihood,
    data=data,
    initial_params=np.array([0.0, 1.0, 5.0, 1.5, 0.6]),
    param_bounds=[(-5, 5), (0.1, 5), (0, 10), (0.1, 5), (0.1, 0.9)],
    n_samples=15000,
    burn_in=3000,
    proposal_std=0.15
)

print("Component 1:")
print(f"  Mean: {np.mean(samples[:, 0]):.2f}")
print(f"  Std: {np.mean(samples[:, 1]):.2f}")
print("\nComponent 2:")
print(f"  Mean: {np.mean(samples[:, 2]):.2f}")
print(f"  Std: {np.mean(samples[:, 3]):.2f}")
print(f"\nMixing weight: {np.mean(samples[:, 4]):.2f}")

3. Time Series Model (AR process)

import numpy as np
from optimizr import mcmc_sample

def log_likelihood(params, data):
    """Autoregressive AR(2) model"""
    phi1, phi2, sigma = params
    
    if sigma <= 0:
        return -np.inf
    
    # Check stationarity conditions
    if abs(phi1) + abs(phi2) >= 1:
        return -np.inf
    
    # Compute residuals
    predictions = phi1 * data[1:-1] + phi2 * data[:-2]
    residuals = (data[2:] - predictions) / sigma
    
    return -0.5 * np.sum(residuals**2) - (len(data) - 2) * np.log(sigma)

# Generate AR(2) process
np.random.seed(42)
n = 500
phi1_true, phi2_true = 0.6, -0.2
sigma_true = 0.5

data = np.zeros(n)
for t in range(2, n):
    data[t] = (phi1_true * data[t-1] + 
               phi2_true * data[t-2] + 
               np.random.normal(0, sigma_true))

samples = mcmc_sample(
    log_likelihood_fn=log_likelihood,
    data=data,
    initial_params=np.array([0.5, -0.1, 0.5]),
    param_bounds=[(-0.99, 0.99), (-0.99, 0.99), (0.01, 5)],
    n_samples=15000,
    burn_in=2000,
    proposal_std=0.05
)

print(f"φ₁: True={phi1_true:.2f}, Est={np.mean(samples[:, 0]):.2f}")
print(f"φ₂: True={phi2_true:.2f}, Est={np.mean(samples[:, 1]):.2f}")
print(f"σ: True={sigma_true:.2f}, Est={np.mean(samples[:, 2]):.2f}")

Diagnostics and Visualization

Trace Plots

import matplotlib.pyplot as plt

fig, axes = plt.subplots(3, 1, figsize=(10, 8))
param_names = ['Mean', 'Std Dev', 'Parameter 3']

for i, (ax, name) in enumerate(zip(axes, param_names)):
    ax.plot(samples[:, i], alpha=0.7)
    ax.set_ylabel(name)
    ax.axhline(np.mean(samples[:, i]), color='r', linestyle='--', 
               label='Mean')
    ax.legend()

axes[-1].set_xlabel('Iteration')
plt.tight_layout()
plt.show()

Posterior Distributions

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

for i, ax in enumerate(axes):
    ax.hist(samples[:, i], bins=50, density=True, alpha=0.7)
    ax.axvline(np.mean(samples[:, i]), color='r', linestyle='--',
               label=f'Mean: {np.mean(samples[:, i]):.3f}')
    ax.set_xlabel(f'Parameter {i+1}')
    ax.set_ylabel('Density')
    ax.legend()

plt.tight_layout()
plt.show()

Autocorrelation

from scipy.stats import acf

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

for i, ax in enumerate(axes):
    autocorr = acf(samples[:, i], nlags=100)
    ax.plot(autocorr)
    ax.axhline(0, color='k', linestyle='--', alpha=0.3)
    ax.set_xlabel('Lag')
    ax.set_ylabel('Autocorrelation')
    ax.set_title(f'Parameter {i+1}')

plt.tight_layout()
plt.show()

Acceptance Rate

# Estimate acceptance rate from consecutive samples
def acceptance_rate(samples):
    changes = np.sum(np.diff(samples, axis=0) != 0, axis=1)
    return np.mean(changes > 0)

rate = acceptance_rate(samples)
print(f"Acceptance rate: {rate:.2%}")

# Ideal range: 20-40% for Metropolis-Hastings
if rate < 0.15:
    print("⚠ Acceptance rate too low. Try decreasing proposal_std.")
elif rate > 0.50:
    print("⚠ Acceptance rate too high. Try increasing proposal_std.")
else:
    print("✓ Acceptance rate is in good range.")

Performance Notes

• Rust Backend: When available, MCMC sampling is 50-100x faster than pure Python implementations.

• Python Fallback: A pure NumPy fallback is automatically used if Rust is unavailable.

• Proposal Tuning: The proposal_std parameter significantly affects convergence:

  • Too small: Slow exploration, high acceptance rate
  • Too large: Poor acceptance rate, slow convergence
  • Optimal: 20-40% acceptance rate

Tips and Best Practices

1. Choosing Burn-in Period

# Run a short chain to visualize convergence
test_samples = mcmc_sample(
    log_likelihood_fn=log_likelihood,
    data=data,
    initial_params=initial_params,
    param_bounds=bounds,
    n_samples=5000,
    burn_in=0,  # Keep all samples for inspection
    proposal_std=0.1
)

# Plot to determine burn-in
plt.plot(test_samples[:, 0])
plt.xlabel('Iteration')
plt.ylabel('Parameter')
plt.title('Determine burn-in period')
plt.show()

2. Multiple Chains

# Run multiple chains with different starting points
n_chains = 4
all_samples = []

for i in range(n_chains):
    # Random initialization
    init = np.random.uniform(
        [b[0] for b in bounds],
        [b[1] for b in bounds]
    )
    
    samples = mcmc_sample(
        log_likelihood_fn=log_likelihood,
        data=data,
        initial_params=init,
        param_bounds=bounds,
        n_samples=5000,
        burn_in=1000
    )
    all_samples.append(samples)

# Check convergence across chains
means = [np.mean(s[:, 0]) for s in all_samples]
print(f"Chain means: {means}")
print(f"Variance: {np.var(means):.6f}")

3. Adaptive Proposal

# Start with exploration, then refine
samples_phase1 = mcmc_sample(
    log_likelihood_fn=log_likelihood,
    data=data,
    initial_params=initial_params,
    param_bounds=bounds,
    n_samples=5000,
    burn_in=1000,
    proposal_std=0.2  # Larger for exploration
)

# Use posterior mean as new starting point
new_init = np.mean(samples_phase1, axis=0)

samples_phase2 = mcmc_sample(
    log_likelihood_fn=log_likelihood,
    data=data,
    initial_params=new_init,
    param_bounds=bounds,
    n_samples=10000,
    burn_in=1000,
    proposal_std=0.05  # Smaller for refinement
)

Common Issues and Solutions

Issue	Cause	Solution
Chains don't converge	Poor initialization	Use multiple chains or better starting values
High autocorrelation	Proposal too small	Increase proposal_std
Low acceptance rate	Proposal too large	Decrease proposal_std
Bimodal posterior	Multiple modes	Use longer chains or parallel tempering
Numerical errors	Overflow in likelihood	Use log-likelihood correctly

See Also

• HMM API - For regime detection and sequence modeling
• MCMC Theory - Mathematical background
• Examples - Complete working examples