How to plot heat map with matplotlib?

Question

How to use python and matplotlib to plot a picture like following? I know how to plot the 2D heat map, but it frustrated me a lot with plotting the bar on top of the heat map, and the bar between the color bar and heat map. How to add those two bars on the picture, and show the number in x axis or y axis belongs to which group?

Thanks very much for all the responses.

Answer 1

A systematic and straightforward approach, although a bit more cumbersome at the start, is to use matplotlib.gridspec.GridSpec .

First set up the grid:

import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

fig = plt.figure()
gs = GridSpec(2, 3, width_ratios=[10, 1, 1], height_ratios=[1, 10])

This gives us a grid of 2 rows and 3 columns, where the lower left axis will be 10x10 and the other axes will be either 10x1 or 1x10 in relative sizes. These ratios can be tweaked to your liking. Note that the top center/right axes will be empty.

big_ax = fig.add_subplot(gs[1,0]) # bottom left
top_ax = fig.add_subplot(gs[0,0]) # top left
right_ax = fig.add_subplot(gs[1,1]) # bottom center
cbar_ax = fig.add_subplot(gs[1,2]) # bottom right

I will use a generic genome picture I found via google for the top and right image:

and will generate a random heatmap. I use imshow(aspect='auto') so that the image objects and heatmap take up the full space of their respective axes (otherwise they will override the height/width ratios set by gridspec).

im = plt.imread('/path/to/image.png')
# Plot your heatmap on big_ax and colorbar on cbar_ax
heatmap = big_ax.imshow(np.random.rand(10, 10), aspect='auto', origin='lower')
cbar = fig.colorbar(heatmap, cax=cbar_ax)

# Show your images on top_ax and right_ax
top_ax.imshow(im, aspect='auto')

# need to rotate my image. 
# you may not have to if you have two different images
from scipy import ndimage 
right_ax.imshow(ndimage.rotate(im, 90), aspect='auto')

# Clean up the image axes (remove ticks, etc.)
right_ax.set_axis_off()
top_ax.set_axis_off()

# remove spacing between axes
fig.subplots_adjust(wspace=0.05, hspace=0.05)

It's not super glamorous (especially with the default jet colormap), but you could easily use this to reproduce the figure your OP.

---

---

Edit: So if you want to generate that genome-like plot on the top and right, you could try something like this for the top bar:

from matplotlib.patches import Rectangle
from matplotlib.collections import PatchCollection

# draw the black line
top_ax.axhline(0, color='k', zorder=-1)

# box x-coords and text labels
boxes = zip(np.arange(0.1, 1, 0.2), np.arange(0.2, 1, 0.2))
box_text = ('A1', 'B1', 'B2', 'A2')
# color indicators for boxes
colors = (0, 1, 1, 0)
# construct Rects
patches = [Rectangle(xy=(x0, -1), width=(x1-x0), height=2) for x0,x1 in boxes]
p = PatchCollection(patches, cmap='jet')
# this maps the colors in [0,1] to the cmap above
p.set_array(np.array(colors))
top_ax.add_collection(p)

# add text
[top_ax.text((x0+x1)/2., 1.2, text, ha='center') 
    for (x0,x1), text in zip(boxes, box_text)]

# adjust ylims
top_ax.set_ylim(-2, 2)

For something the right axis, you can do the same thing but use axvline and swap the x-coords for y-coords.

right_ax.axvline(0, color='k', zorder=-1)

patches = [Rectangle(xy=(-1, y0), width=2, height=(y1-y0)) for y0, y1 in boxes]
p = PatchCollection(patches, cmap='jet')
p.set_array(np.array(colors))
right_ax.add_collection(p)

[right_ax.text(1.2, (y0+y1)/2., text, va='center') 
    for (y0, y1), text in zip(boxes, box_text)]
right_ax.set_xlim(-2,2)

These modifications lead to something like: