我猜您正在寻找numpy.where：

import matplotlib.pyplot as plt
import numpy as np

ampl = 2
wl = 3
xmax = 9

def square_wave(x, wave_length, amplitude):
    return np.where(
        x % wave_length <= wave_length / 2, amplitude, -amplitude
    )

x = np.linspace(0, xmax, 100)
plt.plot(x, square_wave(x=x, wave_length=wl, amplitude=ampl))
plt.axhline(color="black")
plt.show()

你的问题是，你在整个数组x上应用f，一般来说这是个好主意，因为numpy的关键就是通过这样做来加速计算（python是一种非常慢的语言，但是numpy不是用python写的，所以numpy能为你做的每一批计算都快得多）。
因此，您在这里体验到的是numpy的广播和矢量化功能：一段代码，f，很明显，你写的时候考虑到了标量，几乎可以把参数x作为一个数组来使用。
x%T是一个数组，其每个元素都是x[i]%T，其中i在range(len(x))中。
x%T<=T/2是一个布尔数组，其每个元素为真或假，具体取决于是否为x[i]%T<=T/2。
然后它就停止工作了，因为你可以用一个布尔数组作为if的条件，if的条件必须是一个truth值：要么是真，要么不是。你得到的是一个数组，一些是真，一些是假。因此你得到了错误：

ValueError: The truth value of an array with more than one element is ambiguous.

现在您有3个解决方案，从最简单的到需要最多调整的，但也从最慢的到最快的。
1.您可以直接在标量上调用f，因此构建一个数组y（如y[i]=f(x[i])）并将其传递给plot

y = np.zeros_like(x)
for i in range(len(x)):
    y[i]=f(x[i])
plt.plot(x,y)

你的f函数，按原样，运行良好。只是调用出错了。它是用标量调用的，你用vector调用了它。用这个for循环，我更正了一下，用标量调用了它（n次）。
1.你可以让numpy用所谓的vectorize函数来做这件事。
vectorize将一个要在标量（或其他对象，但我感兴趣的是这种情况）上调用的函数转换为可以用任何形状的数组调用的函数，然后返回一个相同形状的数组，其所有元素都是f(z)，因为z是x的任何元素
所以在你的情况下

plt.plot(x, np.vectorize(f)(x))

但这是一个错误的好主意。所有这些它都和我的第一个版本一样，但它为你做了for循环。
尽管如此，它仍然更快。而且仍然相当容易编写（至少在这样一个规范的情况下。使用vectorize有时可能是相当大的挑战）。
1.向量化
但大多数情况下，你也不例外，你可以自己对函数进行向量化，也就是说，重写f来处理一个标量向量，而不是一个标量向量，我们在序言中已经看到，一半的工作已经完成了：你已经计算了一个布尔数组，告诉你返回元素应该是I还是-I。
例如

def f(x):
    return np.where(x%T<=T/2, I, -I)

plt.plot(x, f(x))

x%T<=T/2是我们之前提到的布尔数组，np.where返回一个相同形状的数组，整数值取自第二个或第三个参数。
另一种在数学中更常见的方法是

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

（基于numpy中，用于算术运算的布尔值为0或1这一事实）
时间
在我的计算机上，第一个解决方案（使用for循环）每次运行需要308毫秒。
第二个（使用向量化），每次运行159毫秒，所以快了两倍，因为for循环是用numpy而不是python完成的，但是，for循环中的代码仍然是python代码。
第三个（np.where或算术版本）每次运行需要13 ms。
所以，你看，它不是一个真正的选项。你需要使用第三个选项。一般来说，当使用numpy时，你必须不惜一切代价避免for循环，并尽量把每个操作看作是对数组的操作，而不是对数组中的元素的操作。我提供了3个选项来帮助你理解发生了什么，以及需要做什么。但实际上，只有这3个解是有效的（即使你不关心计算时间，最好在使用numpy时就开始这样想）