{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code is taken from https://matplotlib.org/examples/animation/index.html and extended for the normally distributed case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import beta\n",
    "import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "\n",
    "class UpdateDistbeta(object):\n",
    "    def __init__(self, ax, prob=0.5):\n",
    "        self.success = 0\n",
    "        self.prob = prob\n",
    "        self.line, = ax.plot([], [], 'k-',color='blue')\n",
    "        self.data_point = ax.text(0.02, 0.95, '', transform=ax.transAxes)\n",
    "        self.success_rate = ax.text(0.02, 0.90, '', transform=ax.transAxes)\n",
    "        self.bayes_estimator = ax.text(0.02, 0.85, '', transform=ax.transAxes)\n",
    "        self.x = np.linspace(0, 1, 200)\n",
    "        self.ax = ax\n",
    "\n",
    "        # Set up plot parameters\n",
    "        self.ax.set_xlim(0, 1)\n",
    "        self.ax.set_ylim(0, 15)\n",
    "        self.ax.grid(True)\n",
    "\n",
    "        # This vertical line represents the theoretical value, to\n",
    "        # which the plotted distribution should converge.\n",
    "        self.ax.axvline(prob, linestyle='--', color='red')\n",
    "\n",
    "    def init(self):\n",
    "        self.success = 0\n",
    "        self.line.set_data([], [])\n",
    "        return self.line,\n",
    "\n",
    "    def __call__(self, i):\n",
    "        # This way the plot can continuously run and we just keep\n",
    "        # watching new realizations of the process\n",
    "        if i == 0:\n",
    "            return self.init()\n",
    "\n",
    "        # Choose success based on exceed a threshold with a uniform pick\n",
    "        data = 0\n",
    "        if np.random.rand(1,) < self.prob:\n",
    "            self.success += 1\n",
    "            data = 1\n",
    "        y = beta.pdf(self.x, self.success + 1, (i - self.success) + 1)\n",
    "        self.line.set_data(self.x, y)\n",
    "        rate = self.success/float(i)\n",
    "        bayes_estimator = (self.success+1)/(i+2)\n",
    "        self.data_point.set_text('Data point = %i' % data)\n",
    "        self.success_rate.set_text('MLE for success rate = %.5f' % rate)\n",
    "        self.bayes_estimator.set_text('Bayes estimator = %.5f' % bayes_estimator)\n",
    "        return self.line,\n",
    "\n",
    "# Fixing random state for reproducibility\n",
    "np.random.seed(19680801)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ud = UpdateDistbeta(ax, prob=0.3)\n",
    "animbeta = FuncAnimation(fig, ud, frames=np.arange(100), init_func=ud.init,\n",
    "                     interval=100, blit=True)\n",
    "animbeta.save('moviebeta.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import norm\n",
    "import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "\n",
    "class UpdateDistnormal(object):\n",
    "    def __init__(self, ax, mu0, sigma0, sigma):\n",
    "        self.mu0 = mu0\n",
    "        self.sigma0 = sigma0\n",
    "        self.sigma = sigma \n",
    "        self.data = np.random.normal(self.mu0,self.sigma0,1)\n",
    "        self.cumdata = 0\n",
    "        self.line, = ax.plot([], [], 'k-',color='blue')\n",
    "        self.data_point = ax.text(0.02, 0.95, '', transform=ax.transAxes)\n",
    "        self.cumdata_point = ax.text(0.02, 0.90, '', transform=ax.transAxes)\n",
    "        self.true_value = ax.text(0.02, 0.85, '', transform=ax.transAxes)\n",
    "        self.bayes_estimator = ax.text(0.02,0.80, '',transform=ax.transAxes)\n",
    "        self.x = np.linspace(-10, 10, 1000)\n",
    "        self.ax = ax\n",
    "\n",
    "        # Set up plot parameters\n",
    "        self.ax.set_xlim(-10, 10)\n",
    "        self.ax.set_ylim(0, 5)\n",
    "        self.ax.grid(True)\n",
    "        self.ax.axvline(self.data, linestyle='--', color='red')\n",
    "\n",
    "        # This vertical line represents the theoretical value, to\n",
    "        # which the plotted distribution should converge.\n",
    "        \n",
    "\n",
    "    def init(self):\n",
    "        self.line.set_data([], [])\n",
    "        return self.line,\n",
    "\n",
    "    def __call__(self, i):\n",
    "        # This way the plot can continuously run and we just keep\n",
    "        # watching new realizations of the process\n",
    "        if i == 0:\n",
    "            return self.init()\n",
    "\n",
    "        # Choose success based on exceed a threshold with a uniform pick\n",
    "        data = np.random.normal(self.data,self.sigma,1)\n",
    "        self.cumdata = self.cumdata + data\n",
    "        sigmaupdate = (self.sigma0**(-2)+self.sigma**(-2))**(-0.5)\n",
    "        self.mu0 = (data/self.sigma**2+self.mu0/self.sigma0**2)*sigmaupdate**2\n",
    "        self.sigma0 = sigmaupdate\n",
    "        y = norm.pdf(self.x, self.mu0, self.sigma0)\n",
    "        self.line.set_data(self.x, y)\n",
    "        average = self.cumdata/i\n",
    "        self.data_point.set_text('Data point = %.5f' % data)\n",
    "        self.cumdata_point.set_text('MLE  = %.5f' % average)\n",
    "        self.true_value.set_text('True value = %.5f' % self.data)\n",
    "        self.bayes_estimator.set_text('Bayes estimator = %.5f' % self.mu0)\n",
    "        return self.line,\n",
    "\n",
    "# Fixing random state for reproducibility\n",
    "np.random.seed(19680802)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ud = UpdateDistnormal(ax, mu0 = 5.0, sigma0 = 1.0, sigma = 5.0)\n",
    "animnormal= FuncAnimation(fig, ud, frames=np.arange(100), init_func=ud.init,\n",
    "                     interval=100, blit=True)\n",
    "animnormal.save('movienormal.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}