Обработать nan как ноль в суммировании массива numpy за исключением nan во всех массивах

Question

У меня есть два массива 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 в том же месте.

В результате я ожидаю увидеть это

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] 

, но это очень медленно. Поэтому я ищу более бесчисленное решение для решения этой проблемы.

Благодарим за помощь!

Answer 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 

Answer 2

Я думаю, что мы можем получить немного более кратким, в том же ключе, что и второй подход 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.

Answer 3

На самом деле ваш 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.]]) 

Это будет медленнее, чем @Divakar сек ответ, но я хотел бы отметить, что вы были очень близки уже!:-)