import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt

NX = 20
NY = 40
xvals = np.arange(NX)
yvals = np.arange(NY)
xgrid, ygrid = np.meshgrid(xvals, yvals)

def get_scalarmappable(min=0.0, max=1.0):
    cmap = 'Oranges'
    cmap = plt.get_cmap(cmap)
    norm = mpl.colors.Normalize(vmin=min, vmax=max)
    smap = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    smap._A = []
    return smap

# Construct 2D Function to plot
hist = 1.0 - yvals[np.newaxis, :]/yvals.max() + np.power(xvals[:, np.newaxis]/xvals.max() - 0.5, 3.0)
min = hist.min()
max = hist.max()

# Load colormap
smap = get_scalarmappable(min=min, max=max)

# Construct Figure
fig, ax = plt.subplots(figsize=[8, 6])

# Plot grid
pcm = ax.pcolormesh(xgrid, ygrid, hist.T, norm=smap.norm, cmap=smap.cmap, linewidth=0,
                    rasterized=True, vmin=smap.norm.vmin, vmax=smap.norm.vmax)
pcm.set_edgecolor('face')

# Construct contour line (levels)
num = 3
levels = min + np.arange(1, num+1)*(max-min)/(num+1)

# Create grey backdrop for contour lines
ax.contour(xgrid, ygrid, hist.T, colors='0.50', norm=smap.norm,
           levels=levels, linewidths=6.0, antialiased=True, zorder=10)
# Plot colored contour lines
cs = ax.contour(xgrid, ygrid, hist.T, cmap=smap.cmap, norm=smap.norm,
                levels=levels, linewidths=2.5, antialiased=True, zorder=11, alpha=1.0)

# Plot contour-line labels (at a higher zorder than grey lines)
plt.clabel(cs, zorder=50, inline=True, fontsize=16, colors='0.8')

fig.savefig('clabel_test.png')