If/else statement vs Heaviside function - python-3.x

Question

In my code I have to consider different contributions with respect to different thresholds. In particular I have a function my_index whose output must be compared to the thresholds Z_1, Z_2, Z_3 in order to determine the increment to the variable my_value. In the following MWE, for simplicity sake, the function my_index is just a uniform random generator:
import numpy as np
my_len = 100000
Z_1 = 0.2
Z_2 = 0.4
Z_3 = 0.7
first = 1
second = 2
third = -0.0003
my_value = 0
for i in range(my_len):
my_index = np.random.uniform()
my_value += first*np.heaviside(my_index - Z_1,0)*np.heaviside(Z_2 - my_index,0) + second*np.heaviside(my_index - Z_3,0) + third*np.heaviside(Z_3 - my_index,0)
#if Z_1 < my_index < Z_2 add first
#if my_index > Z_3 add second
#if my_index < Z_3 add third
I have replaced if/else's that could have been used for the thresholds with the Heaviside function see. Keep in mind that, in my original code, this code section has to be iterated up to 10^5 times.
My question is: does this practice make the code faster? Or is the heaviside function (np.heaviside) call better in terms of speed than the if/else control?

Answer 1

In [433]: x=np.arange(-10,10)
In [434]: x
Out[434]:
array([-10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2,
3, 4, 5, 6, 7, 8, 9])
A proper use of heaviside - giving x as array, not a single value:
In [436]: np.heaviside(x,.5)
Out[436]:
array([0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0.5, 1. , 1. ,
1. , 1. , 1. , 1. , 1. , 1. , 1. ])
A list comprehension equivalent:
In [437]: [.5 if i==0 else (0 if i<0 else 1) for i in x]
Out[437]: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.5, 1, 1, 1, 1, 1, 1, 1, 1, 1]
and making an array from that list:
In [438]: np.array([.5 if i==0 else (0 if i<0 else 1) for i in x])
Out[438]:
array([0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0.5, 1. , 1. ,
1. , 1. , 1. , 1. , 1. , 1. , 1. ])
Compare the times:
In [439]: timeit np.heaviside(x,.5)
2.5 µs ± 17.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [440]: timeit np.array([.5 if i==0 else (0 if i<0 else 1) for i in x])
15.1 µs ± 25.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
Iteration on a list is faster (than on an array):
In [441]: timeit np.array([.5 if i==0 else (0 if i<0 else 1) for i in x.tolist()])
6.66 µs ± 195 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
and if we skip the conversion back to a list:
In [442]: timeit [.5 if i==0 else (0 if i<0 else 1) for i in x.tolist()]
2.28 µs ± 3.01 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
For a much larger array, heaviside performance is even better:
In [445]: x=np.arange(-1000,1000)
In [446]: timeit [.5 if i==0 else (0 if i<0 else 1) for i in x.tolist()]
211 µs ± 7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [447]: timeit np.heaviside(x,.5)
13 µs ± 201 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
For the random number generation, taking the whole-array approach is also faster:
In [448]: timeit [np.random.uniform() for _ in range(1000)]
4.62 ms ± 20.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [449]: timeit np.random.uniform(1000)
4.74 µs ± 171 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
I could also time the scalar use of heaviside - that is worse than the if/else in [446]:
In [450]: timeit [np.heaviside(i,.5) for i in x]
8.64 ms ± 44.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In sum:
use whole-array code where possible
when using Python level iteration, use lists and scalar methods instead

Answer 2

Assuming you use the standard CPython interpreter, then performing a Numpy function call like np.heaviside is likely more expensive than doing basic conditionals. However, both are very inefficient. Indeed, conditionals are generally slow and could be replaced with a branchless implementation here (adding/multiplying booleans converted to integers). The most important optimization is to use vectorization because Numpy is design to be efficient on relatively big arrays and not scalar values (mainly due to additional internal checks and function calls). You can generate all the random value in a big array, apply the heaviside function on it multiple times. The resulting code will certainly be 2 or 3 order of magnitude faster!