xarray中idxmax无满足条件值时返回nan的实现方法

问题核心：当布尔数组全为False时，idxmax会返回第一个元素的索引（即0），而非预期的nan。要实现需求，只需在结果中把没有任何符合条件值的行替换为nan即可，具体步骤如下：

1. 计算初始索引结果 #

用(da > 阈值).idxmax(dim='x')得到每行第一个符合条件的索引，全不符合时会返回0。

2. 生成掩码标记无效行 #

通过(da > 阈值).any(dim='x')判断每行是否存在符合条件的值，取反后得到标记无效行的掩码（True表示该行无符合条件的值）。

3. 替换无效行为nan #

用xarray.DataArray.where()方法，将掩码标记的位置替换为nan，同时自动转换数据类型（因为整数类型无法存储nan，会转为浮点型）。

完整代码示例 #

import numpy as np
import xarray as xr

# 构造示例DataArray
data = [[2., 1., 2., 0., -2.],
        [-4., np.nan, 2., np.nan, -2.],
        [np.nan, np.nan, 1., np.nan, np.nan]]
da = xr.DataArray(data, coords={'y': [-1, 0, 1], 'x': [0, 1, 4, 9, 16]}, dims=['y', 'x'])

# 阈值设为4时的处理
threshold = 4
bool_arr = da > threshold
idx_result = bool_arr.idxmax(dim='x')
# 替换无符合条件值的行为nan
final_result = idx_result.where(bool_arr.any(dim='x'), np.nan)

print(final_result)

输出结果 #

<xarray.DataArray 'x' (y: 3)>
array([nan, nan, nan])
Coordinates:
  * y        (y) int64 -1 0 1

也可以写成一行简洁版：

final_result = (da > 4).idxmax(dim='x').where((da > 4).any(dim='x'), np.nan)

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