I am trying to index the maximum windspeed of each grid point at each timestep for a set of data I am working on. I have a variable ws_interp which is the windspeed profile with 4 dimensions (time, height, latitude, longitude). I want to index the maximum windspeed at each grid coordinate at each timestep with an index ws_index. I should therefore output a 3 dimensional array when I index ws_interp however I keep getting a 4 dimensional array the same size as ws_interp.
Here is a reproducible example of my problem:
import xarray as xr
import numpy as np
a = np.arange(5*3*2*1)
c = a.reshape((5,3,2,1))
ws_interp = xr.DataArray(c, dims=("time", "isosurf", "latitude", "longitude"))
ws_interp.shape
(5, 3, 2, 1)
b = np.random.choice([0, 1], size=(10))
d = b.reshape((5,2,1))
ws_index = xr.DataArray(d,dims=("time", "latitude","longitude"))
ws_index.shape
(5, 2, 1)
ws_interp[ws_index].shape
(5, 2, 1, 3)
I don't understand why my indexed xarray is shape (5,2,1,3) when I want it to be (5,2,1) like the shape of the index array. It did what I wanted when the time dimension wasn't involved but now it doesn't seem to work. Any helped would be most appreciated, thank you!