Efficient and Pythonic way to calculate Euclidean distance to the nearest nonzero element, for each nonzero element in NumPy 2D array

user_5495385 0
Viewed 176

I have a bi-dimensional NumPy array of shape M × N with many values set to 0 and others with value ≠ 0.
The following is an example of the aforesaid matrix:

A = np.array([[0, 0, 0, 1, 0, 2, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 1, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0, 3, 0], [1, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 6, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0]])

And that's it nicely formatted:

A = [[0 0 0 1 0 2 0 0]
     [0 0 0 0 0 0 0 0]
     [0 0 0 0 1 0 0 0]
     [0 1 0 0 0 0 0 0]
     [0 0 0 1 0 0 3 0]
     [1 0 0 0 0 0 0 0]
     [0 0 0 0 6 0 0 0]
     [0 0 1 0 0 0 0 0]]

My task is to find, for each nonzero element (e.g. 1, 2, 1, 1, 1, 3, 1, 6 and 1) in the 2D array (A), the distance to the nearest nonzero element (except itself) by means of Euclidean distance, and then create a list (L) with the calculated distances.
The following invariant must hold:

if np.count_nonzero(A) < 2:
    assert len(L) == 0
else
    assert np.count_nonzero(A) == len(L)

Calculations for array A is the following:

  • Nearest nonzero element for A[0, 3] = 1 is A[0, 5] = 2 at distance = 2
  • Nearest nonzero element for A[0, 5] = 2 is A[0, 3] = 1 at distance = 2
  • Nearest nonzero element for A[2, 4] = 1 is A[0, 3] = 1 at distance = 2.24
  • Nearest nonzero element for A[3, 1] = 1 is A[4, 3] = 1 at distance = 2.24
  • Nearest nonzero element for A[4, 3] = 1 is A[2, 4] = 1 at distance = 2.24
  • Nearest nonzero element for A[4, 6] = 3 is A[2, 4] = 1 at distance = 2.83
  • Nearest nonzero element for A[5, 0] = 1 is A[3, 1] = 1 at distance = 2.24
  • Nearest nonzero element for A[6, 4] = 6 is A[4, 3] = 1 at distance = 2.24
  • Nearest nonzero element for A[7, 2] = 1 is A[6, 4] = 6 at distance = 2.24

The list L is then L = [2, 2, 2.24, 2.24, 2.24, 2.83, 2.24, 2.24, 2.24].

I wrote the following code to solve the problem, and I think it works correctly, but it has two problems: it's the naive, brute force and non vectorized solution of (M² × N²) time complexity, and is not very clear, concise and succinct; that is, it's not Pythonic.

def get_distance_list(A):
    L = []
    for (m, n), a_mn in np.ndenumerate(A):
        # skip this element if its value is 0
        if a_mn == 0:
            continue
        d_min = math.inf
        for (k, l), a_kl in np.ndenumerate(A):
            # skip this element if its value is 0 or if it's me
            if a_kl == 0 or (m, n) == (k, l):
                continue
            d = scipy.spatial.distance.euclidean((m, n), (k, l))
            d_min = min(d_min, d)
        # in case there are less than two nonzero values in the matrix,
        # the returned list must be empty, so only add the distance
        # if it's different than the default value of +inf
        if d_min != math.inf:
            L.append(d_min)
    return L

Do you know if there is a built-in function (maybe in NumPy, SciPy, SciKit, etc.) which can replace the one I wrote, or if there is a faster/vectorized and more Pythonic way to solve the problem?

matrix numpy arrays python
2 Answers

I think using scipy.spatial.KDTree is perfect for this.

from scipy.spatial import KDTree

nonzeros = np.transpose(np.nonzero(A))
t = KDTree(nonzeros)
dists, nns = t.query(nonzeros, 2)

for (i, j), d in zip(nns, dists[:,1]):
    print(nonzeros[i], "is closest to", nonzeros[j], "with distance", d)

Result:

[0 3] is closest to [0 5] with distance 2.0
[0 5] is closest to [0 3] with distance 2.0
[2 4] is closest to [0 5] with distance 2.23606797749979
[3 1] is closest to [4 3] with distance 2.23606797749979
[4 3] is closest to [3 1] with distance 2.23606797749979
[4 6] is closest to [2 4] with distance 2.8284271247461903
[5 0] is closest to [3 1] with distance 2.23606797749979
[6 4] is closest to [4 3] with distance 2.23606797749979
[7 2] is closest to [6 4] with distance 2.23606797749979
user_5656350

This uses numpy, there could well be other functions that can streamline this.

import numpy as np

A = np.array([[0, 0, 0, 1, 0, 2, 0, 0], 
              [0, 0, 0, 0, 0, 0, 0, 0], 
              [0, 0, 0, 0, 1, 0, 0, 0], 
              [0, 1, 0, 0, 0, 0, 0, 0], 
              [0, 0, 0, 1, 0, 0, 3, 0],
              [1, 0, 0, 0, 0, 0, 0, 0], 
              [0, 0, 0, 0, 6, 0, 0, 0], 
              [0, 0, 1, 0, 0, 0, 0, 0]])

x, y = np.where( A != 0 )

print( x,'\n', y)                                                      
# [0 0 2 3 4 4 5 6 7]  # x coords 
# [3 5 4 1 3 6 0 4 2]  # y coords

diff_x = np.subtract.outer( x, x )  # differences all x from each x
diff_y = np.subtract.outer( y, y )  # all y from each y

distance  = np.sqrt( diff_x * diff_x + diff_y * diff_y )

# Or using np.complex
# point = x + y * np.complex( 0, 1 )  # x and y in the complex plane
# distance = abs(np.subtract.outer( point, point ))  # Euclidian distances point to all points

distance[ distance == 0 ] = distance.max()  # or use np.diagonal to remove zeroes.

ind = np.argmin( distance, axis = 1 )  # indices of minimums

for i, ix in enumerate( ind ):
    print( x[i], y[i], 'close to', x[ix], y[ix], 'distance = ', distance[ i, ix ] )

This produces:

0 3 close to 0 5 distance =  2.0
0 5 close to 0 3 distance =  2.0
2 4 close to 0 3 distance =  2.23606797749979
3 1 close to 4 3 distance =  2.23606797749979
4 3 close to 2 4 distance =  2.23606797749979
4 6 close to 2 4 distance =  2.8284271247461903
5 0 close to 3 1 distance =  2.23606797749979
6 4 close to 4 3 distance =  2.23606797749979
7 2 close to 6 4 distance =  2.23606797749979
user_71539280
Related