Lesson 2

Since we consider discrete random variables, it is convenient to visualize the probability mass function using bar charts.

The code below generates such bar charts for four random variables:

1. the outcome of a single die roll;
2. the sum of numbers on two dice;
3. the minimum of the numbers on two dice;
4. the number of heads in twenty coin flips.

from itertools import product
import matplotlib.pyplot as plt
from numpy import unique

for description, sample_space, random_variable in (
        ('one-dice-value', list(range(1, 7)), lambda v: v),
        ('two-dice-sum', list(product(range(1, 7), repeat=2)), lambda v: sum(v)),
        ('two-dice-min', list(product(range(1, 7), repeat=2)), lambda v: min(v)),
        ('twenty-coins-heads', list(product(range(2), repeat=20)), lambda v: sum(v)),
):
    values = [random_variable(u) for u in sample_space]
    values, counts = unique(values, return_counts=True)
    plt.clf()
    plt.bar(values, counts / sum(counts))
    plt.savefig(f'histogram-{description}.png')