A/B test with python (Bandit)
In this article, I am going to write a Bayesian Bandit algorithm.
import matplotlib.pyplot as plt import numpy as np from scipy.stats import beta NUM_TRIALS = 2000 BANDIT_PROBABILITIES = [0.2, 0.5, 0.75]
Bandit probabilities are divided into three value just in case. Quick favorite.
Trial number is 2000.
So, I am going to define a class called Bandit.
class Bandit: def __init__(self, p): self.p = p self.a = 1 self.b = 1 def pull(self): # arm of slot machine return np.random.random() < self.p def sample(self): # sample from current beta distribution return np.random.beta(self.a, self.b) def update(self, x): self.a += x self.b += 1 - x
This works like slot machine.
a, b are the beta parameters defined a uniform distribution in the beginning and p is the probability of winning.
This class should have a ability to pull the slot machine defined definition pull() returns the random number and we can get also sample from its current distribution.
Lastly, we have also update function. x is either 0 or 1.
def plot(bandits, trial): x = np.linspace(0, 1, 200) for b in bandits: y = beta.pdf(x, b.a, b.b) plt.plot(x, y, label="real p: %.4f" % b.p) plt.title("Bandit distributions after %s trials" % trial) plt.legend() plt.show()
This is just going to plot PDF of each bandit and we can compare it on the same chart.
OK. Then, we need to run a actual experiment.
def experiment(): bandits = [Bandit(p) for p in BANDIT_PROBABILITIES] sample_points = [5, 10, 20, 50, 100, 200, 500, 1000, 1500, 1999] for i in range(NUM_TRIALS): bestb = None maxsample = -1 allsamples = [] for b in bandits: sample = b.sample() allsamples.append("%.4f" % sample) if sample > maxsample: maxsample = sample bestb = b if i in sample_points: print("Current samples: %s" % allsamples) plot(bandits, i) x = bestb.pull() bestb.update(x) if __name__ == '__main__': experiment()
Bandit should be initialize at first.
Sample points are where we want to reasonably show a plot. Any values are ok.
We also keep track of the maximum sample and best will be.