Convert CSV to heatmap

to_heatmap.py

import numpy as np
import codecs
import matplotlib.pyplot as plt
import sys
import csv
reload(sys)
sys.setdefaultencoding("ISO-8859-1")

in_file = sys.argv[1]
out_file = sys.argv[2]

fig, ax = plt.subplots()

reader = csv.reader(codecs.open(in_file,'rb','utf8'), delimiter="\t")
xlabels = reader.next()[2:]

ylabels = []
matrix = []
for row in reader:
    ylabels.append(row[0] + ' & ' + row[1])
    row_data = [float(item) for item in row[2:]]
    matrix.append(row_data)

heatmap = ax.pcolor(np.array(matrix), cmap=plt.cm.Reds,edgecolors='k')

# n x n matrix
ax.set_xticks(np.arange(0, len(xlabels)) + 0.5)
ax.set_yticks(np.arange(0, len(ylabels)) + 0.5)
# Hide ticks on both axes.
ax.tick_params(axis='both', which='both',length=0)

ax.set_xticklabels(xlabels, minor=False, fontsize=10,
rotation='vertical')
ax.set_yticklabels(ylabels, minor=False, fontsize=10)

plt.title('Title')

for y in range(len(ylabels)):
    for x in range(len(xlabels)-1):
        plt.text(x + 0.5, y + 0.5, '%.3f' % matrix[y][x],
        horizontalalignment='center', verticalalignment='center',
        fontsize=6)

plt.colorbar(heatmap)

print "Saving heatmap to ",out_file
plt.savefig(out_file)