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

Question

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

# 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 无法做到这一点。

Answer 1

这是比@Jeff 的回答更直接的方法，告诉loadtxt 使用将(1.2,0.16) 映射到1.20+0.16j 的辅助函数parse_pair 将其直接加载到复杂数组中：

>>> 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 ]])

或者在熊猫中：

>>> import pandas as pd
>>> f.seek(0)
>>> pd.read_csv(f, delimiter=' ', index_col=0, names=['a', 'b'],
...             converters={1: parse_pair, 2: parse_pair})
             a           b
1  (1.2+0.16j)  (2.8+1.1j)
2  (2.85+6.9j)  (5.8+2.2j)

Answer 2

由于这个问题在 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('()'))))

Answer 3

如果您的文件像您显示的那样只有 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 模块。