Efficient elementwise argmin of matrix-vector difference

Question

Suppose an array a.shape == (N, M) and a vector v.shape == (N,). The goal is to compute argmin of abs of v subtracted from every element of a - that is,

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))

I have a vectorized implementation via np.matlib.repmat, and it's much faster, but takes O(M*N^2) memory, unacceptable in practice. Computation's still done on CPU so best bet seems to be implementing the for-loop in C as an extension, but maybe Numpy already has this logic implemented.

Does it? Any use-ready Numpy functions implementing above efficiently?

@Gulzar The for-loop is, but my vectorized implementation replicates `a` `N` times, a copy for each element in `v`. Added it for reference. — OverLordGoldDragon, Oct 25 '20 at 15:15
What if there are two candidates in `v` that are at equidistant from `a[i,j]`. Does it matter if we choose one over another? — Divakar, Oct 25 '20 at 16:05
@Divakar Nope, pick any. If you plan on implementing this, I doubt you can manage something in Python unless Numpy exposes C tools; I can (and eventually will if nothing is found here) open a Code Review question on this, in C or C++. — OverLordGoldDragon, Oct 25 '20 at 16:11
Yup, here's your suprise - https://stackoverflow.com/a/64526158/. — Divakar, Oct 25 '20 at 16:36

Divakar · Accepted Answer · 2020-10-25T17:01:52.423

4

Inspired by this post, we can leverage 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

Some more perf. boost with numexpr on large datasets :

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

Timings

Setup :

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)

edited Oct 25 '20 at 17:01

answered Oct 25 '20 at 16:36

Divakar

218,885
19
262
358

How would you do it for a general function, not for `np.argmin`? just a search over relevant indices, rather than `np.repeat`? – Gulzar Oct 25 '20 at 16:43
@Gulzar No, its specific to argmin. There's no generic solution. – Divakar Oct 25 '20 at 16:43
I started reading about `np.nditer`, is it irrelevant? I would expect some way to iterate over what gets returned from `np.broadcast_arrays`, is there no such thing? – Gulzar Oct 25 '20 at 16:45
@Gulzar `np.nditer` is for looping, which we need to look to avoid. Or iteration of any kind is what I am looking to avoid. In most cases we can avoid `np.broadcast_arrays` with explicit extension of dims if needed and then let ufuncs operate. – Divakar Oct 25 '20 at 16:46
I don't know how this works (will look into later), but comparison against for-loop doesn't lie - excellent, thank you. I'm guessing the "don't care equidistant" assumption was important? – OverLordGoldDragon Oct 25 '20 at 16:57
1

@OverLordGoldDragon Yup that dont care equidistant was important. – Divakar Oct 25 '20 at 16:58
1

@OverLordGoldDragon You can check out linked post at the start of this post, if you are up for reading on the theory behind using `searchsorted` here. – Divakar Oct 25 '20 at 17:02
1

[Credited](https://github.com/OverLordGoldDragon/ssqueezepy/blob/0.5.0rc/ssqueezepy/algos.py#L7). – OverLordGoldDragon Nov 06 '20 at 02:50

Efficient elementwise argmin of matrix-vector difference

1 Answers1

Linked