Indexing xarray data with variable length DataArray

Question

I am trying to extract data from xarray dataset using DataArray indexing. My goal is to obtain the data along different line segments overlapping the array. For that I have obtained indices of each of the lines (these are of different sizes based on the length).

For example for line 1 : x = [1,2,3], y=[7,8,9] and similarly for line 2 is x=[1,4,5,6,8], y=[0,2,7,9,6] and so on I have some of the lines which are 100x 2. For this I have tried like below :

df=xarray_dataset
indx=xr.DataArray([[1,2,3],[1,4,5,6,8],[2,3]])
indy=xr.DataArray([[7,9,8],[0,2,7,9,6],[4,5]])
dx_sel=df.isel(x=indx,y=indy)

However what I understand that the length of each of the data array index needs to be equal. Is there a way I can handle such issues. Basically these indices represent the x and y coordinates of different segments within the data frame and get the mean of each of the segment, I have 100s of such segments if there are only few I would be able to use a loop for each of the segment indexes however it's not computationally efficient to use a loop for each segment.

This is a similar issue with numpy array as well. Is there a way to pass NaN or something similar in the index so that we could make the equal shape but no data is extracted for that index.

Keisuke FUJII Keisuke FUJII · Accepted Answer · 2020-07-07T02:27:17

You can use set_index -> unstack mechanism, which is based on pd.MultiIndex.

In [4]: df = xr.DataArray(np.arange(110).reshape(10, 11),  
   ...:                   dims=['x', 'y'])  
In [5]: indx=xr.DataArray([1,2,3, 1,4,5,6,8, 2,3], 
   ...:                   dims=['index'],  
   ...:                   coords={'i': ('index', [0,0,0, 1,1,1,1,1, 2,2]), 
   ...:                           'j': ('index', [0,1,2, 0,1,2,3,4, 0,1])}) 
   ...:  
   ...: indy=xr.DataArray([7,9,8, 0,2,7,9,6, 4,5], dims=['index'], 
   ...:                   coords={'i': ('index', [0,0,0, 1,1,1,1,1, 2,2]), 
   ...:                           'j': ('index', [0,1,2, 0,1,2,3,4, 0,1])})       

In [8]: df.isel(x=indx, y=indy).set_index(index=['i', 'j']).unstack('index')                                         
Out[8]: 
<xarray.DataArray (i: 3, j: 5)>
array([[18., 31., 41., nan, nan],
       [11., 46., 62., 75., 94.],
       [26., 38., nan, nan, nan]])
Coordinates:
  * i        (i) int64 0 1 2
  * j        (j) int64 0 1 2 3 4

Here, indx and indy has non-dimensional coordinates, i and j, which are essentially the original position of the index in the 2-dimensional space.

Indexing xarray data with variable length DataArray

1 Answers