Small utility to create a discrete matplotlib colormap

discrete_cmap.py

# By Jake VanderPlas
# License: BSD-style

import matplotlib.pyplot as plt
import numpy as np


def discrete_cmap(N, base_cmap=None):
    """Create an N-bin discrete colormap from the specified input map"""

    # Note that if base_cmap is a string or None, you can simply do
    #    return plt.cm.get_cmap(base_cmap, N)
    # The following works for string, None, or a colormap instance:

    base = plt.cm.get_cmap(base_cmap)
    color_list = base(np.linspace(0, 1, N))
    cmap_name = base.name + str(N)
    return base.from_list(cmap_name, color_list, N)


if __name__ == '__main__':
    N = 5

    x = np.random.randn(40)
    y = np.random.randn(40)
    c = np.random.randint(N, size=40)

    # Edit: don't use the default ('jet') because it makes @mwaskom mad...
    plt.scatter(x, y, c=c, s=50, cmap=discrete_cmap(N, 'cubehelix'))
    plt.colorbar(ticks=range(N))
    plt.clim(-0.5, N - 0.5)
    plt.show()

figure_1.png

Thanks, this is great! I did however notice that it doesn't work for the perceptually uniform colormaps (at least, not in 1.5.1) because they are ListedColormap instances.

Replacing return base.from_list(cmap_name, color_list, N) with return LinearSegmentedColormap.from_list(cmap_name, color_list, N) is sufficient (LinearSegmentedColormap is in matplotlib.colors).

This is brilliant. Thanks a lot for the time saved. :)

Thank you for sharing this! This is great and simple. A function like this should be part of matplotlib.

This is exactly what I was looking for, glad you shared your ideas! Thanks a lot!

Thanks a lot for this guide. But I have a question. How to handle the guide when you want your colorbar from 3 to 6 as an example? Because in your code it always starts with 0.

Not sure why, but it doesn't actually work for me. I had to replace the final line in the function with return ListedColormap(color_list, name=cmap_name).

Thanks!

I received a ValueError complaining that cmap_name was not a valid value for name; supported values are (followed by the list of predefined color map names).
To fix it, I had to register the colormap. Specifically, I replaced return base.from_list(cmap_name, color_list, N) with

custom_cmap = base.from_list(cmap_name, color_list, N)
plt.cm.register_cmap(name=cmap_name, cmap=custom_cmap)
return custom_cmap

This worked for me

def discrete_cmap(N, base_cmap=None):
    """Create an N-bin discrete colormap from the specified input map"""

    # Note that if base_cmap is a string or None, you can simply do
    #    return plt.cm.get_cmap(base_cmap, N)
    # The following works for string, None, or a colormap instance:

    base = plt.cm.get_cmap(base_cmap)
    color_list = base(np.linspace(0, 1, N, 0))
    cmap_name = base.name + str(N)
    return plt.cm.colors.ListedColormap(color_list, color_list, N)

Exactly what I need! Thank you so much for sharing!

Incrediubly great snippet, thanks a lot.
I have a qustion though.
How would you go about inputting your own list of discreet colors?

Thank you so much for this piece of code!!

what if I have values that are negative, e.g., -0.5, -1, -1.5, -3 ... the function does not work in this case. Any tips?

@Yingjie4Science
Should world out of the box with negative values:

def discrete_cmap(N, base_cmap=None):
    import matplotlib as pplt
    """Create an N-bin discrete colormap from the specified input map"""

    # Note that if base_cmap is a string or None, you can simply do
    #    return plt.cm.get_cmap(base_cmap, N)
    # The following works for string, None, or a colormap instance:

    base = pplt.cm.get_cmap(base_cmap)
    color_list = base(np.linspace(0, 1, N, 0))
    cmap_name = base.name + str(N)
    return pplt.cm.colors.ListedColormap(color_list, color_list, N)

N = 10
cmap = discrete_cmap(N)
x = np.random.randint(N, size = (2, 10))*-1

fig, ax = plt.subplots()
h = ax.scatter(*x, c = x[0], cmap = cmap)
fig.colorbar(h)
fig.show()

@cvanelteren Works great. Thanks!