将CSV文件加载到NumPy并按名称访问列

Question

我有一个 csv 文件，里面的表头像这样：

给定这个 test.csv 文件：

"A","B","C","D","E","F","timestamp"
611.88243,9089.5601,5133.0,864.07514,1715.37476,765.22777,1.291111964948E12
611.88243,9089.5601,5133.0,864.07514,1715.37476,765.22777,1.291113113366E12
611.88243,9089.5601,5133.0,864.07514,1715.37476,765.22777,1.291120650486E12

我想把它加载成一个有3行7列的矩阵/数组，并且还想通过给定的 列名 来访问 列向量。如果我使用 genfromtxt（如下所示），我得到的是一个有3行（每行一条数据）但没有列的数组。

r = np.genfromtxt('test.csv',delimiter=',',dtype=None, names=True)
print r
print r.shape

[ (611.88243, 9089.5601000000006, 5133.0, 864.07514000000003, 1715.3747599999999, 765.22776999999996, 1291111964948.0)
 (611.88243, 9089.5601000000006, 5133.0, 864.07514000000003, 1715.3747599999999, 765.22776999999996, 1291113113366.0)
 (611.88243, 9089.5601000000006, 5133.0, 864.07514000000003, 1715.3747599999999, 765.22776999999996, 1291120650486.0)]
(3,)

我可以通过列名来获取列向量，方法是这样的：

print r['A']
  [ 611.88243  611.88243  611.88243]

如果我使用 load.txt，那么我得到的是一个有3行7列的数组，但无法通过 列名 来访问 列（如下所示）。

numpy.loadtxt(open("test.csv","rb"),delimiter=",",skiprows=1)

我得到的是

  [ [611.88243,9089.5601,5133.0,864.07514,1715.37476,765.22777,1.291111964948E12]
    [611.88243,9089.5601,5133.0,864.07514,1715.37476,765.22777,1.291113113366E12]
    [611.88243,9089.5601,5133.0,864.07514,1715.37476,765.22777,1.291120650486E12] ]

有没有什么方法可以在 Python 中同时满足这两个需求（即像 np.genfromtxt 一样通过列名访问列，并且像 np.loadtxt 一样得到一个矩阵）？

Answer 1

因为你的数据是同质的——也就是说，所有的元素都是浮点数——所以你可以用 genfromtxt 返回的数据创建一个二维数组的视图。例如，

In [42]: r = np.genfromtxt("test.csv", delimiter=',', names=True)

创建一个 numpy 数组，它是 r 的一个“视图”。这个数组是一个普通的 numpy 数组，但它是用 r 中的数据创建的：

In [43]: a = r.view(np.float64).reshape(len(r), -1)

In [44]: a.shape
Out[44]: (3, 7)

In [45]: a[:, 0]
Out[45]: array([ 611.88243,  611.88243,  611.88243])

In [46]: r['A']
Out[46]: array([ 611.88243,  611.88243,  611.88243])

r 和 a 指向同一块内存：

In [47]: a[0, 0] = -1

In [48]: r['A']
Out[48]: array([  -1.     ,  611.88243,  611.88243])

Answer 2

如果你只用numpy的话，你的选择就只有这几种。要么用一个形状是(3,7)的同质数据类型的ndarray，要么用一个形状是(3,)的结构化数组，这种数组可以包含不同类型的数据。

如果你真的想要一个有标签的列，并且形状是(3,7)的数据结构（还有很多其他好处），你可以使用pandas的DataFrame：

In [67]: import pandas as pd
In [68]: df = pd.read_csv('data'); df
Out[68]: 
           A          B     C          D           E          F     timestamp
0  611.88243  9089.5601  5133  864.07514  1715.37476  765.22777  1.291112e+12
1  611.88243  9089.5601  5133  864.07514  1715.37476  765.22777  1.291113e+12
2  611.88243  9089.5601  5133  864.07514  1715.37476  765.22777  1.291121e+12    

In [70]: df['A']
Out[70]: 
0    611.88243
1    611.88243
2    611.88243
Name: A, dtype: float64

In [71]: df.shape
Out[71]: (3, 7)

---

如果你想用纯NumPy或Python的方式，可以用一个字典来把列名和索引对应起来：

import numpy as np
import csv
with open(filename) as f:
    reader = csv.reader(f)
    columns = next(reader)
    colmap = dict(zip(columns, range(len(columns))))

arr = np.matrix(np.loadtxt(filename, delimiter=",", skiprows=1))
print(arr[:, colmap['A']])

这样就能得到

[[ 611.88243]
 [ 611.88243]
 [ 611.88243]]

这样，arr就是一个NumPy矩阵，可以通过标签来访问列，使用的语法是

arr[:, colmap[column_name]]