rng.py

import random
from itertools import accumulate

# random number generator that will not generate the same sequence
# of pseudorandom numbers if the `random` module is used elsewhere.
# the advantage is a 60x speedup
class Rng:
    def __init__(self, outputs, probabilities, seed=0):
        # test inputs for consistency
        if len(outputs) != len(probabilities):
            raise ValueError('Number of possible outputs does '
                             'not match number of probabilities')
        # test probabilities for consistency
        if(not (all(0 < p <= 1 for p in probabilities) and
                  abs(sum(probabilities) - 1) < 1E-9 )) :
            raise ValueError('Invalid probabilities')

        # protect against the user changing mutable sequences
        self.outputs = tuple(outputs)
        self.cum_probs = tuple(accumulate(probabilities))

        random.seed(seed)

    def nextNum(self):
        test = random.random()
        for number, cum_prob in zip(self.outputs, self.cum_probs):
            if test < cum_prob:
                return number

# random number generator that is guaranteed to produce the
# same sequence of random numbers, even if the `random` module
# is used elsewhere.
class RngSafe(Rng):
    def __init__(self, outputs, probabilities, seed=0):
        # isolate the state of our RNG from others
        old_state = random.getstate()
        super().__init__(outputs, probabilities, seed)
        self.state = random.getstate()
        random.setstate(old_state)

    def nextNum(self):
        old_state = random.getstate()
        random.setstate(self.state)
        try:
            return super().nextNum()
        finally:
            self.state = random.getstate()
            random.setstate(old_state)