高效序列化numpy布尔数组

Question

我有成千上万的NumPy布尔数组，想把它们当作字典的键来用。（这个字典的值是我们观察到每个数组的次数。）但是，NumPy数组是不能直接用作键的，因为它们不具备哈希值。我想尽可能高效地把这些数组进行序列化。

这里有两个关于效率的定义需要考虑：

1. 内存使用效率；占用空间越小越好
2. 序列化和恢复数组的计算时间效率；耗时越少越好

我希望在这两个相互竞争的需求之间找到一个好的平衡，不过，对我来说，内存使用效率更重要，我愿意牺牲一些计算时间。

我希望有两个特点能让这个任务更简单：

1. 我可以保证所有数组的大小和形状都是一样的
2. 这些数组是布尔类型的，这意味着可以简单地用1和0来表示它们，形成一个比特序列

有没有一种高效的Python（2.7，或者如果可能的话，2.6）数据结构，可以把这些数组序列化成（也许是某种字节结构），并且能给出数组和这个结构之间转换的例子，以及如何从这个结构恢复到原始数组？

注意，不需要存储每个索引是True还是False的信息；只要有一个结构能存储数组为True的索引就足够了，这样就能恢复出数组。

一个足够的解决方案可以处理一维数组，但一个好的解决方案也应该能处理二维数组，而一个优秀的解决方案则应该能处理更高维度的数组。

Answer 1

我会用np.packbits把数组转换成位域。这种方法在内存使用上比较高效，因为它充分利用了一个字节的所有位。而且代码也相对简单。

