prettyplotlib.pcolormesh: Improving heatmaps in matplotlib

Both positive and negative values

The default matplotlib pcolormesh heatmaps use a rainbow colormap, which has been known to mislead data visualization. Specifically, "the rainbow color map is universally inferior to all other color maps". Unfortunately, matplotlib took its default colors from MATLAB, and there the default is also rainbow.


In [17]:
import matplotlib
# Make sure that original paramters are used, and no matplotlibrc files interfere
matplotlib.rcParams = matplotlib.rcParamsDefault

In [18]:
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(1)

np.random.seed(10)

#ax.pcolor(np.random.randn((10,10)))
#ax.pcolor(np.random.randn(10), np.random.randn(10))
p = ax.pcolormesh(np.random.randn(10,10))
fig.colorbar(p)
fig.savefig('pcolormesh_matplotlib_default.png')


Using the same zero-centered randomly distributed gaussian distribution, we can plot it using prettplotlib with a few modifications:

ppl.pcolormesh(fig, ax, np.random.randn(10,10))

You'll notice that the "hot" (large, positive) color is still red, and the "cold" (small, negative) color is still blue, but the in between colors are gradations of red and blue, so it's easier to tell the difference between values.


In [19]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(1)

np.random.seed(10)

ppl.pcolormesh(fig, ax, np.random.randn(10,10))
fig.savefig('pcolormesh_prettyplotlib_default.png')


You may have noticed similar changes as were made in prettyplotlib.scatter, where axis lines were removed, and blacks were changed to almost black.

Only positive (or negative) values

If your data is only positive (or negative), matplotlib does nothing to change the color scale. It's still a rainbow, but look at the colorbar, the range is different (0 to 1 instead of -2 to +2)


In [20]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(1)

np.random.seed(10)

p = ax.pcolormesh(np.random.uniform(size=(10,10)))
fig.colorbar(p)
fig.savefig('pcolormesh_matplotlib_positive_default.png')


If your data is only positive or negative, then prettyplotlib will auto-detect this and use a single-color colormap. The default for positive data is the reds colormap.


In [21]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(1)

np.random.seed(10)

ppl.pcolormesh(fig, ax, np.abs(np.random.randn(10,10)))
fig.savefig('pcolormesh_prettyplotlib_positive.png')


And the default for negative data is the blues colormap.


In [22]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(1)

np.random.seed(10)

ppl.pcolormesh(fig, ax, -np.abs(np.random.randn(10,10)))
# fig.savefig('pcolormesh_prettyplotlib_negative.png')


Out[22]:
<matplotlib.collections.QuadMesh at 0x10c518750>

Plus you can add $x$- and $y$-ticklabels directly!

Normally, when you add $x$- and $y$-ticklables on pcolormesh in matplotlib, they're not centered on the blocks, and you have to do a lot of annoying work just getting a label on each box. You have to specify the xticks explicitly, since you want to label each box.

xticks = range(10)
yticks = range(10)

xticklabels=string.uppercase[:10]
yticklabels=string.lowercase[-10:]

ax.set_xticks(xticks)
ax.set_xticklabels(xticklabels)

ax.set_yticks(yticks)
ax.set_yticklabels(yticklabels)

In [23]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
import numpy as np
import string

fig, ax = plt.subplots(1)

np.random.seed(10)

p = ax.pcolormesh(np.abs(np.random.randn(10,10)))
fig.colorbar(p)

xticks = range(10)
yticks = range(10)

xticklabels=string.uppercase[:10]
yticklabels=string.lowercase[-10:]

ax.set_xticks(xticks)
ax.set_xticklabels(xticklabels)

ax.set_yticks(yticks)
ax.set_yticklabels(yticklabels)


fig.savefig('pcolormesh_matplotlib_positive_labels.png')


But prettyplotlib.pcolormesh assumes that you want the xticklabels and yticklabels on each block, and makes it easy to specify.

ppl.pcolormesh(fig, ax, np.random.uniform(size=(10,10)), 
               xticklabels=string.uppercase[:10], 
               yticklabels=string.lowercase[-10:])

In [24]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
import numpy as np
import string

fig, ax = plt.subplots(1)

