Numpy - 查找矩阵乘法的最近邻

Question

我有一个包含一千个 128 维特征的数据集，形状例如(1000,128)。

我想找到形状为 (128,1) 的 128 维特征的排序最近邻。

通过数据集 (1000,128) 和特征 (128,1) 之间的矩阵乘法计算的距离，这将给出 (1000,1) 形状的相似性数组：

数据集 (1000,128) x 特征 (128,1) = 相似度 (1000,1)

这是通过以下方式完成的：

# features.shape=(1000,128) ; feature.shape=(128,1) ; similarities.shape=(1000,1)
similarities = features.dot(feature)

计算距离（相似度）后，我正在使用以下代码查找最近的邻居：

# The n Nearest Neighbors Indexes (But Not Sorted)
nearest_neighbours_indexes_unsorted = np.argpartition(similarities, kth=-n)[-n:]

# The n Nearest Neighbors (But Not Sorted)
nearest_neighbours_similarities_unsorted = similarities[nearest_neighbours_indexes_unsorted]

# The Indexes of n Nearest Neighbors Sorted
nearest_neighbours_indexes_sorted = np.flip(nearest_neighbours_indexes_unsorted[np.argsort(nearest_neighbours_similarities_unsorted)], axis=0)

此代码对数百万个数据的运行速度非常快（如果有人有使其更快的提示，我很感兴趣）但我希望能够一次性找到多个特征的最近邻居：

数据集 (1000,128) x 特征 (128,n) = 相似度 (1000,n)

一种方法是为循环中的每个功能计算上述代码（这很慢），另一种方法是更改​​代码以适应多维索引，这就是我卡住的地方：我不知道如何为 (128,n) 而不是 (128,1) 形状的特征编写上述代码。

Answer 1

获取最大、最小 n 索引、沿轴的元素的辅助函数

这是一个辅助函数，用于从使用 np.argpartition 和 np.take_along_axis 的通用 ndarray 中沿通用轴选择顶部 n-largest 索引 -

def take_largest_indices_along_axis(ar, n, axis):    
    s = ar.ndim*[slice(None,None,None)]
    s[axis] = slice(-n,None,None)
    idx = np.argpartition(ar, kth=-n, axis=axis)[tuple(s)]
    sidx = np.take_along_axis(ar,idx, axis=axis).argsort(axis=axis)
    return np.flip(np.take_along_axis(idx, sidx, axis=axis),axis=axis)

扩展它以获得 n 最小的索引 -

def take_smallest_indices_along_axis(ar, n, axis):    
    s = ar.ndim*[slice(None,None,None)]
    s[axis] = slice(None,n,None)
    idx = np.argpartition(ar, kth=n, axis=axis)[tuple(s)]
    sidx = np.take_along_axis(ar,idx, axis=axis).argsort(axis=axis)
    return np.take_along_axis(idx, sidx, axis=axis)

并扩展这些以选择最大或最小的n 元素本身，这将是np.take_along_axis 的简单用法，如下所列 -

def take_largest_along_axis(ar, n, axis):
    idx = take_largest_indices_along_axis(ar, n, axis)
    return np.take_along_axis(ar, idx, axis=axis)

def take_smallest_along_axis(ar, n, axis):
    idx = take_smallest_indices_along_axis(ar, n, axis)
    return np.take_along_axis(ar, idx, axis=axis)

样本运行

# Sample setup
In [200]: np.random.seed(0)
     ...: ar = np.random.randint(0,99,(5,5))

In [201]: ar
Out[201]: 
array([[44, 47, 64, 67, 67],
       [ 9, 83, 21, 36, 87],
       [70, 88, 88, 12, 58],
       [65, 39, 87, 46, 88],
       [81, 37, 25, 77, 72]])

取最大的n 索引，沿轴的元素-

In [202]: take_largest_indices_along_axis(ar, n=2, axis=0)
Out[202]: 
array([[4, 2, 2, 4, 3],
       [2, 1, 3, 0, 1]])

In [203]: take_largest_indices_along_axis(ar, n=2, axis=1)
Out[203]: 
array([[4, 3],
       [4, 1],
       [2, 1],
       [4, 2],
       [0, 3]])

In [251]: take_largest_along_axis(ar, n=2, axis=0)
Out[251]: 
array([[81, 88, 88, 77, 88],
       [70, 83, 87, 67, 87]])

In [252]: take_largest_along_axis(ar, n=2, axis=1)
Out[252]: 
array([[67, 67],
       [87, 83],
       [88, 88],
       [88, 87],
       [81, 77]])

取最小的n 索引，沿轴的元素-

In [232]: take_smallest_indices_along_axis(ar, n=2, axis=0)
Out[232]: 
array([[1, 4, 1, 2, 2],
       [0, 3, 4, 1, 0]])

In [233]: take_smallest_indices_along_axis(ar, n=2, axis=1)
Out[233]: 
array([[0, 1],
       [0, 2],
       [3, 4],
       [1, 3],
       [2, 1]])

In [253]: take_smallest_along_axis(ar, n=2, axis=0)
Out[253]: 
array([[ 9, 37, 21, 12, 58],
       [44, 39, 25, 36, 67]])

In [254]: take_smallest_along_axis(ar, n=2, axis=1)
Out[254]: 
array([[44, 47],
       [ 9, 21],
       [12, 58],
       [39, 46],
       [25, 37]])

---

在这里解决我们的问题

对于我们的例子，假设输入是similarities，形状为(1000,128)，代表1000个数据点和128个特征，我们想为每个数据点寻找最大的n=10特征，然后应该是-

take_largest_indices_along_axis(similarities, n=10, axis=1) # indices
take_largest_along_axis(similarities, n=10, axis=1) # elements

最终的索引/值数组的形状为(1000, n)。

使用给定的数据集形状运行示例 -

In [257]: np.random.seed(0)
     ...: similarities = np.random.randint(0,99,(1000,128))

In [263]: take_largest_indices_along_axis(similarities, n=10, axis=1).shape
Out[263]: (1000, 10)

In [264]: take_largest_along_axis(similarities, n=10, axis=1).shape
Out[264]: (1000, 10)

如果您希望为每个特征获得n 的最大数据点，那么最终的索引/值数组的形状将是(n, 128)，然后使用axis=0。