Colorbars and legends

ProPlot includes some useful improvements to the matplotlib API that make working with colorbars and legends much easier.

Axes colorbars and legends

In matplotlib, colorbars are added to the edges of subplots with the figure method matplotlib.figure.Figure.colorbar using e.g. fig.colorbar(m, ax=ax, location='right'). In ProPlot, this is done using the new axes colorbar method proplot.axes.Axes.colorbar with e.g. ax.colorbar(m, loc='r'). The proplot.axes.Axes.colorbar method preserves subplot aspect ratios and visual symmetry between subplots by allocating new space in the figure GridSpec rather than “stealing” space from the parent subplot (see the section on automatic subplot spacing for details).

ProPlot tries to make the usage of proplot.axes.Axes.colorbar consistent with legend, and includes an improved proplot.axes.Axes.legend legend method that tries to do the same:

  • Just like colorbar, proplot.axes.Axes.legend can draw “outer” legends along the edges of subplots when you request a side location for the legend (e.g. loc='right' or loc='r'). If you draw multiple colorbars and legends on one side, they are “stacked” on top of each other.

  • Just like legend, proplot.axes.Axes.colorbar can draw “inset” colorbars when you request an inset location for the colorbar (e.g. loc='upper right' or loc='ur'). Inset colorbars have optional background “frames” that can be configured with various colorbar keywords.

You can also draw colorbars and legends on-the-fly by supplying keyword arguments to various plotting commands. To plot data and draw a colorbar in one go, pass a location (e.g. colorbar='r') to methods that accept a cmap argument (e.g. contourf). To draw a legend or colorbar-legend in one go, pass a location (e.g. legend='r' or colorbar='r') to methods that accept a cycle argument (e.g. plot). Use legend_kw and colorbar_kw to pass keyword arguments to the colorbar and legend functions. This feature is powered by the cmap_changer and cycle_changer wrappers.

  1. [1]:
  1. import proplot as plot
  2. import numpy as np
  3. with plot.rc.context(abc=True):
  4.     fig, axs = plot.subplots(ncols=2, share=0)
  6. # Colorbars
  7. ax = axs[0]
  8. state = np.random.RandomState(51423)
  9. m = ax.heatmap(state.rand(10, 10), colorbar='t', cmap='dusk')
  10. ax.colorbar(m, loc='r')
  11. ax.colorbar(m, loc='ll', label='colorbar label')
  12. ax.format(title='Axes colorbars', suptitle='Axes colorbars and legends demo')
  14. # Legends
  15. ax = axs[1]
  16. ax.format(title='Axes legends', titlepad='0em')
  17. hs = ax.plot(
  18.     (state.rand(10, 5) - 0.5).cumsum(axis=0), linewidth=3,
  19.     cycle='ggplot', legend='t',
  20.     labels=list('abcde'), legend_kw={'ncols': 5, 'frame': False}
  21. )
  22. ax.legend(hs, loc='r', ncols=1, frame=False)
  23. ax.legend(hs, loc='ll', label='legend label')
  24. axs.format(xlabel='xlabel', ylabel='ylabel')

_images/colorbars_legends_2_0.svg

  1. [2]:
  1. import proplot as plot
  2. import numpy as np
  3. fig, axs = plot.subplots(nrows=2, share=0, axwidth='55mm', panelpad='1em')
  4. axs.format(suptitle='Stacked colorbars demo')
  5. state = np.random.RandomState(51423)
  6. N = 10
  7. # Repeat for both axes
  8. for j, ax in enumerate(axs):
  9.     ax.format(
  10.         xlabel='data', xlocator=np.linspace(0, 0.8, 5),
  11.         title=f'Subplot #{j+1}'
  12.     )
  13.     for i, (x0, y0, x1, y1, cmap, scale) in enumerate((
  14.         (0, 0.5, 1, 1, 'grays', 0.5),
  15.         (0, 0, 0.5, 0.5, 'reds', 1),
  16.         (0.5, 0, 1, 0.5, 'blues', 2)
  17.     )):
  18.         if j == 1 and i == 0:
  19.             continue
  20.         data = state.rand(N, N) * scale
  21.         x, y = np.linspace(x0, x1, N + 1), np.linspace(y0, y1, N + 1)
  22.         m = ax.pcolormesh(
  23.             x, y, data, cmap=cmap,
  24.             levels=np.linspace(0, scale, 11)
  25.         )
  26.         ax.colorbar(m, loc='l', label=f'dataset #{i+1}')

