0

I have N points in 2D space (N row, 2 column) and I want to find nearest k points for each point (x_n,y_n) in the point set and then sort them. Here is the codes for this purpose. I want to speed up the following codes:

def nearst_sort(x,y,k):
    N = len(x)    
    A = np.zeros((N,k))
    R = np.zeros((N,N))
    R = (x - x[np.newaxis].transpose())**2 + (y -y[np.newaxis].transpose())**2
    ix = np.argsort(R, kind='stable')
    ix = ix.transpose()
    A=ix[0:k,:].transpose()
    return A

My sample data is as follows:

x   y
0   0
0   0.5
0   1
0.5 0
0.5 0.5
0.5 1
1   0
1   0.5
1   1

I have also tried functions from scipy.spatial.KDTree but could not get good results. Any help would be appreciated.

For Bonder
  • 67
  • 5

1 Answers1

0

For me KDTree is much faster:

>>> import numpy as np
>>> from scipy.spatial import cKDTree as KDTree
>>> from timeit import timeit
>>> 
>>> z = np.random.randn(1000, 2)
>>> k = 5
>>> 
>>> KDTree(z).query(z, k)
(array([[0.        , 0.21130505, 0.22903208, 0.32009477, 0.38000444],
       [0.        , 0.03969915, 0.06698214, 0.08423566, 0.10740011],
       [0.        , 0.04964421, 0.08194808, 0.11576068, 0.12022531],
       ...,
       [0.        , 0.00721785, 0.03346301, 0.03617199, 0.04193239],
       [0.        , 0.05147871, 0.05619545, 0.08028866, 0.08744349],
       [0.        , 0.03733766, 0.06359033, 0.06861222, 0.0698981 ]]), array([[  0, 391, 134, 462, 575],
       [  1,  87, 879, 846, 122],
       [  2, 793, 314, 564, 483],
       ...,
       [997, 390, 432, 165, 952],
       [998, 194, 457, 775, 629],
       [999, 158, 522, 862, 791]]))
>>> nearst_sort(*z.T, k)
array([[  0., 391., 134., 462., 575.],
       [  1.,  87., 879., 846., 122.],
       [  2., 793., 314., 564., 483.],
       ...,
       [997., 390., 432., 165., 952.],
       [998., 194., 457., 775., 629.],
       [999., 158., 522., 862., 791.]])
>>> timeit(lambda: KDTree(z).query(z, k), number=100)
0.12790076900273561
>>> timeit(lambda: nearst_sort(*z.T, k), number=100)
6.5285790269990684

That's a factor of 50. May depend on the example, though.

Paul Panzer
  • 51,835
  • 3
  • 54
  • 99
  • how do you define nearst_sort function? – For Bonder Mar 19 '19 at 03:22
  • Your approach is very fast. But is it possible to modify your approach to get results exactly same as results of nearst_sort function. For example for given data the results are not equal. Namely np.array_equal(ii,jj) returns False for distances,ii=KDTree(z,y).query(z, 3), jj=nearst_sort(*z.T, 3) where z=(x,y) – For Bonder Mar 19 '19 at 15:06
  • That is because on a regular grid such as your example there will be many ties and they happpen to be sorted differently. For example, your code selects `(0,0);(0,0.5);(0.5,0)` while KDTree selects `(0,0);(0.5,0);(0,0.5)`. Obviously, both answers are equally valid. – Paul Panzer Mar 19 '19 at 15:22