xarray四维数组使用三维索引取数后维度不符合预期问题咨询

问题原因 #

你直接使用ws_interp[ws_index]索引时，xarray默认将传入的索引数组作用于第一个维度（即time维度），而非你预期的isosurf（高度）维度。
没有加time维度时，第一个维度就是高度维度，所以索引结果符合预期；加入time维度后，第一个维度变为time，索引逻辑不符合需求，才会输出多余的维度。

解决方案 #

使用xarray自带的isel方法，显式指定要索引的维度为isosurf即可，代码如下：

# 显式指定对isosurf维度使用ws_index做逐位置索引
result = ws_interp.isel(isosurf=ws_index)
print(result.shape) 
# 输出为(5, 2, 1)，和ws_index形状完全一致，符合需求

如果需要用numpy底层操作实现，也可以用take_along_axis方法：

# 给ws_index增加高度维度的轴，对齐后沿高度维度取值
np_result = np.take_along_axis(ws_interp.values, ws_index.values[:, np.newaxis, ...], axis=1).squeeze(axis=1)
print(np_result.shape)
# 输出同样为(5, 2, 1)

内容的提问来源于stack exchange，提问作者Luke Hickinbotham