{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1bba0128",
   "metadata": {},
   "source": [
    "# Bayesian fits\n",
    "\n",
    "This is a notebook showing a visualization of a Byesian interpretation for the line fit.\n",
    "The idea is that we generate all possible line fits and reject them if they are not compatible with the data.\n",
    "Naturally, this is not how most fits are done, but it illustrates the logic behind the Bayes Theorem when applied for fitting lines.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "23feddde",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb1286f0",
   "metadata": {},
   "source": [
    "We start by generating some fake dataset, which is simple enough that we can visualize the results easily. For this reason, the dataset will contain only two  variables.\n",
    "\n",
    "The simulated example data will be $f(x) = 2 x + 1 + \\epsilon$, where $\\epsilon \\sim \\mathcal{N}(\\mu=0, \\sigma=0.5)$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "aef54435-78af-4b13-aee2-2bc576c773cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "true_alpha = 0.3\n",
    "true_beta = 0.5\n",
    "true_eps = 0.5\n",
    "def generate_data(N: int) -> np.ndarray:\n",
    "    x = 2*np.random.randn(N, 1)\n",
    "    epsilon = true_eps*np.random.randn(N, 1)\n",
    "    z = true_beta*x + true_alpha + epsilon\n",
    "    return np.concatenate((x, z), axis=1).astype(np.float32)\n",
    "\n",
    "data = generate_data(N=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04432613-39cd-4209-8833-7af585e7eed9",
   "metadata": {},
   "source": [
    "Given some parameters $\\alpha$ and $\\beta$, we don't know the true $f(x)$, but we assume it is something like $f(x) = \\beta x + \\alpha$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a795239a-03eb-483a-8b07-635e3d9dbc55",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def f(x: np.ndarray, alpha: float, beta: float) -> np.ndarray:\n",
    "    return beta*x + alpha"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2fba887-dc1a-4400-ab40-59fec452dc83",
   "metadata": {},
   "source": [
    "Let's plot it for reference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "fead461d-4d3a-48f2-a525-c518be4f1cf8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 40256 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(8, 6))\n",
    "plt.scatter(data[:,0], data[:,1], label=r\"Data $\\mathcal{D}$\")\n",
    "plt.plot(data[:,0], f(data[:,0], true_alpha, true_beta), c='r', ls='-', alpha=0.5, lw=4, label=r\"True $y=\\beta x + \\alpha$\")\n",
    "xs = np.linspace(-2, 2, 10)\n",
    "plt.plot(xs, f(xs, true_alpha, true_beta) - 2*true_eps, c='r', ls='--', alpha=0.5, lw=4)\n",
    "plt.plot(xs, f(xs, true_alpha, true_beta) + 2*true_eps, c='r', ls='--', alpha=0.5, lw=4)\n",
    "xa = 2\n",
    "ya = f(xa, true_alpha, true_beta)\n",
    "plt.annotate(xy=(xa, ya + 2*true_eps), xytext=(xa, ya - 2*true_eps), arrowprops=dict(arrowstyle='|-|', lw=3, alpha=0.5, color=\"green\"), text=r\"\")\n",
    "plt.text(xa+0.5, ya, s=r\"$\\epsilon$ noise\", size=16)\n",
    "plt.legend(frameon=False)\n",
    "plt.savefig(\"data.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e098e8-0082-4451-927a-073cf242787e",
   "metadata": {},
   "source": [
    "Let's show several random lines and reject the ones that do not agree with the likelihood:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "08d9daaf-25c5-4d06-b4a6-6cecaaa35651",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pylab as pylab\n",
    "params = {'legend.fontsize': 'x-large',\n",
    "          'figure.figsize': (15, 5),\n",
    "         'axes.labelsize': 'x-large',\n",
    "         'axes.titlesize':'x-large',\n",
    "         'xtick.labelsize':'x-large',\n",
    "         'ytick.labelsize':'x-large'}\n",
    "pylab.rcParams.update(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "4c7002b2-5cb5-4929-8e01-964144ee60b6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_log_prior(A, B):\n",
    "    #return 0.25*A/A/1/5\n",
    "    return -0.5*A**2/1**2 - 0.5*B**2/1**2 #- (2*np.pi*p_width**2)\n",
    "def sample_prior() -> Tuple[float, float]:\n",
    "    #return (2*np.random.rand() - 1)*1, (2*np.random.rand() - 1)*5\n",
    "    return 1*np.random.randn(), 1*np.random.randn()\n",
    "def get_log_likelihood(x, y, A, B, sigma):\n",
    "    error = f(x, A, B) - y\n",
    "    return np.sum(-0.5*error**2/sigma**2, axis=0) # - (2*np.pi*sigma**2)\n",
    "def plot(I: int):\n",
    "    # how many lines to show in last column\n",
    "    K = 20\n",
    "    # assumed noise std. dev.\n",
    "    sigma=0.5\n",
    "    # region of alpha and beta to plot\n",
    "    R = 4\n",
    "    # take a subset of data points\n",
    "    D = data[:I,:]\n",
    "    # make plot\n",
    "    fig, ax = plt.subplots(figsize=(4*5, 4), ncols=4)\n",
    "    # create alpha and beta grid\n",
    "    A = np.linspace(-R, R, 50)\n",
    "    B = np.linspace(-R, R, 50)\n",
    "    A, B = np.meshgrid(A, B)\n",
    "    \n",
    "    # calculate the log prior\n",
    "    log_prior = get_log_prior(A, B)\n",
    "    # calculate the log likelihood\n",
    "    log_likelihood = get_log_likelihood(D[:, 0, np.newaxis, np.newaxis],\n",
    "                                        D[:, 1, np.newaxis, np.newaxis],\n",
    "                                        A[np.newaxis, ...],\n",
    "                                        B[np.newaxis, ...],\n",
    "                                        sigma=sigma)\n",
    "    # the log posterior is, except for a constant factor, just the sum of those\n",
    "    log_posterior = log_prior + log_likelihood\n",
    "\n",
    "    m=ax[0].imshow(np.exp(log_prior), extent=(-R, R, -R, R), aspect=\"equal\")\n",
    "    #plt.colorbar(m,ax=ax[0], orientation=\"horizontal\")\n",
    "    ax[0].set(xlabel=r\"$\\alpha$\", ylabel=r\"$\\beta$\", title=r\"$p(\\alpha, \\beta)$\")\n",
    "    \n",
    "    m=ax[1].imshow(np.exp(log_likelihood), extent=(-R, R, -R, R), aspect=\"equal\")\n",
    "    #plt.colorbar(m,ax=ax[1], orientation=\"horizontal\")\n",
    "    \n",
    "    ax[1].set(xlabel=r\"$\\alpha$\", ylabel=r\"$\\beta$\", title=r\"$p(\\mathcal{D}|\\alpha,\\beta)$\")\n",
    "    m=ax[2].imshow(np.exp(log_posterior), extent=(-R, R, -R, R), aspect=\"equal\")\n",
    "    #plt.colorbar(m,ax=ax[2], orientation=\"horizontal\")\n",
    "    ax[2].set(xlabel=r\"$\\alpha$\", ylabel=r\"$\\beta$\", title=r\"$p(\\alpha, \\beta|\\mathcal{D})$\")\n",
    "\n",
    "    if I > 0:\n",
    "        ax[3].scatter(D[:, 0], D[:,1], label=\"Data\")\n",
    "    ax[3].set(xlabel=\"x\", ylabel=\"y\", title=\"Data\", xlim=(-5, 5), ylim=(-5, 5))\n",
    "\n",
    "    # sample alpha and beta\n",
    "    def get_line():\n",
    "        a, b = sample_prior()\n",
    "        if len(D) == 0:\n",
    "            return a, b\n",
    "        for _ in range(1000):\n",
    "            ap, bp = a + 0.1*np.random.randn(), b + 0.1*np.random.randn()\n",
    "            A = min(1, np.exp(get_log_likelihood(D[:, 0], D[:, 1], ap, bp, sigma=sigma) - get_log_likelihood(D[:, 0], D[:, 1], a, b, sigma=sigma) ))\n",
    "            if np.random.rand() <= A:\n",
    "                a = ap\n",
    "                b = bp\n",
    "        return a, b\n",
    "    for _ in range(K):\n",
    "        a, b = get_line()\n",
    "        ax[3].axline((0, a), slope=b, lw=2, ls='-', c='r', alpha=0.5)\n",
    "    #plt.legend(frameon=False)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(f\"bayes_reg_{I}.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "693827d3-c5ac-4b6d-b019-233799cd5d78",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 78028 characters]",
      "text/plain": [
       "<Figure size 2000x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot(I=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1fbd00b-e7f2-4eff-90db-e1d44d104fba",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
