python - 高效地创建面具 - Numpy/Python

Question

假设我有一个形状为 (50, 10000, 10000) 的 NumPy 数组，其中包含 1000 个不同的“簇”。例如，某处会有一个只有 1s 的小体积，另一个有 2s 的小体积，等等。我想遍历每个集群来创建一个掩码，如下所示:

for i in np.unique(arr)[1:]:
    mask = arr == i
    #do other stuff with mask

创建每个掩码大约需要 15 秒，迭代 1000 个集群需要 4 个多小时。有没有一种可能的方法来加速代码，或者这是最好的方法，因为无法避免遍历数组的每个元素？

编辑:数组的数据类型是 uint16

Answer 1

我假设 arr 是稀疏的:

• 你说簇很小，1000 个簇不会平铺那么大的阵列
• 您遍历 np.unique(arr)[1:]，所以我假设第一个唯一值是 0

在这种情况下，我建议利用 scipy.sparse.csr_matrix

from scipy.sparse import csr_matrix
sp_arr = csr_matrix(arr.reshape(1,-1))

这会将您的大密集数组变成单行压缩稀疏行数组。由于稀疏数组不喜欢超过 2 个维度，这会诱使它使用 ravelled 索引。现在 sp_arr 有 data(簇标签)、indices(拼凑的索引)和 indptr(这是这里微不足道，因为我们只有一行)。所以，

for i in np.unique(sp_arr.data):  # as a bonus this `unique` call should be faster too
    x, y, z = np.unravel_index(sp_arr.indices[sp_arr.data == i], arr.shape)

应该更有效地给出等效坐标

for i in np.unique(arr)[:1]:
    x, y, z = np.nonzero(arr == i)

其中 x, y, z 是 mask 中 True 值的索引。从那里您可以重建 mask 或处理索引(推荐)。

您也可以完全使用 numpy 来完成此操作，并且最后仍然有一个 bool 掩码，但内存效率稍低:

all_mask = arr != 0    # points assigned to any cluster
data = arr[all_mask]   # all cluster labels
for i in np.unique(data):
    mask = all_mask.copy()
    mask[mask] = data == i  # now mask is same as before