在 numpy 数组求和中将 nan 视为零，但所有数组中的 nan 除外

Posted 2023-03-11

【中文标题】在 numpy 数组求和中将 nan 视为零，但所有数组中的 nan 除外

【英文标题】：Treat nan as zero in numpy array summation except for nan in all arrays

【发布时间】：2017-07-01 18:43:02

【问题描述】：

我有两个 numpy 数组 NS、EW 来总结。他们每个人在不同的位置都有缺失值，比如

NS = array([[  1.,   2.,  nan],
       [  4.,   5.,  nan],
       [  6.,  nan,  nan]])
EW = array([[  1.,   2.,  nan],
       [  4.,  nan,  nan],
       [  6.,  nan,   9.]]

如何以 numpy 方式执行求和运算，如果一个数组在某个位置有 nan，它将把 nan 视为零，如果两个数组在同一位置有 nan，则保持 nan。

我期望看到的结果是

SUM = array([[  2.,   4.,  nan],
           [  8.,  5.,  nan],
           [  12.,  nan,   9.]])

当我尝试时

SUM=np.add(NS,EW)

它给了我

SUM=array([[  2.,   4.,  nan],
       [  8.,  nan,  nan],
       [ 12.,  nan,  nan]])

当我尝试时

SUM = np.nansum(np.dstack((NS,EW)),2)

它给了我

SUM=array([[  2.,   4.,   0.],
       [  8.,   5.,   0.],
       [ 12.,   0.,   9.]])

当然，我可以通过元素级的操作来实现我的目标，

for i in range(np.size(NS,0)):
    for j in range(np.size(NS,1)):
        if np.isnan(NS[i,j]) and np.isnan(EW[i,j]):
            SUM[i,j] = np.nan
        elif np.isnan(NS[i,j]):
            SUM[i,j] = EW[i,j]
        elif np.isnan(EW[i,j]):
            SUM[i,j] = NS[i,j]
        else:
            SUM[i,j] = NS[i,j]+EW[i,j]

但是速度很慢。所以我正在寻找一个更 numpy 的解决方案来解决这个问题。

提前感谢您的帮助！

【参考方案1】：

方法 #1： 一种使用 np.where 的方法 -

def sum_nan_arrays(a,b):
    ma = np.isnan(a)
    mb = np.isnan(b)
    return np.where(ma&mb, np.nan, np.where(ma,0,a) + np.where(mb,0,b))

示例运行 -

In [43]: NS
Out[43]: 
array([[  1.,   2.,  nan],
       [  4.,   5.,  nan],
       [  6.,  nan,  nan]])

In [44]: EW
Out[44]: 
array([[  1.,   2.,  nan],
       [  4.,  nan,  nan],
       [  6.,  nan,   9.]])

In [45]: sum_nan_arrays(NS, EW)
Out[45]: 
array([[  2.,   4.,  nan],
       [  8.,   5.,  nan],
       [ 12.,  nan,   9.]])

方法 #2： 混合使用boolean-indexing 可能更快 -

def sum_nan_arrays_v2(a,b):
    ma = np.isnan(a)
    mb = np.isnan(b)
    m_keep_a = ~ma & mb
    m_keep_b = ma & ~mb
    out = a + b
    out[m_keep_a] = a[m_keep_a]
    out[m_keep_b] = b[m_keep_b]
    return out

运行时测试-

In [140]: # Setup input arrays with 4/9 ratio of NaNs (same as in the question)
     ...: a = np.random.rand(3000,3000)
     ...: b = np.random.rand(3000,3000)
     ...: a.ravel()[np.random.choice(range(a.size), size=4000000, replace=0)] = np.nan
     ...: b.ravel()[np.random.choice(range(b.size), size=4000000, replace=0)] = np.nan
     ...: 

In [141]: np.nanmax(np.abs(sum_nan_arrays(a, b) - sum_nan_arrays_v2(a, b))) # Verify
Out[141]: 0.0

In [142]: %timeit sum_nan_arrays(a, b)
10 loops, best of 3: 141 ms per loop

In [143]: %timeit sum_nan_arrays_v2(a, b)
10 loops, best of 3: 177 ms per loop

In [144]: # Setup input arrays with lesser NaNs
     ...: a = np.random.rand(3000,3000)
     ...: b = np.random.rand(3000,3000)
     ...: a.ravel()[np.random.choice(range(a.size), size=4000, replace=0)] = np.nan
     ...: b.ravel()[np.random.choice(range(b.size), size=4000, replace=0)] = np.nan
     ...: 

In [145]: np.nanmax(np.abs(sum_nan_arrays(a, b) - sum_nan_arrays_v2(a, b))) # Verify
Out[145]: 0.0

In [146]: %timeit sum_nan_arrays(a, b)
10 loops, best of 3: 69.6 ms per loop

In [147]: %timeit sum_nan_arrays_v2(a, b)
10 loops, best of 3: 38 ms per loop

【讨论】：

完美运行，并且比我使用的元素级操作快约 200 倍。感谢您的帮助！

【参考方案2】：

实际上您的nansum 方法几乎奏效了，您只需再次添加nans：

def add_ignore_nans(a, b):
    stacked = np.array([a, b])
    res = np.nansum(stacked, axis=0)
    res[np.all(np.isnan(stacked), axis=0)] = np.nan
    return res

>>> add_ignore_nans(a, b)
array([[  2.,   4.,  nan],
       [  8.,   5.,  nan],
       [ 12.,  nan,   9.]])

这将比@Divakars 的答案慢，但我想说你已经很接近了！ :-)

【讨论】：

我明白了，我错过了一个额外的逻辑和语句来过滤掉索引。感谢您的帮助！

【参考方案3】：

我认为我们可以更简洁一些，就像 Divakar 的第二种方法一样。与a = NS 和b = EW：

na = numpy.isnan(a)
nb = numpy.isnan(b)
a[na] = 0
b[nb] = 0
a += b
na &= nb
a[na] = numpy.nan

这些操作在可能的情况下就地完成以节省内存，假设这在您的场景中是可行的。最终结果在a。

【讨论】：

是的，内存越少越好，因为计算可以在大矩阵上执行。我将在我的代码中切换到您的解决方案。谢谢！