Search code examples
pythonnumpymatrixgridcoordinates

Radius search in list of coordinates


I'm using python 2.7 and numpy (import numpy as np).
I have a list of x-y coordinates in the following shape:
coords = np.zeros((100, 2), dtype=np.int)
I have a list of values corresponding to these coordinates:
values = np.zeros(100, dtype=np.int)
My program is populating these arrays.

Now, for each coordinate, I want to find neighbours within radius r that have a non-zero value. What's the most efficient way to do that?


Solution

  • Demo:

    import pandas as pd
    from scipy.spatial.distance import pdist, squareform
    

    In [101]: np.random.seed(123)
    
    In [102]: coords = np.random.rand(20, 2)
    
    In [103]: r = 0.3
    
    In [104]: d = pd.DataFrame(squareform(pdist(coords)))
    
    In [105]: d
    Out[105]:
              0         1         2         3         4         5         6         7         8         9         10        11        12  \
    0   0.000000  0.539313  0.138885  0.489671  0.240183  0.566555  0.343214  0.541508  0.525761  0.295906  0.566702  0.326087  0.045059
    1   0.539313  0.000000  0.509028  0.765644  0.299834  0.212418  0.535287  0.253292  0.378472  0.305322  0.504946  0.501173  0.545672
    2   0.138885  0.509028  0.000000  0.369830  0.240542  0.484970  0.459329  0.449965  0.591335  0.217102  0.434730  0.187983  0.100192
    3   0.489671  0.765644  0.369830  0.000000  0.579235  0.639118  0.827519  0.585140  0.946945  0.474554  0.383486  0.266724  0.444612
    4   0.240183  0.299834  0.240542  0.579235  0.000000  0.364005  0.335128  0.355671  0.368796  0.148598  0.482379  0.327450  0.251218
    5   0.566555  0.212418  0.484970  0.639118  0.364005  0.000000  0.676135  0.055591  0.576447  0.272729  0.315123  0.399127  0.555655
    6   0.343214  0.535287  0.459329  0.827519  0.335128  0.676135  0.000000  0.679527  0.281035  0.481218  0.813671  0.621056  0.387169
    7   0.541508  0.253292  0.449965  0.585140  0.355671  0.055591  0.679527  0.000000  0.602427  0.245620  0.261309  0.350237  0.526773
    8   0.525761  0.378472  0.591335  0.946945  0.368796  0.576447  0.281035  0.602427  0.000000  0.498845  0.811462  0.695304  0.559738
    9   0.295906  0.305322  0.217102  0.474554  0.148598  0.272729  0.481218  0.245620  0.498845  0.000000  0.333842  0.208528  0.282959
    10  0.566702  0.504946  0.434730  0.383486  0.482379  0.315123  0.813671  0.261309  0.811462  0.333842  0.000000  0.254850  0.533784
    11  0.326087  0.501173  0.187983  0.266724  0.327450  0.399127  0.621056  0.350237  0.695304  0.208528  0.254850  0.000000  0.288072
    12  0.045059  0.545672  0.100192  0.444612  0.251218  0.555655  0.387169  0.526773  0.559738  0.282959  0.533784  0.288072  0.000000
    13  0.339648  0.350100  0.407307  0.769145  0.202592  0.501132  0.185248  0.511020  0.186913  0.347808  0.678357  0.527288  0.372879
    14  0.530211  0.104003  0.473790  0.689158  0.303486  0.109841  0.589377  0.149459  0.468906  0.257676  0.404710  0.431203  0.527905
    15  0.622118  0.178856  0.627453  0.923461  0.391044  0.387645  0.509836  0.431502  0.273610  0.450269  0.683313  0.656742  0.639993
    16  0.337079  0.211995  0.297111  0.582175  0.113238  0.251168  0.434076  0.246505  0.403684  0.107671  0.409858  0.316172  0.337886
    17  0.271897  0.311029  0.313864  0.668400  0.097022  0.424905  0.252905  0.426640  0.279160  0.243693  0.576241  0.422417  0.296806
    18  0.664617  0.395999  0.554151  0.592343  0.504234  0.184188  0.833801  0.157951  0.758223  0.376555  0.212643  0.410605  0.642698
    19  0.328445  0.719013  0.238085  0.186618  0.476045  0.642499  0.671657  0.594990  0.828653  0.413697  0.465589  0.245340  0.284878
    
              13        14        15        16        17        18        19
    0   0.339648  0.530211  0.622118  0.337079  0.271897  0.664617  0.328445
    1   0.350100  0.104003  0.178856  0.211995  0.311029  0.395999  0.719013
    2   0.407307  0.473790  0.627453  0.297111  0.313864  0.554151  0.238085
    3   0.769145  0.689158  0.923461  0.582175  0.668400  0.592343  0.186618
    4   0.202592  0.303486  0.391044  0.113238  0.097022  0.504234  0.476045
    5   0.501132  0.109841  0.387645  0.251168  0.424905  0.184188  0.642499
    6   0.185248  0.589377  0.509836  0.434076  0.252905  0.833801  0.671657
    7   0.511020  0.149459  0.431502  0.246505  0.426640  0.157951  0.594990
    8   0.186913  0.468906  0.273610  0.403684  0.279160  0.758223  0.828653
    9   0.347808  0.257676  0.450269  0.107671  0.243693  0.376555  0.413697
    10  0.678357  0.404710  0.683313  0.409858  0.576241  0.212643  0.465589
    11  0.527288  0.431203  0.656742  0.316172  0.422417  0.410605  0.245340
    12  0.372879  0.527905  0.639993  0.337886  0.296806  0.642698  0.284878
    13  0.000000  0.408426  0.339019  0.274263  0.105627  0.668252  0.643427
    14  0.408426  0.000000  0.282070  0.194058  0.345013  0.294029  0.663142
    15  0.339019  0.282070  0.000000  0.344028  0.355134  0.568361  0.854775
    16  0.274263  0.194058  0.344028  0.000000  0.181494  0.399730  0.513362
    17  0.105627  0.345013  0.355134  0.181494  0.000000  0.581128  0.551910
    18  0.668252  0.294029  0.568361  0.399730  0.581128  0.000000  0.649183
    19  0.643427  0.663142  0.854775  0.513362  0.551910  0.649183  0.000000
    

    result:

    In [107]: d[(0 < d) & (d < r)].apply(lambda x: x.dropna().index.tolist())
    Out[107]:
    0                       [2, 4, 9, 12, 17]
    1                   [4, 5, 7, 14, 15, 16]
    2               [0, 4, 9, 11, 12, 16, 19]
    3                                [11, 19]
    4            [0, 1, 2, 9, 12, 13, 16, 17]
    5                   [1, 7, 9, 14, 16, 18]
    6                             [8, 13, 17]
    7               [1, 5, 9, 10, 14, 16, 18]
    8                         [6, 13, 15, 17]
    9     [0, 2, 4, 5, 7, 11, 12, 14, 16, 17]
    10                            [7, 11, 18]
    11                  [2, 3, 9, 10, 12, 19]
    12               [0, 2, 4, 9, 11, 17, 19]
    13                      [4, 6, 8, 16, 17]
    14               [1, 5, 7, 9, 15, 16, 18]
    15                             [1, 8, 14]
    16         [1, 2, 4, 5, 7, 9, 13, 14, 17]
    17            [0, 4, 6, 8, 9, 12, 13, 16]
    18                         [5, 7, 10, 14]
    19                         [2, 3, 11, 12]
    dtype: object