np.random.seed(10)

ppl.pcolormesh(fig, ax, np.random.randn(10,10), 
               xticklabels=string.uppercase[:10], 
               yticklabels=string.lowercase[-10:])
fig.savefig('pcolormesh_prettyplotlib_labels.png')


Or pick your own colormap! The diverging colormap PRGn or Purple and Green is pretty nice. I usually use this website to look up the colormaps: Every Colorbrewer Scale


In [25]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
from prettyplotlib import brewer2mpl
import numpy as np
import string

green_purple = brewer2mpl.get_map('PRGn', 'diverging', 11).mpl_colormap

fig, ax = plt.subplots(1)

np.random.seed(10)

ppl.pcolormesh(fig, ax, np.random.randn(10,10), 
               xticklabels=string.uppercase[:10], 
               yticklabels=string.lowercase[-10:],
               cmap=green_purple, center_value=0)
fig.savefig('pcolormesh_prettyplotlib_labels_other_cmap_diverging.png')



In [26]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
from prettyplotlib import brewer2mpl
import numpy as np
import string

green_purple = brewer2mpl.get_map('PRGn', 'diverging', 11).mpl_colormap

fig, ax = plt.subplots(1)

np.random.seed(10)

ppl.pcolormesh(fig, ax, np.random.randn(10,10), 
               xticklabels=string.uppercase[:10], 
               yticklabels=string.lowercase[-10:],
               cmap=green_purple, center_value=2)
fig.savefig('pcolormesh_prettyplotlib_labels_other_cmap_diverging_center_value.png')


Or if you want your own colormap for positive-only data:


In [27]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
from prettyplotlib import brewer2mpl
import numpy as np
import string

red_purple = brewer2mpl.get_map('RdPu', 'Sequential', 9).mpl_colormap

fig, ax = plt.subplots(1)

np.random.seed(10)

ppl.pcolormesh(fig, ax, np.abs(np.random.randn(10,10)),
               xticklabels=string.uppercase[:10], 
               yticklabels=string.lowercase[-10:],
               cmap=red_purple)
fig.savefig('pcolormesh_prettyplotlib_labels_other_cmap_sequential.png')


Plus, this will take the usual parameters of pcolormesh like if you want to rescale your data to log-scale:

from matplotlib.colors import LogNorm
...
ppl.pcolormesh(..., norm=LogNorm(vmin=x.min().min(), vmax=x.max().max()))

In [28]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
from prettyplotlib import brewer2mpl
import numpy as np
import string
from matplotlib.colors import LogNorm

red_purple = brewer2mpl.get_map('RdPu', 'Sequential', 9).mpl_colormap

fig, ax = plt.subplots(1)

np.random.seed(10)

x = np.abs(np.random.randn(10,10))
ppl.pcolormesh(fig, ax, x,
               xticklabels=string.uppercase[:10], 
               yticklabels=string.lowercase[-10:],
               cmap=red_purple, 
               norm=LogNorm(vmin=x.min().min(), vmax=x.max().max()))
fig.savefig('pcolormesh_prettyplotlib_labels_lognorm.png')


And now you can easily make beautiful heatmaps!

Extras


In [31]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
import numpy as np
import string

np.random.seed(10)

ppl.pcolormesh(-np.abs(np.random.randn(10,10)),
               xticklabels=string.uppercase[:10], 
               yticklabels=string.lowercase[-10:])
fig = plt.gcf()
fig.savefig('pcolormesh_prettyplotlib_negative_labels.png')



In [30]:
import prettyplotlib as ppl
import matplotlib.pyplot as plt
import numpy as np
import string

fig, ax = plt.subplots(1)

np.random.seed(10)

p = ax.pcolormesh(-np.abs(np.random.randn(10,10)))
fig.colorbar(p)

xticks = range(10)
yticks = range(10)

xticklabels=string.uppercase[:10]
yticklabels=string.lowercase[-10:]

ax.set_xticks(xticks)
ax.set_xticklabels(xticklabels)

ax.set_yticks(yticks)
ax.set_yticklabels(yticklabels)

fig.savefig('pcolormesh_matplotlib_negative_labels.png')



In [30]:


In [ ]: