Trying to plot a picewise-defined function in Python

Question

Just trying to plot a piecewise defined function in python but having trouble with an error. "The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()". I understand that the issue is python cannot analyze the truth of an entire array at once, however I don't believe a.any() or a.all() resolves the problem. Here are my two attempts at the code:

x = linspace(0,3,1000)
def S(x):
 if 0 <= x <= 1:
 return -1 + x - (5/9)*x**2 + (5/9)*x**3
 elif 1 <= x <= 2:
 return (14/9)*(x-1) + (10/9)*(x-1)**2 - (2/3)*(x-1)**3
 elif 2 <= x <=3:
 return 2 + (16/9)*(x-2) - (8/9)*(x-2)**2 - (1/9)*(x-2)**3
plot(x, S(x), 'k')
show()

x = linspace(0,3,1000)
plot(x, piecewise(x, [x<=1, 1<=x<=2, x>=2], [lambda x: -1 + x - (5/9)*x**2 + (5/9)*x**3, 
 lambda x: (14/9)*(x-1) + (10/9)*(x-1)**2 - (2/3)*(x-1)**3,
 lambda x: 2 + (16/9)*(x-2) - (8/9)*(x-2)**2 - (1/9)*(x-2)**3]))

• Roy Smart: You need to use numpy.where to define a piecewise function like that.

Answer 1

Good morning!

The issue you encounter has to do with typing:

• x = linspace(0,3,1000) defines a np array (of type np.ndarray).
• Your function S takes in a x of type float.
• In plot(x, S(x), 'k'), you provides your function S with an np.ndarray. Thus, the first check fails since it tries to compare the np.ndarray x with 0 and 1: if 0 <= x <= 1 (with x being an array).

To solve that, you can use python list comprehension to apply S for each sample of your array x. It would give you the following:

import numpy as np
from matplotlib.pyplot import plot, show
x: np.ndarray = np.linspace(0,3,1000)
def S(x: float) -> float:
 if 0 <= x <= 1:
 return -1 + x - (5/9)*x**2 + (5/9)*x**3
 elif 1 <= x <= 2:
 return (14/9)*(x-1) + (10/9)*(x-1)**2 - (2/3)*(x-1)**3
 elif 2 <= x <=3:
 return 2 + (16/9)*(x-2) - (8/9)*(x-2)**2 - (1/9)*(x-2)**3
y: list[float] = [S(i) for i in x]
plot(x, y, 'k')
show()

Which outputs:

plot of piecewise function