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

Numpy-查找矩阵乘法的最近邻居 辅助函数可沿轴获取最大，最小的n索引元素

这是一个辅助函数，

n-largest

可使用

np.argpartition

和

np.take_along_axis

-从通用ndarray沿通用轴选择顶部索引。

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 setupIn [200]: np.random.seed(0)     ...: ar = np.random.randint(0,99,(5,5))In [201]: arOut[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) # indicestake_largest_along_axis(similarities, n=10, axis=1) # elements

最终的索引/值数组将为shape

(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).shapeOut[263]: (1000, 10)In [264]: take_largest_along_axis(similarities, n=10, axis=1).shapeOut[264]: (1000, 10)

相反，如果您希望获得

n

每个功能的最大数据点，则最终的index / values数组将为shape

(n, 128)

，请使用

axis=0

。