Python：将浮点数字符串解包为复数

Question

我刚开始学习编程，现在正在读取一个二进制文件里的信号。这个数据是由两个4字节的浮点数组成的，合起来就是一个复数，这种结构会重复出现大约1500次。

我一直在用一个循环来提取这些数据，并把复数添加到一个数组里。

for x in range(dimX):
    for y in range(dimY):
        complexlist=[]
        #2 floats, each 4 bytes, is one complex number
        trace=stream.readBytes(8*dimZ)
        #Unpack as list of floats
        floatlist=struct.unpack("f"*2*dimZ,trace)
        for i in range(0,len(floatlist)-1,2):
            complexlist.append(complex(floatlist[i],floatlist[i+1]))        
        data[x][y]=np.array(complexlist)

其中，dimX可能有几千个，DimY一般小于30，而dimZ小于1500。

但是在处理大文件时，这种方法非常慢。

有没有办法一次性读取整个数据块，然后直接解压到复数数组里呢？

Answer 1

是的，有这个方法。你可以跳过通过Python的复杂类型，因为在内部，numpy把一个包含n个复数的数组表示为一个包含2n个浮点数的数组。

这里有一个简单的例子，展示了这个是怎么工作的：

>>> import numpy as np
>>> a = np.array([1.,2.,3.,4.])
>>> a
array([ 1.,  2.,  3.,  4.])
>>> a.dtype
dtype('float64')
>>> a.dtype = complex
>>> a
array([ 1.+2.j,  3.+4.j])
>>> 

不过要注意，如果最开始的数组的类型不是float，这个方法就不适用了。

>>> a = np.array([1,2,3,4])
>>> a
array([1, 2, 3, 4])
>>> a.dtype
dtype('int64')
>>> a.dtype = complex
>>> a
array([  4.94065646e-324 +9.88131292e-324j,
         1.48219694e-323 +1.97626258e-323j])
>>>

在你的情况下，你需要的类型是np.dtype('complex64')，因为你每个复数占用64位（2*4*8）。

for x in range(dimX):
    for y in range(dimY):
        #2 floats, each 4 bytes, is one complex number
        trace=stream.readBytes(8*dimZ)
        a = np.frombuffer(trace,dtype=np.dtype('complex64'))
        data[x][y] = a

这样应该能让你快很多。这里有一个例子，展示了numpy.frombuffer()是怎么工作的：

>>> binary_string = struct.pack('2f', 1,2)
>>> binary_string
'\x00\x00\x80?\x00\x00\x00@'
>>> numpy.frombuffer(binary_string, dtype=np.dtype('complex64'))
array([ 1.+2.j], dtype=complex64)
>>> 

编辑：我之前不知道有numpy.frombuffer()这个方法。所以我一直是先创建一个字符数组，然后再改变类型来达到同样的效果。谢谢@wim。

编辑2：

至于进一步的速度优化，你可能会发现使用列表推导式比显式的for循环更快。

for x in range(dimX):
    data[x] = [np.frombuffer(stream.readBytes(8*dimZ), dtype=np.dtype('complex64')) for y in range(dimY)]

再往上提升一下：

data = [[np.frombuffer(stream.readBytes(8*dimZ), dtype=np.dtype('complex64'))
         for y in range(dimY)]
         for x in range(dimX)]