Search code examples
pythonmatplotlibcolorbar

create discrete colorbar from colormap in python


I want to use a given colormap (let's say viridis) and create plots and colorbars with discrete colors from that colormap.

I used to use mpl.cm.get_cmap("viridis", 7) for 7 different colors, but this function is deprecated and will be removed. The recommendation is to use matplotlib.colormaps[name] or matplotlib.colormaps.get_cmap(obj) instead, but neither of these allow to specify a number of discrete colors.

So far I have only found complicated workarounds online, is anyone aware of a simple, straightforward way as I had originally? Thank you!

Here is a sample code with simplified data. It took me a while to get the colorbar axis right how I wanted it, so I would prefer to stick with what I already have instead of a different plt.colorbar() solution.

import os.path as op
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.rcParams.update({'font.size': 30})
from mpl_toolkits.axes_grid1 import make_axes_locatable

PLOT = '/tmp/'


def main():
    data = np.random.random((20,20))
    data[5,:] = np.nan
    fig, ax = plt.subplots(figsize=(8.3,12))
    divider = make_axes_locatable(ax)

    cm = mpl.cm.get_cmap('viridis', 7)
    cm.set_bad('darkgrey', alpha=1)

    plt.pcolormesh(data, cmap=cm, vmin=0,vmax=1)
    ax.axis('off')

    cax = divider.append_axes("right", size="5%", pad=0.2)
    cb = plt.colorbar(cax=cax)

    plt.savefig(op.join(PLOT, 'test.png'), bbox_inches='tight', dpi=300)
    plt.clf()
    plt.close()


if __name__ == "__main__":
    main()


Solution

  • A simple option to replace matplotlib.cm.get_cmap (deprecated in 3.7.0) is to use matplotlib.pyplot.get_cmap (that was preserved in the API for backward compatibility).

    Another one woud be to make a resampled colormap (used under the hood) :

    resampled(lutsize) [source]

    Return a new colormap with lutsize entries.

    Note that the lut is a parameter in get_cmap(name=None, lut=None).

        cm = mpl.cm.get_cmap("viridis", 7)   # (-)
        cm = plt.get_cmap("viridis", 7)      # (+)    
    

    enter image description here