#!/usr/bin/env python

# pcolormesh test/example adapted from
# https://matplotlib.org/stable/gallery/images_contours_and_fields/pcolormesh_levels.html#making-levels-using-norms

import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm
from matplotlib.ticker import MaxNLocator
import numpy as np

# make these smaller to increase the resolution
dx, dy = 0.05, 0.05

# generate 2 2d grids for the x & y bounds
y, x = np.mgrid[slice(1, 5 + dy, dy),
                slice(1, 5 + dx, dx)]

z = np.sin(x)**10 + np.cos(10 + y*x) * np.cos(x)

# x and y are bounds, so z should be the value *inside* those bounds.
# Therefore, remove the last value from the z array.
z = z[:-1, :-1]


# pick the desired colormap, sensible levels, and define a normalization
# instance which takes data values and translates those into levels.
cmap_0 = plt.get_cmap('PiYG')
levels_0 = MaxNLocator(nbins=15).tick_values(z.min(), z.max())
norm_0 = BoundaryNorm(levels_0, ncolors=cmap_0.N, clip=True)

fig, (ax_0, ax_1, ax_2, ax_3, ax_4, ax_5) = plt.subplots(nrows=6)
fig.set_size_inches(5, 12)

im_0 = ax_0.pcolormesh(x, y, z, cmap=cmap_0, norm=norm_0)
fig.colorbar(im_0, ax=ax_0)
ax_0.set_title('pcolormesh with levels')


ax_1.set_title('3 intervals/colors covering the\nfull data range ([%.3f, %.3f])' % (z.min(), z.max()))

cmap_1 = ListedColormap(['pink', 'lavender', 'green'])
levels_1 = [-1.1, -0.3, 0.3, 1.1]
norm_1 = BoundaryNorm(levels_1, cmap_1.N)

im_1 = ax_1.pcolormesh(x, y, z,
                       cmap=cmap_1, norm=norm_1,
                       zorder=20)
fig.colorbar(im_1, ax=ax_1)


ax_2.set_title('Top of data range is above the color scale\n...but set_over color is NOT specified')

cmap_2 = ListedColormap(['pink', 'lavender'])
levels_2 = [-1.1, -0.3, 0.3]
norm_2 = BoundaryNorm(levels_2, cmap_2.N)

im_2 = ax_2.pcolormesh(x, y, z,
                       cmap=cmap_2, norm=norm_2,
                     zorder=20)
fig.colorbar(im_2, ax=ax_2)


ax_3.set_title('BLUE color specified for\ndata ABOVE the color scale')

cmap_3 = ListedColormap(['pink', 'lavender'])
cmap_3.set_over('blue')
levels_3 = [-1.1, -0.3, 0.3]
norm_3 = BoundaryNorm(levels_3, cmap_3.N)

im_3 = ax_3.pcolormesh(x, y, z,
                       cmap=cmap_3, norm=norm_3,
                       zorder=20)
text_3 = ax_3.text(0.05, 0.3, 'Test text\nin the BACKGROUND',
                   fontsize='xx-large',
                   horizontalalignment='left',
                   verticalalignment='center', transform=ax_3.transAxes,
                   zorder=10)
fig.colorbar(im_3, ax=ax_3)


ax_4.set_title('Color above the scale is now\nfully TRANSPARENT (alpha=0)')

cmap_4 = ListedColormap(['pink', 'lavender'])
cmap_4.set_over('black', alpha=0)
levels_4 = [-1.1, -0.3, 0.3]
norm_4 = BoundaryNorm(levels_4, cmap_4.N)

im_4 = ax_4.pcolormesh(x, y, z,
                       cmap=cmap_4, norm=norm_4,
                       zorder=20)
text_4 = ax_4.text(0.05, 0.3, 'Test text\nin the BACKGROUND',
                   fontsize='xx-large',
                   horizontalalignment='left',
                   verticalalignment='center', transform=ax_4.transAxes,
                   zorder=10)
fig.colorbar(im_4, ax=ax_4)


ax_5.set_title('Color above the scale is NOT TRANSPARENT\n...but the data above 0. is MASKED')

# MASK the values we do not want, instead of plotting them with a
# fully transparent color
zm = np.ma.masked_greater(z, 0.)

cmap_5 = ListedColormap(['pink', 'lavender'])
# For testing purpose, we explicitly specify a fully opaque color
# (but opaque is the default anyway)
cmap_5.set_over('black', alpha=1)
levels_5 = [-1.1, -0.3, 0.3]
norm_5 = BoundaryNorm(levels_5, cmap_5.N)

im_5 = ax_5.pcolormesh(x, y, zm,
                       cmap=cmap_5, norm=norm_5,
                       zorder=20)
text_5 = ax_5.text(0.05, 0.3, 'Test text\nin the BACKGROUND',
                   fontsize='xx-large',
                   horizontalalignment='left',
                   verticalalignment='center', transform=ax_5.transAxes,
                   zorder=10)
fig.colorbar(im_5, ax=ax_5)


# adjust spacing between subplots so `ax1` title and `ax0` tick labels
# don't overlap
fig.tight_layout()

fig.savefig('test_pcolormesh.png')

plt.show()