matplotlib comes with tons of fantastic examples. I’m not as familiar with matplotlib as I probably should be, so I often find myself wanting to tinker a bit, but needing to refer to those examples. Since matplotlib comes with such wonderful documentation, I though it would be great to just turn those docs into IPython Notebooks for easy tinkering. That’s probably biting off a bit more than I want to chew at the moment, considering that the matplotlib docs are fairly involved and written in reStructuredText instead of markdown (what the IPython Notebook uses).

Luckily, the IPython Notebook format is so mind-bendingly sane that I didn’t even need to read any documentation to understand it. So, instead, I wrote a bit of code that gobbles up matplotlib example scripts and spits out IPython Notebooks. The whole notebook is JSON, but I only want simple things, so I hardcode everything except for the cells. (After Daniel’s comment below, I started to write my own JSONEncoder. Then, I realized that I was right about the “it’s all JSON” thing and rewrote the notebook class). I have a little IPyNB class that knows how to add cells to itself and spit out the results as strings and files:

import os, json, glob
        
class IPyNB():
    def __init__(self,name):
        self.name = name
        self.cells = []
        self.d = {‘metadata’:{‘name’:”},
                  ‘nbformat’:3,
                  ‘nbformat_minor’:0,
                  ‘worksheets’:[{'cells':[],’metadata’:{}}]
                  }
        
    def addcell(self,celltype,content):
        cell = {‘cell_type’:celltype,
                  ‘metadata’:{}
                  }
        if celltype == ‘code’:
            cell['collapsed'] = False
            cell['input'] = content.split(‘\n’)
            cell['language'] = ‘python’
            cell['outputs'] = []
        elif celltype == ‘markdown’:
            cell['source'] = content.split(‘\n’)
        #elif self.celltype = ‘raw’:
        else:
            raise NotImplemented(‘Unknown cell type: {celltype}’.format(celltype=self.celltype))
        self.d['worksheets'][0]['cells'].append(cell)

    def tostring(self):
        return json.dumps(self.d,
                          sort_keys=True,
                          indent=1, separators=(‘,’, ‘: ‘),
            )

    def tofile(self,outdir=’.',overwrite=False):
        fname = os.path.join(outdir,self.name+’.ipynb’)
        if os.path.isfile(fname) and not overwrite:
            raise IOError(“File {fname} already exists”.format(fname=fname))
        f = open(fname,’w')
        f.write(self.tostring())
        f.close()

I’ve defined a few celltypes. It’s easy to add more if there are more.

And then a couple of functions to read in the matplotlib examples and spit out notebooks:

def make_mpl_examples(subdir=’images_contours_and_fields’,basedir=’~/coding/matplotlib/examples’,outdir=’.',overwrite=False):
    n = IPyNB(subdir + ‘ Examples’)
    n.addcell(‘markdown’,”"”#matplotlib examples
The below examples are taken directly from the matplotlib example directory {subdir} and rendered in an IPython Notebook. Before you run them, you should execute one of the first two cells, depending on whether you want inline rendering or not. The inline rendering is usually much nicer for interactive work, but doesn’t support all features (e.g. animation).”"”.format(subdir=subdir)
    n.addcell(‘code’,'%matplotlib inline’)
    n.addcell(‘code’,'%matplotlib’)
    pat = os.path.join(os.path.expanduser(basedir),subdir,’*.py’)
    examples = glob.glob(pat)
    if not examples:
        raise IOError(‘No files found: {pat}’.format(pat=pat))
    for ex in examples:
        n.addcell(‘markdown’,'## {exname}’.format(exname=ex))
        n.addcell(‘code’,open(ex).read())
    n.tofile(overwrite=overwrite,outdir=outdir)

def make_all_mpl_examples(outdir=’output’,basedir=’~/coding/matplotlib/examples’,overwrite=False):
    # There must be a builtin!
    bd = os.path.expanduser(basedir)
    subdirs = [d for d in os.listdir(bd) if os.path.isdir(os.path.join(bd,d))]
    for s in subdirs:
        print “Doing”,s
        make_mpl_examples(subdir=s,basedir=basedir,outdir=outdir,overwrite=overwrite)

It spits out one notebook per example directory, but is easy to change. It does what I want very nicely: gives me an immediate way to tweak the matplotlib examples. I’ll leave the bigger project (turning all of the matplotlib docs into IPython Notebooks) for later, if ever. The first thing I’d need to figure out is how to execute the code in the cells programatically. I’m sure there’s an API.

Here’s an example of some examples. I’ve executed several of the cells so that there’s output. This is significantly less interesting on the web, and more interesting if you run it in a notebook on your own so that you can tweak things.

%matplotlib inline

%matplotlib

In [2]:

"""
Simple demo of the imshow function.
"""
import matplotlib.pyplot as plt
import matplotlib.cbook as cbook

image_file = cbook.get_sample_data('ada.png')
image = plt.imread(image_file)

plt.imshow(image)
plt.axis('off') # clear x- and y-axes
plt.show()

In [3]:

"""
Demo of image that's been clipped by a circular patch.
"""
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.cbook as cbook


image_file = cbook.get_sample_data('grace_hopper.png')
image = plt.imread(image_file)

fig, ax = plt.subplots()
im = ax.imshow(image)
patch = patches.Circle((260, 200), radius=200, transform=ax.transData)
im.set_clip_path(patch)

plt.axis('off')
plt.show()

In [4]:

"""
Shows how to combine Normalization and Colormap instances to draw
"levels" in pcolor, pcolormesh and imshow type plots in a similar
way to the levels keyword argument to contour/contourf.

"""

import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm
from matplotlib.ticker import MaxNLocator
import numpy as np


# make these smaller to increase the resolution
dx, dy = 0.05, 0.05

# generate 2 2d grids for the x & y bounds
y, x = np.mgrid[slice(1, 5 + dy, dy),
                slice(1, 5 + dx, dx)]

z = np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x)

# x and y are bounds, so z should be the value *inside* those bounds.
# Therefore, remove the last value from the z array.
z = z[:-1, :-1]
levels = MaxNLocator(nbins=15).tick_values(z.min(), z.max())


# pick the desired colormap, sensible levels, and define a normalization
# instance which takes data values and translates those into levels.
cmap = plt.get_cmap('PiYG')
norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)

plt.subplot(2, 1, 1)
im = plt.pcolormesh(x, y, z, cmap=cmap, norm=norm)
plt.colorbar()
# set the limits of the plot to the limits of the data
plt.axis([x.min(), x.max(), y.min(), y.max()])
plt.title('pcolormesh with levels')



plt.subplot(2, 1, 2)
# contours are *point* based plots, so convert our bound into point
# centers
plt.contourf(x[:-1, :-1] + dx / 2.,
             y[:-1, :-1] + dy / 2., z, levels=levels,
             cmap=cmap)
plt.colorbar()
plt.title('contourf with levels')


plt.show()

In [5]:

"""
Demo of the `streamplot` function.

A streamplot, or streamline plot, is used to display 2D vector fields. This
example shows a few features of the stream plot function:

    * Varying the color along a streamline.
    * Varying the density of streamlines.
    * Varying the line width along a stream line.
"""
import numpy as np
import matplotlib.pyplot as plt

Y, X = np.mgrid[-3:3:100j, -3:3:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U*U + V*V)

plt.streamplot(X, Y, U, V, color=U, linewidth=2, cmap=plt.cm.autumn)
plt.colorbar()

f, (ax1, ax2) = plt.subplots(ncols=2)
ax1.streamplot(X, Y, U, V, density=[0.5, 1])

lw = 5*speed/speed.max()
ax2.streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw)

plt.show()

In [6]:

"""
Demo of the streamplot function with masking.

This example shows how streamlines created by the streamplot function skips
masked regions and NaN values.
"""
import numpy as np
import matplotlib.pyplot as plt

w = 3
Y, X = np.mgrid[-w:w:100j, -w:w:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U*U + V*V)

mask = np.zeros(U.shape, dtype=bool)
mask[40:60, 40:60] = 1
U = np.ma.array(U, mask=mask)
U[:20, :20] = np.nan

plt.streamplot(X, Y, U, V, color='r')

plt.imshow(~mask, extent=(-w, w, -w, w), alpha=0.5,
           interpolation='nearest', cmap=plt.cm.gray)

plt.show()

Comments from old blog