import numpy as np
import matplotlib.pyplot as plt

"""
Example from lecture 21 sept. For most functions
that take a number x as input we can also pass in
an entire array, and it automatically does the
calculations for all elements in the array and
returns an array of the same size. However, this
does not work for functions that contain if-tests,
since an expression like x < 0 will result in
an array with True/False values, which cannot be used
directly in an if-test. 
"""

#a simple implementation of the Heaviside function
def H(x):
    if x < 0:
        return 0
    else:
	    return 1

"""
This will not work, because of the if-test:
x = np.linspace(-2, 2, 100)
y = H(x)
Below are three alternative ways to solve this problem. 
"""

""" 
Fix 1; use a for-loop. This function takes an
array x as input, loops over it, and calls the
H(x) function for each element in x.
"""
def H_loop(x):
    y = np.zeros_like(x)  
    for i in range(len(x)):
         y[i] = H(x[i])
    return y

"""
Fix 2; use NumPy's vectorize function.
np.vectorize takes our original function as argument, 
and returns a modified function that accepts an array argument.
"""
H_vec = np.vectorize(H)

"""Fix 3; use np.where. This is the most efficient solution. 
np.where(a, v1, v2) takes three arguments. The first is an
array containing True/False values, and the other two are numbers.
It will return a new array with the same size as 'a', where all the
'True' elements are replaced by the value 'v1' and all the 'False'
elements by 'v2'.
"""
def H_pro(x):
    return np.where(x < 0, 0.0, 1.0)




x = np.linspace(-2, 2, 100)

#H_loop, H_vec, H_pro can all be used like this:
y = H_pro(x)

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

"""
Terminal> python heaviside.py
(output is a plot)
"""