如何用NumPy从文件中读取复数?

时间:2013-05-20 23:22:42

标签: python numpy complex-numbers

我需要以下列格式读取复数列:

# index; (real part, imaginary part); (real part, imaginary part) 

  1              (1.2, 0.16)                  (2.8, 1.1)
  2              (2.85, 6.9)                  (5.8, 2.2)

NumPy似乎非常适合只使用一个分隔符读取数据列,但括号似乎破坏了使用numpy.loadtxt()的任何尝试。

有没有一种聪明的方法可以用Python读取文件,或者最好是只读取文件,删除所有括号,然后将其提供给NumPy?

这需要为成千上万的文件完成,所以我想要一种自动化的方式,但是NumPy可能不具备此功能。

3 个答案:

答案 0 :(得分:4)

这是一个比@ Jeff的回答更直接的方法,告诉loadtxt使用将parse_pair映射到{{1}的帮助函数(1.2,0.16)将其直接加载到复杂数组中}:

1.20+0.16j

或者在熊猫中:

>>> import re
>>> import numpy as np

>>> pair = re.compile(r'\(([^,\)]+),([^,\)]+)\)')
>>> def parse_pair(s):
...    return complex(*map(float, pair.match(s).groups()))

>>> s = '''1 (1.2,0.16) (2.8,1.1)
2 (2.85,6.9) (5.8,2.2)'''
>>> from cStringIO import StringIO
>>> f = StringIO(s)

>>> np.loadtxt(f, delimiter=' ', dtype=np.complex,
...            converters={1: parse_pair, 2: parse_pair})
array([[ 1.00+0.j  ,  1.20+0.16j,  2.80+1.1j ],
       [ 2.00+0.j  ,  2.85+6.9j ,  5.80+2.2j ]])

答案 1 :(得分:4)

由于这个问题在pandas中仍然是not resolved,让我添加另一个解决方案。您可以在阅读之后使用单行修改DataFrame

import pandas as pd

df = pd.read_csv('data.csv')
df = df.apply(lambda col: col.apply(lambda val: complex(val.strip('()'))))

答案 2 :(得分:2)

如果您的文件只有5个列,就像您已经显示的那样,您可以将其提供给带有正则表达式进行转换的pandas,并在每行上用逗号替换括号。之后,您可以按this SO answer中的建议合并它们以获得复数。

Pandas让它变得更容易,因为你可以将正则表达式传递给它的read_csv方法,这样你就可以编写更清晰的代码并使用像这样的转换器。 numpy版本的优势在于您可以为分隔符传递正则表达式。

import pandas as pd
from StringIO import StringIO
f_str = "1 (2, 3) (5, 6)\n2 (3, 4) (4, 8)\n3 (0.2, 0.5) (0.6, 0.1)"
f.seek(0)

def complex_converter(txt):
    txt = txt.strip("()").replace(", ", "+").replace("+-", "-") + "j"
    return complex(txt)

df = pd.read_csv(buf, delimiter=r" \(|\) \(", converters = {1: complex_converter, 2: complex_converter}, index_col=0)
编辑:看起来@Dougal在我发布之前提出了这个......真的只取决于你想要如何处理复数。我希望能够避免明确使用re模块。