My Matplotlib Cheatsheet (#1)

I find myself Googling the same matplotlib/pyplot queries over and over again. To make life easy, I started collecting some of them in a series of blog posts.
In this first post, I'll look at subplots and annotations.

Turn off axes

We can remove axes, ticks, borders, etc. with ax.axis('off'). This can be used to remove unnecessary subplots created with the subplots function from pyplot, for instance, when we only want to use the upper triangle of the grid. Here's an example:

	import matplotlib.pyplot as plt
	import numpy as np
	import scipy.stats as sts
	## create some random data
	n = 3
	Sigma = sts.wishart.rvs(scale=np.eye(n), df=n)
	xs = sts.multivariate_normal.rvs(cov=Sigma, size=1000)
	## plot histograms and scatter plots
	fig, axs = plt.subplots(n, n, figsize=(7,7))
	for i in range(n):
	for j in range(n):
	if i > j:
	axs[i, j].axis('off')
	elif i == j:
	axs[i, j].hist(xs[:, i], density=True)
	else:
	axs[i, j].scatter(xs[:, i], xs[:, j], s=5)
	fig.tight_layout()
	fig.savefig("axis-off.png", bbox_inches='tight', dpi=200)

view raw axis-off.py hosted with ❤ by GitHub

The result is:

Share axes between subplots after plotting

An alternative way of making a ragged array of subplots makes use of gridspec. The problem here is that it is a bit more difficult to share x and y axes. Of course add_subplot has keyword arguments sharex and sharey, but then we have to distinguish between the first and subsequent subplots. A better solution is the get_shared_x_axes() method. Here's an example:

	import matplotlib.pyplot as plt
	import numpy as np
	from matplotlib.gridspec import GridSpec
	## we want to make 23 plots using 5 columns
	n = 23
	ncols = 5
	## compute the required number of rows:
	nrows = n // ncols + 1 if n % ncols else 0
	## use gridspec to make the subplots
	fig = plt.figure(figsize=(7,7))
	gs = GridSpec(nrows, ncols)
	## reduce space between the subplots
	gs.update(wspace=0.07, hspace=0.07)
	## list will contain the subplots
	axs = []
	## some parameters for making our data
	u, v = 0.2, 0.3
	ts = np.linspace(0, 2*np.pi, 100)
	for i in range(n):
	## compute the row and column index
	col, row = i % ncols, i // ncols
	## and use gridspec to get a new subplot
	ax = fig.add_subplot(gs[row, col])
	## create some data and plot
	xs = (1+row) * (np.sin((col+1) * ts) + u * np.sin(ts))
	ys = (1+col) * (np.cos((row+1) * ts) + v * np.sin(ts))
	ax.plot(xs, ys, linewidth=3)
	## add the new subplot to the list
	axs.append(ax)
	## make y axis invisible
	if col > 0:
	ax.get_yaxis().set_visible(False)
	## make x axis invisible
	if row > 0:
	ax.get_xaxis().set_visible(False)
	else: ## move the x-ticks to the top
	ax.xaxis.set_ticks_position('top')
	## now we make the x and y axes shared and auto scale them
	axs[0].get_shared_x_axes().join(*axs)
	axs[0].autoscale(axis='x')
	axs[0].get_shared_y_axes().join(*axs)
	axs[0].autoscale(axis='y')
	fig.savefig("distorted-lissajous.png", bbox_inches='tight', dpi=200)

view raw share-axes.py hosted with ❤ by GitHub

The result is:

Hybrid (or blended) transforms for annotations

When you annotate a point in a plot, the location of the text is often relative to the data in one coordinate, but relative to the axis (e.g. in the middle) in the other. I used to do this with inverse transforms, but it turns out that there is a better way: the blended_transform_factory function from matplotlib.transforms.
Suppose that I want to annotate three points in three subplots. The arrows should point to these three points, and I want the text to be located above the point, but the text in the 3 subplots has to be vertically aligned.

Notice that the y-axes are not shared between the subplots! To accomplish the alignment, we have to use the annotate method with a custom transform.

	import matplotlib.pyplot as plt
	import numpy as np
	## import the function we'll need for making the transforms
	from matplotlib.transforms import blended_transform_factory
	fig, axs = plt.subplots(1, 3, figsize=(10,3), sharex=True)
	fs = [
	lambda x: np.sqrt(x**3),
	lambda x: np.sqrt(x**2*(x+1)),
	lambda x: np.sqrt((x+1)*x*(x-1))
	]
	## plot some elliptic curves (the real part at least)
	xs = np.linspace(-1, 2, 1000)
	for f, a in zip(fs, axs):
	a.plot(xs, f(xs), color='k', linewidth=2)
	a.plot(xs, -f(xs), color='k', linewidth=2)
	def hybrid_trans(ax):
	"""
	This function returns a blended/hybrid
	transformation w.r.t. subplot ax.
	x-coord: use the data for positioning
	y-coord: use relative position w.r.t y-axis
	"""
	return blended_transform_factory(ax.transData, ax.transAxes)
	## avoid repetition: common key-word arguments for annotate
	kwargs = {
	"arrowprops" : {"arrowstyle": "->"},
	"ha" : "center",
	"va" : "bottom"
	}
	## give subplots names to avoid indexing
	ax, bx, cx = axs
	## the argument xytext determines where the text is positioned,
	## and we can pass a transformation with textcoords.
	## we'll use our custom "hybrid" transform so that the x-coord
	## corresponds to the data, and the y-coord is relative.
	ax.annotate("cusp", xy=(0,0), xytext=(0, 0.8),
	textcoords=hybrid_trans(ax), **kwargs)
	bx.annotate("singularity", xy=(0,0), xytext=(0, 0.8),
	textcoords=hybrid_trans(bx), **kwargs)
	## this is a bit of a hack: I want two annotate two points
	## with one text box. I'll make the second text box invisible
	## by setting alpha=0
	text = "two real components"
	xcoords = [-0.25, 1.1]
	alphas = [1, 0]
	for x, alpha in zip(xcoords, alphas):
	## use fs[2] to get the correct point
	cx.annotate(text, xy=(x,fs[2](x)), xytext=(0,0.8), alpha=alpha,
	textcoords=hybrid_trans(cx), **kwargs)
	## add some labels and titles...
	bx.set_xlabel("x")
	ax.set_ylabel("y")
	bx.set_title("three elliptic curves")
	fig.savefig("annotations.png", bbox_inches='tight', dpi=200)

view raw transform-annotate.py hosted with ❤ by GitHub