Argmin efficiente per elemento della differenza matrice-vettore

Question

Supponiamo un array a.shape == (N, M)e un vettore v.shape == (N,). L'obiettivo è quello di calcolare argmindi absdella vsottratto ogni elemento di a- che è,

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

Ho un'implementazione vettoriale tramite np.matlib.repmat, ed è molto più veloce, ma richiede O(M*N^2)memoria, inaccettabile nella pratica. Il calcolo è ancora fatto sulla CPU, quindi la soluzione migliore sembra implementare il ciclo for in C come estensione, ma forse Numpy ha già implementato questa logica.

Lo fa? Eventuali funzioni Numpy pronte all'uso che vengono implementate in modo efficiente sopra?

Answer 1

Ispirati da this post, possiamo sfruttare 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

Alcuni più perf. potenziare con numexprgrandi set di dati:

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

Tempistiche

Impostare :

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)