Bayesian update example for DAT-500

visualize.py

import numpy as np
import matplotlib.pyplot as plt

from matplotlib.animation import FuncAnimation
from numpy.random import default_rng
from scipy.stats import beta


class UpdateProbabilityDistribution:

    def __init__(self, ax, prob=0.5):
        self.rng = default_rng()
        self.success = 0
        self.prob = prob
        self.line, = ax.plot([], [], 'k-')
        self.x = np.linspace(0, 1, 200)
        self.ax = ax
    
        self.ax.set_title("Bayesian Modeling the Likelihood of Me Working Out\nOn Any Given Day")
        self.ax.set_xlim(left=0.0, right=1.0)
        self.ax.set_ylim(bottom=0.0, top=10.0)
        self.ax.set_xlabel("Probability")
        self.ax.set_ylabel("Probability Density")
        self.ax.grid(True)

        self.ax.axhline(y=0, color='black', linewidth=2.0)

    def __call__(self, i):
        if i == 0:
            self.success = 1
            self.line.set_data([], [])

            return self.line,
    
        self.success += self.rng.binomial(1, self.prob)

        distribution = beta(self.success + 1, i + 1 - self.success)
        y = distribution.pdf(self.x)

        #self.ax.set_ylim(bottom=0.0, top=y.max()*(1.0 + (1.0 / i)))
        self.line.set_data(self.x, y)

        return self.line,


fig, ax = plt.subplots()

probability_distribution_updater = UpdateProbabilityDistribution(ax, prob=0.28)
animation = FuncAnimation(fig, probability_distribution_updater, frames=90, interval=250, blit=True)
animation.save("bayesian-update.mp4", fps=4)
plt.show()