How to put multiple colormap patches in a matplotlib legend?

Question

Situation at hand:

I have multiple groups of lines, where the lines within the same group vary according to some group specific parameter. I assign each of these lines within the same group a color from a colormap according to this parameter using a different colormap for each group.

Now, I would like to add a legend to the plot with one entry per group of lines.

Solution for only one set of lines:

If I had only one group of lines the best way of labelling would be to add a colorbar as suggested in the answer to: Matplotlib: Add colorbar to non-mappable object.

How best to do this for multiple sets of lines?

As I have multiple such groups of lines, I do not want to add a colorbar for each new parameter. Instead, I would rather put patches filled with the corresponding colormaps in the legend (as a sort of mini colorbar).

Minimal working example:

In the following you can find a minimal working example of the situation at hand. Note, though, that I heavily simplified the calculation of the lines which hides the parameter dependence. Thus, my "parameter" param here is just the index I am iterating over. My actual code calculates the x and y values depending on a model parameter with more complicated functions. Accordingly the maximum param_max here is the same for each group of lines, though actually it would not be.

import numpy as np
import matplotlib.pyplot as plt

x_array = np.linspace(1, 10, 10)
y_array = x_array
param_max = x_array.size
cmaps = [plt.cm.spring, plt.cm.winter]  # set of colormaps 
                                        # (as many as there are groups of lines)
plt.figure()
for param, (x, y) in enumerate(zip(x_array, y_array)):  
    x_line1 = np.linspace(x, 1.5 * x, 10)
    y_line1 = np.linspace(y**2, y**2 - x, 10)
    x_line2 = np.linspace(1.2 * x, 1.5 * x, 10)
    y_line2 = np.linspace(2 * y, 2 * y - x, 10)
    # plot lines with color depending on param using different colormaps:
    plt.plot(x_line1, y_line1, c=cmaps[0](param / param_max))
    plt.plot(x_line2, y_line2, c=cmaps[1](param / param_max))
plt.show()

This produces the plot shown above.

---

Since I could not find anything directly answering this on stackoverflow, I tried finding a solution myself which you can find in the answers section. If there is a more direct/proper way of doing this I would be happy to know.

Answer 1

I adapted the solution of the answer by ImportanceOfBeingErnest to "Create a matplotlib mpatches with a rectangle bi-colored for figure legend" to this case. As linked there, the instructions in the section on Implementing a custom legend handler in the matplotlib legend guide were particularly helpful.

Result:

Solution:

I created the class HandlerColormap derived from matplotlib's base class for legend handlers HandlerBase. HandlerColormap takes a colormap and a number of stripes as arguments.

For the argument cmap a matplotlib colormap instance should be given.

The argument num_stripes determines how (non-)continuous the color gradient in the legend patch will be.

As instructed in HandlerBase I override its create_artist method using the given dimension such that the code should be (automatically) scaled by fontsize. In this new create_artist method I create multiple stripes (slim matplotlib Rectangles) colored according to the input colormap.

Code:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.legend_handler import HandlerBase

class HandlerColormap(HandlerBase):
    def __init__(self, cmap, num_stripes=8, **kw):
        HandlerBase.__init__(self, **kw)
        self.cmap = cmap
        self.num_stripes = num_stripes
    def create_artists(self, legend, orig_handle, 
                       xdescent, ydescent, width, height, fontsize, trans):
        stripes = []
        for i in range(self.num_stripes):
            s = Rectangle([xdescent + i * width / self.num_stripes, ydescent], 
                          width / self.num_stripes, 
                          height, 
                          fc=self.cmap((2 * i + 1) / (2 * self.num_stripes)), 
                          transform=trans)
            stripes.append(s)
        return stripes

x_array = np.linspace(1, 10, 10)
y_array = x_array
param_max = x_array.size
cmaps = [plt.cm.spring, plt.cm.winter]  # set of colormaps 
                                        # (as many as there are groups of lines)
plt.figure()
for param, (x, y) in enumerate(zip(x_array, y_array)):  
    x_line1 = np.linspace(x, 1.5 * x, 10)
    y_line1 = np.linspace(y**2, y**2 - x, 10)
    x_line2 = np.linspace(1.2 * x, 1.5 * x, 10)
    y_line2 = np.linspace(2 * y, 2 * y - x, 10)
    # plot lines with color depending on param using different colormaps:
    plt.plot(x_line1, y_line1, c=cmaps[0](param / param_max))
    plt.plot(x_line2, y_line2, c=cmaps[1](param / param_max))

cmap_labels = ["parameter 1 $\in$ [0, 10]", "parameter 2 $\in$ [-1, 1]"]
# create proxy artists as handles:
cmap_handles = [Rectangle((0, 0), 1, 1) for _ in cmaps]
handler_map = dict(zip(cmap_handles, 
                       [HandlerColormap(cm, num_stripes=8) for cm in cmaps]))
plt.legend(handles=cmap_handles, 
           labels=cmap_labels, 
           handler_map=handler_map, 
           fontsize=12)
plt.show()