Python for Bioinformatics: Yet another heatmapper

Saturday, March 13, 2010

Yet another heatmapper

I'm planning to try a small Python MCMC simulation, utilizing a simple 2D grid of densities. I wrote a new heatmap script as a first step. I'll post the code today and then use it in later updates.

One small thing, the original version asked matplotlib for the xlim and ylim and then calculated fractional square dimensions to use. Since it seemed a bit slow, I changed the script so that we tell the Axes object what dimensions we want (the array's shape, naturally). Then there are no fractions and the increment is +1. However, it was not significantly faster at all. Here is the timeit output:

$ python -m timeit -s 'from HM_Utils import demo' 'demo(10)'
10 loops, best of 3: 225 msec per loop
$ python -m timeit -s 'from HM_Utils_frac import demo' 'demo(10)'
10 loops, best of 3: 229 msec per loop

Here's the listing:

import matplotlib.pyplot as plt
import numpy as np

def init_fig():
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.xaxis.set_visible(False)
    ax.yaxis.set_visible(False)
    return ax,fig
    
def show_array(A):
    for row in A:
        L = [str(round(n,2)).ljust(4) for n in row]
        print '  '.join(L)
    
# origin is lower left
def make_heatmap(A,ax):
    R,C = A.shape
    ax.set_xlim(C)
    ax.set_ylim(R)
    for r in range(R):
        for c in range(C):
            n = A[r][c]
            assert 0.0 <= n <= 1.0
            p = plt.Rectangle((r,c),
                    width=1,height=1,
                    # black is 0.0, w is 1.0
                    facecolor=str(1-n),
                    edgecolor='w')
            ax.add_patch(p)
            
def demo(N=10):
    ax,fig = init_fig()
    A = np.random.uniform(size=N**2)
    A.shape = (N,N)
    if N <= 10 and __name__ == '__main__':  
        show_array(A)
    make_heatmap(A,ax)
    plt.savefig('example.png')
    
if __name__ == '__main__':
    demo(50)