NumPy 2D数组:选择圆圈中的索引

时间:2018-03-16 21:58:43

标签: python arrays numpy indexing

对于某些矩形,我们可以非常有效地选择2D数组中的所有索引:

arr[y:y+height, x:x+width]

...其中(x, y)是矩形的左上角,heightwidth矩形的高度(行数)和宽度(列数)选择。

现在,假设我们要在给定中心坐标(cx, cy)和半径r的某个圆中选择二维数组中的所有索引。有效地实现这一目标是否有一个笨拙的功能?

目前,我通过使用Python循环将索引添加到缓冲区(列表)中来手动预先计算索引。因此,对于大型2D阵列来说,这是非常无效的,因为我需要将位于某个圆圈中的每个整数排队。

# buffer for x & y indices
indices_x = list()
indices_y = list()

# lower and upper index range
x_lower, x_upper = int(max(cx-r, 0)), int(min(cx+r, arr.shape[1]-1))
y_lower, y_upper = int(max(cy-r, 0)), int(min(cy+r, arr.shape[0]-1))
range_x = range(x_lower, x_upper)
range_y = range(y_lower, y_upper)

# loop over all indices
for y, x in product(range_y, range_x):
    # check if point lies within radius r
    if (x-cx)**2 + (y-cy)**2 < r**2:
        indices_y.append(y)
        indices_x.append(x)

# circle indexing
arr[(indices_y, indices_x)]

如前所述,对于较大的阵列/圆圈,此过程效率非常低。有什么想法可以加快速度吗?

如果有更好的方法来索引圆圈,这是否也适用于“任意”2D形状?例如,我可以以某种方式传递一个函数来表示任意形状的点的成员资格,以获得数组的相应numpy索引吗?

1 个答案:

答案 0 :(得分:5)

您可以定义包含圆圈的蒙版。下面,我已经演示了一个圆圈,但你可以在mask赋值中编写任意函数。字段mask的尺寸为arr,如果满足右侧条件,则值True,否则为False。此掩码可以与索引运算符结合使用,以仅分配选择的索引,如行arr[mask] = 123.所示。

import numpy as np
import matplotlib.pyplot as plt

x = np.arange(0, 32)
y = np.arange(0, 32)
arr = np.zeros((y.size, x.size))

cx = 12.
cy = 16.
r = 5.

# The two lines below could be merged, but I stored the mask
# for code clarity.
mask = (x[np.newaxis,:]-cx)**2 + (y[:,np.newaxis]-cy)**2 < r**2
arr[mask] = 123.

# This plot shows that only within the circle the value is set to 123.
plt.figure()
plt.pcolormesh(x, y, arr)
plt.colorbar()
plt.show()