Tuesday, December 15, 2009

Matplotlib in OS X (3)


I've been playing with matplotlib, and I have to say that the figures it produces are beautiful. Although the documentation is extensive, I do find it hard to navigate to the right pages, or find helpful examples by direct search. But of course, GIYF.

Here is a proof of principle for a heat map in a few short lines of code. I still have to work on placing text, getting the right colors, etc. The fact that it's all Python is terrific. We use a scatterplot and adjust the size of the plot symbol (a square) to be as large as we need.

Probably it's obvious, but the call plt.axis('equal') is crucial because by default the axes are not typically equal and the "squares" are therefore not square. The call yL = sum(yL,[]) is a clever trick I found on Stack Overflow for flattening a list.

[UPDATE: As is not uncommon, I found a simpler API after making this post. The result is shown below. The code is the second of the two samples.]



import matplotlib.pyplot as plt
import random

xL = range(1,11) * 10
yL = [[j]*10 for j in range(1,11)]
yL = sum(yL,[])

R = xrange(100000)
cL = [random.choice(R) for i in range(len(xL))]

fig = plt.figure()
ax = fig.add_subplot(111,axisbg='0.9') # light gray
plt.axis('equal')
ax.scatter(xL, yL,c=cL,
marker='s',edgecolor='w',s=1250)

plt.savefig('example.pdf')


import matplotlib.pyplot as plt
import numpy as np
import random

R = xrange(10000)
data = [random.choice(R) for i in range(100)]
data = np.array(data)
data.shape = (10,10)

fig = plt.figure()
plt.hot()
plt.pcolormesh(data)
plt.colorbar()
plt.savefig('example.png')