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()