更快地矢量化 np.unique 以获得多维索引矩阵中的项目计数

Question

我对计算索引在形状为 (size, 3) 的矩阵中出现的次数感兴趣，其中 size 是一个很大的数字。出于说明目的，在一个循环中，它看起来像：

max_ = 30
size = 300000
idx_matrix = np.random.randint(0, max_, (size, 3))
result = np.zeros((max_, max_, max_), dtype=int)
for i in range(len(idx_matrix)): 
    result[tuple(idx_matrix[i, :])] += 1 
result = result / np.sum(result) # Normalization

我使用 np.unique 和 return_counts=True 进行了以下矢量化实现，但是对于我的目的来说速度不够快，因为我将执行此操作数百次。

import numpy as np

def foo():
    max_ = 30
    size = 300000
    idx_matrix = np.random.randint(0, max_, (size, 3))
    result = np.zeros((max_, max_, max_), dtype=int)

    unique, counts = np.unique(idx_matrix, axis=0, return_counts=True)
    result[unique[:, 0], unique[:, 1], unique[:, 2]] = counts
    return result / np.sum(result)

>>> %timeit foo()
    464 ms ± 22.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

我怎样才能进一步改进这段代码，使它花费的时间比半秒少得多？

Answer 1

您可以使用 at method of the numpy.add ufunc 执行此操作。下面我把foo()的生成去掉了idx_matrix，所以你可以验证foo()和foo2()return是一样的结果，等等时间不包括 idx_matrix.

的生成

def foo(idx_matrix, max_):
    result = np.zeros((max_, max_, max_), dtype=int)
    unique, counts = np.unique(idx_matrix, axis=0, return_counts=True)
    result[unique[:, 0], unique[:, 1], unique[:, 2]] = counts
    return result / len(idx_matrix)


def foo2(idx_matrix, max_):
    result = np.zeros((max_, max_, max_), dtype=int)
    np.add.at(result, tuple(idx_matrix.T), 1)
    return result / len(idx_matrix)


max_ = 30
size = 300000
idx_matrix = np.random.randint(0, max_, (size, 3))

这是 foo(idx_matrix, max_) 和 foo2(idx_matrix, max_) 的时间：

In [64]: %timeit foo(idx_matrix, max_)
349 ms ± 5.12 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [65]: %timeit foo2(idx_matrix, max_)
26.3 ms ± 456 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

更快地矢量化 np.unique 以获得多维索引矩阵中的项目计数

Faster vectorization of np.unique to get the item count in a matrix of multidimensional indices

numpy

vectorization

python-3.x