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

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!