import numpy as np
array=np.array([True,False]*20)
Hash=np.packbits(array).tostring()
dict={}
dict[Hash]=10
print(np.unpackbits(np.fromstring(Hash,np.uint8)).astype(np.bool)[:len((array)])

要注意可变长度的布尔数组，代码无法区分比如说6个或7个元素的全为假（False）的数组。对于多维数组，你可能需要进行一些形状调整。

如果这样还是不够高效，而且你的数组很大，你可能可以通过打包来进一步减少内存使用：

import bz2
Hash_compressed=bz2.compress(Hash,1)

不过，这种方法不适用于随机的、无法压缩的数据。

Answer 2

我有三个建议。第一个建议是直接借鉴aix的内容。问题在于，bitarray对象是可变的，它们的hash值和内容无关（也就是说，对于bitarray b，hash(b) == id(b)）。虽然可以通过一些方法来解决这个问题，正如aix的回答所示，但实际上你根本不需要bitarray，你可以直接使用tostring！

In [1]: import numpy
In [2]: a = numpy.arange(25).reshape((5, 5))
In [3]: (a > 10).tostring()
Out[3]: '\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x01
         \x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01'

这样我们就得到了一个不可变的字节字符串，非常适合用作字典的键。需要说明的是，这些转义字符没有被转义，所以这是在不使用bitstring风格序列化的情况下，最紧凑的表示方式。

In [4]: len((a > 10).tostring())
Out[4]: 25

转换回去也很简单且快速：

In [5]: numpy.fromstring((a > 10).tostring(), dtype=bool).reshape(5, 5)
Out[5]: 
array([[False, False, False, False, False],
       [False, False, False, False, False],
       [False,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True]], dtype=bool)
In [6]: %timeit numpy.fromstring((a > 10).tostring(), dtype=bool).reshape(5, 5)
100000 loops, best of 3: 5.75 us per loop

和aix一样，我也没能找到简单的方法来存储维度信息。如果你必须要这个信息，那你可能得忍受更长的键。不过，cPickle似乎是个不错的选择。尽管如此，它的输出还是大了10倍...

In [7]: import cPickle
In [8]: len(cPickle.dumps(a > 10))
Out[8]: 255

而且速度也更慢：

In [9]: cPickle.loads(cPickle.dumps(a > 10))
Out[9]: 
array([[False, False, False, False, False],
       [False, False, False, False, False],
       [False,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True]], dtype=bool)
In [10]: %timeit cPickle.loads(cPickle.dumps(a > 10))
10000 loops, best of 3: 45.8 us per loop

我的第三个建议是使用bitstring，特别是bitstring.ConstBitArray。它的思路和的解决方案相似，但ConstBitArray是不可变的，所以在hash方面能满足你的需求。

In [11]: import bitstring

你需要显式地将numpy数组展平：

In [12]: b = bitstring.ConstBitArray((a > 10).flat)
In [13]: b.bin
Out[13]: '0b0000000000011111111111111'

它是不可变的，所以哈希效果很好：

In [14]: hash(b)
Out[14]: 12144

转换回数组也超级简单，但同样，形状信息会丢失。

In [15]: numpy.array(b).reshape(5, 5)
Out[15]: 
array([[False, False, False, False, False],
       [False, False, False, False, False],
       [False,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True]], dtype=bool)

而且这是最慢的选项：

In [16]: %timeit numpy.array(b).reshape(5, 5)
1000 loops, best of 3: 240 us per loop

---

这里有更多信息。我一直在尝试和测试，得出了以下结论。首先，bitarray在正确使用时比bitstring快得多：

In [1]: %timeit numpy.array(bitstring.ConstBitArray(a.flat)).reshape(5, 5)
1000 loops, best of 3: 283 us per loop

In [2]: %timeit numpy.array(bitarray.bitarray(a.flat)).reshape(5, 5)
10000 loops, best of 3: 19.9 us per loop

其次，从上面可以看出，所有的tostring操作都是不必要的；你也可以直接将numpy数组展平。但实际上，aix的方法更快，所以下面修订后的数据就是基于这个。

下面是结果的完整总结。首先，定义：

small_nda = numpy.arange(25).reshape(5, 5) > 10
big_nda = numpy.arange(10000).reshape(100, 100) > 5000
small_barray = bitarray.bitarray(small_nda.flat)
big_barray = bitarray.bitarray(big_nda.flat)
small_bstr = bitstring.ConstBitArray(small_nda.flat)
big_bstr = bitstring.ConstBitArray(big_nda.flat)

keysize是sys.getsizeof({small|big}_nda.tostring())、sys.getsizeof({small|big}_barray) + sys.getsizeof({small|big}_barray.tostring())或sys.getsizeof({small|big}_bstr) + sys.getsizeof({small|big}_bstr.tobytes())的结果——后两种方法返回的是打包成字节的位字符串，所以应该能很好地估算每种方法占用的空间。

speed是从{small|big}_nda转换为键再转换回来所需的时间，加上将bitarray对象转换为字符串以进行哈希的时间，这个时间如果你缓存字符串就是一次性成本，如果不缓存就是每次字典操作的成本。

         small_nda   big_nda   small_barray   big_barray   small_bstr   big_bstr

keysize  64          10040     148            1394         100          1346

speed    2.05 us     3.15 us   3.81 us        96.3 us      277 us       92.2ms  
                             + 161 ns       + 257 ns 

从结果来看，bitarray的速度非常快，而aix建议的bitarray子类应该也能很好地工作。它确实比bitstring快多了。很高兴看到你接受了那个答案。

另一方面，我仍然对numpy.array.tostring()方法有些依恋。它生成的键在渐进上大约是8倍，但对于大数组来说，速度提升依然很可观——在我机器上对于大数组大约快了30倍。这是个不错的权衡。不过，可能在成为瓶颈之前不值得去麻烦它。

Answer 3

最开始，我建议使用bitarray这个库。不过，正如@senderle所说的，由于bitarray是可变的，所以不能直接用它作为字典（dict）的键。

这里有一个修改后的解决方案（内部仍然基于bitarray）：

import bitarray

class BoolArray(object):

  # create from an ndarray
  def __init__(self, array):
    ba = bitarray.bitarray()
    ba.pack(array.tostring())
    self.arr = ba.tostring()
    self.shape = array.shape
    self.size = array.size

  # convert back to an ndarray
  def to_array(self):
    ba = bitarray.bitarray()
    ba.fromstring(self.arr)
    ret = np.fromstring(ba.unpack(), dtype=np.bool)[:self.size]
    return ret.reshape(self.shape)

  def __cmp__(self, other):
    return cmp(self.arr, other.arr)

  def __hash__(self):
    return hash(self.arr)

import numpy as np

x = (np.random.random((2,3,2))>0.5)
b1 = BoolArray(x)
b2 = BoolArray(x)
d = {b1: 12}
d[b2] += 1
print d
print b1.to_array()

这个方法适用于Python 2.5及以上版本，每个数组元素需要一个比特，并且支持任何形状或维度的数组。

补充说明：在最近的版本中，你需要把ba.tostring和ba.fromstring替换成ba.tobytes和ba.frombytes（从0.4.0版本开始不再推荐使用）。