python - 使用 matplotlib 绘制分段函数会导致 ValueError : The truth value of an array with more than one element is ambiguous

Question

我想使用 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))

我收到以下错误:

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

但是，在使用 (x < 0).any() 时，我只是得到 True ，在 (x < 0).all() 上我得到 False ，这两者都不适合我的情况。

我在这里找到了答案: Python Error: Using matplotlib: Truth value of an array with more than one element is ambiguous. Use a.any() or a.all() ，但是问题和答案都充满了不必要的代码，因此我决定发布一个简洁的解决方案。

Answer 1

您可以使用

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

或者，对于适用于 0 和 1 以外的数字的解决方案，

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

此类问题的一般解决方案是将函数向量化:

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()