Don't show zero values on 2D heat map

Question

I want to plot a 2D map of a sillicon wafer dies. Hence only the center portion have values and corners have the value 0. I'm using matplotlib's plt.imshow to obtain a simple map as follows:

data = np.array([[ 0. ,  0. ,  1. ,  1. ,  0. ,  0. ],
       [ 0. ,  1. ,  1. ,  1. ,  1. ,  0. ],
       [ 1. ,  2. ,  0.1,  2. ,  2. ,  1. ],
       [ 1. ,  2. ,  2. ,  0.1,  2. ,  1. ],
       [ 0. ,  1. ,  1. ,  1. ,  1. ,  0. ],
       [ 0. ,  0. ,  1. ,  1. ,  0. ,  0. ]])

plt.figure(1)
plt.imshow(data ,interpolation='none')
plt.colorbar()

And I obtain the following map:

Is there any way to remove the dark blue areas where the values are zeros while retaining the shape of the 'wafer' (the green, red and lighter blue areas)? Meaning the corners would be whitespaces while the remainder retains the color configuration.

Or is there a better function I could use to obtain this?

Answer 1

There are two ways to get rid of the dark blue corners:

You can flag the data with zero values:

data[data == 0] = np.nan
plt.imshow(data, interpolation = 'none', vmin = 0)

Or you can create a masked array for imshow:

data_masked = np.ma.masked_where(data == 0, data)
plt.imshow(data_masked, interpolation = 'none', vmin = 0)

The two solutions above both solve your problem, although the use of masks is a bit more general.

If you want to retain the exact color configuration you need to manually set the vmin/vmax arguments for plotting the image. Passing vmin = 0 to plt.imshow above makes sure that the discarded zeros still show up on the color bar.