think_bayes_notes.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "id": "c24e3db6",
   "metadata": {},
   "source": [
    "---\n",
    "title: \"Think Bayes - Notes\"\n",
    "format:\n",
    "  html:\n",
    "    code-fold: true\n",
    "    code-tools: true\n",
    "    number-sections: true\n",
    "    toc: true\n",
    "    toc-location: left\n",
    "    toc-depth: 5\n",
    "    toc-expand: 2\n",
    "    callout-icon: false\n",
    "    highlight-style: tango\n",
    "    code-line-numbers: true\n",
    "    embed-resources: true\n",
    "    grid:\n",
    "        body-width: 1100px\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f7b0207",
   "metadata": {},
   "source": [
    "Notes related to the book [Think Bayes v2](https://allendowney.github.io/ThinkBayes2/) which introduces Bayesian statistics using computational methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9f4ceb7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:07.879698Z",
     "iopub.status.busy": "2023-08-25T06:16:07.878878Z",
     "iopub.status.idle": "2023-08-25T06:16:09.109533Z",
     "shell.execute_reply": "2023-08-25T06:16:09.109087Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "# import dependencies\n",
    "\n",
    "from itertools import product\n",
    "from typing import Optional\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scipy\n",
    "from empiricaldist import Pmf, Cdf\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import pymc as pm\n",
    "import bambi as bmb\n",
    "import arviz as az"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91526dc8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.112789Z",
     "iopub.status.busy": "2023-08-25T06:16:09.112374Z",
     "iopub.status.idle": "2023-08-25T06:16:09.115140Z",
     "shell.execute_reply": "2023-08-25T06:16:09.114752Z"
    }
   },
   "outputs": [],
   "source": [
    "sns.set_style(\"whitegrid\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea8454a4",
   "metadata": {},
   "source": [
    "## Bayes's theorem\n",
    "\n",
    "::: {.callout-tip title=\"Theory\"}\n",
    "\n",
    "Given\n",
    "\n",
    "$$ P(A|B) = \\frac{P(A \\cap B)}{P(B)} $$\n",
    "\n",
    "and \n",
    "\n",
    "$$ P(A \\cap B) = P(B \\cap A) $$\n",
    "\n",
    "We can prove the Bayes's Theorem:\n",
    "\n",
    "$$ P(A|B) = \\frac{P(A) * P(B|A)}{P(B)} $$\n",
    "\n",
    "Which can be written as \n",
    "\n",
    "$$ P(H|Obs) = \\frac{P(H) * P(Obs|H)}{P(Obs)} $$\n",
    "\n",
    "\n",
    "Where:\n",
    "\n",
    "\\begin{align}\n",
    "    P(H|Obs)) & : \\text{Posterior} \\\\\n",
    "    P(H) & : \\text{Prior} \\\\\n",
    "    P(Obs|H)   & : \\text{Likelihood } A \\\\\n",
    "    P(Obs)   & : \\text{Total probability of observation}\\\\\n",
    "\\end{align}\n",
    "\n",
    "___\n",
    "\n",
    "Usually, to compute the posterior with the grid method, we specify the prior in order it includes all possible values of Obs with null intersection:\n",
    "\n",
    "$$ P(Obs) = \\sum_{i=1}^n P(H) * P(Obs_i | H) $$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e95e3a75",
   "metadata": {},
   "source": [
    "## Pmf / Cdf\n",
    "\n",
    "The followings are class implemented in the empiricaldist library which extands pd.Series class.\n",
    "\n",
    "- **Pmf:** Probablity mass function  \n",
    "- **Cdf:** Cumulative density function\n",
    "\n",
    "### Creation and usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f21960dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.118098Z",
     "iopub.status.busy": "2023-08-25T06:16:09.117785Z",
     "iopub.status.idle": "2023-08-25T06:16:09.124486Z",
     "shell.execute_reply": "2023-08-25T06:16:09.124174Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "\n",
    "# ----------\n",
    "# Create Pmf\n",
    "\n",
    "# from sequence\n",
    "pmf = Pmf.from_seq([1, 0, 0, 1, 1, 1])\n",
    "\n",
    "# specify range\n",
    "range_qs = np.linspace(0, 10, 101)\n",
    "ps = scipy.stats.norm(5, 1).pdf(range_qs)  # normal distribution\n",
    "pmf = Pmf(ps, range_qs)\n",
    "pmf.normalize()\n",
    "\n",
    "# ------------------\n",
    "# Specific functions\n",
    "\n",
    "pmf.mean()\n",
    "pmf.credible_interval(0.9)\n",
    "pmf.make_cdf()\n",
    "pmf.idxmax()\n",
    "pmf.max_prob()\n",
    "\n",
    "\n",
    "# ----------\n",
    "# Create Cdf\n",
    "\n",
    "# from pmf\n",
    "cdf = pmf.make_cdf()\n",
    "pmf = cdf.make_pmf()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa5aa4c0",
   "metadata": {},
   "source": [
    "### Mixture\n",
    "\n",
    "**TODO**: Add corresponding matrix operations  \n",
    "\n",
    "We can compute a distribution as mixture of other distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "537bc3d5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.126084Z",
     "iopub.status.busy": "2023-08-25T06:16:09.125971Z",
     "iopub.status.idle": "2023-08-25T06:16:09.128526Z",
     "shell.execute_reply": "2023-08-25T06:16:09.128226Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "\n",
    "def make_mixture(pmf: Pmf, pmf_seq: list[pmf]) -> Pmf:\n",
    "    \"\"\"\n",
    "    Make a mixture of distributions\n",
    "    \n",
    "    pmf: map each hypothesis to a probability\n",
    "    pmf_seq: sequence pmf, for each hypothesis\n",
    "    \"\"\"\n",
    "\n",
    "    df = pd.DataFrame(pmf_seq).fillna(0).transpose()\n",
    "    df *= np.array(pmf)\n",
    "    total = df.sum(axis=1)\n",
    "\n",
    "    return Pmf(total)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46061a8c",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}\n",
    "\n",
    "We choose a dice at random on a bag and roll it 1 time, we don't know the nb of sides, but we know the corresponding probability (e.g. the number of dices on the bag):\n",
    "\\begin{align}\n",
    "    P(nb\\_sides=4) = 1 / 6  \\\\\n",
    "    P(nb\\_sides=6) = 2 / 6 \\\\\n",
    "    P(nb\\_sides=8) = 3 / 6 \\\\\n",
    "\\end{align}\n",
    "\n",
    "Given this uncertainty, we can compute the Pmf of the outcome\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b203978",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.130196Z",
     "iopub.status.busy": "2023-08-25T06:16:09.130001Z",
     "iopub.status.idle": "2023-08-25T06:16:09.138519Z",
     "shell.execute_reply": "2023-08-25T06:16:09.138147Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | tbl-cap: Pmf of distributions mixture\n",
    "\n",
    "pmf_dice = Pmf([1, 2, 3], [4, 6, 8])\n",
    "pmf_dice.normalize()\n",
    "\n",
    "dice_4 = Pmf(1 / 4, range(1, 4 + 1))\n",
    "dice_6 = Pmf(1 / 6, range(1, 6 + 1))\n",
    "dice_8 = Pmf(1 / 8, range(1, 8 + 1))\n",
    "\n",
    "make_mixture(pmf=pmf_dice, pmf_seq=[dice_4,dice_6,dice_8])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cc01cba",
   "metadata": {},
   "source": [
    "### Utils Pmf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adcb1ebd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.140036Z",
     "iopub.status.busy": "2023-08-25T06:16:09.139927Z",
     "iopub.status.idle": "2023-08-25T06:16:09.144370Z",
     "shell.execute_reply": "2023-08-25T06:16:09.144062Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "\n",
    "def pmf_from_dist(dist, qs: np.array) -> Pmf:\n",
    "    \"\"\"\n",
    "    Make a discrete approximation of a scipy distribution.\n",
    "    dist: SciPy distribution object\n",
    "    qs: quantities\n",
    "    \"\"\"\n",
    "\n",
    "    ps = dist.pdf(qs)\n",
    "    pmf = Pmf(ps, qs)\n",
    "    pmf.normalize()\n",
    "\n",
    "    return pmf\n",
    "\n",
    "\n",
    "def pmf_kde_from_sample(sample: np.array, qs: np.array, **options) -> Pmf:\n",
    "    \"\"\"\n",
    "    Make a kernel density estimate from a sample\n",
    "\n",
    "    sample: sequence of values\n",
    "    qs: quantities where we should evaluate the KDE\n",
    "    \"\"\"\n",
    "\n",
    "    kde = scipy.stats.gaussian_kde(sample)\n",
    "    ps = kde(qs)\n",
    "    pmf = Pmf(ps, qs, **options)\n",
    "    pmf.normalize()\n",
    "\n",
    "    return pmf\n",
    "\n",
    "\n",
    "def kde_from_pmf(pmf: Pmf, n: int = 101, qs: Optional[np.array] = None, **options) -> Pmf:\n",
    "    \"\"\"\n",
    "    Make a kernel density estimate from a Pmf.\n",
    "\n",
    "    pmf: Pmf object\n",
    "    n: number of points\n",
    "    \"\"\"\n",
    "\n",
    "    if not qs:\n",
    "        qs = np.linspace(pmf.qs.min(), pmf.qs.max(), n)\n",
    "\n",
    "    kde = scipy.stats.gaussian_kde(pmf.qs, weights=pmf.ps)\n",
    "    ps = kde.evaluate(qs)\n",
    "    pmf = Pmf(ps, qs, **options)\n",
    "    pmf.normalize()\n",
    "\n",
    "    return pmf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34a24a50",
   "metadata": {},
   "source": [
    "## Grid method\n",
    "\n",
    "Given a prior and likelihood, we can compute the posterior estimation for the entire range of hypothesis\n",
    "\n",
    "### 1-parameter grid\n",
    "\n",
    "**Steps:**\n",
    "\n",
    "(1) Define prior ($P(H)$)\n",
    "\n",
    "\n",
    "(2) Define likelihood ($P(Obs|H)$)\n",
    "\n",
    "\n",
    "(3) Multiply prior and likelihood ($P(Obs|H) * P(Obs|H)$)\n",
    "\n",
    "\n",
    "(4) Normalize (Divide by $P(Obs) =  \\sum_{i=1}^n P(H) * P(Obs_i | H)$)\n",
    "\n",
    "\n",
    "(5) Assert the range of your parameters covers all plausible values\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4eca753",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "Throwing a coin 10 times, we observe 3 tails. What is the posterior rate of tails ?  \n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11cca960",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.169914Z",
     "iopub.status.busy": "2023-08-25T06:16:09.169452Z",
     "iopub.status.idle": "2023-08-25T06:16:09.175653Z",
     "shell.execute_reply": "2023-08-25T06:16:09.175283Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "# Define prior\n",
    "range_qs = np.linspace(0, 1, 101)\n",
    "pmf_prior = Pmf(1, range_qs)\n",
    "pmf_prior.normalize()\n",
    "\n",
    "# Define likelihood\n",
    "likelihood = scipy.stats.binom.pmf(3, 10, pmf_prior.qs)\n",
    "\n",
    "# Multiply prior and likelihood\n",
    "pmf_posterior = pmf_prior * likelihood\n",
    "\n",
    "# Normalize\n",
    "pmf_posterior.normalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a9d1886",
   "metadata": {
    "code_folding": [
     0
    ],
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.177264Z",
     "iopub.status.busy": "2023-08-25T06:16:09.177081Z",
     "iopub.status.idle": "2023-08-25T06:16:09.416879Z",
     "shell.execute_reply": "2023-08-25T06:16:09.416503Z"
    }
   },
   "outputs": [],
   "source": [
    "print(f\"\"\"Parameter value:\n",
    "    - Mean : {pmf_posterior.mean():.3f}\n",
    "    - 90% Credible interval: {pmf_posterior.credible_interval(0.9)}\n",
    "    - Most probable value: {pmf_posterior.idxmax()}\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "pmf_posterior.plot(\n",
    "    **{\n",
    "        \"title\": \"Posterior probability of parameter given observations\",\n",
    "        \"xlabel\": \"Parameter possible values\",\n",
    "        \"ylabel\": \"Probability\",\n",
    "    }\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b816aa2",
   "metadata": {},
   "source": [
    "### 2-parameters grid\n",
    "\n",
    "We can use the 1-parameter grid approach with 2 parameters\n",
    "\n",
    "**Functions:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "079edd03",
   "metadata": {
    "code_folding": [
     47
    ],
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.419242Z",
     "iopub.status.busy": "2023-08-25T06:16:09.418883Z",
     "iopub.status.idle": "2023-08-25T06:16:09.424091Z",
     "shell.execute_reply": "2023-08-25T06:16:09.423744Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "\n",
    "def make_joint(pmf1: Pmf, pmf2: Pmf) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Compute the outer product of two Pmfs, usually priors of two parameters.\n",
    "    Output dataframe: (columns = pmf1.qs / index = pmf2.qs)\n",
    "    \"\"\"\n",
    "\n",
    "    X, Y = np.meshgrid(pmf1, pmf2)\n",
    "    return pd.DataFrame(X * Y, columns=pmf1.qs, index=pmf2.qs)\n",
    "\n",
    "\n",
    "def plot_joint(\n",
    "    joint: pd.DataFrame,\n",
    "    title: str = \"Plot joint\",\n",
    "    xlabel: str = \"x\",\n",
    "    ylabel: str = \"y\",\n",
    "    cmap: str = \"Blues\",\n",
    "    add_contour: bool = True,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Plot a joint distribution with a color mesh.\n",
    "    x-axis: joint DataFrame columns\n",
    "    y-axis: joint DataFrame index\n",
    "    \"\"\"\n",
    "\n",
    "    vmax = joint.to_numpy().max() * 1.1\n",
    "\n",
    "    plt.pcolormesh(\n",
    "        joint.columns,\n",
    "        joint.index,\n",
    "        joint,\n",
    "        cmap=cmap,\n",
    "        vmax=vmax,\n",
    "        shading=\"nearest\",\n",
    "    )\n",
    "\n",
    "    plt.colorbar()\n",
    "\n",
    "    if add_contour:\n",
    "        plt.contour(joint.columns, joint.index, joint, linewidths=1.5, cmap=\"gray_r\")\n",
    "    \n",
    "    plt.gca().set(title=title, xlabel=xlabel, ylabel=ylabel)\n",
    "\n",
    "\n",
    "def joint_to_marginal(joint: pd.DataFrame) -> tuple[Pmf, Pmf]:\n",
    "    \"\"\"\n",
    "    Compute the marginal distribution for a joint distribution\n",
    "    Outuput tuple:\n",
    "        - First distribution: marginal of columns\n",
    "        - Second distribution: marginal of index\n",
    "    \"\"\"\n",
    "\n",
    "    marginal_columns = joint.sum(axis=0)\n",
    "    marginal_index = joint.sum(axis=1)\n",
    "\n",
    "    pmf_marginal_columns = Pmf(marginal_columns)\n",
    "    pmf_marginal_index = Pmf(marginal_index)\n",
    "\n",
    "    return pmf_marginal_columns, pmf_marginal_index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "948c2dbb",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "We draw at random a man and a woman in the population.  \n",
    "The woman is taller than the man, but we don't know how much.  \n",
    "&rarr; Compute the posterior size of each one.\n",
    "\n",
    "- Man prior size: Normal distribution with $(\\mu, \\sigma²)$ = (178, 7.7)\n",
    "- Woman prior size: Normal distribution with $(\\mu, \\sigma²)$ = (163, 7.3)\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd664122",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.425748Z",
     "iopub.status.busy": "2023-08-25T06:16:09.425560Z",
     "iopub.status.idle": "2023-08-25T06:16:09.432638Z",
     "shell.execute_reply": "2023-08-25T06:16:09.432321Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "\n",
    "# Prior\n",
    "\n",
    "range_size = np.arange(140, 205, 0.5)\n",
    "pmf_prior_size_man = Pmf(scipy.stats.norm(178, 7.7).pdf(range_size), range_size + 10)\n",
    "pmf_prior_size_man.normalize()\n",
    "\n",
    "pmf_prior_size_woman = Pmf(scipy.stats.norm(163, 7.3).pdf(range_size), range_size)\n",
    "pmf_prior_size_woman.normalize()\n",
    "\n",
    "joint_prior_size = make_joint(pmf_prior_size_man, pmf_prior_size_woman)\n",
    "\n",
    "\n",
    "# Likelihood\n",
    "\n",
    "# Observed: woman > man\n",
    "col_vals, index_vals = np.meshgrid(joint_prior_size.columns, joint_prior_size.index)\n",
    "likelihood = col_vals < index_vals\n",
    "\n",
    "\n",
    "# Bayes theorem\n",
    "\n",
    "joint_posterior_size = joint_prior_size * likelihood\n",
    "joint_posterior_size /= joint_posterior_size.sum().sum()\n",
    "\n",
    "\n",
    "# Marginal\n",
    "pmf_posterior_size_man, pmf_posterior_size_woman = joint_to_marginal(joint_posterior_size) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a0a2554",
   "metadata": {
    "code_folding": [],
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.434509Z",
     "iopub.status.busy": "2023-08-25T06:16:09.434397Z",
     "iopub.status.idle": "2023-08-25T06:16:09.814539Z",
     "shell.execute_reply": "2023-08-25T06:16:09.814135Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot results\n",
    "\n",
    "plt.subplot(311)\n",
    "plot_joint(joint_prior_size, title=\"Prior size distribution\", xlabel=\"Man\", ylabel=\"Woman\")\n",
    "plt.show()\n",
    "\n",
    "plt.subplot(312)\n",
    "plot_joint(\n",
    "    pd.DataFrame(likelihood, columns=joint_prior_size.columns, index=joint_prior_size.index),\n",
    "    title=\"likelihood\",\n",
    "    xlabel=\"Man\",\n",
    "    ylabel=\"Woman\",\n",
    ")\n",
    "plt.show()\n",
    "\n",
    "plt.subplot(313)\n",
    "pmf_prior_size_man.plot(label=\"Prior size man (cm)\", linestyle=\"--\", color=\"orangered\")\n",
    "pmf_prior_size_woman.plot(label=\"Prior size woman (cm)\", linestyle=\"--\", color=\"navy\")\n",
    "pmf_posterior_size_man.plot(label=\"Posterior size man (cm)\", color=\"orangered\")\n",
    "pmf_posterior_size_woman.plot(label=\"Posterior size woman (cm)\", color=\"navy\")\n",
    "plt.gca().set(title=\"Posterior marginal distributions\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c140955",
   "metadata": {},
   "source": [
    "## Probability laws\n",
    "\n",
    "The following section introduces several probability laws and their usage in bayesin statistics\n",
    "\n",
    "### Binomial\n",
    "\n",
    "::: {.callout-tip title=\"Theory\"}\n",
    "\n",
    "\n",
    "Given $n$ independant Bernouilli trials with a given conversion rate $p$, the number of success is described by the binomial law\n",
    "\n",
    "$$ X \\sim B(n,p) \\implies P(X=k) =  \\binom{N}{k}p^{k}q^{n-k} $$\n",
    "\n",
    ":::\n",
    "\n",
    "#### Estimating proportion p\n",
    "\n",
    "##### Grid Method\n",
    "\n",
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "We flip a coin 10 times, we get 3 tails, what is the probability of a tail ?\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc1153e4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.817002Z",
     "iopub.status.busy": "2023-08-25T06:16:09.816669Z",
     "iopub.status.idle": "2023-08-25T06:16:09.821868Z",
     "shell.execute_reply": "2023-08-25T06:16:09.821535Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "\n",
    "nb_conversion, nb_trials = 3, 10\n",
    "\n",
    "# Prior\n",
    "range_qs = np.linspace(0, 1, 101)\n",
    "pmf_prior = Pmf(1, range_qs)\n",
    "pmf_prior.normalize()\n",
    "\n",
    "# Likelihood\n",
    "likelihood = scipy.stats.binom.pmf(nb_conversion, nb_trials, pmf_prior.qs)  # binom.pmf(k, n, p)\n",
    "\n",
    "# Bayes formula\n",
    "pmf_posterior = pmf_prior * likelihood\n",
    "pmf_posterior.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b03063df",
   "metadata": {},
   "source": [
    "##### Conjugate priors\n",
    "\n",
    "::: {.callout-tip title=\"Theory\"}\n",
    "\n",
    "\n",
    "The Beta distribution is the conjugate distribution of the Binomial distribution.\n",
    "\n",
    "Given:\n",
    "\n",
    "$$ prior \\sim \\beta(a, b) $$\n",
    "$$ nb\\_conversions: k $$\n",
    "$$ nb\\_trials: n $$\n",
    "\n",
    "\n",
    "The posterior can be computed as:\n",
    "\n",
    "$$ posterior \\sim \\beta(a + k, b + n - k) $$\n",
    "\n",
    "To get a prior with uniform distribution between 0 and 1, we can use _a = b = 1_\n",
    "\n",
    "$$ U(0, 1) = \\beta(1, 1) $$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8962e71",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.823488Z",
     "iopub.status.busy": "2023-08-25T06:16:09.823371Z",
     "iopub.status.idle": "2023-08-25T06:16:09.828362Z",
     "shell.execute_reply": "2023-08-25T06:16:09.828041Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "\n",
    "nb_conversion, nb_trials = 3, 10\n",
    "range_qs = np.linspace(0, 1, 101)\n",
    "\n",
    "pmf_posterior = Pmf(scipy.stats.beta(1 + nb_conversion, 1 + nb_trials - nb_conversion).pdf(range_qs), range_qs)\n",
    "\n",
    "pmf_posterior.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a6d284b",
   "metadata": {},
   "source": [
    "##### MCMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f521d05",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:09.830270Z",
     "iopub.status.busy": "2023-08-25T06:16:09.830161Z",
     "iopub.status.idle": "2023-08-25T06:16:12.882194Z",
     "shell.execute_reply": "2023-08-25T06:16:12.881729Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "\n",
    "nb_conversion, nb_trials = 3, 10\n",
    "\n",
    "\n",
    "with pm.Model() as model:\n",
    "    conversion_rates = pm.Uniform(\"conversion_rates\", 0, 1)\n",
    "    nb_conversion = pm.Binomial(\"nb_conversion\", n=nb_trials, observed=nb_conversion, p=conversion_rates)\n",
    "    trace_binom = pm.sample(1000, return_inferencedata=False)\n",
    "\n",
    "posterior_samples = trace_binom.get_values(\"conversion_rates\")\n",
    "\n",
    "pd.Series(posterior_samples).quantile([0.05, .95]).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c5d0180",
   "metadata": {},
   "source": [
    "#### Estimating count\n",
    "\n",
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "On an audience of 0 to 2000 people:\n",
    "\n",
    "* 2 are born the 11th of may\n",
    "* 1 is born the 23th of may\n",
    "* 0 is born the 1st of August\n",
    "\n",
    "&rarr; How many people is there ?\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51a581a6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:12.885391Z",
     "iopub.status.busy": "2023-08-25T06:16:12.885046Z",
     "iopub.status.idle": "2023-08-25T06:16:12.890400Z",
     "shell.execute_reply": "2023-08-25T06:16:12.890107Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "# Prior\n",
    "range_qs = np.arange(0, 2000)\n",
    "pmf_prior = Pmf(1, range_qs)\n",
    "pmf_prior.normalize()\n",
    "\n",
    "# Likelihood\n",
    "likelihood = np.ones(len(pmf_prior))\n",
    "\n",
    "\n",
    "for nb_conversion in [2, 1, 0]:\n",
    "    likelihood *= scipy.stats.binom.pmf(nb_conversion, pmf_prior.qs, 1 / 365)\n",
    "\n",
    "# Bayes formula\n",
    "pmf_posterior = pmf_prior * likelihood\n",
    "pmf_posterior.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4d107e3",
   "metadata": {},
   "source": [
    "### Multinomial\n",
    "\n",
    "Modelize n independant trials with k possible outcomes at each trials, with a given probability for each outcome"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dc8e704",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "Can be used in mark and recapture problem:\n",
    "\n",
    "* A reviewer detects 20 errors errors on a coding's review\n",
    "* A second reviewer detects 15 errors, among which 3 errors were detected by the first reviewer\n",
    "* Given that each reviewer has a probability p1, p2 of detecting an error, every error fall into one of those categories:\n",
    "\n",
    "\n",
    "$$ P(\\bar{R_1} \\cap \\bar{R_2}) = (1 - p_1) * (1 - p_2) $$\n",
    "$$ P(\\bar{R_1} \\cap R_2) = (1 - p_1) * p_2 $$\n",
    "$$ P(R_1 \\cap \\bar{R_2}) = p_1 * (1 - p_2) $$\n",
    "$$ P(R_1 \\cap R_2) = p_1 * p_2 $$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77843db3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:12.892009Z",
     "iopub.status.busy": "2023-08-25T06:16:12.891902Z",
     "iopub.status.idle": "2023-08-25T06:16:23.655648Z",
     "shell.execute_reply": "2023-08-25T06:16:23.655273Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "\n",
    "# Priors\n",
    "range_p = np.linspace(0, 1, 51)\n",
    "range_nb_errors = np.arange(32, 350, 5)\n",
    "\n",
    "cols = [\"p1\", \"p2\", \"nb_e\"]\n",
    "\n",
    "prior = pd.DataFrame(list(product(range_p, range_p, range_nb_errors)), columns=cols)\n",
    "prior[\"p\"] = 1 / len(prior)\n",
    "\n",
    "prior = Pmf(prior.set_index(cols).p)\n",
    "\n",
    "# Likelihood\n",
    "\n",
    "k10 = 20 - 3  # nb seen by first reviewer only\n",
    "k01 = 15 - 3  # nb seen by second reviewer only\n",
    "k11 = 3  # nb seen by both reviewers\n",
    "\n",
    "\n",
    "likelihoods = prior.copy()\n",
    "\n",
    "for p1, p2, nb_e in prior.index:\n",
    "    k00 = nb_e - (k10 + k01 + k11)\n",
    "\n",
    "    obs = [k00, k10, k01, k11]\n",
    "    probs = [\n",
    "        (1 - p1) * (1 - p2),  # prob k00\n",
    "        p1 * (1 - p2),        # prob k10\n",
    "        (1 - p1) * p2,        # prob k01\n",
    "        p1 * p2,              # prob k11\n",
    "    ]\n",
    "\n",
    "    likelihood = scipy.stats.multinomial.pmf(obs, nb_e, probs)\n",
    "\n",
    "    likelihoods[p1, p2, nb_e] = likelihood\n",
    "\n",
    "# Bayes formula \n",
    "posterior = prior * likelihoods\n",
    "posterior.normalize()\n",
    "\n",
    "# Marginal \n",
    "marginal_nb_e = Pmf(posterior.rename(\"p\").reset_index().groupby(\"nb_e\").p.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbc2a51c",
   "metadata": {},
   "source": [
    "### Poisson \n",
    "\n",
    "::: {.callout-tip title=\"Theory\"}\n",
    "\n",
    "\n",
    "Given:\n",
    "\n",
    "- An event occur at frequency $\\lambda$ on a given period\n",
    "- The chances of occuring is uniform on a given period\n",
    "- The occurence of an event doesn't impact the occurence of other events\n",
    "\n",
    "The number of event which occurend on a given period follow a Poisson distribution:\n",
    "\n",
    "$$ P(X=k) = \\frac{\\lambda^k * e^{-\\lambda}}{k!} $$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3cb94c7",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "Suppose we can model the number of goals of a footbal team by a poisson process. Our prior follow a gamma law of parameter 1.4.  \n",
    "&rarr; Compute the posterior distribution of $\\lambda$ if a team score 4 goals\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2d1ad9",
   "metadata": {},
   "source": [
    "#### Grid method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d508ddbe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:23.657723Z",
     "iopub.status.busy": "2023-08-25T06:16:23.657602Z",
     "iopub.status.idle": "2023-08-25T06:16:23.662670Z",
     "shell.execute_reply": "2023-08-25T06:16:23.662394Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "nb_events_obs = 4\n",
    "\n",
    "# Prior\n",
    "range_qs = np.linspace(0, 10, 101)\n",
    "pmf_prior = Pmf(scipy.stats.gamma(1.4).pdf(range_qs), range_qs)\n",
    "pmf_prior.normalize()\n",
    "\n",
    "# Likelihood\n",
    "likelihood = scipy.stats.poisson(pmf_prior.qs).pmf(nb_events_obs)\n",
    "\n",
    "# Bayes formula\n",
    "pmf_posterior = pmf_prior * likelihood\n",
    "pmf_posterior.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72390f6b",
   "metadata": {},
   "source": [
    "#### Conjugate priors\n",
    "\n",
    "::: {.callout-tip title=\"Theory\"}\n",
    "\n",
    "\n",
    "The Gamma distribution is the conjugate distribution of the Poisson distribution.\n",
    "\n",
    "Given:\n",
    "\n",
    "$$ prior \\sim Gamma(\\alpha, \\beta) $$\n",
    "$$ nb\\_event: k  $$\n",
    "$$ Duration: t $$\n",
    "\n",
    "\n",
    "The posterior can be computed as:\n",
    "\n",
    "$$ posterior \\sim Gamma(\\alpha + k, \\beta + t) $$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de175917",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:23.664222Z",
     "iopub.status.busy": "2023-08-25T06:16:23.664111Z",
     "iopub.status.idle": "2023-08-25T06:16:23.668099Z",
     "shell.execute_reply": "2023-08-25T06:16:23.667789Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "alpha, beta = 1.4, 1\n",
    "obs_k, obs_t = 4, 1\n",
    "\n",
    "alpha_updated, beta_updated = alpha + obs_k, (beta + obs_t)\n",
    "\n",
    "gamma_dist = scipy.stats.gamma(alpha_updated, scale=1 / beta_updated) # scale factor equal to 1 / b\n",
    "                               \n",
    "range_qs = np.linspace(0, 10, 101)\n",
    "pmf_posterior_conjugate = Pmf(gamma_dist.pdf(range_qs), range_qs)\n",
    "pmf_posterior_conjugate.normalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "796c4b49",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:23.669505Z",
     "iopub.status.busy": "2023-08-25T06:16:23.669368Z",
     "iopub.status.idle": "2023-08-25T06:16:23.891338Z",
     "shell.execute_reply": "2023-08-25T06:16:23.890946Z"
    }
   },
   "outputs": [],
   "source": [
    "pmf_prior.plot(label=\"Prior\")\n",
    "pmf_posterior.plot(label=\"Posterior grid method\")\n",
    "pmf_posterior_conjugate.plot(label='Posterior conjugate method', linestyle=\"dotted\")\n",
    "\n",
    "plt.legend()\n",
    "plt.title = \"Distribution of lambda\"\n",
    "plt.xlabel = \"Lambda\"\n",
    "plt.ylabel = \"Probability\"\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea751708",
   "metadata": {},
   "source": [
    "#### MCMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9edd7544",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:23.893200Z",
     "iopub.status.busy": "2023-08-25T06:16:23.893069Z",
     "iopub.status.idle": "2023-08-25T06:16:26.501731Z",
     "shell.execute_reply": "2023-08-25T06:16:26.501316Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "obs_k = 4\n",
    "\n",
    "with pm.Model() as model:\n",
    "    lam = pm.Gamma(\"lambda\", alpha=1.4, beta=1)\n",
    "    nb_events = pm.Poisson(\"nb_event\", mu=lam, observed=obs_k)\n",
    "    trace = pm.sample(1000, return_inferencedata=False)\n",
    "\n",
    "posterior_samples = trace.get_values(\"lambda\")\n",
    "\n",
    "pd.Series(posterior_samples).quantile([0.05, .95]).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e386f9e4",
   "metadata": {},
   "source": [
    "### Exponential\n",
    "\n",
    "::: {.callout-tip title=\"Theory\"}\n",
    "\n",
    "\n",
    "Given the hypothesis of the Poisson distribution, the repartition of duration to the next event follows an exponential distribution:\n",
    "\n",
    "$$ P(X=t) = \\lambda * e^{-\\lambda  t} $$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b93bbc5e",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "Suppose we can model the number of goals of a footbal team by a poisson process. Our prior follow a gamma law of parameter 1.4.  \n",
    "&rarr; Compute the posterior distribution of $\\lambda$ if a team scores a goal at 11th and 23th minute\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7fbbc8b",
   "metadata": {},
   "source": [
    "#### Grid method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aacb1e5b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:26.503799Z",
     "iopub.status.busy": "2023-08-25T06:16:26.503654Z",
     "iopub.status.idle": "2023-08-25T06:16:26.509360Z",
     "shell.execute_reply": "2023-08-25T06:16:26.509073Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "elapsed_time_1 = 11 / 90\n",
    "elapsed_time_2 = (23 - 11) / 90\n",
    "\n",
    "# Prior: range of lambda\n",
    "range_qs = np.linspace(0, 10, 101)\n",
    "pmf_prior = Pmf(scipy.stats.gamma(1.4).pdf(range_qs), range_qs)\n",
    "pmf_prior.normalize()\n",
    "\n",
    "\n",
    "# Likelihood\n",
    "lams = range_qs\n",
    "likelihood_1 = lams * np.exp(-lams * elapsed_time_1)\n",
    "likelihood_2 = lams * np.exp(-lams * elapsed_time_2)\n",
    "\n",
    "# Bayes formula\n",
    "pmf_posterior = pmf_prior * likelihood_1 * likelihood_2\n",
    "pmf_posterior.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "975e8e76",
   "metadata": {},
   "source": [
    "### Normal\n",
    "\n",
    "The aim is usually to estimate the parameters $(\\mu, \\sigma²)$ of the normal distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3f4d524",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "We compare the results of a given value accross a control and test group. We assume the results follow a normal distribution.  \n",
    "&rarr; We want to know if the difference of mean between groups is significant\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b7d881",
   "metadata": {},
   "source": [
    "#### Grid method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eddce0b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:26.511508Z",
     "iopub.status.busy": "2023-08-25T06:16:26.511383Z",
     "iopub.status.idle": "2023-08-25T06:16:26.597378Z",
     "shell.execute_reply": "2023-08-25T06:16:26.596992Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "\n",
    "control_obs = [24, 43, 58, 71, 43, 49, 61, 44, 67, 49, 53, 56, 59, 52, 62, 54, 57, 33, 46, 43, 57]\n",
    "test_obs = [42, 43, 55, 26, 62, 37, 33, 41, 19, 54, 20, 85, 46, 10, 17, 60, 53, 42, 37, 42, 55, 28, 48]\n",
    "\n",
    "\n",
    "# Define priors\n",
    "\n",
    "range_mu = np.linspace(30, 60, 201)\n",
    "pmf_prior_mu = Pmf(1, range_mu)\n",
    "pmf_prior_mu.normalize()\n",
    "\n",
    "range_sigma = np.linspace(7, 30, 201)\n",
    "pmf_prior_sigma = Pmf(1, range_sigma)\n",
    "pmf_prior_sigma.normalize()\n",
    "\n",
    "prior_joint = make_joint(pmf_prior_mu, pmf_prior_sigma)\n",
    "\n",
    "\n",
    "# Likelihood\n",
    "\n",
    "# Control group\n",
    "mu_mesh, sigma_mesh, obs_mesh = np.meshgrid(pmf_prior_mu.qs, pmf_prior_sigma.qs, control_obs)\n",
    "likelihood_control = scipy.stats.norm(mu_mesh, sigma_mesh).pdf(obs_mesh)\n",
    "likelihood_control = likelihood_control.prod(axis=-1)\n",
    "\n",
    "# Test group\n",
    "mu_mesh, sigma_mesh, obs_mesh = np.meshgrid(pmf_prior_mu.qs, pmf_prior_sigma.qs, test_obs)\n",
    "likelihood_test = scipy.stats.norm(mu_mesh, sigma_mesh).pdf(obs_mesh)\n",
    "likelihood_test = likelihood_test.prod(axis=-1)\n",
    "\n",
    "\n",
    "# Bayes formula\n",
    "posterior_joint_control = prior_joint * likelihood_control\n",
    "posterior_joint_control /= posterior_joint_control.sum().sum()\n",
    "\n",
    "posterior_joint_test = prior_joint * likelihood_test\n",
    "posterior_joint_test /= posterior_joint_test.sum().sum()\n",
    "\n",
    "\n",
    "# Marginal\n",
    "pmf_posterior_mu_control, pmf_posterior_sigma_control = joint_to_marginal(posterior_joint_control)\n",
    "pmf_posterior_mu_test, pmf_posterior_sigma_test = joint_to_marginal(posterior_joint_test)\n",
    "\n",
    "# Diff of mu\n",
    "posterior_mu_diff = pmf_posterior_mu_control.sub_dist(pmf_posterior_mu_test)\n",
    "posterior_mu_diff = kde_from_pmf(posterior_mu_diff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7c2f518",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:26.600158Z",
     "iopub.status.busy": "2023-08-25T06:16:26.599928Z",
     "iopub.status.idle": "2023-08-25T06:16:27.072011Z",
     "shell.execute_reply": "2023-08-25T06:16:27.071578Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.subplot(221)\n",
    "plt.contour(posterior_joint_control)\n",
    "plt.contour(posterior_joint_test)\n",
    "plt.gca().set(title=\"Joint posterior density\")\n",
    "\n",
    "plt.subplot(222)\n",
    "pmf_posterior_mu_control.plot(label=\"Control mu\")\n",
    "pmf_posterior_mu_test.plot(label=\"Test mu\")\n",
    "plt.legend()\n",
    "plt.gca().set(title=\"Marginal posterior density of mu\")\n",
    "\n",
    "plt.subplot(223)\n",
    "pmf_posterior_sigma_control.plot(label=\"Control sigma\")\n",
    "pmf_posterior_sigma_test.plot(label=\"Test sigma\")\n",
    "plt.legend()\n",
    "plt.gca().set(title=\"Marginal posterior density of sigma\")\n",
    "\n",
    "plt.subplot(224)\n",
    "posterior_mu_diff.plot(label=\"Difference mu\")\n",
    "plt.legend()\n",
    "plt.gca().set(title=\"Posterior mu difference\")\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98d45e3c",
   "metadata": {},
   "source": [
    "#### MCMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20ad154d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:27.073687Z",
     "iopub.status.busy": "2023-08-25T06:16:27.073534Z",
     "iopub.status.idle": "2023-08-25T06:16:31.688178Z",
     "shell.execute_reply": "2023-08-25T06:16:31.687804Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "with pm.Model() as model:\n",
    "    mu = pm.Uniform(\"mu\", 30, 60)\n",
    "    sigma = pm.Uniform(\"sigma\", 7, 30)\n",
    "    result = pm.Normal(\"result\", mu=mu, sigma=sigma, observed=control_obs)\n",
    "    trace_binom = pm.sample(1000, return_inferencedata=False)\n",
    "\n",
    "posterior_samples = trace_binom.get_values(\"mu\")\n",
    "\n",
    "pd.Series(posterior_samples).quantile([0.05, .95]).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c52f7ae",
   "metadata": {},
   "source": [
    "### Weibull\n",
    "\n",
    "The Weibull distribution is used in survival analysis to predict time until an event (e.g. death)   \n",
    "Here we'll use the _weibull_min_ implementation of scipy. Given a parameter $\\lambda$ and _k_ its pdf describe the time until an event\n",
    "\n",
    "####  Parameters' influence\n",
    "\n",
    "- $\\lambda$: Mean of lifetime\n",
    "- $k$: Shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c26ae3a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:31.690009Z",
     "iopub.status.busy": "2023-08-25T06:16:31.689875Z",
     "iopub.status.idle": "2023-08-25T06:16:32.003199Z",
     "shell.execute_reply": "2023-08-25T06:16:32.002820Z"
    }
   },
   "outputs": [],
   "source": [
    "range_wb = np.linspace(0, 10, 101)\n",
    "\n",
    "fig, axs = plt.subplots(nrows=1, ncols=2)\n",
    "\n",
    "for lam in [1, 3, 5]:\n",
    "    for k in [3]:\n",
    "        sns.lineplot(scipy.stats.weibull_min(c=k, scale=lam).pdf(range_wb), label=f\"(lam;k) = ({lam};{k})\", ax=axs[0])\n",
    "\n",
    "for lam in [3]:\n",
    "    for k in [1, 3, 5]:\n",
    "        sns.lineplot(scipy.stats.weibull_min(c=k, scale=lam).pdf(range_wb), label=f\"(lam;k) = ({lam};{k})\", ax=axs[1])\n",
    "\n",
    "\n",
    "axs[0].title.set_text('Variation of lambda')\n",
    "axs[1].title.set_text('Variation of k')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "992f2aa3",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "We stop an experiment 8 years after the beginning  \n",
    "Some subjects arrived during the experiment, some died and some survived  \n",
    "Suppose we can model their lifetime by the Weibull distribution  \n",
    "&rarr; Compute the distribution of $k$ and $\\lambda$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c94be32d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:32.004922Z",
     "iopub.status.busy": "2023-08-25T06:16:32.004799Z",
     "iopub.status.idle": "2023-08-25T06:16:32.010084Z",
     "shell.execute_reply": "2023-08-25T06:16:32.009784Z"
    }
   },
   "outputs": [],
   "source": [
    "# Generate distribution and observations\n",
    "\n",
    "np.random.seed(1)\n",
    "\n",
    "# Distribution\n",
    "k, lam = 0.8, 3\n",
    "actual_dist = scipy.stats.weibull_min(c=k, scale=lam)\n",
    "\n",
    "\n",
    "# Observations\n",
    "obs_time = 8\n",
    "nb_obs = 20\n",
    "\n",
    "start = np.random.uniform(0, 8, size=nb_obs)\n",
    "duration = actual_dist.rvs(size=nb_obs, random_state=42)\n",
    "\n",
    "df_obs = pd.DataFrame({\"start\": start,\"duration\": duration,}).round(2)\n",
    "df_obs[\"end\"] = df_obs[\"start\"] + df_obs[\"duration\"]\n",
    "\n",
    "df_obs[\"censored\"] = df_obs[\"end\"] > obs_time\n",
    "df_obs[\"obs_t\"] = df_obs[\"end\"].clip(upper=8) - df_obs[\"start\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cd907aa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:32.011575Z",
     "iopub.status.busy": "2023-08-25T06:16:32.011461Z",
     "iopub.status.idle": "2023-08-25T06:16:32.143540Z",
     "shell.execute_reply": "2023-08-25T06:16:32.143063Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.hlines(\n",
    "    y=df_obs.index,\n",
    "    xmin=df_obs.start,\n",
    "    xmax=df_obs.end,\n",
    "    color=df_obs.censored.apply(lambda c: 'blue' if not c else \"red\"),\n",
    "    alpha=0.4)\n",
    "\n",
    "plt.scatter(df_obs.end, df_obs.index, color='blue')\n",
    "plt.axvline(x=obs_time)\n",
    "plt.gca().set(title=\"Subjects' lifetime\", xlabel=\"Lifetime\", ylabel=\"Subject n°\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "961f1554",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:32.145512Z",
     "iopub.status.busy": "2023-08-25T06:16:32.145239Z",
     "iopub.status.idle": "2023-08-25T06:16:32.151377Z",
     "shell.execute_reply": "2023-08-25T06:16:32.151000Z"
    }
   },
   "outputs": [],
   "source": [
    "#| tbl-cap: Example of subject's data\n",
    "\n",
    "df_obs.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fa0243d",
   "metadata": {},
   "source": [
    "#### Grid method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67a0a90f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:32.152916Z",
     "iopub.status.busy": "2023-08-25T06:16:32.152732Z",
     "iopub.status.idle": "2023-08-25T06:16:32.169598Z",
     "shell.execute_reply": "2023-08-25T06:16:32.169229Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "\n",
    "# Define priors\n",
    "range_k = np.linspace(0.1, 2, 101)\n",
    "pmf_prior_k = Pmf(1, range_k)\n",
    "pmf_prior_k.normalize()\n",
    "\n",
    "range_lam = np.linspace(0.1, 6, 101)\n",
    "pmf_prior_lam = Pmf(1, range_lam)\n",
    "pmf_prior_lam.normalize()\n",
    "\n",
    "joint_prior = make_joint(pmf_prior_k, pmf_prior_lam)\n",
    "\n",
    "# Likelihood\n",
    "\n",
    "k_mesh, lam_mesh, time_mesh = np.meshgrid(range_k, range_lam, df_obs[~df_obs.censored].obs_t)\n",
    "likelihood_uncensored = scipy.stats.weibull_min(c=k_mesh, scale=lam_mesh).pdf(time_mesh)\n",
    "likelihood_uncensored = likelihood_uncensored.prod(axis=-1)\n",
    "\n",
    "k_mesh, lam_mesh, time_mesh = np.meshgrid(range_k, range_lam, df_obs[df_obs.censored].obs_t)\n",
    "likelihood_censored = scipy.stats.weibull_min(c=k_mesh, scale=lam_mesh).sf(time_mesh)\n",
    "likelihood_censored = likelihood_censored.prod(axis=-1)\n",
    "\n",
    "\n",
    "# Bayes formula\n",
    "joint_posterior = joint_prior * likelihood_uncensored * likelihood_censored\n",
    "joint_posterior /= joint_posterior.sum().sum()\n",
    "\n",
    "\n",
    "# Marginal\n",
    "posterior_k, posterior_lam = joint_to_marginal(joint_posterior)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4aa6d66",
   "metadata": {},
   "source": [
    "**TODO:**\n",
    "\n",
    "* Add MCMC with censor and uncensored data\n",
    "* Add Hypergeometric distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad023de6",
   "metadata": {},
   "source": [
    "## Regression\n",
    "\n",
    "### Logistic regression\n",
    "\n",
    "::: {.callout-tip title=\"Theory\"}\n",
    "\n",
    "\n",
    "To demonstrate the logisic regression formula, we can use the following:\n",
    "\n",
    "$$ O(H|obs) = O(H) * \\frac{P(Obs|H)}{P(Obs|\\bar{H})} $$\n",
    "\n",
    "$$\n",
    "O(x) = \\frac{P(x)}{1 - P(x)}\n",
    "$$\n",
    "\n",
    "If observations are independant, then,\n",
    "\n",
    "$$\n",
    "O(H|obs^n) = O(H) * (\\frac{P(Obs|H)}{P(Obs|\\bar{H})})^n \n",
    "$$\n",
    "\n",
    "$$\n",
    "\\implies log(O(H|obs^n)) = log(O(H)) + n *log(\\frac{P(Obs|H)}{P(Obs|\\bar{H})})\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\implies O(H|obs^n) = e^{log(O(H)) + n *log(\\frac{P(Obs|H)}{P(Obs|\\bar{H})})}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\implies O(H|obs^n) = e^{a * x + b}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\implies P(H|obs^n) = \\frac{1}{1 + e^{-(a*x+b)}}\n",
    "$$\n",
    "\n",
    "With:  \n",
    "$$ a =log (\\frac{P(Obs|H)}{P(Obs|\\bar{H})}): Log \\: likelihood \\:ratio $$\n",
    "$$ b = log(O(H)): prior $$\n",
    "\n",
    ":::\n",
    "\n",
    "The aim is to define the value of $a$ and $b$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45011ab4",
   "metadata": {},
   "source": [
    "#### Grid method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b3bb97",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "We look at the risk of accident given a temperature.  \n",
    "We assume this can be modelled by a logistic regression.  \n",
    "In the following we have: $(a: slope)$ and $(b: intercept)$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98acb891",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:32.171686Z",
     "iopub.status.busy": "2023-08-25T06:16:32.171571Z",
     "iopub.status.idle": "2023-08-25T06:16:32.175190Z",
     "shell.execute_reply": "2023-08-25T06:16:32.174920Z"
    }
   },
   "outputs": [],
   "source": [
    "temp = [66, 70, 69, 68, 67, 72, 73, 70, 57, 63, 70, 78, 67, 53, 67, 75, 70, 81, 76, 79, 75, 76, 58]\n",
    "converted = [0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1]\n",
    "\n",
    "df = pd.DataFrame({\"temp\": temp, \"conv\": converted})\n",
    "\n",
    "# Normalize data\n",
    "offset = df[\"temp\"].mean()\n",
    "df[\"x\"] = df[\"temp\"] - offset\n",
    "df[\"y\"] = df[\"conv\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edd0f3e6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:32.176629Z",
     "iopub.status.busy": "2023-08-25T06:16:32.176439Z",
     "iopub.status.idle": "2023-08-25T06:16:34.135963Z",
     "shell.execute_reply": "2023-08-25T06:16:34.135584Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "\n",
    "# Priors\n",
    "range_a = np.linspace(-0.8, 0.1, 51)\n",
    "range_b = np.linspace(-5, 1, 51)\n",
    "\n",
    "cols = [\"a\", \"b\"]\n",
    "\n",
    "prior = pd.DataFrame(list(product(range_a, range_b)), columns=cols)\n",
    "prior[\"p\"] = 1 / len(prior)\n",
    "\n",
    "prior = Pmf(prior.set_index(cols).p)\n",
    "\n",
    "\n",
    "# Likelihood\n",
    "likelihood = prior.copy()\n",
    "\n",
    "for a, b in prior.index:\n",
    "    likelihoods_converted = scipy.special.expit(a * df[df.y == 1].x + b)\n",
    "    likelihoods_not_converted = 1 - scipy.special.expit((a * df[df.y == 0].x + b))\n",
    "\n",
    "    likelihood[a, b] = likelihoods_converted.prod() * likelihoods_not_converted.prod()\n",
    "\n",
    "\n",
    "# Bayes formula\n",
    "posterior = prior * likelihood\n",
    "posterior /= posterior.sum()\n",
    "\n",
    "# Posterior marginal\n",
    "posterior_a = Pmf(posterior.rename(\"p\").reset_index().groupby(\"a\").p.sum())\n",
    "posterior_b = Pmf(posterior.rename(\"p\").reset_index().groupby(\"b\").p.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fddfa6a",
   "metadata": {},
   "source": [
    "#### MCMC - PyMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5b3901f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:34.137989Z",
     "iopub.status.busy": "2023-08-25T06:16:34.137850Z",
     "iopub.status.idle": "2023-08-25T06:16:39.280993Z",
     "shell.execute_reply": "2023-08-25T06:16:39.280645Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "with pm.Model() as model:\n",
    "    a = pm.Uniform(\"a\", -0.8, 0.1)\n",
    "    b = pm.Uniform(\"b\", -5, 1)\n",
    "    \n",
    "    est_expit = 1 / (1 + np.exp(- (a * df.x + b)))\n",
    "    bernouilli = pm.Bernoulli('trials', p=est_expit, observed=df.y)\n",
    "    \n",
    "    trace_binom = pm.sample(1000, return_inferencedata=False)\n",
    "\n",
    "a_posterior_samples = trace_binom.get_values(\"a\")\n",
    "b_posterior_samples = trace_binom.get_values(\"b\")\n",
    "\n",
    "posterior_samples =pd.DataFrame({\"a\": a_posterior_samples, \"b\":b_posterior_samples})\n",
    "posterior_samples.mean().round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02fc6e09",
   "metadata": {},
   "source": [
    "#### MCMC - Bambi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e851143",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:39.282711Z",
     "iopub.status.busy": "2023-08-25T06:16:39.282599Z",
     "iopub.status.idle": "2023-08-25T06:16:43.234863Z",
     "shell.execute_reply": "2023-08-25T06:16:43.234432Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "model = bmb.Model(\"y ~ x\", df, family=\"bernoulli\")\n",
    "fitted = model.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ea1df7a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.237080Z",
     "iopub.status.busy": "2023-08-25T06:16:43.236858Z",
     "iopub.status.idle": "2023-08-25T06:16:43.259436Z",
     "shell.execute_reply": "2023-08-25T06:16:43.259062Z"
    }
   },
   "outputs": [],
   "source": [
    "az.summary(fitted)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cd4554e",
   "metadata": {},
   "source": [
    "#### Interpretation\n",
    "\n",
    "##### Interpretation of intercept\n",
    "\n",
    "We can use the posterior distribution of the intercept to get the probability when x = 0:\n",
    "\n",
    "$b = log(O(H)) \\implies P(H) = \\frac{1}{1 + e^{-b}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "130df142",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.261281Z",
     "iopub.status.busy": "2023-08-25T06:16:43.261091Z",
     "iopub.status.idle": "2023-08-25T06:16:43.263720Z",
     "shell.execute_reply": "2023-08-25T06:16:43.263455Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "qs = posterior_b.qs\n",
    "ps = posterior_b.ps\n",
    "\n",
    "\n",
    "pmf_at_0 = Pmf(ps, scipy.special.expit(qs))\n",
    "\n",
    "# identical as :\n",
    "pmf_at_0 = posterior_b.transform(scipy.special.expit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9db53d4e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.265239Z",
     "iopub.status.busy": "2023-08-25T06:16:43.265121Z",
     "iopub.status.idle": "2023-08-25T06:16:43.422488Z",
     "shell.execute_reply": "2023-08-25T06:16:43.421841Z"
    }
   },
   "outputs": [],
   "source": [
    "pmf_at_0.plot()\n",
    "plt.gca().set(title=\"P(y=1|x=0),  Here T° = Offset = 70°\", xlabel=\"P(y=1) at 70°\", ylabel=\"PDF\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28b22d0d",
   "metadata": {},
   "source": [
    "##### Interpretation of slope\n",
    "\n",
    "The distribution of $a$ gives us the log likelihood ratio so:\n",
    "\n",
    "$$ a = log (\\frac{P(Obs|H)}{P(Obs|\\bar{H})}) \n",
    "\\implies\n",
    "likelihood \\: ratio = \\frac{P(Obs|H)}{P(Obs|\\bar{H})} = e^a\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fcb67b4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.424063Z",
     "iopub.status.busy": "2023-08-25T06:16:43.423926Z",
     "iopub.status.idle": "2023-08-25T06:16:43.426703Z",
     "shell.execute_reply": "2023-08-25T06:16:43.426432Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "ps = posterior_a.ps\n",
    "qs = posterior_a.qs\n",
    "\n",
    "pmf_likelihood_ratio = Pmf(ps, np.exp(qs))\n",
    "\n",
    "# identical as:\n",
    "pmf_likelihood_ratio = posterior_a.transform(np.exp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7580f02a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.427999Z",
     "iopub.status.busy": "2023-08-25T06:16:43.427886Z",
     "iopub.status.idle": "2023-08-25T06:16:43.546576Z",
     "shell.execute_reply": "2023-08-25T06:16:43.546146Z"
    }
   },
   "outputs": [],
   "source": [
    "pmf_likelihood_ratio.plot()\n",
    "plt.gca().set(title=\"Posterior marginal distribution off likelihood ratio\", xlabel=\"Likelihood ratio of 1°\", ylabel=\"PDF\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "874b8204",
   "metadata": {},
   "source": [
    "##### Uncertainty on logistic curve\n",
    "\n",
    "By leveraging simulations on posteriors, we can define credible interval on our logistic regression:  \n",
    "\n",
    "* For each value of $x$:\n",
    "    + Compute probabilities by sampling from the posterior distribution of $a$ and $b$\n",
    "    + Compute the related credible interval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b4d08a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.549113Z",
     "iopub.status.busy": "2023-08-25T06:16:43.548992Z",
     "iopub.status.idle": "2023-08-25T06:16:43.680104Z",
     "shell.execute_reply": "2023-08-25T06:16:43.679705Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "# Simulation parameters\n",
    "RANGE_T = np.linspace(30, 85, 101)\n",
    "RANGE_X = RANGE_T - offset\n",
    "\n",
    "NB_SAMPLE_PER_X = 10_000\n",
    "QUANTS = [0.05, 0.5, 0.95]\n",
    "\n",
    "x_to_probs_thresh = {} # Will contain the q5, q50 and q95 of the distribution of probs for each x\n",
    "\n",
    "\n",
    "# For each value of x\n",
    "for x in RANGE_X:\n",
    "\n",
    "    # compute n probs by sampling from posteriors\n",
    "    log_odds = posterior_a.choice(NB_SAMPLE_PER_X) * x + posterior_b.choice(NB_SAMPLE_PER_X) \n",
    "    probs = scipy.special.expit(log_odds)\n",
    "    x_to_probs_thresh[x] = pd.Series(probs).quantile(QUANTS)\n",
    "\n",
    "# Format the data\n",
    "x_to_probs_thresh = pd.DataFrame(x_to_probs_thresh).T.rename_axis(\"x\").reset_index()\n",
    "x_to_probs_thresh[\"temp\"] = RANGE_T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85f79395",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.681927Z",
     "iopub.status.busy": "2023-08-25T06:16:43.681813Z",
     "iopub.status.idle": "2023-08-25T06:16:43.821732Z",
     "shell.execute_reply": "2023-08-25T06:16:43.821376Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.fill_between(\n",
    "    x_to_probs_thresh[\"temp\"],\n",
    "    x_to_probs_thresh[0.05],\n",
    "    x_to_probs_thresh[0.95],\n",
    "    color=\"cornflowerblue\",\n",
    "    alpha=0.2,\n",
    ")\n",
    "\n",
    "sns.lineplot(data=x_to_probs_thresh, x=\"temp\", y=0.5, label=\"Logistic model\")\n",
    "sns.scatterplot(\n",
    "    data=df, x=\"temp\", y=\"y\", alpha=0.2, color=\"orangered\", label=\"Observed data\"\n",
    ")\n",
    "\n",
    "plt.gca().set(\n",
    "    title=\"Median P(Y = 1 | T°) with 90% credible interval\",\n",
    "    xlabel=\"T°\",\n",
    "    ylabel=\"P(Y=1)\",\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1017fc45",
   "metadata": {},
   "source": [
    "### Linear regression\n",
    "\n",
    "::: {.callout-tip title=\"Theory\"}\n",
    "\n",
    "\n",
    "The assumption of the linear regression is that you can model $y$ in the following way:\n",
    "$$ y = (\\sum_{i=1}^{n} w_i * x_i) + \\epsilon $$\n",
    "\n",
    "With:\n",
    "$$\\epsilon \\sim \\mathcal{N}(0,\\,\\sigma^{2})  $$\n",
    "\n",
    "\n",
    ":::\n",
    "\n",
    "The aim is to estimate the values of $w_i$ and $\\sigma$\n",
    "\n",
    "#### Grid method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60f650dd",
   "metadata": {},
   "source": [
    ":::{.callout-note title=\"Example\"}  \n",
    "\n",
    "We observe the quantity of snow each year for 54 years, does the snow quantity tends to increase ?  \n",
    "The following notations are used: $$y = a * x + b + \\mathcal{N}(0,\\,\\sigma^{2})$$\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fcbac47",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.823725Z",
     "iopub.status.busy": "2023-08-25T06:16:43.823581Z",
     "iopub.status.idle": "2023-08-25T06:16:43.828654Z",
     "shell.execute_reply": "2023-08-25T06:16:43.828360Z"
    }
   },
   "outputs": [],
   "source": [
    "year = [1967, 1968, 1969, 1970, 1971, 1972, 1973, 1974, 1975, 1976, 1977, 1978, 1979, 1980, 1981, 1982, 1983, 1984, 1985, 1986, 1987, 1988, 1989, 1990, 1991, 1992, 1993, 1994, 1995, 1996, 1997, 1998, 1999, 2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020]\n",
    "amount_fall = [28.6, 44.7, 99.2, 66.8, 54.6, 68.7, 12.5, 46.5, 61.9, 57.7, 65.5, 100.6, 34.4, 27.7, 54.2, 64.1, 50.2, 56.3, 49.7, 38.3, 89.3, 53.5, 38.5, 52.5, 30.7, 39.0, 80.1, 88.7, 46.0, 124.2, 77.1, 24.1, 56.6, 43.7, 90.8, 61.5, 112.0, 51.8, 110.7, 40.6, 54.2, 59.1, 85.8, 57.6, 82.3, 31.1, 99.1, 80.6, 141.1, 64.0, 68.0, 90.4, 59.6, 12.2]\n",
    "\n",
    "df_snow = pd.DataFrame({\"year\": year, \"amount_fall\": amount_fall})\n",
    "\n",
    "offset = df_snow.year.mean()\n",
    "df_snow[\"x\"] = df_snow.year - offset\n",
    "df_snow[\"y\"] = df_snow.amount_fall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2b5f538",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:16:43.830321Z",
     "iopub.status.busy": "2023-08-25T06:16:43.830180Z",
     "iopub.status.idle": "2023-08-25T06:17:03.229505Z",
     "shell.execute_reply": "2023-08-25T06:17:03.229138Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "\n",
    "# Priors\n",
    "\n",
    "range_a = np.linspace(-0.5, 1.5, 51)\n",
    "range_b = np.linspace(54, 75, 21)\n",
    "range_sigma = np.linspace(20, 35, 31)  # epsilon sigma\n",
    "\n",
    "cols = [\"a\", \"b\", \"sigma\"]\n",
    "\n",
    "prior = pd.DataFrame(list(product(range_a, range_b, range_sigma)), columns=cols)\n",
    "prior[\"p\"] = 1 / len(prior)\n",
    "\n",
    "prior = Pmf(prior.set_index(cols).p)\n",
    "\n",
    "\n",
    "# Likelihood\n",
    "\n",
    "likelihood = prior.copy()\n",
    "\n",
    "for a, b, sigma in prior.index:\n",
    "\n",
    "    #likelihood = scipy.stats.multinomial.pmf(obs, nb_e, probs)\n",
    "    \n",
    "    y_est = a * df_snow.x + b\n",
    "    residuals = df_snow.y - y_est\n",
    "    likelihoods = scipy.stats.norm(0, sigma).pdf(residuals)\n",
    "\n",
    "    likelihood[a, b, sigma] = likelihoods.prod()\n",
    "\n",
    "\n",
    "# Bayes formula\n",
    "\n",
    "posterior = prior * likelihood\n",
    "posterior.normalize()\n",
    "\n",
    "\n",
    "# Marginal\n",
    "\n",
    "marginal_a = Pmf(posterior.rename(\"p\").reset_index().groupby(\"a\").p.sum())\n",
    "marginal_b = Pmf(posterior.rename(\"p\").reset_index().groupby(\"b\").p.sum())\n",
    "marginal_sigma = Pmf(posterior.rename(\"p\").reset_index().groupby(\"sigma\").p.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bae1145c",
   "metadata": {},
   "source": [
    "#### MCMC - PyMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42c53a3c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:17:03.231662Z",
     "iopub.status.busy": "2023-08-25T06:17:03.231537Z",
     "iopub.status.idle": "2023-08-25T06:17:08.381002Z",
     "shell.execute_reply": "2023-08-25T06:17:08.380644Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "with pm.Model() as model:\n",
    "    a = pm.Uniform(\"a\", -0.5, 1.5)\n",
    "    b = pm.Uniform(\"b\", 54, 75)\n",
    "    sigma = pm.Uniform(\"sigma\", 20, 35)\n",
    "\n",
    "    est_y = a * df_snow.x + b\n",
    "    y = pm.Normal(\"y\", mu=y_est, sigma=sigma, observed=df_snow.y)\n",
    "    trace_binom = pm.sample(1000, return_inferencedata=False)\n",
    "\n",
    "a_posterior_samples = trace_binom.get_values(\"a\")\n",
    "b_posterior_samples = trace_binom.get_values(\"b\")\n",
    "sigma_posterior_samples = trace_binom.get_values(\"sigma\")\n",
    "\n",
    "posterior_samples = pd.DataFrame(\n",
    "    {\n",
    "        \"a\": a_posterior_samples,\n",
    "        \"b\": b_posterior_samples,\n",
    "        \"sigma\": sigma_posterior_samples,\n",
    "    }\n",
    ")\n",
    "posterior_samples.mean().round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bd4f221",
   "metadata": {},
   "source": [
    "#### MCMC - Bambi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35897385",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:17:08.383138Z",
     "iopub.status.busy": "2023-08-25T06:17:08.382759Z",
     "iopub.status.idle": "2023-08-25T06:17:13.004424Z",
     "shell.execute_reply": "2023-08-25T06:17:13.004012Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "model = bmb.Model(\"y ~ x\", df_snow)\n",
    "fitted = model.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d1798c7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:17:13.006300Z",
     "iopub.status.busy": "2023-08-25T06:17:13.006175Z",
     "iopub.status.idle": "2023-08-25T06:17:13.026679Z",
     "shell.execute_reply": "2023-08-25T06:17:13.026341Z"
    }
   },
   "outputs": [],
   "source": [
    "az.summary(fitted)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "434f287d",
   "metadata": {},
   "source": [
    "#### Uncertainty on regression curve\n",
    "\n",
    "By leveraging simulations on posteriors, we can define credible interval on our linear regression:  \n",
    "\n",
    "* For each value of $x$:\n",
    "    + Compute probabilities by sampling from the posterior distribution of $a$, $b$ and $\\sigma$\n",
    "    + Compute the related credible interval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9900e2d7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:17:13.028253Z",
     "iopub.status.busy": "2023-08-25T06:17:13.028144Z",
     "iopub.status.idle": "2023-08-25T06:17:13.256111Z",
     "shell.execute_reply": "2023-08-25T06:17:13.255714Z"
    }
   },
   "outputs": [],
   "source": [
    "# | code-fold: show\n",
    "# | output: false\n",
    "\n",
    "# Simulation parameters\n",
    "RANGE_YEAR = np.linspace(df_snow.year.min(), df_snow.year.max(), 101)\n",
    "RANGE_X = RANGE_YEAR - offset\n",
    "\n",
    "NB_SAMPLE_PER_X = 10_000\n",
    "QUANTS = [0.05, 0.5, 0.95]\n",
    "\n",
    "x_to_y_thresh = {} # Will contain the q5, q50 and q95 of the distribution of probs for each x\n",
    "\n",
    "# For each value of x\n",
    "for x in RANGE_X:\n",
    "    \n",
    "    samples_a = marginal_a.choice(NB_SAMPLE_PER_X) \n",
    "    samples_b = marginal_b.choice(NB_SAMPLE_PER_X)\n",
    "    samples_sigma = marginal_sigma.choice(NB_SAMPLE_PER_X)\n",
    "\n",
    "    # compute n probs by sampling from posteriors\n",
    "    ys = samples_a * x + samples_b + scipy.stats.norm(0, samples_sigma).rvs(NB_SAMPLE_PER_X) \n",
    "    x_to_y_thresh[x] = pd.Series(ys).quantile(QUANTS)\n",
    "\n",
    "# Format the data\n",
    "x_to_y_thresh = pd.DataFrame(x_to_y_thresh).T.rename_axis(\"x\").reset_index()\n",
    "x_to_y_thresh[\"year\"] = RANGE_YEAR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "095f5f54",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-25T06:17:13.257960Z",
     "iopub.status.busy": "2023-08-25T06:17:13.257827Z",
     "iopub.status.idle": "2023-08-25T06:17:13.427892Z",
     "shell.execute_reply": "2023-08-25T06:17:13.427563Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.fill_between(\n",
    "    x_to_y_thresh[\"year\"],\n",
    "    x_to_y_thresh[0.05],\n",
    "    x_to_y_thresh[0.95],\n",
    "    color=\"cornflowerblue\",\n",
    "    alpha=0.2,\n",
    ")\n",
    "\n",
    "sns.lineplot(data=x_to_y_thresh, x=\"year\", y=0.5, label=\"Linear model\")\n",
    "sns.scatterplot(\n",
    "    data=df_snow, x=\"year\", y=\"y\", alpha=0.2, color=\"orangered\", label=\"Observed data\"\n",
    ")\n",
    "\n",
    "plt.gca().set(\n",
    "    title=\"Median Snowfall with 90% credible interval\",\n",
    "    xlabel=\"Year\",\n",
    "    ylabel=\"Snowfall (cm)\",\n",
    ")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "184.333px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}