xarray, indexing with multidimensional coordinates

Question

I'm not sure if it is possible, but I would like to add a new coordinate to a xarray which is the same size as the data that its indexing, but has a different number of dimensions.

For example, take a 2D array, one dimension is a date, and the other is time of day:

data = np.random.rand(5, 24)
date_idx = pd.date_range("20200921", "20200925")
time_of_day_idx = np.arange(0, np.timedelta64(1, "D"), np.timedelta64(1, "h"))

da = xr.DataArray(data,
                  dims=("date", "time_of_day"),
                  coords={"date" : date_idx,
                          "time_of_day": time_of_day_idx})

I can add a new 2D coordinate with the times for each value in the array.

time_idx = np.add.outer(date_idx, time_of_day_idx)
da.assign_coords(time=(("date", "time_of_day"), time_idx))

But what I'd really like to be able to do is to stack the dimensions together to create a new one dimensional coordinate for a time-series in the same array, something like this:

da.assign_coords(time=(("date", "time_of_day"), time_idx.flatten()))

However, doing this raises:

ValueError: dimensions ('date', 'time_of_day') must have the same length as the number of data dimensions, ndim=1

I know I can reshape or flatten a 2D xarray to 1D and create the new index for the time-series but it would be really handy to be able to have a 1D coordinate that can index into a 2D xarray. I've played a bit with stacking and MultiIndex, but I found the same problem, each index has separate values access each dimension, rather than a single value that can be used to index the location across both dimensions.

Does anyone know if this is possible?

Answer 1

This is not currently possible with xarrays.