Pairwise distance between objects (Xarray) - scipy

Question

I have 3 cars travelling in space (x,y) at 10 time steps.
For each time step I want to calculate the pairwise Euclidean distance between cars.
import numpy as np
from scipy.spatial.distance import pdist
import xarray as xr
data = np.random.rand(3,2,10)
times = pd.date_range('2000-01-01', periods=10)
space = ['x','y']
cars = ['a','b','c']
foo = xr.DataArray(data, coords=[cars,space,times], dims = ['cars','space','time'])
The for loop iteration below works fine, each input is 3*2 array , and pdist is happily calculating a condensed distance matrix for all the pairwise distances between cars
for label,group in foo.groupby('time'):
print(group.shape, type(group), pdist(group) )
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.45389929 0.96104589 0.51489773]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.87532985 0.49758256 0.4418555 ]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.44036486 0.17947479 0.39842543]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.52294711 0.26278261 0.78106623]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.30004324 0.62807379 0.40601505]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.48351623 0.38331324 0.30677522]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.83682031 0.38409803 0.455275 ]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.33614753 0.50814237 0.49033016]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.17365559 0.33567641 0.30382769]
(3, 2) <class 'xarray.core.dataarray.DataArray'> [0.76981095 0.18099241 0.91187884]
but this simple call ( which should do the identical operation as I understand it ) is failing.
foo.groupby('time').apply(pdist)
AttributeError: 'numpy.ndarray' object has no attribute 'dims'
It seems to be having trouble with the return shape ? do I need a u_func here ?
BTW all these calls work fine and returns as expected with a variety of shapes:
foo.groupby('time').apply(np.mean)
foo.groupby('time').apply(np.mean,axis=0)
foo.groupby('time').apply(np.mean,axis=1)
thanks in advance for any pointers...

Answer 1

pdist changes the array size and therefore xarray can not find its coordinates.
How about the following?
In [12]: np.sqrt(((foo - foo.rename(cars='cars1'))**2).sum('space'))
Out[12]:
<xarray.DataArray (cars: 3, time: 10, cars1: 3)>
array([[[0. , 0.131342, 0.352521],
[0. , 0.329914, 0.859899],
[0. , 0.933117, 0.351842],
[0. , 0.802514, 0.426005],
[0. , 0.167081, 0.563704],
[0. , 0.9822 , 0.145496],
[0. , 0.894892, 0.457217],
[0. , 0.333222, 0.505805],
[0. , 0.377352, 0.604625],
[0. , 0.467771, 0.62544 ]],
[[0.131342, 0. , 0.243476],
[0.329914, 0. , 0.813076],
[0.933117, 0. , 0.847525],
[0.802514, 0. , 0.390665],
[0.167081, 0. , 0.562188],
[0.9822 , 0. , 0.957067],
[0.894892, 0. , 0.525863],
[0.333222, 0. , 0.835241],
[0.377352, 0. , 0.894856],
[0.467771, 0. , 0.594124]],
[[0.352521, 0.243476, 0. ],
[0.859899, 0.813076, 0. ],
[0.351842, 0.847525, 0. ],
[0.426005, 0.390665, 0. ],
[0.563704, 0.562188, 0. ],
[0.145496, 0.957067, 0. ],
[0.457217, 0.525863, 0. ],
[0.505805, 0.835241, 0. ],
[0.604625, 0.894856, 0. ],
[0.62544 , 0.594124, 0. ]]])
Coordinates:
* cars (cars) <U1 'a' 'b' 'c'
* time (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-10
* cars1 (cars1) <U1 'a' 'b' 'c'
If you like to have a similar output to pdist, apply_ufunc can be used,
In [21]:xr.apply_ufunc(pdist, foo, input_core_dims=[['cars', 'space']],
...: output_core_dims=[['cars_pair']], vectorize=True)
...:
Out[21]:
<xarray.DataArray (time: 10, cars_pair: 3)>
array([[0.131342, 0.352521, 0.243476],
[0.329914, 0.859899, 0.813076],
[0.933117, 0.351842, 0.847525],
[0.802514, 0.426005, 0.390665],
[0.167081, 0.563704, 0.562188],
[0.9822 , 0.145496, 0.957067],
[0.894892, 0.457217, 0.525863],
[0.333222, 0.505805, 0.835241],
[0.377352, 0.604625, 0.894856],
[0.467771, 0.62544 , 0.594124]])
Coordinates:
* time (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-10
Dimensions without coordinates: cars_pair