Cointhrowing_converging_to_wrong_posterior.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import comb\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def average_likelihood(p,q,n=10):\n",
    "    l_ps = []\n",
    "    l_qs = []\n",
    "    for n_head in range(n+1):\n",
    "        l_p = p**n_head * (1-p)**(n-n_head) * comb(n,n_head)\n",
    "        l_q = q**n_head * (1-q)**(n-n_head) * comb(n,n_head)\n",
    "        l_ps += [l_p]\n",
    "        l_qs += [l_q]\n",
    "    l_ps = torch.tensor(l_ps)\n",
    "    l_qs = torch.tensor(l_qs)\n",
    "\n",
    "    likelihood = l_ps @ l_qs\n",
    "    return likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = .5\n",
    "qs = [.3,.6]\n",
    "ns = [1,2,4,8,16,32,64,128,256,512,1024]\n",
    "likelihoods = []\n",
    "\n",
    "for q in qs:\n",
    "    for n in ns:\n",
    "        avg_l = average_likelihood(p,q,n=n).item()\n",
    "        print(q,n,avg_l)\n",
    "        likelihoods+=[{'q': q, 'n': n, 'avg_l': avg_l}]\n",
    "likelihoods = pd.DataFrame(likelihoods)\n",
    "\n",
    "likelihoods_per_n = likelihoods.set_index(['n','q']).unstack(1).to_numpy()\n",
    "ppd = (likelihoods_per_n / likelihoods_per_n.sum(1)[:,None]) @ np.array(qs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "palette = sns.color_palette(\"muted\")\n",
    "sns.set_palette(palette)\n",
    "\n",
    "\n",
    "fig, axs = plt.subplots(1,2, figsize=(10,5))\n",
    "\n",
    "axs[0].plot(ns, ppd, color=palette[1], marker='o', label='Bayesian Prediction')\n",
    "axs[0].axhline(y=p, color=palette[0], linestyle='--', label='True Probability')\n",
    "\n",
    "axs[0].axhline(y=.3, color='gray', linestyle='--')\n",
    "axs[0].axhline(y=.6, color='gray', linestyle='--')\n",
    "\n",
    "axs[0].set_xlabel(\"Number of In-Context Examples (Coin Flips)\")\n",
    "axs[0].set_ylabel(\"Predicted Probability of Heads\")\n",
    "\n",
    "\n",
    "\n",
    "#axs[0].set_xscale('log')\n",
    "axs[0].legend()\n",
    "#plt.xlim(ns[0],ns[-1])\n",
    "\n",
    "axs[1].plot(ns, likelihoods_per_n[:,0], color=palette[2], marker='o', label=f'Coin with p={qs[0]}')\n",
    "axs[1].plot(ns, likelihoods_per_n[:,1], color=palette[3], marker='o', label=f'Coin with p={qs[1]}')\n",
    "\n",
    "axs[1].set_yscale('log')\n",
    "#axs[1].set_xscale('log')\n",
    "axs[1].set_ylabel(f'Likelihood of Observed Data')\n",
    "axs[1].set_xlabel(\"Number of In-Context Examples (Coin Flips)\")\n",
    "\n",
    "axs[1].legend()\n",
    "\n",
    "# utils.tikzplotlib_save(f'figures/bad_prior_coin_flip_failure.tex', axis_width=r\".4\\textwidth\", axis_height=r\".4\\textwidth\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}