Trying to plot the triangular function in python using arrays

Question

My question is: Write a vectorized version of the triangle function. The function should take the NumPy array as input and returns the vectorized array. You should also display the triangle using matplot.

I'm struggling to vectorize the function

t(x) = 0, x < 0; x, 0 <= x <1; 2-x, 1 <= x < 2, 0, x >= 2 t(x) = 0, x < 0; x, 0 <= x <1; 2-x, 1 <= x < 2, 0, x >= 2 .

Not even sure how to put that in an array tbh

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
def triangle (x):
    return np.where(x<0, 0, 1)
plt.ylim(0,1)
plt.xlim(0,8)
plt.show(x)

Answer 1

As I understand, the idea is to not use if statement.

import numpy as np
import matplotlib.pyplot as plt

def my_fun(x):
    y = np.zeros_like(x)
    mask = (x >= 0) & (x < 1)
    y[mask] = x[mask]
    mask = (x >= 1) & (x < 2)
    y[mask] = 2. - x[mask]
    return y

x = np.random.rand(1000) * 5 - 1.
y = my_fun(x)

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

I'm using masking to select values which are modified. Also x values have to be masked, so that the calculations are done properly. Note that plt.show() normally does not take arguments and is only showing what you have plotted before.