# 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.