Dee Jay
Dee Jay

Reputation: 11

Add a custom color map to legend using matplotlib in python

I want to create a custom color map (between two greens[0,150,0] to [0,250,0]) and then use it in the legend in the plot. This is an image plot using RGBA.

Code to create the color map is as follows:

from matplotlib.colors import ListedColormap

N = 256
vals = np.ones((N, 4))
vals[:, 0] = np.linspace(0/256, 0, N)
vals[:, 1] = np.linspace(250/256, 150/255, N)
vals[:, 2] = np.linspace(0/256, 0, N)
newcmp = ListedColormap(vals)

Code to create the legend is as follows:

cmap = {1: [155/255, 118/255, 83/255, 1],
            2: [168/255, 210/255, 151/255, 1],
            3: newcmp,
            4: [175/255, 206/255, 208/255, 1],
            5: [249/255, 231/255, 157/255, 1],
            6: [209/255, 217/255, 208/255, 1],
            7: [225/255, 166/255, 49/255, 1],
            8: [128/255, 128/255, 0, 1],
            9: [120/255, 89/255, 86/255, 1],
            10: [60/255, 45/255, 21/255, 1],
            11: [230/255, 167/255, 125/255, 1],
            12: [170/255, 157/255, 132/255, 1],
            13: [84/255, 96/255, 76/255, 1]
            }
    labels = {1: 'A',
              2: 'B',
              3: 'C',
              4: 'D',
              5: 'E',
              6: 'F',
              7: 'G',
              8: 'H',
              9: 'I',
              10: 'J',
              11: 'K',
              12: 'L',
              13: 'M'
              }
    patches = [mpatches.Patch(color=cmap[i], label=labels[i]) for i in cmap]
    plt.legend(handles=patches, loc='center left', bbox_to_anchor=(1, 0.5))

But since the objects are of 2 different types (one is Patch and the other a ListedColormap. How should i go about this to make it happen?

Upvotes: 1

Views: 730

Answers (1)

just_another_profile
just_another_profile

Reputation: 499

I understand your approach with Patch - that was my original approach too (I never got it to work). Instead, you can use "tuple legend handler" to create a colorbar in the legend. Following the answer in Display matplotlib legend element as 2D line of colormap, here is how you can go about it:

import numpy as np
import matplotlib.pyplot as plt

from matplotlib.colors import ListedColormap
from matplotlib.legend_handler import HandlerTuple

if __name__ == "__main__":

    # Custom colormap
    N = 256
    vals = np.ones((N, 4))
    vals[:, 0] = np.linspace(0/256, 0, N)
    vals[:, 1] = np.linspace(250/256, 150/255, N)
    vals[:, 2] = np.linspace(0/256, 0, N)
    newcmp = ListedColormap(vals)

    # Start figure
    fig, ax = plt.subplots()

    lines = []  # list of lines to be used for the legend
    for i in range(5):
        # Create 5 sine waves
        t = 20
        fs = 10 + i
        samples = np.linspace(0, t, int(fs*t), endpoint=False)
        wave = np.sin(samples)
        # Plot - the i * 30 + 50 is to get different colors in the cmap that are far away from each other,
        # otherwise all greens look the same
        line, = ax.plot(wave, color=newcmp(i * 40 + 50))
        lines.append(line)

    ax.legend(handles=[tuple(lines)],
              labels=['Radius'],
              handlelength=5, handler_map={tuple: HandlerTuple(ndivide=None, pad=0)})

    plt.show()

Sine waves with colorbar in legend

Cheers!

Upvotes: 1

Related Questions