1

I want to plot a piece wise function using matplotlib:

import numpy as np
import matplotlib.pyplot as plt

def pwf(x):
    return 0 if x < 0 else 1
x = np.linspace(-1, 1, 100)
plt.plot(x, pwf(x))

and I get the following error:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

However, on using (x < 0).any() I simply get True and on (x < 0).all() I get False, neither of which is suitable for my case.

I found the answer here: Python Error: Using matplotlib: Truth value of an array with more than one element is ambiguous. Use a.any() or a.all(), but the question and answer are so bloated with unnecessary code that I decided to post a concise solution out of it.

banan3'14
  • 3,810
  • 3
  • 24
  • 47
  • Add a `print(x)` to `pwf`. And a `print(x<0)`. If `x` has more than one element, that won't work in a `if/else` clause. – hpaulj Jan 10 '20 at 00:44
  • related?: https://stackoverflow.com/questions/10062954/valueerror-the-truth-value-of-an-array-with-more-than-one-element-is-ambiguous – David Cary Sep 15 '20 at 20:42

2 Answers2

1

You can use

def pwf(x):
    return (x < 0).astype(float)

or, for a solution that would work with numbers other than 0 and 1,

def pwf(x):
    return np.array([1,0])[(x < 0).astype(int)]

The general solution to such kind of problem would be to vectorize the function:

import matplotlib.pyplot as plt
import numpy as np

def pwf(x):
    return 0 if x < 0 else 1

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

plt.plot(x, np.vectorize(pwf)(x))
plt.show()
ImportanceOfBeingErnest
  • 321,279
  • 53
  • 665
  • 712
0

A simple solution is to use Python's list comprehension:

x = np.linspace(-1, 1, 100)
y = [pwf(x) for x in x]
plt.plot(x, y)

Credit goes to @Mike Müller for his answer https://stackoverflow.com/a/34565170/7769052.

banan3'14
  • 3,810
  • 3
  • 24
  • 47