矩阵向量差的有效元素 argmin

Posted 2023-02-24

【中文标题】矩阵向量差的有效元素 argmin

【英文标题】：Efficient elementwise argmin of matrix-vector difference

【发布时间】：2021-02-07 23:48:39

【问题描述】：

假设一个数组a.shape == (N, M) 和一个向量v.shape == (N,)。目标是从a 的每个元素 中减去abs 的abs 和v 的argmin - 也就是说，

out = np.zeros(N, M)
for i in range(N):
    for j in range(M):
        out[i, j] = np.argmin(np.abs(a[i, j] - v))

我通过np.matlib.repmat 有一个vectorized implementation，它更快，但占用O(M*N^2) 内存，在实践中是不可接受的。计算仍然在 CPU 上完成，所以最好的办法似乎是在 C 中实现 for 循环作为扩展，但也许 Numpy 已经实现了这个逻辑。

是吗？是否有任何可用的 Numpy 函数有效地实现上述？

【问题讨论】：

@Gulzar for 循环是，但我的矢量化实现复制了aN 次，为v 中的每个元素复制一个副本。添加它以供参考。

如果v 中有两个候选人与a[i,j] 等距怎么办。如果我们选择一个而不是另一个有关系吗？

@Divakar 不，随便选一个。如果你打算实现这个，我怀疑你可以用 Python 管理一些东西，除非 Numpy 公开了 C 工具；我可以（如果在此处找不到任何内容，最终会）在 C 或 C++ 中打开一个关于此的代码审查问题。

@Divakar 好吧，200k 代表 - 我可能会大吃一惊。

是的，这是你的惊喜 - ***.com/a/64526158。

【参考方案1】：

受this post 启发，我们可以利用np.searchsorted -

def find_closest(a, v):
    sidx = v.argsort()
    v_s = v[sidx]
    idx = np.searchsorted(v_s, a)
    idx[idx==len(v)] = len(v)-1
    idx0 = (idx-1).clip(min=0)
    
    m = np.abs(a-v_s[idx]) >= np.abs(v_s[idx0]-a)
    m[idx==0] = 0
    idx[m] -= 1
    out = sidx[idx]
    return out

更多性能。在大型数据集上使用 numexpr 提升：

import numexpr as ne

def find_closest_v2(a, v):
    sidx = v.argsort()
    v_s = v[sidx]
    idx = np.searchsorted(v_s, a)
    idx[idx==len(v)] = len(v)-1
    idx0 = (idx-1).clip(min=0)
    
    p1 = v_s[idx]
    p2 = v_s[idx0]
    m = ne.evaluate('(idx!=0) & (abs(a-p1) >= abs(p2-a))', 'p1':p1, 'p2':p2, 'idx':idx)
    idx[m] -= 1
    out = sidx[idx]
    return out

时间

设置：

N,M = 500,100000
a = np.random.rand(N,M)
v = np.random.rand(N)

In [22]: %timeit find_closest_v2(a, v)
4.35 s ± 21.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [23]: %timeit find_closest(a, v)
4.69 s ± 173 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

【讨论】：

对于一般功能，而不是np.argmin，您将如何做？只是搜索相关索引，而不是np.repeat?

@Gulzar 不，它特定于 argmin。没有通用的解决方案。

@OverLordGoldDragon 是的，不关心等距很重要。

Credited.