基于多个数组的多个条件为numpy数组赋值

Question

我有一个包含海洋和大气数据的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) 

这会产生以下结果：

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

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

接下来，我想要在上面两个数组结果之间取元素-wise的最大值。因此，我使用np.fmax来获取元素-wise的最大值。

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

这会产生下面的结果：

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

上面的结果几乎是我需要的。唯一的问题是缺少了nan值。我希望最终的结果是：

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

有没有人能帮我把原始海洋数组中相同索引的值替换成nan？

Answer 1

你可以使用数学运算符：

out = final_cat * np.nan ** np.isnan(ocean + atmos)
out
array([[ 1.,  2.,  2.],
       [ 3., nan,  1.],
       [ 3.,  1., nan]])

Answer 2

一个简单的方法是用 numpy.where 来处理输出：

bin1 = [2, 4, 6]
bin2 = [4, 6, 8]
ocean_cat = np.digitize(ocean, bin1)
atmos_cat = np.digitize(atmos, bin2) 
final_cat = np.where(np.isnan(ocean), np.nan,
                     np.fmax(ocean_cat, atmos_cat))

如果两个数组都可能包含 NaN（缺失值）：

final_cat = np.where(np.isnan(ocean)|np.isnan(atmos),
                     np.nan,
                     np.fmax(ocean_cat, atmos_cat))

或者使用 np.isnan(ocean)&np.isnan(atmos)，这样只有当两个输入都是 NaN 时，结果才会是 NaN。

输出结果：

array([[ 1.,  2.,  2.],
       [ 3., nan,  1.],
       [ 3.,  1., nan]])

对于任意数量的输入数组，通用的方法是：

arrays = [ocean, atmos]
bins = [bin1, bin2]

out = np.where(np.logical_or.reduce([np.isnan(a) for a in arrays]),
               np.nan,
               np.fmax.reduce([np.digitize(a, b) for a,b in zip(arrays, bins)])
               )