时间戳

首页 > 解决方案 > 向量化是否可以应用于由两个或多个函数定义的函数 Z= f(X,Y)

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

问题描述

我正在使用 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 倍，具体取决于网格大小。在这种情况下是否可以应用矢量化？如何？谢谢！

标签： pythonnumpy

解决方案

一般使用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 的方法

推荐阅读