Python np.nansum() 数组，但在对两个 nan 求和时不返回 0

Question

上下文：我有一个 3D numpy 数组 (3,10,10)：

arr

。我有一个列表，收集

arr

第 0 维的索引：

grouped_indices

。我想计算

arr

这些分组索引的总和，并将它们存储在主机数组中：

host_arr

。

问题：我正在使用

np.nansum()

，但是两个 NaN 的和给我一个 0，我希望它返回一个 NaN。计算总和后，我不想将所有零设置为 NaN。

问题： 如何计算 n 2D 数组（相同形状）的 nansum，但将所有数组都具有 NaN 的任何单元格设置为 NaN ？

示例：

import numpy as np
import matplotlib.pyplot as plt

# Generate example data
np.random.seed(0)
arr_shape = (10, 10)
num_arrays = 3

# Create a 3D numpy array with random values
arr = np.random.rand(num_arrays, *arr_shape)

# Introduce NaNs
arr[0, :5, :5] = np.nan
arr[1, 2:7, 2:7] = np.nan
arr[2] = np.nan
arr[2, :2, :2] = 10

# Generate a list of arrays containing indices of the 0th dimension of arr
grouped_indices = [np.array([0,1]), np.array([0,1,2])]

# Create a host array that is the sum of grouped_indices slices
host_arr = np.array([np.nansum(arr[indices], axis=0) for indices in grouped_indices])

# Plot the nansums
plt.figure()
plt.imshow(host_arr[0]) # indices [2:5, 2:5] should be NaNs
plt.colorbar()
plt.figure()
plt.imshow(host_arr[1]) # indices [2:5, 2:5] should be NaNs too
plt.colorbar()

Answer 1

IIUC，您可以使用

np.all

+

np.isnan

来检测所有

NaN

的位置，并在

NaN

之后明确将这些值设置为

np.nansum

host_arr = np.array([np.nansum(arr[indices], axis=0) for indices in grouped_indices])

host_arr[0][np.all(np.isnan(arr[grouped_indices[0]]), axis=0)] = np.nan
host_arr[1][np.all(np.isnan(arr[grouped_indices[1]]), axis=0)] = np.nan

那么结果就是：