如何计算沿某一维度的 3d 体积的非零元素的平均值

Posted 2023-03-24

【中文标题】如何计算沿某一维度的 3d 体积的非零元素的平均值

【英文标题】：How to compute mean for non zero elements of a 3d volume along one of the dimensions

【发布时间】：2021-03-27 13:31:15

【问题描述】：

我有一个 3d 体积 DxHxW，我只想计算沿 D 维度的非零值的平均值，以获得每个元素的 2d HxW 平均值。这是我当前的实现，将所有内容设置为nan，然后计算nanmean。但是，它很慢，在 Python 中是否有更优化的方法来执行此操作？

volume_3d[volume_3d == 0] = np.nan
mean_2d = np.nanmean(volume_3d, 0))
mean_2d[np.isnan(mean_2d)] = 0

【问题讨论】：

尝试将沿轴的总和除以非零计数。我相信它会更快

【参考方案1】：

我将测试数组创建为：

D = 2; H = 3; W = 4
volume_3d = np.arange(1, D * H * W + 1, dtype=float).reshape(D, W, H)
volume_3d[0, 1, 0] = volume_3d[0, 2, 1] = volume_3d[1, 2, 1] = 0

其实，为了不改变源数组，你原来的代码 可以改写成如下函数：

def f1(arr):
    wrk = np.copy(arr)
    wrk[wrk == 0] = np.nan
    mean_2d = np.nanmean(wrk, 0)
    mean_2d[np.isnan(mean_2d)] = 0
    return mean_2d

要得到结果，你可以运行res1 = f1(volume_3d)，结果是：

array([[ 7.,  8.,  9.],
       [16., 11., 12.],
       [13.,  0., 15.],
       [16., 17., 18.]])

我的建议如何更快地计算这个结果是：

def f2(arr):
    cnt = np.sum(arr != 0, axis=0)
    tot = np.sum(arr, axis=0)
    return np.divide(tot, cnt, out=np.zeros_like(tot), where=cnt!=0)

使用 %timeit 我比较了执行时间：

对于您的代码：330 µs, 对于我的代码：160 µs,

即我的代码运行速度提高了 2 倍以上。

另一个值得一提的细节是您的代码引发了RuntimeWarning： 空切片的平均值和我的运行没有警告。

我想对于更大的数组，差异应该更明显。 自己试试吧。