python - 高效重复 numpy.where

Question

我有一个代码，我想在其中检查成对的坐标是否落入某些矩形。但是，有很多矩形，我不确定如何将以下代码推广到许多矩形。我只能在循环中使用 eval 来完成它，但这非常难看。

这是一个代码，用于检查由坐标组成的 DataFrame 的每个条目对应的矩形。如果它属于第一个，则分配 0，第二个分配 1，否则分配 nan。我想要这样的代码，它会产生类似的结果，假设我们有一个大的 Rectangle 对象列表，而不在最后一行应用 eval 或循环。非常感谢。

from matplotlib.patches import Rectangle

rec1 = Rectangle((0,0), 100, 100)
rec2 = Rectangle((100,0), 100, 100)
x = np.random.poisson(100, size=200)
y = np.random.poisson(80, size=200)
xy = pd.DataFrame({"x" : x, "y" : y}).values
e1 = np.asarray(rec1.get_extents())
e2 = np.asarray(rec2.get_extents())
r1m1, r1m2 = np.min(e1), np.max(e1)
r2m1, r2m2 = np.min(e2), np.max(e2)
out = np.where(((xy >= r1m1) & (xy <= r1m2)).all(axis=1), 0, 
               np.where(((xy >= r2m1) & (xy <= r2m2)).all(axis=1), 1, np.nan))

编辑这是一个有 3 个矩形的版本

rec1 = Rectangle((0,0), 100, 100)
rec2 = Rectangle((0,100), 100, 100)
rec3 = Rectangle((100,100), 100, 100)
x = np.random.poisson(100, size=200)
y = np.random.poisson(100, size=200)
xy = pd.DataFrame({"x" : x, "y" : y}).values
e1 = np.asarray(rec1.get_extents())
e2 = np.asarray(rec2.get_extents())
e3 = np.asarray(rec3.get_extents())
r1m1, r1m2 = np.min(e1), np.max(e1)
r2m1, r2m2 = np.min(e2), np.max(e2)
r3m1, r3m2 = np.min(e3), np.max(e3)
out = np.where(((xy >= r1m1) & (xy <= r1m2)).all(axis=1), 0, 
           np.where(((xy >= r2m1) & (xy <= r2m2)).all(axis=1), 1, 
           np.where(((xy >= r3m1) & (xy <= r3m2)).all(axis=1), 2, np.nan)))

我喜欢得到的值是 0、1、2 或 np.nan。但输出仅由 0 和 1 组成。

Answer 1

这是一个使用 NumPy broadcasting 的向量化方法-

# Store extents in a 3D array
e = np.dstack((e1,e2,e3))

# Get a valid mask for the X's and Y's and then the combined one
x_valid_mask = (xy[:,0] >= e[0,0,:,None]) & (xy[:,0] <= e[1,0,:,None])
y_valid_mask = (xy[:,1] >= e[0,1,:,None]) & (xy[:,1] <= e[1,1,:,None])
valid_mask = x_valid_mask & y_valid_mask

# Finally use argmax() to choose the rectangle each pt belongs. We can use
# argmax to choose the first matching one and that works here because
# we are guaranteed to have the recatnagles mutually exclusive
out = np.where(valid_mask.any(0), valid_mask.argmax(0), np.nan)

让我们运行一个示例来验证这里的东西 -

1) 设置随机输入:

In [315]: rec1 = Rectangle((0,0), 100, 100)
     ...: rec2 = Rectangle((0,100), 100, 100)
     ...: rec3 = Rectangle((100,100), 100, 100)
     ...: 

In [316]: e1 = np.asarray(rec1.get_extents())
     ...: e2 = np.asarray(rec2.get_extents())
     ...: e3 = np.asarray(rec3.get_extents())
     ...: 

2) 查看 rec3 的范围:

In [317]: e3
Out[317]: 
array([[ 100.,  100.],
       [ 200.,  200.]])

3) xy 随机获得 5 分:

In [319]: x = np.random.poisson(100, size=5)
     ...: y = np.random.poisson(100, size=5)
     ...: xy = pd.DataFrame({"x" : x, "y" : y}).values
     ...: 

4) 让我们设置 pt[1] 使其位于 rec3 中。所以，这个 pt 的 o/p 应该是 2。

In [320]: xy[1] = [150,175]

5) 让我们设置 pt[3] 使其位于所有矩形之外。所以，对应的 o/p 应该是 NaN。

In [321]: xy[3] = [400,400]

6) 运行发布的代码并打印输出:

In [323]: out
Out[323]: array([ nan,   2.,   2.,  nan,   2.])

可以看出out[1]是2，out[3]是NaN，这是预期的早些。