从大型掩码数组中计算Numpy平均值

Question

怎样从一个很大的被屏蔽数组中获取平均值呢？通常我会直接使用 .mean()，但对于非常大的数组，这个方法对我来说是失败的。

想象一下，创建一个包含一百万个元素的数组，所有元素的值都是500，像这样：

a = np.ones(1000000, dtype=np.int16) * 500

然后创建一个随机的屏蔽，并把这两个结合成一个新的 masked array：

mask = np.random.randint(0, 2, a.size)
b = np.ma.masked_array(a, mask=mask)

这个数组 b 继承了 dtype，并且也是 int16 类型。

从 b 中获取平均值可以用不同的方法，但结果都是一样的。不过，non-ma 的函数会忽略屏蔽，不应该使用。

print(np.average(b), np.mean(b), np.ma.average(b), np.ma.mean(b))
(500.0, 500.0, 500.0, 500.0)

但是如果我把原始数组的大小从一百万增加到一千万，结果就变成了：

print(np.average(b), np.mean(b), np.ma.average(b), np.ma.mean(b))
(500.0, -359.19365132075774, -359.19365132075774, -359.19365132075774)

现在只有 np.average 看起来 是正确的，但如前所述，它忽略了屏蔽，计算的是整个数组的平均值。这可以通过将一些被屏蔽的值改为 b[b.mask] = 1000 来证明。我本以为 np.mean 也会这样做。

把被屏蔽的数组 b 转换为 float32 会得到：

b = b.astype(np.float32)
print(np.average(b), np.mean(b), np.ma.average(b), np.ma.mean(b))
(511.18945, 510.37895680000003, 510.37895680000003, 510.37895680000003)

而把被屏蔽的数组 b 转换为 float64（根据文档，这应该是默认的）会得到：

b = b.astype(np.float64)
print(np.average(b), np.mean(b), np.ma.average(b), np.ma.mean(b))
(500.0, 500.0, 500.0, 500.0)

所以转换为 float64 似乎有效，但我更希望避免这样，因为这会大大增加内存占用。

在测试这些内容时，我还注意到，对非屏蔽数组（a）调用 np.ma.average 在大小为一百万时结果正确，而在一千万时结果错误，而 np.ma.mean 对这两种大小都是正确的。

有人能解释一下在这个例子中 dtype 和数组的 size 之间的关系吗？我对这种情况发生的原因和如何正确处理它感到有点困惑。

所有操作都是在64位的Win 7机器上使用Numpy 1.8.1完成的。通过conda安装。

这里有一个笔记本，复现了我所做的内容：

http://nbviewer.ipython.org/gist/RutgerK/69b60da73f464900310a

Answer 1

当你用 b[b.mask] = 1000 这样的方式改变一些被屏蔽的值时，就能看到这个问题。我本来以为 np.mean 也会有同样的效果。

其实这样理解是不对的，b.mask 是在有被屏蔽值的地方返回 True。当你给这些被屏蔽的值赋新值时，实际上是把它们“解屏蔽”了，这样数组里的所有值就都变成有效的了。你可以用 b[np.invert(b.mask)] 来代替。

所以这样做应该是可以的：

import numpy as np

a = np.ones(10000000, dtype=np.int64) * 500

mask = np.random.randint(0, 2, a.size)
b = np.ma.masked_array(a, mask=mask)

b[np.invert(b.mask)] = 1000
print(np.average(b), np.mean(b), np.ma.average(b), np.ma.mean(b))

这样会给你正确的值，除了 np.average 可能不行。

另外，当你得到负数或不正确的值时，通常是因为整数溢出。用 dtype=np.int64 代替应该能解决这个问题。

编辑： 还有一个选择是用 Python 的整数，设置 dtype=object，而不是固定宽度的整数，但这样会慢一些，这个改动会导致 np.average 崩溃，但其他方法都能正常工作。

编辑 2： 正如评论中提到的，在这种情况下其实不需要增加数组元素的大小，我们可以直接调用 np.mean(b, dtype=np.float64)，这样 np.mean 就会使用更大的累加器来避免溢出。