optimiz-rs/examples/notebooks/02_mcmc_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dc5d5825",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OptimizR MCMC Module Loaded!\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import stats\n",
    "from optimizr import mcmc_sample\n",
    "\n",
    "np.random.seed(42)\n",
    "print(\"OptimizR MCMC Module Loaded!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddfbd617",
   "metadata": {},
   "source": [
    "## Example 1: Inferring Parameters of a Normal Distribution\n",
    "\n",
    "Given observed data $\\{x_1, \\ldots, x_n\\}$, infer $\\mu$ and $\\sigma$.\n",
    "\n",
    "### Likelihood\n",
    "$$L(\\mu, \\sigma | \\mathbf{x}) = \\prod_{i=1}^n \\frac{1}{\\sqrt{2\\pi\\sigma^2}} \\exp\\left(-\\frac{(x_i - \\mu)^2}{2\\sigma^2}\\right)$$\n",
    "\n",
    "### Log-Likelihood\n",
    "$$\\log L(\\mu, \\sigma | \\mathbf{x}) = -\\frac{n}{2}\\log(2\\pi) - n\\log(\\sigma) - \\frac{1}{2\\sigma^2}\\sum_{i=1}^n (x_i - \\mu)^2$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f929c18e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate synthetic data\n",
    "true_mu = 5.0\n",
    "true_sigma = 2.0\n",
    "n_obs = 100\n",
    "\n",
    "observed_data = np.random.normal(true_mu, true_sigma, n_obs)\n",
    "\n",
    "print(f\"True parameters: μ={true_mu}, σ={true_sigma}\")\n",
    "print(f\"Sample mean: {observed_data.mean():.3f}\")\n",
    "print(f\"Sample std: {observed_data.std():.3f}\")\n",
    "\n",
    "# Plot data\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.hist(observed_data, bins=20, density=True, alpha=0.6, color='skyblue', edgecolor='black')\n",
    "x_range = np.linspace(observed_data.min(), observed_data.max(), 100)\n",
    "plt.plot(x_range, stats.norm.pdf(x_range, true_mu, true_sigma), \n",
    "         'r-', linewidth=2, label=f'True: N({true_mu}, {true_sigma}²)')\n",
    "plt.xlabel('Value', fontsize=12)\n",
    "plt.ylabel('Density', fontsize=12)\n",
    "plt.title('Observed Data Distribution', fontsize=14, fontweight='bold')\n",
    "plt.legend()\n",
    "plt.grid(alpha=0.3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5626ba97",
   "metadata": {},
   "source": [
    "### Define Log-Likelihood Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37ac93a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_likelihood_normal(params, data):\n",
    "    \"\"\"\n",
    "    Log-likelihood for Normal(μ, σ²) given data.\n",
    "    \n",
    "    Args:\n",
    "        params: [μ, σ]\n",
    "        data: observed data points\n",
    "    \"\"\"\n",
    "    mu, sigma = params\n",
    "    \n",
    "    # Ensure sigma is positive\n",
    "    if sigma <= 0:\n",
    "        return -np.inf\n",
    "    \n",
    "    n = len(data)\n",
    "    residuals = (data - mu) / sigma\n",
    "    \n",
    "    log_lik = -0.5 * n * np.log(2 * np.pi)\n",
    "    log_lik -= n * np.log(sigma)\n",
    "    log_lik -= 0.5 * np.sum(residuals**2)\n",
    "    \n",
    "    return log_lik\n",
    "\n",
    "# Test the function\n",
    "test_params = [5.0, 2.0]\n",
    "print(f\"Log-likelihood at true params: {log_likelihood_normal(test_params, observed_data):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a36ddef7",
   "metadata": {},
   "source": [
    "### Run MCMC Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60e64783",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MCMC parameters\n",
    "initial_params = [0.0, 1.0]  # Start far from true values\n",
    "param_bounds = [(-10, 10), (0.1, 10)]  # μ ∈ [-10, 10], σ ∈ [0.1, 10]\n",
    "proposal_std = [0.5, 0.2]  # Proposal step sizes\n",
    "n_samples = 20000\n",
    "burn_in = 2000\n",
    "\n",
    "print(\"Running MCMC sampling...\")\n",
    "samples, acceptance_rate = mcmc_sample(\n",
    "    log_likelihood_fn=log_likelihood_normal,\n",
    "    data=observed_data,\n",
    "    initial_params=initial_params,\n",
    "    param_bounds=param_bounds,\n",
    "    proposal_std=proposal_std,\n",
    "    n_samples=n_samples,\n",
    "    burn_in=burn_in\n",
    ")\n",
    "\n",
    "print(f\"\\nAcceptance rate: {acceptance_rate:.2%}\")\n",
    "print(f\"Generated {len(samples)} samples after burn-in\")\n",
    "print(f\"\\nPosterior estimates:\")\n",
    "print(f\"μ: {samples[:, 0].mean():.3f} ± {samples[:, 0].std():.3f}\")\n",
    "print(f\"σ: {samples[:, 1].mean():.3f} ± {samples[:, 1].std():.3f}\")\n",
    "print(f\"\\nTrue values: μ={true_mu}, σ={true_sigma}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2306bc9b",
   "metadata": {},
   "source": [
    "## Visualize MCMC Results\n",
    "\n",
    "### Trace Plots - Check Convergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "856b7ca8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, figsize=(14, 8))\n",
    "\n",
    "# Trace plots\n",
    "axes[0, 0].plot(samples[:, 0], linewidth=0.5, alpha=0.7)\n",
    "axes[0, 0].axhline(true_mu, color='red', linestyle='--', linewidth=2, label='True μ')\n",
    "axes[0, 0].set_xlabel('Sample', fontsize=11)\n",
    "axes[0, 0].set_ylabel('μ', fontsize=11)\n",
    "axes[0, 0].set_title('Trace Plot: μ', fontsize=13, fontweight='bold')\n",
    "axes[0, 0].legend()\n",
    "axes[0, 0].grid(alpha=0.3)\n",
    "\n",
    "axes[0, 1].plot(samples[:, 1], linewidth=0.5, alpha=0.7, color='orange')\n",
    "axes[0, 1].axhline(true_sigma, color='red', linestyle='--', linewidth=2, label='True σ')\n",
    "axes[0, 1].set_xlabel('Sample', fontsize=11)\n",
    "axes[0, 1].set_ylabel('σ', fontsize=11)\n",
    "axes[0, 1].set_title('Trace Plot: σ', fontsize=13, fontweight='bold')\n",
    "axes[0, 1].legend()\n",
    "axes[0, 1].grid(alpha=0.3)\n",
    "\n",
    "# Posterior distributions\n",
    "axes[1, 0].hist(samples[:, 0], bins=50, density=True, alpha=0.6, color='skyblue', edgecolor='black')\n",
    "axes[1, 0].axvline(true_mu, color='red', linestyle='--', linewidth=2, label='True μ')\n",
    "axes[1, 0].axvline(samples[:, 0].mean(), color='green', linestyle='-', linewidth=2, label='Posterior mean')\n",
    "axes[1, 0].set_xlabel('μ', fontsize=11)\n",
    "axes[1, 0].set_ylabel('Density', fontsize=11)\n",
    "axes[1, 0].set_title('Posterior Distribution: μ', fontsize=13, fontweight='bold')\n",
    "axes[1, 0].legend()\n",
    "axes[1, 0].grid(alpha=0.3)\n",
    "\n",
    "axes[1, 1].hist(samples[:, 1], bins=50, density=True, alpha=0.6, color='orange', edgecolor='black')\n",
    "axes[1, 1].axvline(true_sigma, color='red', linestyle='--', linewidth=2, label='True σ')\n",
    "axes[1, 1].axvline(samples[:, 1].mean(), color='green', linestyle='-', linewidth=2, label='Posterior mean')\n",
    "axes[1, 1].set_xlabel('σ', fontsize=11)\n",
    "axes[1, 1].set_ylabel('Density', fontsize=11)\n",
    "axes[1, 1].set_title('Posterior Distribution: σ', fontsize=13, fontweight='bold')\n",
    "axes[1, 1].legend()\n",
    "axes[1, 1].grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b20c718",
   "metadata": {},
   "source": [
    "### Joint Posterior Distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad2eebc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "# 2D histogram\n",
    "plt.hist2d(samples[:, 0], samples[:, 1], bins=50, cmap='Blues')\n",
    "plt.colorbar(label='Sample Density')\n",
    "\n",
    "# Mark true values\n",
    "plt.scatter([true_mu], [true_sigma], c='red', s=200, marker='*', \n",
    "           edgecolors='black', linewidths=2, label='True values', zorder=5)\n",
    "\n",
    "# Mark posterior mean\n",
    "plt.scatter([samples[:, 0].mean()], [samples[:, 1].mean()], \n",
    "           c='green', s=200, marker='o', edgecolors='black', \n",
    "           linewidths=2, label='Posterior mean', zorder=5)\n",
    "\n",
    "plt.xlabel('μ', fontsize=12)\n",
    "plt.ylabel('σ', fontsize=12)\n",
    "plt.title('Joint Posterior Distribution', fontsize=14, fontweight='bold')\n",
    "plt.legend(fontsize=11)\n",
    "plt.grid(alpha=0.3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "049df073",
   "metadata": {},
   "source": [
    "## Example 2: Logistic Regression with MCMC\n",
    "\n",
    "Bayesian inference for binary classification.\n",
    "\n",
    "### Model\n",
    "$$P(y=1 | \\mathbf{x}, \\boldsymbol{\\beta}) = \\frac{1}{1 + \\exp(-\\boldsymbol{\\beta}^T \\mathbf{x})}$$\n",
    "\n",
    "### Log-Likelihood\n",
    "$$\\log L(\\boldsymbol{\\beta}) = \\sum_{i=1}^n \\left[y_i \\log p_i + (1-y_i) \\log(1-p_i)\\right]$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072639a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate synthetic classification data\n",
    "from sklearn.datasets import make_classification\n",
    "\n",
    "X, y = make_classification(n_samples=200, n_features=2, n_redundant=0,\n",
    "                          n_informative=2, random_state=42, n_clusters_per_class=1)\n",
    "\n",
    "# Add intercept\n",
    "X_with_intercept = np.column_stack([np.ones(len(X)), X])\n",
    "\n",
    "print(f\"Features shape: {X_with_intercept.shape}\")\n",
    "print(f\"Class distribution: {np.bincount(y)}\")\n",
    "\n",
    "# Visualize data\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.scatter(X[y == 0, 0], X[y == 0, 1], c='blue', label='Class 0', alpha=0.6, s=50)\n",
    "plt.scatter(X[y == 1, 0], X[y == 1, 1], c='red', label='Class 1', alpha=0.6, s=50)\n",
    "plt.xlabel('Feature 1', fontsize=12)\n",
    "plt.ylabel('Feature 2', fontsize=12)\n",
    "plt.title('Binary Classification Data', fontsize=14, fontweight='bold')\n",
    "plt.legend()\n",
    "plt.grid(alpha=0.3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6163cb52",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_likelihood_logistic(beta, X, y):\n",
    "    \"\"\"\n",
    "    Log-likelihood for logistic regression.\n",
    "    \"\"\"\n",
    "    z = X @ beta\n",
    "    # Numerically stable sigmoid\n",
    "    p = 1 / (1 + np.exp(-np.clip(z, -500, 500)))\n",
    "    p = np.clip(p, 1e-10, 1 - 1e-10)  # Avoid log(0)\n",
    "    \n",
    "    log_lik = np.sum(y * np.log(p) + (1 - y) * np.log(1 - p))\n",
    "    \n",
    "    # Add weak prior: beta ~ N(0, 10²)\n",
    "    log_prior = -0.5 * np.sum(beta**2) / 100\n",
    "    \n",
    "    return log_lik + log_prior\n",
    "\n",
    "# Prepare data tuple\n",
    "logistic_data = (X_with_intercept, y)\n",
    "\n",
    "# MCMC for logistic regression\n",
    "initial_beta = np.zeros(3)  # [intercept, coef1, coef2]\n",
    "beta_bounds = [(-10, 10)] * 3\n",
    "beta_proposal_std = [0.1] * 3\n",
    "\n",
    "print(\"Running MCMC for logistic regression...\")\n",
    "beta_samples, beta_acceptance = mcmc_sample(\n",
    "    log_likelihood_fn=log_likelihood_logistic,\n",
    "    data=logistic_data,\n",
    "    initial_params=initial_beta,\n",
    "    param_bounds=beta_bounds,\n",
    "    proposal_std=beta_proposal_std,\n",
    "    n_samples=15000,\n",
    "    burn_in=1500\n",
    ")\n",
    "\n",
    "print(f\"\\nAcceptance rate: {beta_acceptance:.2%}\")\n",
    "print(f\"\\nPosterior estimates:\")\n",
    "print(f\"β₀ (intercept): {beta_samples[:, 0].mean():.3f} ± {beta_samples[:, 0].std():.3f}\")\n",
    "print(f\"β₁: {beta_samples[:, 1].mean():.3f} ± {beta_samples[:, 1].std():.3f}\")\n",
    "print(f\"β₂: {beta_samples[:, 2].mean():.3f} ± {beta_samples[:, 2].std():.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cc9ff7c",
   "metadata": {},
   "source": [
    "### Visualize Decision Boundary with Uncertainty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eb5d0d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot decision boundaries from posterior samples\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "# Plot data\n",
    "plt.scatter(X[y == 0, 0], X[y == 0, 1], c='blue', label='Class 0', alpha=0.6, s=50, zorder=3)\n",
    "plt.scatter(X[y == 1, 0], X[y == 1, 1], c='red', label='Class 1', alpha=0.6, s=50, zorder=3)\n",
    "\n",
    "# Create grid\n",
    "x1_min, x1_max = X[:, 0].min() - 1, X[:, 0].max() + 1\n",
    "x2_min, x2_max = X[:, 1].min() - 1, X[:, 1].max() + 1\n",
    "\n",
    "# Plot decision boundaries from random posterior samples\n",
    "n_boundary_samples = 100\n",
    "indices = np.random.choice(len(beta_samples), n_boundary_samples, replace=False)\n",
    "\n",
    "for idx in indices:\n",
    "    beta = beta_samples[idx]\n",
    "    # Decision boundary: β₀ + β₁x₁ + β₂x₂ = 0\n",
    "    # => x₂ = -(β₀ + β₁x₁) / β₂\n",
    "    if abs(beta[2]) > 0.01:  # Avoid division by zero\n",
    "        x1_line = np.array([x1_min, x1_max])\n",
    "        x2_line = -(beta[0] + beta[1] * x1_line) / beta[2]\n",
    "        plt.plot(x1_line, x2_line, 'gray', alpha=0.02, linewidth=0.5, zorder=1)\n",
    "\n",
    "# Plot mean decision boundary\n",
    "beta_mean = beta_samples.mean(axis=0)\n",
    "if abs(beta_mean[2]) > 0.01:\n",
    "    x1_line = np.array([x1_min, x1_max])\n",
    "    x2_line = -(beta_mean[0] + beta_mean[1] * x1_line) / beta_mean[2]\n",
    "    plt.plot(x1_line, x2_line, 'black', linewidth=3, label='Mean boundary', zorder=2)\n",
    "\n",
    "plt.xlim(x1_min, x1_max)\n",
    "plt.ylim(x2_min, x2_max)\n",
    "plt.xlabel('Feature 1', fontsize=12)\n",
    "plt.ylabel('Feature 2', fontsize=12)\n",
    "plt.title('Logistic Regression: Decision Boundary with Uncertainty', fontsize=14, fontweight='bold')\n",
    "plt.legend()\n",
    "plt.grid(alpha=0.3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c4facb3",
   "metadata": {},
   "source": [
    "## MCMC Diagnostics\n",
    "\n",
    "### Autocorrelation - Check Mixing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab8cd82",
   "metadata": {},
   "outputs": [],
   "source": [
    "from statsmodels.graphics.tsaplots import plot_acf\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "plot_acf(samples[:, 0], lags=100, ax=axes[0], alpha=0.05)\n",
    "axes[0].set_title('Autocorrelation: μ', fontsize=13, fontweight='bold')\n",
    "axes[0].set_xlabel('Lag', fontsize=11)\n",
    "\n",
    "plot_acf(samples[:, 1], lags=100, ax=axes[1], alpha=0.05)\n",
    "axes[1].set_title('Autocorrelation: σ', fontsize=13, fontweight='bold')\n",
    "axes[1].set_xlabel('Lag', fontsize=11)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"Low autocorrelation at large lags indicates good mixing!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2ed2ccd",
   "metadata": {},
   "source": [
    "## Key Takeaways\n",
    "\n",
    "1. **MCMC samples from complex distributions** using only likelihood evaluations\n",
    "2. **Metropolis-Hastings** uses proposal distribution and accept/reject steps\n",
    "3. **Burn-in period** allows chain to converge to target distribution\n",
    "4. **Diagnostics** (trace plots, autocorrelation) verify convergence and mixing\n",
    "5. **OptimizR provides 50-100x speedup** for likelihood evaluations\n",
    "6. **Bayesian inference** naturally quantifies uncertainty in parameters\n",
    "\n",
    "## Further Reading\n",
    "\n",
    "- Gelman, A., et al. (2013). \"Bayesian Data Analysis\" - Chapter 11\n",
    "- Brooks, S., et al. (2011). \"Handbook of Markov Chain Monte Carlo\"\n",
    "- Betancourt, M. (2017). \"A Conceptual Introduction to Hamiltonian Monte Carlo\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rhftlab",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}