3D数组分组求和：如何让np.nansum全NaN单元格返回NaN而非0？

问题说明 #

现有形状为(3,10,10)的3D numpy数组arr，以及按第0维分组的索引列表grouped_indices。需要计算每组索引对应的切片的和存入host_arr，但使用np.nansum()时，当对应位置所有元素都是NaN，求和结果会变成0，需求是让这种情况返回NaN，且不能直接把求和结果里的所有0替换成NaN（避免误改非全NaN但和为0的情况）。

解决方案 #

核心思路是先计算分组的nansum，再精准定位该分组内所有数组对应位置全为NaN的区域，将这些区域的求和结果替换为NaN：

1. 对每个分组索引，提取对应的arr切片
2. 计算该切片的nansum结果
3. 生成掩码：判断切片中每个位置是否在所有子数组里都是NaN
4. 将nansum结果中掩码为True的位置设为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])]

# 修正后的host_arr生成逻辑
host_arr = []
for indices in grouped_indices:
    # 提取当前分组的切片
    slice_arr = arr[indices]
    # 计算nansum
    sum_result = np.nansum(slice_arr, axis=0)
    # 生成全NaN位置的掩码
    all_nan_mask = np.isnan(slice_arr).all(axis=0)
    # 将全NaN位置的sum结果设为NaN
    sum_result[all_nan_mask] = np.nan
    host_arr.append(sum_result)
host_arr = np.array(host_arr)

# Plot the nansums
plt.figure()
plt.imshow(host_arr[0]) # 现在[2:5, 2:5]区域是NaN
plt.colorbar()
plt.figure()
plt.imshow(host_arr[1]) # 现在[2:5, 2:5]区域也是NaN
plt.colorbar()
plt.show()

关键代码解释 #

• np.isnan(slice_arr).all(axis=0)：对分组切片的每个位置，检查所有第0维的元素是否都是NaN，返回和2D子数组同形状的布尔掩码，axis=0表示沿分组维度做判断。
• sum_result[all_nan_mask] = np.nan：仅替换全NaN的位置，不会影响那些非全NaN但求和结果为0的位置，完全匹配需求。

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