heatmap.py

#! /usr/bin/python3

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.patches import Rectangle
import numpy as np
import math
import sys

def verbose_usage_and_exit():
    sys.stderr.write('\nheatmap.py - Draw a heatmap from a .CSV file\n')
    sys.stderr.write('\n')
    sys.stderr.write('Usage:\n')
    sys.stderr.write('\t%s [CSV file] [output PNG file]\n' % sys.argv[0])
    sys.stderr.write('\n')
    exit(0)

if (len(sys.argv) != 3):
    verbose_usage_and_exit()

csvfile = sys.argv[1]
pngfile = sys.argv[2]

with open(csvfile, 'r') as mapfile:
    line = mapfile.readline()
    elems = line[:-1].split(',')
    # print(elems)
    i = 0
    nume = len(elems)
    refs = []
    total_len = 0
    while i < nume:
        refs.append((elems[i], int(elems[i+1])))
        total_len += int(elems[i+1])
        i += 2
    # print(refs)
    
    lines = mapfile.readlines()
    numrows = len(lines)
    numcols = len(lines[0].split(','))
    if numrows != numcols:
        print('Numbers of rows and columns do not match!')
        exit(1)
    
    # Using only numrows from now, since numrows=numcols
    arr = []    
    for i in range(numrows):
        tmp = [float(e) for e in lines[i][:-1].split(',')]
        arr.append([math.log(x+1, 1.5) for x in tmp])
        # arr.append([math.sqrt(x) for x in tmp])
    
    heatmap = np.array(arr)
    fig,ax1 = plt.subplots(1,1, figsize=(16,12))
    fig.set_facecolor("white")
    
    # Calculating ticks
    ticks = []
    cur_len = 0
    for ref in refs:
        cur_len += ref[1]
        ticks.append(cur_len*numrows/total_len)
    
    labels = range(1,len(refs)+1)
    minorticks = []
    prev_tick = 0
    for tick in ticks:
        minorticks.append((prev_tick + tick)//2)
        prev_tick = tick
    
    # X axis ticks and labels
    ax1.xaxis.set_major_locator(ticker.FixedLocator(ticks))
    ax1.xaxis.set_minor_locator(ticker.FixedLocator(minorticks))
    ax1.xaxis.set_major_formatter(ticker.NullFormatter())
    ax1.xaxis.set_minor_formatter(ticker.FixedFormatter(labels))
    
    # Y axis ticks and labels
    ax1.yaxis.set_major_locator(ticker.FixedLocator(ticks))
    ax1.yaxis.set_minor_locator(ticker.FixedLocator(minorticks))
    ax1.yaxis.set_major_formatter(ticker.NullFormatter())
    ax1.yaxis.set_minor_formatter(ticker.FixedFormatter(labels))
        
    # ax1.tick_params(length=10, width=2, color="red", labelbottom=False, labelleft=False)
    ax1.tick_params(which="major", length=10, width=2, colors="red")
    ax1.tick_params(which="minor", length=0, width=0, colors="red")
    
    # Drawing the legend from references
    dummylegend = Rectangle((0, 0), 1, 1, fc="w", fill=False, edgecolor='none', linewidth=0)
    legendtext = []
    for i in range(len(refs)):
        if i > 30:
            legendtext.append("...")
            break
        ref = refs[i];
        text = "{0} - {1}".format(i+1, ref[0])
        legendtext.append(text)
    ax1.legend([dummylegend]*len(legendtext), legendtext, title="Legend:", bbox_to_anchor=(1.05, 1), loc='upper left')
    ax1.imshow(heatmap, cmap='hot', interpolation='nearest')
    fig.savefig(pngfile, bbox_inches='tight')