binary-sum.py

import numpy as np
import timeit
from matplotlib import pyplot as plt

n = 1000
t1s, t2s = [], []
for p in np.linspace(0, 1, num=20):
    y = np.random.binomial(1, p, n)
    x = np.random.rand(n)
    t1 = timeit.timeit('np.sum(x * y)', 'import numpy as np; from __main__ import x, y', number=100000)
    t2 = timeit.timeit('np.sum(x[y > 0])', 'import numpy as np; from __main__ import x, y', number=100000)
    t1s.append(t1)
    t2s.append(t2)

plt.plot(np.linspace(0, 1, num=20), t1s, label='np.sum(x * y)')
plt.plot(np.linspace(0, 1, num=20), t2s, label='np.sum(x[y > 0])')
plt.xlabel('p')
plt.ylabel('t')
plt.legend()
plt.show()