如何保留原Ocean数组NaN值并完成多条件Numpy数组赋值？

我有NetCDF格式的海洋与大气数据集，海洋数据在陆地区域为NaN或-999（示例使用NaN）。示例数据如下：

import numpy as np
ocean = np.array([[2, 4, 5], [6, np.nan, 2], [9, 3, np.nan]])
atmos = np.array([[4, 2, 5], [6, 7, 3], [8, 3, 2]])

需要对海洋和大气数据按多条件生成取值为1-8的新数组：海洋数据中2-4的值赋值为1，4-6的值赋值为2；大气数据同理。

我通过定义分箱列表，使用np.digitize生成分类数组：

bin1 = [2, 4, 6]
bin2 = [4, 6, 8]
ocean_cat = np.digitize(ocean, bin1)
atmos_cat = np.digitize(atmos, bin2) 

得到结果：

# ocean_cat
[[1 2 2]
 [3 3 1]
 [3 1 3]]

# atmos_cat
[[1 0 1]
 [2 2 0]
 [3 0 0]]

随后使用np.fmax获取元素级最大值：

final_cat = np.fmax(ocean_cat, atmos_cat)
print(final_cat)

结果为：

[[1 2 2]
 [3 3 1]
 [3 1 3]]

该结果接近需求，但丢失了原ocean数组中的NaN值，期望的最终结果为：

[[1 2 2]
 [3 nan 1]
 [3 1 nan]]

需要将原ocean数组对应索引的NaN值同步到final_cat数组中。

解决方案 #

可以通过np.isnan()定位原ocean数组中的NaN位置，再将这些位置的值在final_cat中替换为NaN即可。因为np.digitize会把NaN映射为最大分箱索引，所以需要覆盖这些位置的错误值。

完整代码示例 #

import numpy as np

# 原始数据
ocean = np.array([[2, 4, 5], [6, np.nan, 2], [9, 3, np.nan]])
atmos = np.array([[4, 2, 5], [6, 7, 3], [8, 3, 2]])

# 分箱生成分类数组
bin1 = [2, 4, 6]
bin2 = [4, 6, 8]
ocean_cat = np.digitize(ocean, bin1)
atmos_cat = np.digitize(atmos, bin2) 

# 获取元素级最大值
final_cat = np.fmax(ocean_cat, atmos_cat)

# 同步原ocean数组中的NaN值
final_cat[np.isnan(ocean)] = np.nan

print(final_cat)

运行后输出：

[[ 1.  2.  2.]
 [ 3. nan  1.]
 [ 3.  1. nan]]

针对-999缺失值的适配 #

如果海洋数据中用-999表示缺失值，只需将np.isnan(ocean)替换为ocean == -999即可：

final_cat[ocean == -999] = np.nan  # 也可根据需求保持-999作为缺失标记

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