python - 向量化可以应用于由两个或多个函数定义的函数 Z= f(X,Y)

Question

我正在使用 np.meshgrid() 绘制一个标量函数:Z= f(X,Y)。

当 f(X,Y) 由单个方程定义时，我没有问题。但是，当 f(X,Y) 拆分为两个或多个方程时，例如:

X, Y = np.meshgrid(x, y, sparse = True, indexing = 'ij')
Z = f1(X,Y), y1 < y <= y2
Z = f2(X,Y), y > y2

我不知道如何从 meshgrid 数组中提取 'Y' 的值，以便在条件语句中根据 y1 和 y2 进行检查。我尝试了几种方法，如下所示，无济于事:

row1, col1 = np.where(Y == y1)
row2, col2 = np.where(Y == y2)
if( Y > Y[row1][0] and Y <= Y[row2][0] ):
    Z = f1(X, Y)
elif ( Y > Y[row2][0] ):
    Z = f2(X, Y)

当我使用数组计算 Z 时，与矢量化相比，执行速度非常慢，至少慢 10 到 20 倍，具体取决于网格大小。
在这种情况下是否可以应用矢量化？如何？
谢谢!

Answer 1

使用 np.where一般来说。

z = np.where(Y <= y2, f1(X, Y), f2(X, Y))

如果你有一个空情况(就像 Y < y1 一样)你可能需要用 np.zeros_like 初始化

Z = np.zeros_like(Y)  # same shape and dtype as `Y`
np.where(y1 < Y <= y2, f1(X, Y), out = Z)
np.where(y< y2, f2(X,Y), out = Z)

请注意，如果您的函数之一引发超出其范围的错误，您将收到自 np.where 起的错误。计算整个 X, Y 的函数范围。在这种情况下，您将要尝试 @Stef 的方法