将复杂数据读取到numpy数组中

Question

我需要从一个文本文件中读取复数，并把它们放进一个numpy数组里。我的问题和这个关于使用numpy.savetxt和numpy.loadtxt读写复数的问题有点像，不过这里的解决办法是改变数据保存的格式。但我没有这个选择，因为这个文本文件是由其他软件生成的，我无法去修改它。下面是这个文本文件的一个示例：

25 + 0i
8.43818 + -4.94194i
4.46817 + -5.08305i
4.55764 + -3.02201i
2.69138 + -5.43104i
-0.151334 + -4.83717i
1.98336 + -1.3339i
3.59002 + -0.932973i
1.42727 + -0.617317i
1.02005 + -1.14214i
-0.15564 + 2.74564i

我尝试过以下方法：

np.loadtxt('file.txt',delimiter='\n',dtype=np.complex128)

但是我遇到了这个错误：

ValueError: complex() arg is a malformed string

我看到的一些帖子提到，这可能是某些行中+ -表示法的问题，不过即使去掉多余的+符号，我还是会遇到同样的错误。

Answer 1

@abarnert的回答很好，但别忘了loadtxt有一个参数converters，可以用来定制如何处理某个字段。下面有几个例子，展示了如何处理这个文件。

在第一个版本中，默认的分隔符（空格）被保留，所以有三列数据。usecols用来忽略中间那一列（'+'）。一个转换器被用来把最后一列（第二列）转换成浮点数。它通过切片来去掉最后一个字符，也就是'i'。使用这些参数后，loadtxt返回一个形状为(11, 2)的浮点数数组。通过调用view方法并指定类型np.complex128，这个数组被转换成一个形状为(11, 1)的浮点数数组。最后，用[:,0]索引可以得到一个一维的复数值数组。（也可以用ravel、squeeze或reshape来替代。）

In [24]: loadtxt('file.txt', converters={2:lambda f: float(f[:-1])}, usecols=(0,2)).view(np.complex128)[:,0]
Out[24]: 
array([ 25.000000+0.j      ,   8.438180-4.94194j ,   4.468170-5.08305j ,
         4.557640-3.02201j ,   2.691380-5.43104j ,  -0.151334-4.83717j ,
         1.983360-1.3339j  ,   3.590020-0.932973j,   1.427270-0.617317j,
         1.020050-1.14214j ,  -0.155640+2.74564j ])

下一个版本把整行当作一个字段来处理。下面的函数将文件中使用的字符串格式转换为复数：

In [36]: def to_complex(field):
   ....:     return complex(field.replace(' ', '').replace('+-', '-').replace('i', 'j'))
   ....: 

例如

In [39]: to_complex('8.43818 + -4.94194i')
Out[39]: (8.43818-4.94194j)

这个对loadtxt的调用把每一行当作一个字段，并使用转换器把每个字段转换为复数：

In [37]: loadtxt('file.txt', converters={0: to_complex}, dtype=np.complex128, delimiter=';')
Out[37]: 
array([ 25.000000+0.j      ,   8.438180-4.94194j ,   4.468170-5.08305j ,
         4.557640-3.02201j ,   2.691380-5.43104j ,  -0.151334-4.83717j ,
         1.983360-1.3339j  ,   3.590020-0.932973j,   1.427270-0.617317j,
         1.020050-1.14214j ,  -0.155640+2.74564j ])

（分隔符设置为';'，但也可以是文件中没有出现的任何字符。）

Answer 2

不过，这里的解决办法是改变数据保存的格式。

好消息是，你不需要这么做！

numpy.loadtxt 可以接受任何可迭代的行，不仅仅是文件对象。

所以，你可以把文件对象放进一个简单的生成器里，这个生成器可以实时转换行，然后把它传给 loadtxt，这样大家都会满意。

像这样：

def transform_complex(line):
    # insert your code here

with open('file.txt', 'rb') as f:
    lines = map(transform_complex, f)
    arr = np.loadtxt(lines, dtype=np.complex128)

（如果你在用 Python 2.x，并且文件很大，可能想用 itertools.imap 而不是 map。）

在“在这里插入你的代码”部分，你可以填入那个有效但不被接受的解决方案的代码，因为它需要修改文件。由于我在你的链接中没有看到这样的答案，我不太确定那是什么，但举个例子，可能是这个：

def transform_complex(line):
    return line.replace(b'+ -', b'- ')

---

在本地测试时，实际上你的输入有三个问题。

你可以用 savetxt 来测试输出应该是什么样子的。例如：

>>> arr = np.array([1-2j])
>>> f = io.BytesIO()
>>> np.savetxt(f, arr)
>>> f.getvalue()
b' (1.000000000000000000e+00-2.000000000000000000e+00j)\n'

（在 Python 2.x 中，你不会看到 b 前缀。）

并不是所有的差异都很重要——你不需要使用科学计数法，不需要括号等等——但看起来这三个是必须注意的：

• 复数中的 + 周围不能有空格。
• 虚数单位必须用 j，而不是 i。
• 不允许出现 +-。

所以：

def transform_complex(line):
    return line.replace(b' ', b'').replace(b'+-', b'-').replace(b'i', b'j')