如何使用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无法胜任此任务。

- f4hy

Answer 1

以下是比@Jeff的回答更为直接的方法，使用辅助函数parse_pair将(1.2,0.16)映射到1.20+0.16j，以复数数组的形式直接加载到loadtxt中：

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

或者在pandas中：

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

- Danica

---

我以前从未听说过pandas。看起来非常有用。谢谢！ - f4hy

在我的电脑上运行这个确切的示例会返回一个错误："ValueError: complex() arg is a malformed string"。有什么解释吗？ - Joey Dumont

@JoeyDumont 我刚刚再次尝试了一下，使用的是Python 2.7.6 / NumPy 1.7.1，依然可以正常工作。你使用的是哪个版本的Python / NumPy？ - Danica

@JoeyDumont 实际上，我刚刚尝试了Python 3.3，并得到了类似的错误。稍微追踪一下，似乎是因为NumPy将其输入读取为“bytes”，例如complex(b'1')失败并引发了TypeError，而float(b'1')却可以正常工作。这应该是使用loadtxt(..., dtype=complex)读取任何内容时都会遇到的问题；您可以通过添加一个转换器来解决其他所有问题，例如lambda x: complex(x.decode() if isinstance(x, bytes) else x)。parse_pair还需要对其输入进行解码，以便将其与正则表达式匹配。真是令人烦恼... - Danica

我使用 Python 2.7.6 / numpy 1.8.0。我不确定如何更改 parse_pair 来尝试您的解决方案。 - Joey Dumont

啊，numpy 1.8可能正在进行相同类型的字节处理。我手头没有安装来尝试它，但在2.7中，你只需要用s.decode()替换pair.match(s)即可。不过这样做不会给你带来错误；尝试将{0: lambda x: complex(x.decode())}添加到converters字典中？ - Danica

Answer 2

由于这个问题在pandas中仍然未解决，让我再提供另一种解决方案。您可以在读取DataFrame后使用一行代码进行修改：

import pandas as pd

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

- jakosz

Answer 3

如果您的文件只有像您展示的那样的5列，您可以使用正则表达式将括号替换为每行上的逗号，并将其提供给pandas进行转换。之后，您可以像这个SO答案中建议的那样将它们组合起来得到复数。

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模块。

- Jeff Tratner

---

我之前从未听说过pandas。看起来相当有用。谢谢！ - f4hy

只是一个注意事项，如果其他人看这个寻求帮助，你的 complex_converter 如果虚部为负数是无法工作的。complex 方法似乎不适用于“1+-3j”。 - f4hy

@f4hy 已经修复了，将 "+-" 替换为 "-"。现在可以处理负虚数 :) - Jeff Tratner