Trova in modo efficiente gli indici dei punti più vicini sulla griglia 2D non rettangolare

Question

Ho una griglia lon/lat irregolare (non rettangolare) e un gruppo di punti in coordinate lon/lat, che dovrebbero corrispondere ai punti sulla griglia (sebbene essi potrebbe essere leggermente fuori per ragioni numeriche). Ora ho bisogno degli indici dei corrispondenti punti lon/lat.

Ho scritto una funzione che esegue questa operazione, ma è VERAMENTE lenta.

def find_indices(lon,lat,x,y): 
    lonlat = np.dstack([lon,lat]) 
    delta = np.abs(lonlat-[x,y]) 
    ij_1d = np.linalg.norm(delta,axis=2).argmin() 
    i,j = np.unravel_index(ij_1d,lon.shape) 
    return i,j 

ind = [find_indices(lon,lat,p*) for p in points] 

Sono abbastanza sicuro che ci sia una soluzione migliore (e più veloce) in numpy/scipy. Ho già fatto ricerche su Google, ma la risposta mi ha eluso finora.

Qualche suggerimento su come trovare in modo efficiente gli indici dei punti corrispondenti (più vicini)?

PS: questa domanda è emersa da un'altra (click).

Edit: Soluzione

in base alla risposta del @Cong Ma, ho trovato la seguente soluzione:

def find_indices(points,lon,lat,tree=None): 
    if tree is None: 
     lon,lat = lon.T,lat.T 
     lonlat = np.column_stack((lon.ravel(),lat.ravel())) 
     tree = sp.spatial.cKDTree(lonlat) 
    dist,idx = tree.query(points,k=1) 
    ind = np.column_stack(np.unravel_index(idx,lon.shape)) 
    return [(i,j) for i,j in ind] 

Per mettere questa soluzione e anche quella dalla risposta di Divakar in prospettiva, qui ci sono alcune temporizzazioni della funzione in cui sto usando find_indices (e dove è il collo di bottiglia in termini di velocità) (vedi link sopra):

spatial_contour_frequency/pil0    : 331.9553 
spatial_contour_frequency/pil1    : 104.5771 
spatial_contour_frequency/pil2    :  2.3629 
spatial_contour_frequency/pil3    :  0.3287 

pil0 è il mio approccio iniziale, pil1 Divakar e pil2/pil3 la soluzione finale sopra, in cui l'albero viene creato al volo in pil2 (ad es. per ogni iterazione del ciclo in cui viene chiamato find_indices e una sola volta in pil3 (vedere altri thread per i dettagli). Anche se la raffinatezza di Divakar del mio approccio iniziale mi dà un 3x di accelerazione, cKDTree porta questo ad un livello completamente nuovo con un altro 50x di accelerazione! E spostare la creazione dell'albero fuori dalla funzione rende le cose ancora più veloci.

Answer 1

Se i punti sono sufficientemente localizzate, si può provare direttamente s' scipy.spatialcKDTree implementazione, come discusso da me in another post. Quel post riguardava l'interpolazione ma puoi ignorarlo e usare solo la parte della query.

tl; dr versione:

Leggi la documentazione di scipy.sptial.cKDTree. Creare l'albero passando un (n, m) -dhaped numpy ndarray object all'inizializzatore e l'albero verrà creato dalle coordinate nm -dimensionali.

tree = scipy.spatial.cKDTree(array_of_coordinates) 

Dopo di che, utilizzare tree.query() per recuperare il k esimo vicino più prossimo (possibilmente con approssimazione e la parallelizzazione, vedi documentazione), oppure utilizzare tree.query_ball_point() per trovare tutti i vicini all'interno di data tolleranza distanza.

Se i punti non sono ben localizzati e la curvatura sferica/topologia non banale entra in azione, è possibile provare a suddividere il collettore in più parti, ciascuna abbastanza piccola da essere considerata locale.

Answer 2

Ecco un approccio vectorized generico utilizzando scipy.spatial.distance.cdist -

import scipy 

# Stack lon and lat arrays as columns to form a Nx2 array, where is N is grid**2 
lonlat = np.column_stack((lon.ravel(),lat.ravel())) 

# Get the distances and get the argmin across the entire N length 
idx = scipy.spatial.distance.cdist(lonlat,points).argmin(0) 

# Get the indices corresponding to grid's shape as the final output 
ind = np.column_stack((np.unravel_index(idx,lon.shape))).tolist() 

esempio del suo funzionamento -

In [161]: lon 
Out[161]: 
array([[-11. , -7.82 , -4.52 , -1.18 , 2.19 ], 
     [-12. , -8.65 , -5.21 , -1.71 , 1.81 ], 
     [-13. , -9.53 , -5.94 , -2.29 , 1.41 ], 
     [-14.1 , -0.04 , -6.74 , -2.91 , 0.976]]) 

In [162]: lat 
Out[162]: 
array([[-11.2 , -7.82 , -4.51 , -1.18 , 2.19 ], 
     [-12. , -8.63 , -5.27 , -1.71 , 1.81 ], 
     [-13.2 , -9.52 , -5.96 , -2.29 , 1.41 ], 
     [-14.3 , -0.06 , -6.75 , -2.91 , 0.973]]) 

In [163]: lonlat = np.column_stack((lon.ravel(),lat.ravel())) 

In [164]: idx = scipy.spatial.distance.cdist(lonlat,points).argmin(0) 

In [165]: np.column_stack((np.unravel_index(idx,lon.shape))).tolist() 
Out[165]: [[0, 4], [0, 4], [0, 4], [0, 4], [0, 4], [0, 4], [3, 3]] 

test runtime -

definire funzioni:

def find_indices(lon,lat,x,y): 
    lonlat = np.dstack([lon,lat]) 
    delta = np.abs(lonlat-[x,y]) 
    ij_1d = np.linalg.norm(delta,axis=2).argmin() 
    i,j = np.unravel_index(ij_1d,lon.shape) 
    return i,j 

def loopy_app(lon,lat,pts): 
    return [find_indices(lon,lat,pts[i,0],pts[i,1]) for i in range(pts.shape[0])] 

def vectorized_app(lon,lat,points): 
    lonlat = np.column_stack((lon.ravel(),lat.ravel())) 
    idx = scipy.spatial.distance.cdist(lonlat,points).argmin(0) 
    return np.column_stack((np.unravel_index(idx,lon.shape))).tolist() 

Tempi:

In [179]: lon = np.random.rand(100,100) 

In [180]: lat = np.random.rand(100,100) 

In [181]: points = np.random.rand(50,2) 

In [182]: %timeit loopy_app(lon,lat,points) 
10 loops, best of 3: 47 ms per loop 

In [183]: %timeit vectorized_app(lon,lat,points) 
10 loops, best of 3: 16.6 ms per loop 

Per spremere più prestazioni, np.concatenate potrebbero essere utilizzati al posto di np.column_stack.