使用xr.where获取3D索引(time,lat,lon)时出现索引越界问题咨询

我希望从给定的xarray数据集中，使用xr.where获取满足t<22条件的3D索引(time, lat, lon)。数据集代码如下：

import xarray as xr
import numpy as np
import pandas as pd
np.random.seed(0)
t = 15 + 8 * np.random.randn(2, 2, 3)
lon = [[-99.83, -99.32], [-99.79, -99.23]]
lat = [[42.25, 42.21], [42.63, 42.59]]
time = pd.date_range("2014-09-06", periods=3)
reference_time = pd.Timestamp("2014-09-05")
ds = xr.Dataset(
    data_vars=dict(
        t=(['x', 'y', 'time'], t),
    ),
    coords=dict(
        lon=(['x', 'y'], lon),
        lat=(['x', 'y'], lat),
        time=time,
        reference_time=reference_time,
    ),
    attrs=dict(description="Weather related data."),
)

注：原代码中t变量赋值存在笔误，已修正为t=(['x', 'y', 'time'], t)。

我尝试了以下代码：

time, lat_i, lon_i = np.nonzero(xr.where(ds.t<22, 1, 0).data)
lats = ds.x.isel(x=lat_i).data
lons = ds.y.isel(y=lon_i).data
time = ds.time.isel(time=time).data
pairs = list(zip(time, lats, lons))

但出现错误：IndexError: index 2 is out of bounds for axis 0 with size 2，打印索引结果为：

(array([0, 0, 1, 1, 1, 1], dtype=int64),
 array([0, 1, 0, 0, 1, 1], dtype=int64),
 array([1, 2, 1, 2, 0, 1], dtype=int64))

请问为何lon_i中出现了预期外的数值2？

问题根源 #

你搞反了np.nonzero返回索引的维度对应关系：

• 你的t变量维度是['x', 'y', 'time']，对应数组形状是(2,2,3)
• np.nonzero返回的三个索引数组，顺序是x轴索引、y轴索引、time轴索引，而非你假设的time, lat_i, lon_i
• 你把第三个数组（time轴索引，范围0-2）错误当成了lon_i（y轴索引，范围仅0-1），自然会出现超出边界的数值2

修正后的代码 #

先正确对应维度索引，再提取对应的坐标值：

# 获取满足条件的x、y、time索引
x_idx, y_idx, time_idx = np.nonzero(ds.t < 22)

# 提取对应位置的lat和lon值（lat和lon是x,y二维坐标）
lats = ds.lat.isel(x=x_idx, y=y_idx).data
lons = ds.lon.isel(x=x_idx, y=y_idx).data

# 提取对应的时间值
times = ds.time.isel(time=time_idx).data

# 组合成(time, lat, lon)元组列表
pairs = list(zip(times, lats, lons))

补充说明 #

• 无需用xr.where转成1和0，直接用ds.t < 22生成布尔数组传给np.nonzero即可
• lat和lon是x,y二维坐标，需同时用x_idx和y_idx索引才能拿到对应位置的经纬度值，不能单独用x或y索引

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