_images/colorbars_legends_3_0.svg

Figure colorbars and legends

In ProPlot, colorbars and legends can be added to the edge of figures with the proplot.figure.Figure.colorbar and proplot.figure.Figure.legend methods. Figure colorbars and legends are aligned between the edges of the subplot grid, rather than the figure bounds. As with axes colorbars and legends, if you draw multiple colorbars or legends on the same side, they are stacked on top of each other.

To draw a colorbar or legend alongside particular row(s) or column(s) of the subplot grid, use the row, rows, col, or cols keyword arguments. Pass an integer to draw the colorbar or legend beside a single row or column, or pass a tuple to draw it beside a range of rows or columns.

  1. [3]:
  1. import proplot as plot
  2. import numpy as np
  3. fig, axs = plot.subplots(ncols=3, nrows=3, axwidth=1.4)
  4. state = np.random.RandomState(51423)
  5. m = axs.pcolormesh(
  6.     state.rand(20, 20), cmap='grays',
  7.     levels=np.linspace(0, 1, 11), extend='both'
  8. )[0]
  9. axs.format(
  10.     suptitle='Figure colorbars and legends demo', abc=True,
  11.     abcloc='l', abcstyle='a.', xlabel='xlabel', ylabel='ylabel'
  12. )
  13. fig.colorbar(m, label='column 1', ticks=0.5, loc='b', col=1)
  14. fig.colorbar(m, label='columns 2-3', ticks=0.2, loc='b', cols=(2, 3))
  15. fig.colorbar(m, label='stacked colorbar', ticks=0.1, loc='b', minorticks=0.05)
  16. fig.colorbar(m, label='colorbar with length <1', ticks=0.1, loc='r', length=0.7)
  1. [3]:
  1. <matplotlib.colorbar.Colorbar at 0x7fdd0aaeb880>

_images/colorbars_legends_5_1.svg

  1. [4]:
  1. import proplot as plot
  2. import numpy as np
  3. fig, axs = plot.subplots(
  4.     ncols=2, nrows=2, axwidth=1.7,
  5.     share=0, wspace=0.3, order='F'
  6. )
  8. # Plot data
  9. data = (np.random.rand(50, 50) - 0.1).cumsum(axis=0)
  10. m = axs[:2].contourf(data, cmap='grays', extend='both')
  11. colors = plot.Colors('grays', 5)
  12. hs = []
  13. state = np.random.RandomState(51423)
  14. for abc, color in zip('ABCDEF', colors):
  15.     h = axs[2:].plot(state.rand(10), lw=3, color=color, label=f'line {abc}')
  16.     hs.extend(h[0])
  18. # Add colorbars and legends
  19. fig.colorbar(m[0], length=0.8, label='colorbar label', loc='b', col=1, locator=5)
  20. fig.colorbar(m[0], label='colorbar label', loc='l')
  21. fig.legend(hs, ncols=2, center=True, frame=False, loc='b', col=2)
  22. fig.legend(hs, ncols=1, label='legend label', frame=False, loc='r')
  23. axs.format(
  24.     suptitle='Figure colorbars and legends demo',
  25.     abc=True, abcloc='ul', abcstyle='A'
  26. )
  27. for ax, title in zip(
  28.     axs, ['2D dataset #1', '2D dataset #2', 'Line set #1', 'Line set #2']
  29. ):
  30.     ax.format(xlabel='xlabel', title=title)

_images/colorbars_legends_6_0.svg

New colorbar features

