{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Informatics 2 - Foundations of Data Science \n",
    "\n",
    "# Maximum likelihood\n",
    "\n",
    "David Sterratt, 2023-2026"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from scipy.stats import binom\n",
    "from scipy.stats import expon\n",
    "from scipy.stats import norm\n",
    "from scipy.stats import uniform\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import axes3d\n",
    "import seaborn as sns\n",
    "import statsmodels.formula.api as smf\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['figure.max_open_warning'] = 100\n",
    "%matplotlib ipympl\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum likelihood of normally-distributed single-variable data\n",
    "\n",
    "We'll first create some synthetic data, which we will then try to fit a model to using maximum likelihood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "y = np.random.normal(size=10)\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now imagine that we don't know how this data was generated, but that we make the assumptions that each data point was drawn independently from a normal distribution with mean $\\mu$ and variance $\\sigma^2$. The likelihood of the data for any value of $\\mu$ and $\\sigma^2$ is:\n",
    "$$\n",
    "\\begin{align}\n",
    "p(\\mathbf{Y}= y_1, \\dots y_{10}|\\mu, \\sigma^2) & =  p(Y=y_1 |\\mu, \\sigma^2) \\times \\dots \\times p(Y=y_{10} |\\mu, \\sigma^2)\\\\\n",
    "& = \\frac{1}{\\sqrt{2\\pi}\\sigma}\\exp\\left(-\\frac{(y_1 - \\mu)^2}{2\\sigma^2}\\right) \\times \\dots \\times \\frac{1}{\\sqrt{2\\pi}\\sigma}\\exp\\left(-\\frac{(y_{10} - \\mu)^2}{2\\sigma^2}\\right) \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "We'll now define these functions in code. First the likelihood of an individual point $p(Y=y_1 |\\mu, \\sigma^2)$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normal(y, mu, sigma2):\n",
    "    return(1/np.sqrt(2*np.pi*sigma2)*np.exp(-0.5*pow(y-mu,2)/sigma2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then the likelihood of a set of points (which we represent as a python list), i.e.  $p(\\mathbf{Y}= y_1, \\dots y_{10}|\\mu, \\sigma^2)$ above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def likelihood(y, mu, sigma2):\n",
    "    p = normal(y[0], mu, sigma2)\n",
    "    for i in np.arange(1, len(y)):\n",
    "        p = p * normal(y[i], mu, sigma2)\n",
    "    return(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's the probability (likelihood) of generating just the first point $y_1$, given particular values of $\\mu$ and $\\sigma^2$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(normal(y[0], mu=0, sigma2=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What happens if we change $\\mu$ and $\\sigma^2$?\n",
    "\n",
    "Now, what's the likelihood of all the data? Notice anything different?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(likelihood(y, mu=0.0, sigma2=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The likelihood as a function of the parameters $\\mu$ and $\\sigma^2$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is to print the table in the lecture notes\n",
    "y_df = pd.DataFrame(y, index=['$y_{' + str(n) + '}$' for n in np.arange(1, 11)], columns=['Data'])\n",
    "y_df.style.to_latex('y-data.tex')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mus, sigmas = np.meshgrid(np.arange(-1, 1, 0.1), np.arange(0.3, 10, 0.1))\n",
    "ps = likelihood(y, mus, sigmas)\n",
    "\n",
    "ax = plt.figure(figsize=(6,6)).add_subplot(projection='3d')\n",
    "ax.plot_surface(mus, sigmas, ps*1e7, edgecolor='royalblue', lw=0.5, alpha=0.3)\n",
    "ax.set(xlabel='$\\\\mu$', ylabel='$\\\\sigma^2$', zlabel='$P(Y=y_1,\\\\dots,y_n,\\\\mu,\\\\sigma^2)\\\\times 10^7$')\n",
    "ax.view_init(25, -30, 0)\n",
    "plt.show()\n",
    "plt.savefig('normal-likelihood.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What happens as we increase $n$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_orig = y.copy()\n",
    "y = np.random.normal(size=500)\n",
    "ps = likelihood(y, mus, sigmas)\n",
    "\n",
    "ax = plt.figure(figsize=(6,6)).add_subplot(projection='3d')\n",
    "ax.plot_surface(mus, sigmas, ps, edgecolor='royalblue', lw=0.5, alpha=0.3)\n",
    "ax.set(xlabel='$\\\\mu$', ylabel='$\\\\sigma^2$', zlabel='$p$')\n",
    "ax.view_init(25, -30, 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Uhhh? Why do we have a flat plot?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(likelihood(y, 0.1, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What if we plot the log Likelihood?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lognormal(y, mu, sigma2):\n",
    "    return(-0.5*np.log(2*np.pi*sigma2) -0.5*pow(y-mu,2)/sigma2)\n",
    "\n",
    "def loglikelihood(y, mu, sigma2):\n",
    "    logp = lognormal(y[0], mu, sigma2)\n",
    "    for i in np.arange(1, len(y)):\n",
    "        logp = logp + lognormal(y[i], mu, sigma2)\n",
    "    return(logp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.random.normal(size=1000)\n",
    "\n",
    "logps = loglikelihood(y, mus, sigmas)\n",
    "\n",
    "ax = plt.figure(figsize=(6,6)).add_subplot(projection='3d')\n",
    "ax.plot_surface(mus, sigmas, logps, edgecolor='royalblue', lw=0.5, alpha=0.3)\n",
    "ax.set(xlabel='$\\\\mu$', ylabel='$\\\\sigma^2$', zlabel='$\\\\log p$')\n",
    "ax.view_init(25, -30, 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting against $\\log\\sigma^2$\n",
    "\n",
    "This makes the landscape look smoother."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.figure(figsize=(6,6)).add_subplot(projection='3d')\n",
    "ax.plot_surface(mus, np.log(sigmas), logps, edgecolor='royalblue', lw=0.5, alpha=0.3)\n",
    "ax.set(xlabel='$\\\\mu$', ylabel='$\\\\log \\\\sigma^2$', zlabel='$\\\\log p$')\n",
    "ax.view_init(25, -30, 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What happens as we increase $n$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.random.normal(size=500)\n",
    "print(likelihood(y, 0.1, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The log likelihood of the original small (10-value) dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logps = loglikelihood(y_orig, mus, sigmas)\n",
    "\n",
    "ax = plt.figure(figsize=(6,4)).add_subplot(projection='3d')\n",
    "ax.plot_surface(mus, sigmas, logps, edgecolor='royalblue', lw=0.5, alpha=0.3)\n",
    "# ax.contour(mus, np.log(sigmas), logps, zdir='z', offset=-50)\n",
    "ax.scatter([np.mean(y_orig)], [np.var(y_orig)], [-50], color='red')\n",
    "ax.set(xlabel='$\\\\mu$',\n",
    "       ylabel='$\\\\sigma^2$',\n",
    "       zlabel='$\\\\log P(Y=y_1,\\\\dots,y_n,\\\\mu,\\\\sigma^2)$',\n",
    "       zlim=[-50, -15])\n",
    "ax.text(-0.5, 5, -50,\n",
    "        'Max at\\n$\\\\mu=%0.2f$\\n$\\\\sigma^2=%0.2f$'%(np.mean(y_orig), np.var(y_orig)), color='red')\n",
    "ax.view_init(25, -30, 0)\n",
    "plt.savefig('normal-log-likelihood.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.figure(figsize=(6,4)).add_subplot(projection='3d')\n",
    "ax.plot_surface(mus, np.log(sigmas), logps, edgecolor='royalblue', lw=0.5, alpha=0.3)\n",
    "ax.contour(mus, np.log(sigmas), logps, zdir='z', offset=-50)\n",
    "ax.scatter([np.mean(y_orig)], [np.log(np.var(y_orig))], [-50], color='red')\n",
    "ax.set(xlabel='$\\\\mu$',\n",
    "       ylabel='$\\\\log \\\\sigma^2$',\n",
    "       zlabel='$\\\\log P(Y=y_1,\\\\dots,y_n,\\\\mu,\\\\sigma^2)$',\n",
    "       zlim=[-50, -15])\n",
    "ax.view_init(25, -30, 0)\n",
    "plt.savefig('normal-log-likelihood-logsigma2.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.var(y_orig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Application of max likelihood to regression\n",
    "\n",
    "We will look at the squirrel data again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = 'squirrel-clean.csv'\n",
    "if not os.path.exists(filename):\n",
    "    filename = os.path.join('../../data', filename)\n",
    "dat = pd.read_csv(filename)\n",
    "dat[\"Year\"].value_counts()\n",
    "datf = dat[dat['Sex'] == 'F']\n",
    "datm = dat[dat['Sex'] == 'M']\n",
    "datf.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "ax = sns.scatterplot(data=datf, x='Length (mm)', y='Weight (g)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write a function for the likelihood of the data under a linear model parameterised by $\\beta_0$, $\\beta_1$ and variance of the noise $\\sigma^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loglikelihood_regression(x, y, beta0, beta1, sigma2):\n",
    "    logp = lognormal(y[0], beta0 + beta1*x[0], sigma2)\n",
    "    for i in np.arange(1, len(y)):\n",
    "        logp = logp + lognormal(y[i], beta0 + beta1*x[i], sigma2)\n",
    "    return(logp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Demonstration that we can compute the likelihood for any values of the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.array(datf['Length (mm)'])\n",
    "y = np.array(datf['Weight (g)'])\n",
    "loglikelihood_regression(x, y, 0, 10, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll compute the maximum likelihood estimates using the formula (not the function we've just defined). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta1hat = np.sum((x - np.mean(x))*(y - np.mean(y)))/np.sum((x - np.mean(x))*(x - np.mean(x)))\n",
    "beta0hat = np.mean(y) - beta1hat*np.mean(x)\n",
    "sigma2hat = 1/len(x)*np.sum(pow(y - beta0hat - beta1hat*x, 2))\n",
    "print('Max likelihood estimates\\nbeta0 = %0.2f\\nbeta1 = %0.2f\\nsigma2= %0.2f'%(beta0hat, beta1hat, sigma2hat))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will compute the log likelihood function for the optimal value of $\\sigma^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta0s, beta1s = np.meshgrid(np.arange(-700, 0, 1), np.arange(1, 5, 0.002))\n",
    "logps = loglikelihood_regression(x, y, beta0s, beta1s, 187.95)\n",
    "print('Range of log likelihood: (%0.2f, %0.2f)'%(np.min(logps), np.max(logps)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.figure(figsize=(6,3.5)).add_subplot(projection='3d')\n",
    "ax.plot_surface(beta0s, beta1s, logps/1000, edgecolor='royalblue', lw=0.5, alpha=0.3)\n",
    "ax.contour(beta0s, beta1s, logps, zdir='z', offset=-70, levels=[np.log(0.05077) + np.max(logps)])\n",
    "ax.scatter(beta0hat, beta1hat, -70, color='red')\n",
    "ax.set(xlabel='$\\\\beta_0$',\n",
    "       ylabel='$\\\\beta_1$',\n",
    "       zlabel='$\\\\ell(\\\\beta_0,\\\\beta_1,\\\\sigma^2)\\\\times 10^{-3}$', zlim=[-70, 0])\n",
    "plt.savefig('regression-log-likelihood.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(0.05077) + np.max(logps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the probability level that encloses 95% of the mass of the likelihood distribution. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ps = np.exp(logps - np.max(logps))\n",
    "pscum = pd.DataFrame({'ps': ps.flatten(), 'beta0': beta0s.flatten(), 'beta1': beta1s.flatten()}).sort_values('ps')\n",
    "# pscum['ps'] = pscum['ps'].div(pscum['ps'].sum())\n",
    "pscum['pscum'] = pscum['ps'].cumsum()\n",
    "\n",
    "def get_level(pscum, centile):\n",
    "    ind = (pscum['pscum'] - pscum['pscum'].max() * (1-centile)).abs().argsort().iloc[0]\n",
    "    return(pscum.iloc[ind].loc['ps'])\n",
    "\n",
    "levels = [get_level(pscum, centile) for centile in [0.95]]\n",
    "levels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate some sample fits drawn in proportion to the likelihood at each set of parameter values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_fits = pscum.sample(5, weights=pscum['ps'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ax = plt.figure(figsize=(6,6)).add_subplot(projection='3d')\n",
    "fig, ax = plt.subplots(2, 2, figsize=(6, 4))\n",
    "# \\label{5-03-max-likelihood-glm:fig:normal-log-likelihood-both}ax.plot_surface(beta0s, beta1s, ps, edgecolor='royalblue', lw=0.5, alpha=0.3)\n",
    "\n",
    "plt.sca(ax[0, 0])\n",
    "sns.scatterplot(x='Length (mm)', y='Weight (g)', data=datf)\n",
    "for i in range(5):\n",
    "    plt.plot(datf['Length (mm)'],\n",
    "             sample_fits.iloc[i]['beta0'] + sample_fits.iloc[i]['beta1']*datf['Length (mm)'],\n",
    "             color='orange')\n",
    "plt.plot(datf['Length (mm)'],\n",
    "         beta0hat + beta1hat*datf['Length (mm)'],\n",
    "         color='red')\n",
    "\n",
    "ax[1,1].set(xlabel='$\\\\beta_1$', ylabel='$\\\\beta_0$') # zlabel='$\\log p$')\n",
    "ax[1,1].contour(beta1s.transpose(), beta0s.transpose(), ps.transpose(), levels=levels)#  zdir='z', offset=0)\n",
    "ax[1,1].scatter(sample_fits['beta1'], sample_fits['beta0'], color='orange')\n",
    "ax[1,1].scatter([beta1hat], [beta0hat], color='red')\n",
    "\n",
    "ax[1,0].plot(beta0s[0,:], ps.sum(0)/ps.sum())\n",
    "ax[1,0].set(xlabel='$\\\\beta_0$', ylabel='$P(\\\\beta_0|\\\\sigma=\\\\hat\\\\sigma)$')\n",
    "\n",
    "ax[0,1].plot(beta1s[:,0], ps.sum(1)/ps.sum()/0.002)\n",
    "ax[0,1].set(xlabel='$\\\\beta_1$', ylabel='$P(\\\\beta_1|\\\\sigma=\\\\hat\\\\sigma)$')\n",
    "plt.tight_layout()\n",
    "plt.savefig('regression-uncertainty.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datfr = datf.copy()\n",
    "results = smf.ols('Weight ~ Length', data=datf.rename(columns={'Weight (g)': 'Weight', 'Length (mm)': 'Length'})).fit()\n",
    "results.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "2*129.18 + 2*3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum likelihood of Poisson-distributed single-variable data\n",
    "\n",
    "Bortkewitsch, (1898) collated how many deaths there were due to horse kicks in each of 14 Prussian army regiments in the 20 years 1875-1894.\n",
    "\n",
    "The raw data is a set of 280 counts, one for each combination of regiment and year. We denote these counts $y_1,\\dots,y_{n}$, where here $n=280$.\n",
    "\n",
    "We can summarise the counts using the histogram of how many regiment-year pairs had 0 deaths, 1 deaths etc. We denote the number of regiment-year pairs with $k$ deaths as $n_k$.\n",
    "\n",
    "Here is the summarised data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dat = pd.DataFrame({'k':[0, 1, 2, 3, 4,], 'nk':[144, 91, 32, 11, 2]})\n",
    "dat.rename(columns={'k': '$k$', 'nk': '$n_k$'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The log likelihood is of this data under the Poisson distribution is\n",
    "$$\\ell(\\lambda) = \\ln\\lambda \\sum_{i=1}^n y_i -n\\lambda -\\sum_{i=1}^n\\ln y_i!$$\n",
    "\n",
    "We can rewrite  in terms of the summarised data:\n",
    "$$\\ell(\\lambda) = \\ln\\lambda \\sum_{k} n_k k -n\\lambda -\\sum_{k}n_k\\ln k! = \\sum_k n_k (\\ln\\lambda - \\lambda -\\ln k!)$$\n",
    "Here the sum over $k$ is from 0 to the maximum value of $k$, which is 4 in the case of the horse kick data.\n",
    "\n",
    "We now write functions to compute the likelihood of the data given the parameter $\\lambda$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function\n",
    "def array_factorial(k):\n",
    "    return(np.array([math.factorial(int(ki)) for ki in k]))\n",
    "\n",
    "# Log likelihood of data\n",
    "def loglikelihood_poisson(nk, Lambda):\n",
    "    k = np.arange(len(nk))*1.0\n",
    "    return(np.sum(nk*(k*np.log(Lambda) - Lambda - array_factorial(k))))\n",
    "\n",
    "# Poisson distribution: i.e. P(y=k|\\lambda)\n",
    "def poisson(Lambda, k):\n",
    "    return(np.power(Lambda, k)*np.exp(-Lambda)/math.factorial(k))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot the log likelihood against $\\lambda$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Lambdas = np.arange(0.1, 3, 0.01)\n",
    "logps_poisson = np.array([loglikelihood_poisson(dat['nk'], Lambda) for Lambda in Lambdas])\n",
    "plt.figure(figsize=(4,3))\n",
    "plt.plot(Lambdas, logps_poisson)\n",
    "plt.xlabel('$\\\\lambda$')\n",
    "plt.ylabel('$\\\\log p$')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks like the maximum likelihood occurs when $\\lambda$ is around 0.7. We can derive a maximum likelihood estimate of $\\lambda$ analyticallly by finding the point where $d\\ell/d \\lambda=0$. This estimate is\n",
    "$$\\hat\\lambda_{\\text{MLE}}= 1/n\\sum_{i=1}^n y_i = 1/n \\sum_k n_k k$$\n",
    "\n",
    "Let's check this estimate is about 0.7:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mle_poisson(n):\n",
    "    k = np.arange(len(n))*1.0\n",
    "    return(np.sum(n*k)/np.sum(n))\n",
    "\n",
    "mle_poisson(dat['nk'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Yes! We can also experiment to see what would happen if we changed $\\lambda$, computing the predicted number of kicks. Try changing the value `Lambda` below and re-running the cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Lambda = 0.71\n",
    "dat['ppred'] = dat['k'].apply(lambda x: poisson(Lambda, x))\n",
    "dat['npred'] = dat['ppred']*np.sum(dat['nk'])\n",
    "print('Log likelihood with lambda = %1.3f is %4.2f'%(Lambda, loglikelihood_poisson(np.array(dat['nk']), Lambda)))\n",
    "dat.rename(columns={'k': '$k$', 'nk': '$n_k$', 'ppred': 'Predicted $p_k$', 'npred': 'Predicted $n_k$'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## To Bayesian stats\n",
    "\n",
    "Finally, we can convert the log likelihood into the posterior distribution of the parameter $\\lambda$ given the data. Here we have normalised the area under the distribution to 1 approximately - this is the reason for the 681 constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(4,3))\n",
    "plt.plot(Lambdas, np.exp(logps_poisson+681))\n",
    "plt.xlabel('$\\\\lambda$')\n",
    "plt.ylabel('$p$')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}