Python / Numpy - Сохранить массив с заголовками столбцов и строк - PullRequest
Новая
       13

Python / Numpy - Сохранить массив с заголовками столбцов и строк

7 голосов
/ 28 марта 2012

Я хочу сохранить 2D-массив в CSV-файл с информацией о заголовке строки и столбца (например, таблицы).Я знаю, что мог бы использовать аргумент заголовка для numpy.savetxt, чтобы сохранить имена столбцов, но есть ли простой способ также включить какой-либо другой массив (или список) в качестве первого столбца данных (например, заголовки строк)?

Ниже приведен пример того, как я сейчас это делаю.Есть ли лучший способ включить заголовки строк, возможно, какой-то трюк с сохраненным текстом, о котором я не знаю? 

import csv
import numpy as np

data = np.arange(12).reshape(3,4)
# Add a '' for the first column because the row titles go there...
cols = ['', 'col1', 'col2', 'col3', 'col4']
rows = ['row1', 'row2', 'row3']

with open('test.csv', 'wb') as f:
   writer = csv.writer(f)
   writer.writerow(cols)
   for row_title, data_row in zip(rows, data):
      writer.writerow([row_title] + data_row.tolist())

1 Ответ

6 голосов
/ 29 марта 2012

Может быть, вы бы предпочли сделать что-то вроде этого: 

# Column of row titles
rows = np.array(['row1', 'row2', 'row3'], dtype='|S20')[:, np.newaxis]
with open('test.csv', 'w') as f:
    np.savetxt(f, np.hstack((rows, data)), delimiter=', ', fmt='%s')

Это неявное преобразование data в массив строк и занимает около 200 мс на каждый миллион элементов в моем компьютере.

dtype '|S20' означает строки из двадцати символов .Если оно слишком низкое, ваши номера будут разделены: 

>>> np.asarray([123], dtype='|S2')
array(['12'], 
  dtype='|S2')

Другой вариант, который из моего ограниченного тестирования медленнее, но дает вам гораздо больший контроль и не имеет проблемы с прерыванием, будет использовать np.char.mod, как 

# Column of row titles
rows = np.array(['row1', 'row2', 'row3'])[:, np.newaxis]
str_data = np.char.mod("%10.6f", data)
with open('test.csv', 'w') as f:
    np.savetxt(f, np.hstack((rows, str_data)), delimiter=', ', fmt='%s')
...