{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import json\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import scipy.stats as st\n",
    "import scipy.integrate as integrate\n",
    "import matplotlib.tri as mtri\n",
    "from scipy.special import gammaln\n",
    "from sklearn import linear_model\n",
    "from scipy.stats import multivariate_normal\n",
    "# from sklearn.utils.testing import ignore_warnings\n",
    "from sklearn.exceptions import ConvergenceWarning\n",
    "import statsmodels.api as sm\n",
    "\n",
    "sns.set_style(\"whitegrid\")\n",
    "sns.set_palette(\"colorblind\")\n",
    "palette = sns.color_palette()\n",
    "figsize = (12,6)\n",
    "legend_fontsize = 14\n",
    "\n",
    "from matplotlib import rc\n",
    "rc('font',**{'family':'sans-serif'})\n",
    "rc('figure', **{'dpi': 300})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bayesian inference for a coin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = np.arange(-0.05, 1.05, 0.0025)\n",
    "\n",
    "prior_params = (50, 50)\n",
    "experimental_data = (0, 100)\n",
    "\n",
    "## Prior\n",
    "pri = st.beta(prior_params[0], prior_params[1]).pdf\n",
    "norm_prior = 1\n",
    "ys_prior = [ pri(x) / norm_prior for x in xs ]\n",
    "\n",
    "## Likelihood\n",
    "n_heads, n_tails = experimental_data\n",
    "lk = lambda x : x ** n_heads * (1 - x) ** n_tails\n",
    "ys_like = [ lk(x) for x in xs ]\n",
    "\n",
    "## Posterior\n",
    "post = lambda x : lk(x) * pri(x)\n",
    "norm_post = integrate.quad(lambda x: post(x), 0, 1)[0]\n",
    "ys_post = [ post(x) / norm_post if x > 0 and x < 1 else 0 for x in xs ]\n",
    "\n",
    "## Plotting\n",
    "fig = plt.figure(figsize=figsize)\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(xs, ys_prior, linewidth=4, label=r\"Prior\")\n",
    "ax.plot(xs, ys_like, linewidth=4, label=r\"Likelihood\")\n",
    "ax.plot(xs, ys_post, linewidth=4, label=r\"Posterior\")\n",
    "ax.set_xlim((-0.05, 1.05))\n",
    "ax.set_ylim((-0.05, 10.5))\n",
    "ax.set_xlabel(r\"Heads probability $\\theta$\", fontsize=legend_fontsize)\n",
    "ax.legend(loc=\"upper left\", fontsize=legend_fontsize)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Beta priors: the conjugate family"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = np.linspace(0, 1, 500)\n",
    "beta_params = [(0.5, 0.5), (1, 1), (2, 2), (2, 5), (8, 4), (20, 20)]\n",
    "fig = plt.figure(figsize=figsize); ax = fig.add_subplot(111)\n",
    "for a, b in beta_params:\n",
    "    ax.plot(xs, st.beta(a, b).pdf(xs), linewidth=3, label=r\"$\\mathrm{Beta}(%g, %g)$\" % (a, b))\n",
    "ax.set_xlim((0, 1)); ax.set_ylim((0, 5))\n",
    "ax.set_xlabel(r\"Coin bias $\\theta$\", fontsize=legend_fontsize)\n",
    "ax.legend(loc=\"upper left\", fontsize=legend_fontsize)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sequential Bayesian updating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_theta, seed = 0.6, 7\n",
    "flips = np.random.RandomState(seed).rand(50) < true_theta\n",
    "xs = np.linspace(0, 1, 500)\n",
    "fig = plt.figure(figsize=figsize); ax = fig.add_subplot(111)\n",
    "for n in [0, 1, 2, 5, 10, 50]:\n",
    "    h = int(flips[:n].sum()); t = n - h\n",
    "    ax.plot(xs, st.beta(1 + h, 1 + t).pdf(xs), linewidth=3, label=r\"$N=%d$ ($%d$H, $%d$T)\" % (n, h, t))\n",
    "ax.axvline(true_theta, color=\"black\", linestyle=\"dashed\", linewidth=1.5)\n",
    "ax.set_xlim((0, 1)); ax.set_xlabel(r\"Coin bias $\\theta$\", fontsize=legend_fontsize)\n",
    "ax.legend(loc=\"upper left\", fontsize=legend_fontsize)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Laplace's rule vs maximum likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flips = (np.random.RandomState(36).rand(300) < 0.5).astype(int)\n",
    "Ns = np.arange(1, len(flips) + 1); h = np.cumsum(flips)\n",
    "fig = plt.figure(figsize=figsize); ax = fig.add_subplot(111)\n",
    "ax.axhline(0.5, color=\"black\", linestyle=\"dashed\", linewidth=1.5, label=r\"True $\\theta$\")\n",
    "ax.plot(Ns, h / Ns, linewidth=2, label=r\"Maximum likelihood $\\frac{n}{n+m}$\")\n",
    "ax.plot(Ns, (h + 1) / (Ns + 2), linewidth=2, label=r\"Laplace's rule $\\frac{n+1}{n+m+2}$\")\n",
    "ax.set_xscale(\"log\"); ax.set_ylim((-0.05, 1.05))\n",
    "ax.set_xlabel(r\"Number of flips $N$\", fontsize=legend_fontsize)\n",
    "ax.set_ylabel(r\"Predicted $p(\\mathrm{heads})$\", fontsize=legend_fontsize)\n",
    "ax.legend(loc=\"upper right\", fontsize=legend_fontsize)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Dirichlet distribution: concentration by $\\alpha$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corners = np.array([[0, 0], [1, 0], [0.5, np.sqrt(3) / 2]])\n",
    "def bary_grid(n=160):\n",
    "    p1, p2 = [], []\n",
    "    for i in range(n + 1):\n",
    "        for j in range(n + 1 - i):\n",
    "            p1.append(i / n); p2.append(j / n)\n",
    "    p1, p2 = np.array(p1), np.array(p2)\n",
    "    return np.vstack([p1, p2, 1 - p1 - p2]).T\n",
    "def dir_pdf(P, alpha):\n",
    "    alpha = np.array(alpha, float)\n",
    "    logB = np.sum(gammaln(alpha)) - gammaln(alpha.sum())\n",
    "    return np.exp(np.sum((alpha - 1) * np.log(np.clip(P, 1e-12, 1.0)), axis=1) - logB)\n",
    "P = bary_grid(160); cart = P @ corners; tri = mtri.Triangulation(cart[:, 0], cart[:, 1])\n",
    "\n",
    "np.random.seed(0)\n",
    "fig, axes = plt.subplots(1, 4, figsize=(17, 4.6))\n",
    "for ax, a in zip(axes, [(8, 8, 8), (1, 1, 1), (0.5, 0.5, 0.5), (0.1, 0.1, 0.1)]):\n",
    "    xy = np.random.dirichlet(a, 4000) @ corners\n",
    "    ax.scatter(xy[:, 0], xy[:, 1], s=4, alpha=0.15, color=palette[0])\n",
    "    ax.plot(*np.vstack([corners, corners[0]]).T, color=\"0.2\", lw=1)\n",
    "    ax.set_title(r\"$\\alpha=(%g,%g,%g)$\" % a, fontsize=15); ax.axis(\"equal\"); ax.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dirichlet density as a surface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(15, 6))\n",
    "for k, a in enumerate([(6, 6, 6), (2, 6, 12)]):\n",
    "    z = dir_pdf(P, a)\n",
    "    ax = fig.add_subplot(1, 2, k + 1, projection=\"3d\")\n",
    "    ax.plot_trisurf(cart[:, 0], cart[:, 1], z, triangles=tri.triangles, cmap=\"viridis\",\n",
    "                    linewidth=0, antialiased=True)\n",
    "    ax.set_title(r\"$\\alpha=(%g,%g,%g)$\" % a, fontsize=15)\n",
    "    ax.set_zlabel(\"density\"); ax.set_xticks([]); ax.set_yticks([]); ax.view_init(elev=35, azim=-60)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}