Sampling with weighted probabilities In an exam question I need to output some numbers self.random_nums with a certain probability self.probabilities:

I have written the following program that successfully returns the correct answer and also a test at the bottom which confirms that everything is working well.

Here my questions:


Are there any better ways that this problem could be solved? Does my
program go against certain conventions that I should regard?
Why is the unittest ignoring the rounding?
Are there any style issues in my program?


Any other comments are appreciated

from __future__ import division  # in case of python 2
import numpy as np
from matplotlib import pyplot as plot


class RandomGen(object):
    def __init__(self):
        self.random_nums = np.array([-1, 0, 1, 2, 3])
        self.probabilities = np.array([0.01, 0.3, 0.58, 0.1, 0.01])

        # generate cumulative probabilities and add a 0 to the beginning and remove the 1 at the end
        self.cum_probabiliies = np.cumsum(self.probabilities)
        self.cum_probabiliies = np.delete(self.cum_probabiliies, self.cum_probabiliies.size - 1)
        self.cum_probabiliies = np.append([0], self.cum_probabiliies)

    def next_num(self, i):
        num = 0
        while num &lt; i:
            rnd = np.random.uniform(low=0.00, high=1)  # between 0 and 0.99 (less than the high)
            x = (self.cum_probabiliies[self.cum_probabiliies &lt;= rnd]).size - 1
            number = self.random_nums[x]
            yield number
            num += 1

    def get_expected_probabilities(self):
        print("Expected probabilities:")
        target = dict(zip(self.random_nums, self.probabilities))
        target = {k: round(v, 2) for k, v in target.items()}
        print(target)
        return target

    def get_actual_probabilities(self, n):
        # create histogram with corresponding buckets
        hist = np.histogram(n, bins=np.append(self.random_nums, max(self.random_nums + 1)))

        s = sum(hist[0])
        print("Actual Probabilities:")
        effective_probabilities = hist[0] / s
        actual = dict(zip(hist[1], effective_probabilities))
        actual = {k: round(v, 2) for k, v in actual.items()}
        print(actual)
        return actual


if __name__ == "__main__":
    r = RandomGen()
    n = list(r.next_num(10000))
    print("Output:")
    print(n)
    print()
    r.get_expected_probabilities()
    print()
    r.get_actual_probabilities(n)

    plot.hist(n)
    plot.show()


As test I suggest the following:

import unittest
import man_exercise


class TestRandomGen(unittest.TestCase):
    def test_next_num(self):
        iterations = 10000
        delta = iterations / 50000

        r = man_exercise.RandomGen()
        n = list(r.next_num(iterations))

        target = r.get_expected_probabilities() # dictionary with numbers and their expected probabilities
        effective = r.get_actual_probabilities(n) # dictionary with numbers and their actual probabilities

        for k,_ in target.items():
            self.assertAlmostEqual(target[k], effective[k], delta=delta)
Allow native iteration

It is reasonable to expect to be able to iterate over a RandomGen[erator], but the built-in for loop fails right now.

All you need to allow looping with a standard for loop is renaming next_num to __iter__

Use native iteration

    num = 0
    while num &lt; i:
        # ... snip ...
        num += 1


Is not nice, you are manually incrementing a counter in a high level language. Just use (x)range (depending on the Python version):

for num in range(i):
    # ... snip ...