How to let user pick colormap?

Question

Hi everyone I want to define a function to plot map, and here is a simple example:

def PlotMap(df, fig = plt.figure(), size = 111, loc_ix = 0):
    ax = fig.add_subplot(size + loc_ix)
    color = matplotlib.cm.spectral(np.linspace(0,1,100))
    for s in df.index:
#this is for plotting process
#extract polygon from data frame df
        poly = Polygon(df.polygon[s])
#find its color based on the partition
        c = color[df.partition[s]][0:3]
        ax.add_patch(PolygonPatch(poly, fc = c, ec = 'k', alpha = 0.7, zorder = 2))
    ax.axis('scaled')

You see here the color is fixed with 'spectral', and I was wondering how to modify this code, so the user can pick their preferred colormap? The ideal scenario would be add an additional input argument (let's call it cmap) and then we can simply call

PlotMap(df, cmap = 'hot')

to plot a heatmap. (with 'spectral' being the default setting.)

Thanks a lot!

Answer 1

Something like this?

color_map_name = "spectral"
color_func = getattr(matplotlib.cm, color_map_name)
color = color_func(np.linspace(0,1,100))

EDIT: Even better:

color_func = matplotlib.cm.get_cmap(color_map_name)