How to calculate the euclidian distance between a specified coordinate and 2-d xarray data at once

How to calculate the euclidian distance between a specified coordinate and 2-d xarray data at once (without for loop)?

I wrote the following code. But if the data-size become larger, this script appears to be slow. How can I do the same thing without for loop ?

import math
import xarray as xr 
import numpy as np

#set the specified lat, lon
specific_lat = 15
specific_lon = 65

#create sample xarry data
lat = [0, 10, 20]
lon = [50, 60, 70, 80]

#sample data
test_data = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12]])

#to xarray
data_xarray = xr.DataArray(test_data, dims=("lat","lon"), coords={"lat":lat, "lon":lon})

#calculate distance
xarray_distance = data_xarray #copy
xarray_distance.data[:,:] = 0.0 #zero-reset

for lat in data_xarray.lat.data:
    for lon in data_xarray.lon.data:
        xarray_distance.loc[{"lat":lat,"lon":lon}] = math.sqrt((lat- spec_lat)**2 + (lon - spec_lon)**2)

print(xarray_distance)

#<xarray.DataArray (lat: 3, lon: 4)>
#array([[21, 15, 15, 21],
#       [15,  7,  7, 15],
#       [15,  7,  7, 15]])
#Coordinates:
#  * lat      (lat) int64 0 10 20
#  * lon      (lon) int64 50 60 70 80

Answer

First things first, xarray_distance = data_xarray does not copy the object. Both names will see the same object and any changes that occur to it equally. For a copy, you want to do data_xarray.copy().

Now, instead of using loops, you want to take advantage of numpy’s vectorisation and broadcasting to create the 2D array in one go and set that to your xarray’s object data.

lat = np.array(lat)
lon = np.array(lon)
xarray_distance.data = np.sqrt((lat - spec_lat)[:,None]**2 + ((lon - spec_lon)**2))

Result

<xarray.DataArray (lat: 3, lon: 4)>
array([[21.21320344, 15.8113883 , 15.8113883 , 21.21320344],
       [15.8113883 ,  7.07106781,  7.07106781, 15.8113883 ],
       [15.8113883 ,  7.07106781,  7.07106781, 15.8113883 ]])
Coordinates:
  * lat      (lat) int32 0 10 20
  * lon      (lon) int32 50 60 70 80

You’ll notice the data are floats, as they should probably be, considering what they represent. If you want them in the original type of the data, substitute

xarray_distance.data = np.sqrt(...)

with

xarray_distance.data = np.sqrt(...).astype(xarray_distance.data.dtype)