Python绘制方波遇阻：无法在函数中判断Numpy数组x的值，如何解决？

你的代码核心问题在于：自定义函数f(x)用了普通的if条件判断，但传入的x是Numpy数组类型。对数组执行x % T <= T/2会返回一个布尔数组，而普通if语句无法处理这种数组，因此会报错。

下面提供三种可行的解决方法：

方法一：用np.vectorize()包装函数 #

把普通函数转换成能批量处理Numpy数组的矢量化函数，直接适配数组输入：

import matplotlib.pyplot as plt
import numpy as np

I = float(input())
t = float(input())
T = float(input())

def f(x):
    if x % T <= (T / 2):
        return I
    else:
        return -I

# 将普通函数矢量化，支持Numpy数组调用
f_vec = np.vectorize(f)
x = np.linspace(0, t, 1000)
plt.plot(x, f_vec(x))
plt.xlim(0, t)
plt.ylim(-10, 10)
plt.axhline(color="black")
plt.show()

方法二：用np.where()直接做批量判断 #

利用Numpy原生向量化操作，无需循环或if，效率更高：

import matplotlib.pyplot as plt
import numpy as np

I = float(input())
t = float(input())
T = float(input())

x = np.linspace(0, t, 1000)
# 批量判断：满足条件返回I，否则返回-I
y = np.where(x % T <= T/2, I, -I)
plt.plot(x, y)
plt.xlim(0, t)
plt.ylim(-10, 10)
plt.axhline(color="black")
plt.show()

方法三：用布尔索引赋值 #

先创建全为-I的数组，再将满足条件的位置替换为I：

import matplotlib.pyplot as plt
import numpy as np

I = float(input())
t = float(input())
T = float(input())

x = np.linspace(0, t, 1000)
y = np.full_like(x, -I)
# 布尔索引定位目标元素并赋值
y[x % T <= T/2] = I
plt.plot(x, y)
plt.xlim(0, t)
plt.ylim(-10, 10)
plt.axhline(color="black")
plt.show()

提示：方法二和方法三是原生Numpy向量化操作，比方法一的np.vectorize()效率更高（np.vectorize()本质是底层循环），处理大规模数据时更推荐后两种。

内容的提问来源于stack exchange，提问作者Kikibord_2009