Как поместить несколько патчей для цветовой карты в легенду Matplotlib?

Question

Ситуация под рукой:

У меня есть несколько групп линий, где строки в одной и той же группе меняются в зависимости от параметра, определенного для конкретной группы.Я присваиваю каждой из этих строк в одной группе цвет из цветовой карты в соответствии с этим параметром, используя разные цветовые карты для каждой группы.

Теперь я хотел бы добавить легенду к сюжету с одной записью на группу линий.

Решение только для одного набора строк:

Если бы у меня была только одна группа строк, лучшим способом маркировки было быдобавить цветовую шкалу, как предложено в ответе: Matplotlib: добавить цветную полосу к не отображаемому объекту .

Как лучше всего это сделать для нескольких наборов линий?

Поскольку у меня несколько таких групп линий, я не хочу добавлять цветную полосу для каждого нового параметра.Вместо этого я бы предпочел поместить в легенду патчи, заполненные соответствующими цветными картами (как своего рода мини-панель цветов).

Минимальный рабочий пример:

В следующемМожно найти минимальный рабочий пример ситуации под рукой.Обратите внимание, что я сильно упростил вычисление линий, что скрывает зависимость параметра.Таким образом, мой «параметр» param - это просто индекс, по которому я перебираю.Мой реальный код вычисляет значения x и y в зависимости от параметра модели с более сложными функциями.Соответственно, максимум param_max здесь одинаков для каждой группы линий, хотя на самом деле это не так.

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()

Это дает график, показанный выше.

---

Так как яНе удалось найти ничего, что напрямую отвечало бы на это в stackoverflow, я попытался найти решение, которое вы можете найти в разделе ответов.Если есть более прямой / правильный способ сделать это, я был бы рад узнать.

Answer 1

Я адаптировал решение ответа с помощью ImportanceOfBeingErnest к "Создайте для этого случая образцы matplotlib с двухцветным прямоугольником для легенды фигуры" . Как указано там, инструкции в разделе Реализация пользовательского обработчика легенды в руководстве по легендам matplotlib были особенно полезны.

Результат:

Решение:

Я создал класс HandlerColormap, полученный из базового класса matplotlib для обработчиков легенды HandlerBase. HandlerColormap принимает в качестве аргументов цветовую карту и количество полос.

Для аргумента cmap должен быть задан экземпляр matplotlib colormap .

Аргумент num_stripes определяет, насколько (не) непрерывным будет цветовой градиент в патче легенды.

Как указано в HandlerBase, я переопределяю его метод create_artist, используя заданное измерение, так что код должен (автоматически) масштабироваться по размеру шрифта. В этом новом методе create_artist я создаю несколько полос (slim matplotlib Rectangles), окрашенных в соответствии с входной цветовой картой.

Код:

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()