The proplot.figure.Figure.colorbar and proplot.axes.Axes.colorbar methods are wrapped by colorbar_wrapper, which adds several new features.

You can now draw colorbars from lists of colors or lists of artists by passing a list instead of a mappable object. Colorbar minor ticks are now much more robust, and the tick location and formatter arguments are passed through Locator and Formatter. The colorbar width and length can be changed with the width and length keyword args. Colorbar widths are now specified in physical units, which helps avoid colorbars that look “too skinny” or “too fat” and preserves the look of the figure when the figure size changes.

  1. [5]:
  1. import proplot as plot
  2. import numpy as np
  3. fig, axs = plot.subplots(share=0, ncols=2, axwidth=2)
  5. # Colorbars from lines
  6. ax = axs[0]
  7. state = np.random.RandomState(51423)
  8. data = 1 + (state.rand(12, 10) - 0.45).cumsum(axis=0)
  9. cycle = plot.Cycle('algae')
  10. hs = ax.plot(
  11.     data, lw=4, cycle=cycle, colorbar='lr',
  12.     colorbar_kw={'length': '8em', 'label': 'from lines'}
  13. )
  14. axs.colorbar(
  15.     hs, loc='t', values=np.arange(0, 10),
  16.     label='from lines', ticks=2
  17. )
  19. # Colorbars from a mappable
  20. ax = axs[1]
  21. m = ax.contourf(
  22.     data.T, extend='both', cmap='algae',
  23.     levels=plot.arange(0, 3, 0.5)
  24. )
  25. fig.colorbar(
  26.     m, length=1, loc='r', label='inside ticks',
  27.     tickloc='left'
  28. )
  29. ax.colorbar(
  30.     m, loc='ul', length=1, tickminor=True,
  31.     label='inset colorbar', alpha=0.5
  32. )
  33. axs.format(
  34.     suptitle='Colorbar formatting demo',
  35.     xlabel='xlabel', ylabel='ylabel', abovetop=False
  36. )

_images/colorbars_legends_8_0.svg

New legend features

The proplot.figure.Figure.legend and proplot.axes.Axes.legend methods are wrapped by legend_wrapper, which adds several new features.

You can draw legends with centered legend rows, either by passing center=True or by passing list of lists of plot handles. This is accomplished by stacking multiple single-row, horizontally centered legends, then manually adding an encompassing legend frame. You can also modify legend text and handle properties with several keyword args, and switch between row-major and column-major order for legend entries with the order keyword arg (default is row-major).

  1. [6]:
  1. import proplot as plot
  2. import numpy as np
  3. plot.rc.cycle = '538'
  4. labels = ['a', 'bb', 'ccc', 'dddd', 'eeeee']
  5. fig, axs = plot.subplots(ncols=2, span=False, share=1, axwidth=2.3)
  6. hs1, hs2 = [], []
  8. # On-the-fly legends
  9. state = np.random.RandomState(51423)
  10. for i, label in enumerate(labels):
  11.     data = (state.rand(20) - 0.45).cumsum(axis=0)
  12.     h1 = axs[0].plot(
  13.         data, lw=4, label=label, legend='ul',
  14.         legend_kw={'order': 'F', 'title': 'column major'}
  15.     )
  16.     hs1.extend(h1)
  17.     h2 = axs[1].plot(
  18.         data, lw=4, label=label, legend='r', cycle='Set3',
  19.         legend_kw={'ncols': 1, 'frame': False, 'title': 'no frame'}
  20.     )
  21.     hs2.extend(h2)
  23. # Outer legends
  24. ax = axs[0]
  25. ax.legend(
  26.     hs1, loc='b', ncols=3, title='row major', order='C',
  27.     facecolor='gray2'
  28. )
  29. ax = axs[1]
  30. ax.legend(hs2, loc='b', ncols=3, center=True, title='centered rows')
  31. axs.format(xlabel='xlabel', ylabel='ylabel', suptitle='Legend formatting demo')

_images/colorbars_legends_10_0.svg