0

I'm having trouble giving colorbars to a grid of line plots in Matplotlib.

I have a grid of plots, which each shows 64 lines. The lines depict the penalty value vs time when optimizing the same system under 64 different values of a certain hyperparameter h.

Since there are so many lines, instead of using a standard legend, I'd like to use a colorbar, and color the lines by the value of h. In other words, I'd like something that looks like this:

line plot with color bar

The above was done by adding a new axis to hold the colorbar, by calling figure.add_axes([0.95, 0.2, 0.02, 0.6]), passing in the axis position explicitly as parameters to that method. The colorbar was then created as in the example code here, by instantiating a ColorbarBase(). That's fine for single plots, but I'd like to make a grid of plots like the one above.

To do this, I tried doubling the number of subplots, and using every other subplot axis for the colorbar. Unfortunately, this led to the colorbars having the same size/shape as the plots:

awkward square-shaped colorbar

Is there a way to shrink just the colorbar subplots in a grid of subplots like the 1x2 grid above?

Ideally, it'd be great if the colorbar just shared the same axis as the line plot it describes. I saw that the colorbar.colorbar() function has an ax parameter:

ax

parent axes object from which space for a new colorbar axes will be stolen.

That sounds great, except that colorbar.colorbar() requires you to pass in a imshow image, or a ContourSet, but my plot is neither an image nor a contour plot. Can I achieve the same (axis-sharing) effect using ColorbarBase?

SuperElectric
  • 17,548
  • 10
  • 52
  • 69

1 Answers1

0

It turns out you can have different-shaped subplots, so long as all the plots in a given row have the same height, and all the plots in a given column have the same width.

You can do this using gridspec.GridSpec, as described in this answer.

So I set the columns with line plots to be 20x wider than the columns with color bars. The code looks like: 

    grid_spec = gridspec.GridSpec(num_rows, 
                                  num_columns * 2, 
                                  width_ratios=[20, 1] * num_columns)
    colormap_type = cm.cool

    for (x_vec_list, 
         y_vec_list,
         color_hyperparam_vec,
         plot_index) in izip(x_vec_lists,
                             y_vec_lists,
                             color_hyperparam_vecs,
                             range(len(x_vecs))):

        line_axis = plt.subplot(grid_spec[grid_index * 2])
        colorbar_axis = plt.subplot(grid_spec[grid_index * 2 + 1])

        colormap_normalizer = mpl.colors.Normalize(vmin=color_hyperparam_vec.min(), 
                                                   vmax=color_hyperparam_vec.max())
        scalar_to_color_map = mpl.cm.ScalarMappable(norm=colormap_normalizer, 
                                                    cmap=colormap_type)

        colorbar.ColorbarBase(colorbar_axis, 
                              cmap=colormap_type, 
                              norm=colormap_normalizer)

        for (line_index, 
             x_vec, 
             y_vec) in zip(range(len(x_vec_list)),
                           x_vec_list,
                           y_vec_list):
            hyperparam = color_hyperparam_vec[line_index]
            line_color = scalar_to_color_map.to_rgba(hyperparam)

            line_axis.plot(x_vec, y_vec, color=line_color, alpha=0.5)

For num_rows=1 and num_columns=1, this looks like:

enter image description here

Community
  • 1
  • 1
SuperElectric
  • 17,548
  • 10
  • 52
  • 69
  • 1
    There is no need to use colorbarbase. To have a colorbar on axes `ax` you would call `fig.colorbar(scalar_to_color_map, ax=ax)` and get by default a colorbar on the right side of the axes with the same height as the axes (as long as it is not some special axes with equal aspect or so). – ImportanceOfBeingErnest Oct 02 '17 at 23:06