calculate distance from all points in numpy array to a single point on the basis of index

Question

Suppose a 2d array is given as:

arr = array([[1, 1, 1],
             [4, 5, 8],
             [2, 6, 9]])

if point=array([1,1]) is given then I want to calculate the euclidean distance from all indices of arr to point (1,1). The result should be

array([[1.41  , 1.        , 1.41],
       [1.    , 0.        , 1.  ],
       [1.41  , 1.        , 1.41]])

For loop is too slow to do these computations. Is there any faster method to achieve this using numpy or scipy?

Thanks!!!

Answer 1

Approach #1

You can use scipy.ndimage.morphology.distance_transform_edt -

def distmat(a, index):
    mask = np.ones(a.shape, dtype=bool)
    mask[index[0],index[1]] = False
    return distance_transform_edt(mask)

Approach #2

Another with NumPy-native tools -

def distmat_v2(a, index):
    i,j = np.indices(a.shape, sparse=True)
    return np.sqrt((i-index[0])**2 + (j-index[1])**2)

Sample run -

In [60]: a
Out[60]: 
array([[1, 1, 1],
       [4, 5, 8],
       [2, 6, 9]])

In [61]: distmat(a, index=[1,1])
Out[61]: 
array([[1.41421356, 1.        , 1.41421356],
       [1.        , 0.        , 1.        ],
       [1.41421356, 1.        , 1.41421356]])

In [62]: distmat_v2(a, index=[1,1])
Out[62]: 
array([[1.41421356, 1.        , 1.41421356],
       [1.        , 0.        , 1.        ],
       [1.41421356, 1.        , 1.41421356]])

Benchmarking

Other proposed solution(s) :

# https://stackoverflow.com/a/61629292/3293881 @Ehsan
def norm_method(arr, point):
    point = np.asarray(point)
    return np.linalg.norm(np.indices(arr.shape, sparse=True)-point)

Using benchit package (few benchmarking tools packaged together; disclaimer: I am its author) to benchmark proposed solutions.

In [66]: import benchit

In [76]: funcs = [distmat, distmat_v2, norm_method]

In [77]: inputs = {n:(np.random.rand(n,n),[1,1]) for n in [3,10,50,100,500,1000,2000,5000]}

In [83]: T = benchit.timings(funcs, inputs, multivar=True, input_name='Length')

In [84]: In [33]: T.plot(logx=True, colormap='Dark2', savepath='plot.png')

So, distmat_v2 seems to be doing really well, We can further improve on it, by leveraging numexpr.

---

Extend to array of indices

We could extend the listed solutions to cover for the generic/bigger case of list/array of indices w.r.t. whom we need to get euclidean distances at rest of the positions, like so -

def distmat_indices(a, indices):
    indices = np.atleast_2d(indices)
    mask = np.ones(a.shape, dtype=bool)
    mask[indices[:,0],indices[:,1]] = False
    return distance_transform_edt(mask)

def distmat_indices_v2(a, indices):
    indices = np.atleast_2d(indices)
    i,j = np.indices(a.shape, sparse=True)
    return np.sqrt(((i-indices[:,0])[...,None])**2 + (j-indices[:,1,None])**2).min(1)

Sample run -

In [143]: a = np.random.rand(4,5)

In [144]: distmat_indices(a, indices=[[2,2],[0,3]])
Out[144]: 
array([[2.82842712, 2.        , 1.        , 0.        , 1.        ],
       [2.23606798, 1.41421356, 1.        , 1.        , 1.41421356],
       [2.        , 1.        , 0.        , 1.        , 2.        ],
       [2.23606798, 1.41421356, 1.        , 1.41421356, 2.23606798]])

Answer 2

On top of @Divakar's good solutions, if you are looking for something abstract, you can use:

np.linalg.norm(np.indices(arr.shape, sparse=True)-point)

Note that it works with numpy 1.17+ (argument sparse is added on the versions 1.17+ of numpy). Upgrade your numpy and enjoy. In case you have older than 1.17 version of numpy , you can add dimensions to your point by using this:

np.linalg.norm(np.indices(arr.shape)-point[:,None,None], axis=0)

output for point=np.array([1,1]) and given array in question:

[[1.41421356 1.         1.41421356]
 [1.         0.         1.        ]
 [1.41421356 1.         1.41421356]]