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 pltimport numpy as npdef 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 leftdef 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)`