Equivalent of np.take for xarray DataArray

[peendebak]

With numpy one can take the slice of the kth dimension with

X = np.random.rand( (2,3,4,5,6))
k=2
index=1
np.take(X, index, axis=k)

Is there an xarray equivalent of the np.take method?

[TomNicholas]

Yes, it's just .isel().

From the numpy docs:

A call such as np.take(arr, indices, axis=3) is equivalent to arr[:,:,:,indices,...].

In xarray we use dimension names in place of axis numbers, so if if you have a DataArray da with a dimension named time that corresponds to your axis=2, you could achieve the same thing with

da.isel(time=2)

Or in your example, let's say we have created a dataarray with data X and 5 dimensions:

da_X = xr.DataArray(data=X, dims=["a", "b", "c", "d", "e"])

The equivalent isel call would be

index=1
da_X.isel(c=index)

because dimension "c" now corresponds to axis number 2 of your data.

[peendebak]

I would rather select my dimensions by axis (even though in xarray it is perhaps more common to select by name). With isel it is possible though to write a custom take method:

def take(data_array : xarray.DataArray, indices, axis : int):
    dimension_name = list(data_array.dims)[axis]
    return data_array .isel({dimension_name: indices})

An alternative without the isel is

def take(data_array : xarray.DataArray, indices, axis : int):
    m = {axis: indices}
    ix = tuple(m.get(dim, slice(None)) for dim in range(data_array.ndim))
    return data_array[ix]