"""
This example includes a function with an if-test, which
is one of the (few) mathematical functions
which will not automatically work if it is passed
a numpy array as argument. We show two possible solutions
to this problem."""

import numpy as np
import matplotlib.pyplot as plt

#The Heaviside function (piecewise constant)
def H(x):
    if x < 0:
        return 0
    else:
        return 1

x = np.linspace(-5,5,101)

"""
Now, using the standard
y = H(x)
will fail, since the comparison x < 0 returns an array
True/False values, which cannot be used with if.
"""

"""
Remedy nr 1, use the builtin numpy function vectorize.
This is simple, but the code is slow for large arrays:
"""
Hv1 = np.vectorize(H)
y1 = Hv1(x)

"""
Remedy nr 2, change the if-test into using the
numpy function 'where'. This solution will be
much faster for large arrays:
"""
def Hv2(x):
    return np.where(x<0,0,1)

y2 = Hv2(x)

plt.plot(x,y2)
plt.show()