如何高效从xarray DataArray提取阈值极值及对应坐标

原方案问题原因 #

da.where(da>2, drop=True)返回大量NaN的核心原因是：drop=True参数的逻辑是仅删除沿维度全为NaN的整段切片，不会移除切片内部的单个NaN值。只要某段经纬度/时间切片里存在1个符合条件的有效值，整个切片都会被保留，自然会携带大量冗余NaN，符合条件的点越稀疏，冗余问题越明显。
而手写三层嵌套循环的方案，是在Python层逐点做索引取值、判断，存在大量重复的对象构造开销，数据规模大时运行效率会极低。另外该方案基于裁剪后的DataArray遍历，返回的经纬度是裁剪后的相对坐标，并非原始数据集的真实坐标，结果本身存在偏差。

最优实现方案 #

以下两种方案均为向量化实现，无Python层显式循环，运行效率比手写嵌套循环高数百到数千倍，不会产生冗余NaN，且直接返回原始数据集的正确坐标。

方案1：Numpy原生索引（性能最高，适合超大规模数据集） #

直接通过布尔掩码拿到符合条件值的索引，再向量化批量取值构造结果，内存开销最小、速度最快：

import xarray as xr
import numpy as np
import pandas as pd

# 测试数据集构造
values = np.array(
    [[[3, 1, 1],
      [1, 1, 1], 
      [1, 1, 1]],

     [[1, 1, 1],
      [1, 1, 1],
      [1, 1, 4]],

     [[1, 1, 1],
      [1, 1, 1],
      [1, 1, 5]]]
)
da = xr.DataArray(values, dims=('time', 'lat', 'lon'), 
    coords={'time': list(range(3)), 'lat': list(range(3)), 'lon':list(range(3))})

# 生成过滤掩码
mask = da > 2
# 直接获取所有符合条件点的维度索引，无Python层循环
time_idx, lat_idx, lon_idx = mask.values.nonzero()
# 批量取值构造DataFrame
res = pd.DataFrame({
    'Time': da.time.values[time_idx],
    'Latitude': da.lat.values[lat_idx],
    'Longitude': da.lon.values[lon_idx],
    'Value': da.values[time_idx, lat_idx, lon_idx]
})

运行后返回的结果为原始坐标下的正确值：

Time  Latitude  Longitude  Value
0     0         0          0      3
1     1         2          2      4
2     2         2          2      5

方案2：Xarray原生堆叠写法（代码简洁，适合常规规模数据） #

通过stack将多维数组压缩为单维度序列，再直接过滤掉NaN，代码更易读，性能同样远优于手写循环：

# 过滤后堆叠多维度为单点维度，删除空值
stacked_da = da.where(da>2).stack(point=('time', 'lat', 'lon')).dropna('point')
# 转换格式得到目标DataFrame
res = stacked_da.to_dataframe('Value').reset_index()[['time', 'lat', 'lon', 'Value']]
res.columns = ['Time', 'Latitude', 'Longitude', 'Value']

选型提示 #

• 如果数据集网格点规模在千万级以上，优先选择方案1，numpy的nonzero是C层实现的索引逻辑，内存占用和运行速度都达到最优
• 常规十万到百万级网格点的场景，方案2写法更简洁，维护成本更低
• 两种方案都可以完全替代手写三层循环的实现，避免不必要的性能损耗

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