blr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "271e07eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import functools\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from jax.scipy.special import expit\n",
    "from jax.scipy.stats import bernoulli, gamma, norm\n",
    "from sklearn.datasets import load_breast_cancer\n",
    "\n",
    "from samplers import get_min_ess, make_samplers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "688c743e",
   "metadata": {},
   "source": [
    "# Basic setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "42728836",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = load_breast_cancer()\n",
    "X = (data.data - np.mean(data.data, axis=0, keepdims=True)) / np.std(\n",
    "    data.data, axis=0, keepdims=True\n",
    ")\n",
    "X = np.concatenate([X, np.ones((X.shape[0], 1))], axis=1)\n",
    "y = data.target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0e8f9e0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYAAAAD4CAYAAADlwTGnAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8/fFQqAAAACXBIWXMAAAsTAAALEwEAmpwYAAAfBUlEQVR4nO3deXRU553m8e+vqrSvSIhNEgiwMAYMGGSMt7aT2D3geEx6Oouddpss3YTY7pNlutNkPJMzOT1nsnROJ/bEbbcTuydkczuOExM3bsdL4iVjbAQ2mB1ZLBYSIBBIAiG0vfNHXUgha7lCy5XqPp9z6lTd975v1e+VsR7dpe415xwiIhI+kaALEBGRYCgARERCSgEgIhJSCgARkZBSAIiIhFQs6AIGYvz48a6srCzoMkRExpRNmzYdc84VdW8fUwFQVlZGZWVl0GWIiIwpZnagp3btAhIRCSkFgIhISCkARERCSgEgIhJSCgARkZDyFQBmtszMdptZlZmt6WG9mdkD3vqtZrYoYd1jZnbUzLZ1G1NgZs+b2V7vedzgpyMiIn71GwBmFgUeBJYDc4A7zGxOt27LgXLvsQp4KGHd/wWW9fDWa4AXnXPlwIvesoiIjBA/WwBLgCrnXLVzrg14HFjRrc8KYK2L2wDkm9lkAOfcK0BDD++7AviR9/pHwEcuon5fXtp1hH/+fdVwvb2IyJjkJwCKgfcSlmu8toH26W6ic64OwHue0FMnM1tlZpVmVllfX++j3Pf7Q9VxHnhxL11duveBiMg5fgLAemjr/pvUT5+L4px7xDlX4ZyrKCp63zeZfSmfkE1rexeHTp4ZipJERJKCnwCoAUoTlkuA2ovo092Rc7uJvOejPmq5KJdMyAag6uip4foIEZExx08AbATKzWy6maUCtwPruvVZB9zlnQ20FGg8t3unD+uAld7rlcDTA6h7QM4FwN6jzcP1ESIiY06/AeCc6wDuBZ4DdgJPOOe2m9lqM1vtdVsPVANVwA+Au8+NN7OfA68Dl5pZjZl91lv1TeBmM9sL3OwtD4v8zFTGZ6dqC0BEJIGvq4E659YT/yWf2PZwwmsH3NPL2Dt6aT8OfMh3pYN0yYRsBYCISILQfBP4XADEs0pERMITAEXZNLV2UH/qbNCliIiMCuEJgAk5AFQd0W4gEREIUQCUT/ROBa1XAIiIQIgCYEJOGjlpMR0IFhHxhCYAzIyZOhNIROS80AQA6FRQEZFEoQuAo81naTzTHnQpIiKBC1UAlOuaQCIi54UqAM5dE+hdBYCISLgCoGRcJqmxiE4FFREhZAEQjRgzi7LZc0RXBRURCVUAAFw6MZvdhxUAIiKhC4DZk3Opa2ylsUVnAolIuIUuAC6dFL8m0G7tBhKRkAtdAMw+FwCHmwKuREQkWKELgEm56eSmx9il4wAiEnKhCwAzY/akXB0IFpHQC10AQPw4wO7Dzbo7mIiEWmgDoPlsB4dOngm6FBGRwIQyAP54IFi7gUQkvEIZALO8ANCBYBEJs1AGQG56CsX5GdoCEJFQC2UAwB8PBIuIhFWoA+Dd+lO0dXQFXYqISCBCGwCzJ+XQ0eV4V5eGFpGQCnEA5AI6E0hEwiu0ATCjKIvUaIQddbomkIiEU2gDICUa4dJJOWyvbQy6FBGRQIQ2AADmTslle22TLgkhIqEU7gAozuNkS7suCSEioeQrAMxsmZntNrMqM1vTw3ozswe89VvNbFF/Y81soZltMLO3zazSzJYMzZT8mzslfiB4e62OA4hI+PQbAGYWBR4ElgNzgDvMbE63bsuBcu+xCnjIx9hvA193zi0EvuYtj6jLJuUSMQWAiISTny2AJUCVc67aOdcGPA6s6NZnBbDWxW0A8s1scj9jHZDrvc4Dagc5lwHLSI0ysyibHToQLCIhFPPRpxh4L2G5BrjKR5/ifsZ+EXjOzL5DPIiu6enDzWwV8a0Kpk6d6qPcgZk7JZcN1Q1D/r4iIqOdny0A66Gt+2kzvfXpa+zngS8550qBLwGP9vThzrlHnHMVzrmKoqIiH+UOzNwpeRxuauX4qbND/t4iIqOZnwCoAUoTlkt4/+6a3vr0NXYl8JT3+hfEdxeNuLnFOhAsIuHkJwA2AuVmNt3MUoHbgXXd+qwD7vLOBloKNDrn6voZWwvc4L3+ILB3kHO5KHMn5wEKABEJn36PATjnOszsXuA5IAo85pzbbmarvfUPA+uBW4AqoAX4dF9jvbf+a+B+M4sBrXj7+UdaXmYKJeMy2KYDwSISMn4OAuOcW0/8l3xi28MJrx1wj9+xXvtrwOKBFDtc5k7JZYe2AEQkZEL9TeBz5k3JY9+x0zS3tgddiojIiFEAAPNL8wF455B2A4lIeCgAgPnF8QPBW95TAIhIeCgAgHFZqUwrzGRrzcmgSxERGTEKAM+Ckny2vHcy6DJEREaMAsCzoDSf2sZWjja1Bl2KiMiIUAB4FpR4xwFqdBxARMJBAeCZOyWPaMS0G0hEQkMB4MlIjXLpxBy26ECwiISEAiDBgtL4gWDdI1hEwkABkGBBSR5NrR3sP94SdCkiIsNOAZBggfeNYB0HEJEwUAAkKJ+QTUZKlLcVACISAgqABLFohMtL8nhLASAiIaAA6GbxtHFsP9TImbbOoEsRERlWCoBuKqaNo6PL6bpAIpL0FADdLJo6DoDKAycCrkREZHgpALoZl5XKzKIsNikARCTJKQB6sHjaODYfPEFXl74QJiLJSwHQg4ppBZxsaaf62KmgSxERGTYKgB4sLvOOA+zXbiARSV4KgB7MGJ/FuMwUHQcQkaSmAOiBmbF42jgFgIgkNQVALxZNG0f1sdM0nG4LuhQRkWGhAOhFxbQCACr3NwRciYjI8FAA9GJ+SR5psQhv7FMAiEhyUgD0Ij0lyhVT89lQfTzoUkREhoUCoA9LZxSyo66Jxpb2oEsRERlyCoA+LJ1RiHPwpo4DiEgSUgD0YWFpPqmxCG9oN5CIJCFfAWBmy8xst5lVmdmaHtabmT3grd9qZov8jDWzv/HWbTezbw9+OkMrPSXKoqn5bNinABCR5NNvAJhZFHgQWA7MAe4wszndui0Hyr3HKuCh/saa2QeAFcB859xc4DtDMaGhtnRGIdtrm2g8o+MAIpJc/GwBLAGqnHPVzrk24HHiv7gTrQDWurgNQL6ZTe5n7OeBbzrnzgI4544OwXyG3FXT48cBNup0UBFJMn4CoBh4L2G5xmvz06evsbOA683sDTN72cyu7OnDzWyVmVWaWWV9fb2PcofWFVPjxwF0OqiIJBs/AWA9tHW/UH5vffoaGwPGAUuBvwOeMLP39XfOPeKcq3DOVRQVFfkod2ilp0S5olTHAUQk+fgJgBqgNGG5BKj12aevsTXAU95uozeBLmC8/9JHztUz48cBTui6QCKSRPwEwEag3Mymm1kqcDuwrlufdcBd3tlAS4FG51xdP2N/DXwQwMxmAanAscFOaDhcX16Ec/CHd0dleSIiF6XfAHDOdQD3As8BO4EnnHPbzWy1ma32uq0HqoEq4AfA3X2N9cY8Bswws23EDw6vdM6NynswLijJIyc9xqt7FAAikjxifjo559YT/yWf2PZwwmsH3ON3rNfeBtw5kGKDEotGuHbmeF6rOoZzjh4OVYiIjDn6JrBP15WP59DJM1QfOx10KSIiQ0IB4NOflMfPQHp1z8ifiioiMhwUAD5NLcxkWmEmr+7VcQARSQ4KgAG4vnw8r1cfp62jK+hSREQGTQEwANeXF9HS1snmg7pZvIiMfQqAAbh6ZiHRiPGadgOJSBJQAAxAbnoKV5Tm87vdo/K6dSIiA6IAGKAPzJ7A9tomDje2Bl2KiMigKAAG6KbLJgLw0i5tBYjI2KYAGKBZE7Mpzs/gpV1Hgi5FRGRQFAADZGbcdNkEXqs6Rmt7Z9DliIhcNAXARfjgZRNpbe/i/+nqoCIyhikALsLSGQVkpkZ5YaeOA4jI2KUAuAhpsSjXl4/npZ1HGaVXsBYR6ZcC4CJ96LKJHG5qZUddU9CliIhcFAXARfrApRMwg+d36GwgERmbFAAXqSgnjYpp4/iPbYeDLkVE5KIoAAZh+bzJ7DrcTHX9qaBLEREZMAXAICybNwmAZ7UVICJjkAJgEKbkZ7CwNF+7gURkTFIADNItl0/inUONvNfQEnQpIiIDogAYpOXzJgPw7La6gCsRERkYBcAglRZkMq84V8cBRGTMUQAMgeXzJvPWwZMcOnkm6FJERHxTAAyBW+fHdwP9ZkttwJWIiPinABgC0wqzuGJqPr9+61DQpYiI+KYAGCIfWVjMrsPN7D7cHHQpIiK+KACGyK3zJxONGL9+W1sBIjI2KACGSGF2Gn9SPp6n3zpEV5cuES0io58CYAh95Ipiahtb2bi/IehSRET6pQAYQjfPmUhmapRfv62zgURk9PMVAGa2zMx2m1mVma3pYb2Z2QPe+q1mtmgAY//WzJyZjR/cVIKXmRpj2dxJPLO1VjeMF5FRr98AMLMo8CCwHJgD3GFmc7p1Ww6Ue49VwEN+xppZKXAzcHDQMxklPn5lKc2tHbo0hIiMen62AJYAVc65audcG/A4sKJbnxXAWhe3Acg3s8k+xn4X+AqQNEdNr5peQFlhJo+/+V7QpYiI9MlPABQDib/Narw2P316HWtmtwGHnHNb+vpwM1tlZpVmVllfX++j3GCZGZ+4cipv7GvQjWJEZFTzEwDWQ1v3v9h769Nju5llAvcBX+vvw51zjzjnKpxzFUVFRf0WOxr8+eJiohHjicqaoEsREemVnwCoAUoTlkuA7qe59Nant/aZwHRgi5nt99o3m9mkgRQ/Wk3ISedDsyfw5KYa2ju7gi5HRKRHfgJgI1BuZtPNLBW4HVjXrc864C7vbKClQKNzrq63sc65d5xzE5xzZc65MuJBscg5lzTXVP7ElaUcO3WWl3YdDboUEZEe9RsAzrkO4F7gOWAn8IRzbruZrTaz1V639UA1UAX8ALi7r7FDPotR6IZZRUzKTednbyTNCU4ikmRifjo559YT/yWf2PZwwmsH3ON3bA99yvzUMZbEohE+edVU/un5PVTXn2JGUXbQJYmIXEDfBB5GdyyZSmo0wtrXDwRdiojI+ygAhlFRThq3zp/Mk5tqaG5tD7ocEZELKACG2cpryjh1toNfbtIpoSIyuigAhtmC0nwWluaz9vUDuky0iIwqCoAR8Olry6g+dpqX947+bzKLSHgoAEbA8nmTmZyXzr+8/G7QpYiInKcAGAGpsQifvW46G6obeOvgiaDLEREBFAAj5o4lU8nLSOFhbQWIyCihABghWWkxVl49jd/uOELVUV0lVESCpwAYQSuvKSMtFuGRV7QVICLBUwCMoMLsND5RUcqv3jrEoZNngi5HREJOATDCVt0wE8P4/kt7gy5FREJOATDCivMzuH1JKb+orOHg8ZagyxGREFMABOCeD1xCNGI8oK0AEQmQAiAAE3PTuXPpNJ7aXKP7BotIYBQAAfn8jTNJi0X53gvaChCRYCgAAjI+O43PXFfGui21bK05GXQ5IhJCCoAArb5hJuOzU/lfz+wkflM1EZGRowAIUE56Cl+++VLe3N/Ac9sPB12OiISMAiBgH68oYdbEbL7x7C7aOrqCLkdEQkQBELBYNMJ9H57DgeMtrH19f9DliEiIKABGgRtmFXHDrCK+98JejjS1Bl2OiISEAmCU+Pptc2nr7OIfntkRdCkiEhIKgFGibHwW99x4Cc9sreOVPbp1pIgMPwXAKLL6xhlMH5/F/3h6G63tnUGXIyJJTgEwiqTFovzDinkcON7C/9F1gkRkmCkARpnrysfz54tKePjlara8dzLockQkiSkARqGv/ec5FGWn8V9/sUW7gkRk2CgARqG8jBS+9dH5VB09xXef3xN0OSKSpBQAo9QNs4q4Y8lUHnm1mo37G4IuR0SSkK8AMLNlZrbbzKrMbE0P683MHvDWbzWzRf2NNbN/NLNdXv9fmVn+kMwoidz34cuYWpDJF37+FidOtwVdjogkmX4DwMyiwIPAcmAOcIeZzenWbTlQ7j1WAQ/5GPs8MM85Nx/YA3x10LNJMtlpMR785CKOnWrj757coiuGisiQ8rMFsASocs5VO+fagMeBFd36rADWurgNQL6ZTe5rrHPut865Dm/8BqBkCOaTdOYV5/HfbpnNCzuP8uhr+4IuR0SSiJ8AKAbeS1iu8dr89PEzFuAzwLM9fbiZrTKzSjOrrK8P5zdkV15Txp/Omci3/mMXmw6cCLocEUkSfgLAemjrvi+itz79jjWz+4AO4Kc9fbhz7hHnXIVzrqKoqMhHucnHzPjHjy5gSn4Gn/vxJuoazwRdkogkAT8BUAOUJiyXALU++/Q51sxWArcCf+G0g7tPeZkp/OCuClrbO1m1dpO+HyAig+YnADYC5WY23cxSgduBdd36rAPu8s4GWgo0Oufq+hprZsuAvwduc861DNF8ktqsiTl87xML2VbbyFee3KqDwiIyKP0GgHeg9l7gOWAn8IRzbruZrTaz1V639UA1UAX8ALi7r7HemO8DOcDzZva2mT08dNNKXjfNmcjf/umlrNtSy3df0PWCROTixfx0cs6tJ/5LPrHt4YTXDrjH71iv/ZIBVSrn3X3jTPYfO80DL+6lKCeNv1w6LeiSRGQM8hUAMrqYGd/4L5fTcLqNrz29jcKsVG65fHLQZYnIGKNLQYxRsWiE739yEVeU5vPFx9/mtb3Hgi5JRMYYBcAYlpEa5bFPXcmMoiw++6ONCgERGRAFwBiXn5nKz/56KdPHKwREZGAUAEmgIOvCEPj97qNBlyQiY4ACIEmcC4GZRdn81Y8qeWpzTdAlicgopwBIIgVZqfzb55ayZHoBX35iC//y8rv6spiI9EoBkGRy0lP4109fya3zJ/ONZ3fxP9dtp72zK+iyRGQU0vcAklBaLMoDt1/BpNx0fvjaPvYcOcWDf7GIgqzUoEsTkVFEWwBJKhIx/vutc/jOxxaw6eAJbvv+a+ysawq6LBEZRRQASe6ji0t44nNX097ZxZ/98x/4+ZsHdVxARAAFQCgsLM3nN39zHRXTCvjqU+9w908309jSHnRZIhIwBUBITMhJZ+1nlrBm+Wye33GE5fe/oi+NiYScAiBEIhFj9Q0zefLz15CWEuXOR9/gK09u0daASEgpAEJoYWk+z37helbfMJNfbj7ETd99mWffqdOxAZGQUQCEVHpKlDXLZ/P0PddSlJ3G53+6mTsffYNdh3WmkEhYKABCbl5xHk/fey1fv20u22ubuOX+V7nvV+9w7NTZoEsTkWFmY2mzv6KiwlVWVgZdRtI62dLG917Yy483HCAtFmHlNWWsun4G4/QFMpExzcw2Oecq3teuAJDu3q0/xf0v7OU3W2vJSo3x6WvL+My10xUEImOUAkAGbM+RZu5/YS///k4d6SkRPrq4hM9cO50ZRdlBlyYiA6AAkIu250gzj766j1+9dYi2zi5uumwCf3l1GdddMp5oxIIuT0T6oQCQQatvPstPNhzgJxsOcPx0G8X5GXysooSPVZRSnJ8RdHki0gsFgAyZsx2dvLDjKI9vPMhrVfFvE189o5Bb509h2bxJuuqoyCijAJBhUXOihV9U1rBuSy37jp0mGjGumVnIrfMn88HZEynKSQu6RJHQUwDIsHLOsaOuiX/fWsczW+s42NACwPySPG6cVcSNsyewoCRfxwxEAqAAkBHjnGN7bRO/332U3++uZ/PBE3Q5yM9MYUlZAVfNKOSq6QVcNjlXgSAyAhQAEpiTLW28svcYr+yp5819Dee3DnLSY1xZVsDC0nwuL8ljfnEehdnaZSQy1HoLAN0SUoZdfmYqty2Ywm0LpgBQe/IMG/c3sKG6gTf3Hed3u49y7u+Q4vwMFpTmMXdKHpdOzGHWxBxKxmUQ0ZaCyJDTFoAErrm1nW2Hmnjn0Em21jTyzqFGDhxvOb8+IyXKJROyKZ+YzayJOZQVZjG1IJNphZlkpelvGJH+aAtARq2c9BSunlnI1TMLz7c1tbaz98gp9h5pZveRZvYeOcVre4/x1OZDF4wdn51KaUEm0woymVqQyeT8DCblpjMpL51JuenkZ6Zgpq0HkZ4oAGRUyk1PYfG0cSyeNu6C9saWdg42tHCg4TQHG1o4eLyFgw0tbNx/gnVbaunqtkGbFoswKS+dibnpFOWkUZiVSkFWKoXZCa+95/zMVB2UllDxFQBmtgy4H4gCP3TOfbPbevPW3wK0AJ9yzm3ua6yZFQD/BpQB+4GPO+dODH5KkszyMlO4PDOPy0vy3reuvbOL+uaz1DW2cqSplcONrRw+99zYys7aJo6fbqPxTM93QDOD7LQYuekp5KTHn3MzYuSkp5Cb7j17y5mpUTJTY2SmRklPiXrLUTK89oyUqMJERr1+A8DMosCDwM1ADbDRzNY553YkdFsOlHuPq4CHgKv6GbsGeNE5900zW+Mt//3QTU3CJiUaYUp+BlP6uSxFe2cXJ063cfx0Gw3nnk+dpaGlnaYz7TS3dtDU2k5zazu1J1tpam2mubWD5tb2921h9CU1FokHQ0qUtJQoqdEIqTHvEY2Q4j2nXdBmpEaj5/ulxSKkRI1YJEIsakTMiEWMaMSIRY1oJEIsktAe/eP6qF3YJxqJr4t468wg4j2bGZHEZbotd1vfVz8ZO/xsASwBqpxz1QBm9jiwAkgMgBXAWhc/orzBzPLNbDLxv+57G7sCuNEb/yPg9ygAZASkRCNMyE1nQm76gMY55zjd1knTmXZa2jo409ZFS1sHLe2dnGmLP+KvO2hp6+SM197S1klreyftnV20dXTR5j23nOmML3d0nm9r73ReW7zfWNRTcJwLinMBcT4m7IKnP67vrT3hc6zbm7x/zLnlvt+z+/v5Hdd9Dn3xE4v9hef//rPLWTK9wMc7+ecnAIqB9xKWa4j/ld9fn+J+xk50ztUBOOfqzGxCTx9uZquAVQBTp071Ua7I8DAzstNiZI/QmUfOufPB0NUFHV1ddHY5Orocnd6j4/xzz306uhxd55e7zrd3OUdXFzigyzmcczgHXS5hGejqcnS5eD/nvHEOr28f4xL6JY6Lz8t7xl2wnDjv+Pqe+1/YduEy3d6zt8/q9b3ft763Wi5c3xdfG40+OmWlRf2804D4+ZfcUyx1L7e3Pn7G9sk59wjwCMRPAx3IWJGxzMxIi0VJiw39//gi4O+ewDVAacJyCVDrs09fY494u4nwno/6L1tERAbLTwBsBMrNbLqZpQK3A+u69VkH3GVxS4FGb/dOX2PXASu91yuBpwc5FxERGYB+dwE55zrM7F7gOeKncj7mnNtuZqu99Q8D64mfAlpF/DTQT/c11nvrbwJPmNlngYPAx4Z0ZiIi0iddCkJEJMn1dikIP7uAREQkCSkARERCSgEgIhJSCgARkZAaUweBzaweOHCRw8cDx4awnLFAcw4HzTkcBjPnac65ou6NYyoABsPMKns6Cp7MNOdw0JzDYTjmrF1AIiIhpQAQEQmpMAXAI0EXEADNORw053AY8jmH5hiAiIhcKExbACIikkABICISUqEIADNbZma7zazKu//wmGdmpWb2OzPbaWbbzewLXnuBmT1vZnu953EJY77q/Qx2m9l/Cq76wTGzqJm9ZWbPeMtJPWfvFqtPmtku77/31SGY85e8f9fbzOznZpaebHM2s8fM7KiZbUtoG/AczWyxmb3jrXvABnJjZnf+tm7J+SB+Gep3gRlAKrAFmBN0XUMwr8nAIu91DrAHmAN8G1jjta8BvuW9nuPNPQ2Y7v1MokHP4yLn/mXgZ8Az3nJSz5n4PbP/ynudCuQn85yJ30p2H5DhLT8BfCrZ5gz8CbAI2JbQNuA5Am8CVxO/A+OzwHK/NYRhC+D8Te2dc23AuRvTj2nOuTrn3GbvdTOwk/j/OCuI/8LAe/6I93oF8Lhz7qxzbh/xezcsGdGih4CZlQAfBn6Y0Jy0czazXOK/KB4FcM61OedOksRz9sSADDOLAZnE7ySYVHN2zr0CNHRrHtAcvbsp5jrnXnfxNFibMKZfYQiA3m5YnzTMrAy4AngDmOjid2PDe57gdUuWn8P3gK8AXQltyTznGUA98K/ebq8fmlkWSTxn59wh4DvEbxRVR/wOg78lieecYKBzLPZed2/3JQwBMOgb049mZpYN/BL4onOuqa+uPbSNqZ+Dmd0KHHXObfI7pIe2MTVn4n8JLwIecs5dAZwmvmugN2N+zt5+7xXEd3VMAbLM7M6+hvTQNqbm7ENvcxzU3MMQAH5uaj8mmVkK8V/+P3XOPeU1H/E2C/Gej3rtyfBzuBa4zcz2E9+V90Ez+wnJPecaoMY594a3/CTxQEjmOd8E7HPO1Tvn2oGngGtI7jmfM9A51nivu7f7EoYA8HNT+zHHO9L/KLDTOfdPCavWASu91yuBpxPabzezNDObDpQTP3g0ZjjnvuqcK3HOlRH/7/iSc+5OknvOh4H3zOxSr+lDwA6SeM7Ed/0sNbNM79/5h4gf40rmOZ8zoDl6u4mazWyp97O6K2FM/4I+Ej5CR9tvIX6WzLvAfUHXM0Rzuo74pt5W4G3vcQtQCLwI7PWeCxLG3Of9DHYzgDMFRuMDuJE/ngWU1HMGFgKV3n/rXwPjQjDnrwO7gG3Aj4mf/ZJUcwZ+TvwYRzvxv+Q/ezFzBCq8n9O7wPfxrvDg56FLQYiIhFQYdgGJiEgPFAAiIiGlABARCSkFgIhISCkARERCSgEgIhJSCgARkZD6/9E2u/BjgGvRAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "a = 1.0\n",
    "beta = 0.01\n",
    "ndim = X.shape[1]\n",
    "x = np.linspace(0.01, 1000, int(1e5))\n",
    "plt.plot(x, gamma.pdf(x, a, scale=1 / beta))\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def joint_energy(q_hmc, q_other):\n",
    "    return (\n",
    "        -gamma.logpdf(q_other, a, scale=1 / beta)\n",
    "        - jnp.sum(norm.logpdf(q_hmc, scale=1 / jnp.sqrt(q_other)))\n",
    "        - jnp.sum(bernoulli.logpmf(y, p=expit(jnp.dot(X, q_hmc))))\n",
    "    )\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def sample_q_other(q_hmc, key):\n",
    "    key, subkey = jax.random.split(key)\n",
    "    q_other = jax.random.gamma(subkey, a + 0.5 * q_hmc.shape[0]) / (\n",
    "        beta + 0.5 * jnp.sum(q_hmc ** 2)\n",
    "    )\n",
    "    return q_other, key\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def get_step_size(q_other, epsilon):\n",
    "    return epsilon / jnp.sqrt(q_other)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "66ad39ef",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "(\n",
    "    mala_within_gibbs,\n",
    "    hmc_within_gibbs,\n",
    "    mahmc_within_gibbs,\n",
    "    mala_persistent_within_gibbs,\n",
    "    mala_persistent_nonreversible_within_gibbs,\n",
    ") = make_samplers(joint_energy, sample_q_other, get_step_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7f2c9216",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_samples, n_warm_up_samples = int(1e5), int(1e4)\n",
    "L = 5\n",
    "n_chains = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "738c2a5f",
   "metadata": {},
   "source": [
    "# MALA within Gibbs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0dee2f7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 6978.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "### MALA-within-Gibbs, epsilon: 0.11\n",
      "\n",
      "energies_ess: 0.006790891102627784\n",
      "acceptance: 0.8209933638572693\n",
      "accuracy: 0.9876977205276489\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "epsilon = 0.11\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "@functools.partial(jax.vmap, in_axes=(0, None), out_axes=0)\n",
    "def get_mala_within_gibbs_samples(key, epsilon):\n",
    "    def scan_f(carry, ii):\n",
    "        beta, tau, key = carry\n",
    "        beta, tau, key, accept = mala_within_gibbs(beta, tau, key, epsilon, L)\n",
    "        energy = joint_energy(beta, tau)\n",
    "        return (beta, tau, key), (beta, tau, energy, accept)\n",
    "\n",
    "    tau = 150.0\n",
    "    key, subkey = jax.random.split(key)\n",
    "    beta = jax.random.normal(subkey, shape=(ndim,)) / jnp.sqrt(tau)\n",
    "    _, samples = jax.lax.scan(scan_f, (beta, tau, key), jnp.arange(n_samples))\n",
    "    samples = jax.tree_util.tree_map(lambda x: x[n_warm_up_samples:], samples)\n",
    "    return samples\n",
    "\n",
    "\n",
    "key = jax.random.PRNGKey(np.random.randint(int(1e5)))\n",
    "keys = jax.random.split(key, n_chains)\n",
    "samples = get_mala_within_gibbs_samples(keys, epsilon)\n",
    "print(\n",
    "    f\"\"\"\n",
    "### MALA-within-Gibbs, epsilon: {epsilon}\n",
    "\n",
    "energies_ess: {get_min_ess(samples[2][..., None].copy())}\n",
    "acceptance: {np.mean(samples[3])}\n",
    "accuracy: {\n",
    "    np.sum(\n",
    "        y == (\n",
    "            np.mean(\n",
    "                expit(jnp.dot(samples[0][0], X.T)) >= 0.5,\n",
    "                axis=0\n",
    "            ) >= 0.5\n",
    "        ).astype(int)\n",
    "    ) / y.shape[0]\n",
    "}\n",
    "\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4930e8b0",
   "metadata": {},
   "source": [
    "# HMC-within-Gibbs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "772f4db1",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 7503.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "### HMC-within-Gibbs, N: 2, epsilon: 0.09\n",
      "\n",
      "energies_ess: 0.03945755561291712\n",
      "acceptance: 0.881588876247406\n",
      "accuracy: 0.9876977205276489\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "N = 2\n",
    "epsilon = 0.09\n",
    "\n",
    "\n",
    "@functools.partial(jax.jit, static_argnames=\"N\")\n",
    "@functools.partial(jax.vmap, in_axes=(0, None, None), out_axes=0)\n",
    "def get_hmc_within_gibbs_samples(key, N, epsilon):\n",
    "    def scan_f(carry, ii):\n",
    "        beta, tau, key = carry\n",
    "        beta, tau, key, accept = hmc_within_gibbs(beta, tau, key, epsilon, L * N)\n",
    "        energy = joint_energy(beta, tau)\n",
    "        return (beta, tau, key), (beta, tau, energy, accept)\n",
    "\n",
    "    tau = 150.0\n",
    "    key, subkey = jax.random.split(key)\n",
    "    beta = jax.random.normal(subkey, shape=(ndim,)) / jnp.sqrt(tau)\n",
    "    _, samples = jax.lax.scan(scan_f, (beta, tau, key), jnp.arange(n_samples))\n",
    "    samples = jax.tree_util.tree_map(lambda x: x[n_warm_up_samples:], samples)\n",
    "    return samples\n",
    "\n",
    "\n",
    "key = jax.random.PRNGKey(np.random.randint(int(1e5)))\n",
    "keys = jax.random.split(key, n_chains)\n",
    "samples = get_hmc_within_gibbs_samples(keys, N, epsilon)\n",
    "print(\n",
    "    f\"\"\"\n",
    "### HMC-within-Gibbs, N: {N}, epsilon: {epsilon}\n",
    "\n",
    "energies_ess: {get_min_ess(samples[2][..., None].copy()) / N}\n",
    "acceptance: {np.mean(samples[3])}\n",
    "accuracy: {\n",
    "    np.sum(\n",
    "        y == (\n",
    "            np.mean(\n",
    "                expit(jnp.dot(samples[0][0], X.T)) >= 0.5,\n",
    "                axis=0\n",
    "            ) >= 0.5\n",
    "        ).astype(int)\n",
    "    ) / y.shape[0]\n",
    "}\n",
    "\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "123b94d2",
   "metadata": {},
   "source": [
    "# MAHMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "11246c0b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 5440.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "### MAHMC-within-Gibbs, N: 2, epsilon: 0.1\n",
      "\n",
      "energies_ess: 0.047817248469230286\n",
      "acceptance: 0.8319000005722046\n",
      "accuracy: 0.9876977205276489\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "N = 2\n",
    "epsilon = 0.10\n",
    "\n",
    "\n",
    "@functools.partial(jax.jit, static_argnames=\"N\")\n",
    "@functools.partial(jax.vmap, in_axes=(0, None, None), out_axes=0)\n",
    "def get_mahmc_within_gibbs_samples(key, N, epsilon):\n",
    "    def scan_f(carry, ii):\n",
    "        beta, tau, key = carry\n",
    "        beta, tau, key, accept = mahmc_within_gibbs(beta, tau, key, epsilon, L, N)\n",
    "        energy = joint_energy(beta, tau)\n",
    "        return (beta, tau, key), (beta, tau, energy, accept)\n",
    "\n",
    "    tau = 150.0\n",
    "    key, subkey = jax.random.split(key)\n",
    "    beta = jax.random.normal(subkey, shape=(ndim,)) / jnp.sqrt(tau)\n",
    "    _, samples = jax.lax.scan(scan_f, (beta, tau, key), jnp.arange(n_samples))\n",
    "    samples = jax.tree_util.tree_map(lambda x: x[n_warm_up_samples:], samples)\n",
    "    return samples\n",
    "\n",
    "\n",
    "key = jax.random.PRNGKey(np.random.randint(int(1e5)))\n",
    "keys = jax.random.split(key, n_chains)\n",
    "samples = get_mahmc_within_gibbs_samples(keys, N, epsilon)\n",
    "accuracy = (\n",
    "    np.sum(\n",
    "        y\n",
    "        == (np.mean(expit(jnp.dot(samples[0][0], X.T)) >= 0.5, axis=0) >= 0.5).astype(\n",
    "            int\n",
    "        )\n",
    "    )\n",
    "    / y.shape[0]\n",
    ")\n",
    "print(\n",
    "    f\"\"\"\n",
    "### MAHMC-within-Gibbs, N: {N}, epsilon: {epsilon}\n",
    "\n",
    "energies_ess: {get_min_ess(samples[2][..., None].copy()) / N}\n",
    "acceptance: {np.mean(samples[3])}\n",
    "accuracy: {\n",
    "    np.sum(\n",
    "        y == (\n",
    "            np.mean(\n",
    "                expit(jnp.dot(samples[0][0], X.T)) >= 0.5,\n",
    "                axis=0\n",
    "            ) >= 0.5\n",
    "        ).astype(int)\n",
    "    ) / y.shape[0]\n",
    "}\n",
    "\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaecf545",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "# MALA with persistent momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b092386d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 7371.36it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "### MALA-P-within-Gibbs, epsilon: 0.09, alpha: 0.9\n",
      "\n",
      "energies_ess: 0.03631890030070464\n",
      "acceptance: 0.9016111493110657\n",
      "accuracy: 0.9876977205276489\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "epsilon = 0.09\n",
    "alpha = 0.9\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "@functools.partial(jax.vmap, in_axes=(0, None, None), out_axes=0)\n",
    "def get_mala_persistent_within_gibbs_samples(key, epsilon, alpha):\n",
    "    def scan_f(carry, ii):\n",
    "        beta, p, tau, key = carry\n",
    "        beta, p, tau, key, accept = mala_persistent_within_gibbs(\n",
    "            beta, p, tau, key, epsilon, L, alpha\n",
    "        )\n",
    "        energy = joint_energy(beta, tau)\n",
    "        return (beta, p, tau, key), (beta, tau, energy, accept)\n",
    "\n",
    "    tau = 150.0\n",
    "    key, subkey = jax.random.split(key)\n",
    "    beta = jax.random.normal(subkey, shape=(ndim,)) / jnp.sqrt(tau)\n",
    "    key, subkey = jax.random.split(key)\n",
    "    p = jax.random.normal(subkey, shape=beta.shape)\n",
    "    _, samples = jax.lax.scan(scan_f, (beta, p, tau, key), jnp.arange(n_samples))\n",
    "    samples = jax.tree_util.tree_map(lambda x: x[n_warm_up_samples:], samples)\n",
    "    return samples\n",
    "\n",
    "\n",
    "key = jax.random.PRNGKey(np.random.randint(int(1e5)))\n",
    "keys = jax.random.split(key, n_chains)\n",
    "samples = get_mala_persistent_within_gibbs_samples(keys, epsilon, alpha)\n",
    "print(\n",
    "    f\"\"\"\n",
    "### MALA-P-within-Gibbs, epsilon: {epsilon}, alpha: {alpha}\n",
    "\n",
    "energies_ess: {get_min_ess(samples[2][..., None].copy())}\n",
    "acceptance: {np.mean(samples[3])}\n",
    "accuracy: {\n",
    "    np.sum(\n",
    "        y == (\n",
    "            np.mean(\n",
    "                expit(jnp.dot(samples[0][0], X.T)) >= 0.5,\n",
    "                axis=0\n",
    "            ) >= 0.5\n",
    "        ).astype(int)\n",
    "    ) / y.shape[0]\n",
    "}\n",
    "\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3a825c0",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "# MALA with persistent momentum and non-reversible Metropolis accept/reject"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b85694d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 6864.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "### MALA-P-N-within-Gibbs, epsilon: 0.1, alpha: 0.9, delta: 0.015\n",
      "\n",
      "energies_ess: 0.04087781447592402\n",
      "acceptance: 0.8686622381210327\n",
      "accuracy: 0.9876977205276489\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "epsilon = 0.1\n",
    "alpha = 0.9\n",
    "delta = 0.015\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "@functools.partial(jax.vmap, in_axes=(0, None, None, None), out_axes=0)\n",
    "def get_mala_persistent_nonreversible_within_gibbs_samples(key, epsilon, alpha, delta):\n",
    "    def scan_f(carry, ii):\n",
    "        beta, p, tau, v, key = carry\n",
    "        beta, p, tau, v, key, accept = mala_persistent_nonreversible_within_gibbs(\n",
    "            beta, p, tau, v, key, epsilon, L, alpha, delta\n",
    "        )\n",
    "        energy = joint_energy(beta, tau)\n",
    "        return (beta, p, tau, v, key), (beta, tau, energy, accept)\n",
    "\n",
    "    tau = 150.0\n",
    "    key, subkey = jax.random.split(key)\n",
    "    beta = jax.random.normal(subkey, shape=(ndim,)) / jnp.sqrt(tau)\n",
    "    key, subkey = jax.random.split(key)\n",
    "    p = jax.random.normal(subkey, shape=beta.shape)\n",
    "    key, subkey = jax.random.split(key)\n",
    "    v = jax.random.uniform(subkey) * 2 - 1\n",
    "    _, samples = jax.lax.scan(scan_f, (beta, p, tau, v, key), jnp.arange(n_samples))\n",
    "    samples = jax.tree_util.tree_map(lambda x: x[n_warm_up_samples:], samples)\n",
    "    return samples\n",
    "\n",
    "\n",
    "key = jax.random.PRNGKey(np.random.randint(int(1e5)))\n",
    "keys = jax.random.split(key, n_chains)\n",
    "samples = get_mala_persistent_nonreversible_within_gibbs_samples(\n",
    "    keys, epsilon, alpha, delta\n",
    ")\n",
    "print(\n",
    "    f\"\"\"\n",
    "### MALA-P-N-within-Gibbs, epsilon: {epsilon}, alpha: {alpha}, delta: {delta}\n",
    "\n",
    "energies_ess: {get_min_ess(samples[2][..., None].copy())}\n",
    "acceptance: {np.mean(samples[3])}\n",
    "accuracy: {\n",
    "    np.sum(\n",
    "        y == (\n",
    "            np.mean(\n",
    "                expit(jnp.dot(samples[0][0], X.T)) >= 0.5,\n",
    "                axis=0\n",
    "            ) >= 0.5\n",
    "        ).astype(int)\n",
    "    ) / y.shape[0]\n",
    "}\n",
    "\n",
    "\"\"\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}