Sample a discrete distribution (i.e. simulate a loaded die) efficiently via the "alias sampling" method

discreternd.py

"""
Alias sampler
=============

See : <http://www.keithschwarz.com/darts-dice-coins/> by Keith Schwarz.

References
----------

+ Vose, Michael D. 
  "A linear algorithm for generating random numbers with a given distribution." 
  IEEE Transactions on software engineering 17.9 (1991): 972-975.

"""
from __future__ import division, print_function
import numpy as np
unifrnd = np.random.random_sample # uniform in [0, 1)
randint = np.random.randint       # samples zero-based, half-open range


class LoadedDie(object):
    def __init__(self, weights):
        weights = np.asarray(weights, dtype=float)
        self.prob, self.alias = self._build(weights)

    def _build(self, x):
        assert np.all(x >= 0) and np.any(x)

        # normalize
        x_total = np.sum(x)
        if not np.isclose(x_total, 1.0):
            x /= x_total

        N = len(x)
        prob = np.zeros(N, dtype=float)
        alias = np.zeros(N, dtype=int)

        # rescale probabilities x and fill the stacks
        x = N*x
        small = []
        large = []
        for i, p in enumerate(x):
            if p < 1:
                small.append(i)
            else:
                large.append(i)

        # normally, small empties first
        while small and large:
            # assign an alias to l
            l = small.pop()
            g = large.pop()
            prob[l] = x[l]
            alias[l] = g

            # trim off g's donated area and put it back in one of the stacks
            x[g] = (x[g] + x[l]) - 1.0 # more stable than x[g] - (1 - x[l])
            if x[g] < 1.0:
                # possible false assignment if x[g]/N should be 1.0, but falls below
                # due to numerical instability
                small.append(g)
            else:
                large.append(g)

        # for remaining bins, no need to alias
        while large:
            g = large.pop()
            prob[g] = 1.0

        # check in case large emptied before small due to numerical problems
        while small:
            l = small.pop()
            prob[l] = 1.0

        return prob, alias

    def roll(self, n=None):
        if n is None:
            outcome = randint(len(self.prob))
            return self.alias[outcome] if unifrnd() > self.prob[outcome] else outcome
        else:
            outcomes = randint(0, len(self.prob), n)
            missed = unifrnd(n) > self.prob[outcomes]
            outcomes[missed] = self.alias[outcomes[missed]]
            return outcomes