Small utility to create a discrete matplotlib colormap

discrete_colorbar.py

# Original by Jake VanderPlas 
# https://gist.github.com/jakevdp/91077b0cae40f8f8244a
# License: BSD-style
# cvanelteren: small edit since the original was not compatible with latest mpl
import matplotlib.pyplot as plt, 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, 0))
    cmap_name = base.name + str(N)
    # edit cvanelteren
    return plt.cm.colors.ListedColormap(color_list, 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 cvanelteren object based matplotlib
    cmap =  discrete_cmap(N, 'cubehelix')
    fig, ax = plt.subplots()
    h = ax.scatter(x, y, c=c, s=50, cmap=cmap)
    # adjust ticks accordingly
    cbr = fig.colorbar(h, ticks = [i + .5 for i in range(N)])
    cbr.set_ticklabels([i + .5 for i in cbr.get_ticks()])
    fig.show()