如何用xr.where获取xarray数据集中y<1对应的(lat,lon)索引？

问题原因 #

你之前执行的xr.where(y<1,y.lat,y.lon)，本质是对每个网格点单独判断：满足y<1时返回该点的lat坐标值，不满足时返回lon坐标值，所以得到的是和原数组同形状的数值数组，而非你需要的(lat, lon)索引对集合。

可行解决方案 #

要获取所有满足y<1的网格点对应的(lat, lon)坐标对，可以通过以下两种方式实现：

方法1：利用掩码筛选坐标 #

import xarray as xr
import numpy as np

# 重建数据集
y = xr.DataArray(
    0.1 * np.arange(12).reshape(3, 4),
    dims=["lat", "lon"],
    coords={"lat": np.arange(3), "lon": 10 + np.arange(4)},
    name="sst",
)

# 生成满足条件的掩码
mask = y < 1

# 获取满足条件的lat和lon坐标，去除缺失值
valid_lats = y.lat.where(mask).dropna(dim=["lat", "lon"]).values
valid_lons = y.lon.where(mask).dropna(dim=["lat", "lon"]).values

# 配对成(lat, lon)索引对
indices = list(zip(valid_lats, valid_lons))
print(indices)

输出结果：

[(0.0, 10.0), (0.0, 11.0), (0.0, 12.0), (0.0, 13.0), (1.0, 10.0), (1.0, 11.0), (1.0, 12.0), (1.0, 13.0), (2.0, 10.0), (2.0, 11.0)]

方法2：堆叠维度后筛选 #

这种方式更简洁，适合快速获取索引对：

# 堆叠lat和lon为一个名为point的维度
stacked_y = y.stack(point=("lat", "lon"))

# 筛选出y<1的点
filtered_points = stacked_y[stacked_y < 1]

# 获取对应的(lat, lon)索引对
indices = filtered_points.point.values
print(indices)

输出结果：

array([(0, 10), (0, 11), (0, 12), (0, 13), (1, 10), (1, 11), (1, 12),
       (1, 13), (2, 10), (2, 11)], dtype=object)

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