Filled contour using class labels

Question

I'm having a 2D grid of points where each of the points have a corresponding label, which is in the range [0.0, 5.0]. Now I want to do the following:

Plot all points in the grid and color them according to their label.

However, I don't want to do this using a scatter plot. I've tried to do it with a contourf and pcolormesh plot:

import matplotlib.pyplot as plt

np.random.seed(1234)

x = np.linspace(-1.0, 1.0, num=5)
xx, yy = np.meshgrid(x, x)
z = np.random.randint(low=0, high=6, size=xx.shape)
levels = np.arange(0, 6)

fig, axes = plt.subplots(nrows=2, ncols=2)
axes[0, 0].contourf(xx, yy, z)
axes[0, 1].contour(xx, yy, z, colors='k')
axes[1, 0].scatter(xx, yy, marker='.', c=z)
axes[1, 1].pcolormesh(xx, yy, z)
plt.show()

How should I specify the levels of the contourf plot such that I get contour lines separating the labels. (Similar to the pcolormesh plot)

Additionally, how can I fix the color for every label, i.e label 4 should always have color red?

EDIT: This is an example of a contourf plot which produces too many coloured areas:

Actually, there are only two labels in the grid. However, at the border between the two areas, several additional contour lines are drawn.

For the example above, there should be a single contour line separating the two areas (cyan and blue)

I appreciate any help.

Answer 1

If I understand you correctly, you want something like the pcolormesh plot, but with only the outlines. One way to achieve this is to extend (or widen) your array such that it contains the same value many times in x and y direction. This basically means that your z consists of many plateaus with very steep gradients in between. You can do this easily with np.repeat. Below I show an example where each point in the original data is expanded to a 20x20 plateau.

The colours of the plots can be fixed by creating a custom colormap. In your case, using a ListedColormap should be enough. When using contour, you also have to specify the levels at which the contours should be drawn in order to make it work correctly.

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colors

cmap = colors.ListedColormap([
    'royalblue', 'cyan','yellow', 'orange', 'red', 'green'
])

np.random.seed(1234)

num = 5
x = np.linspace(-1.0, 1.0, num=num)
xx, yy = np.meshgrid(x, x)
z = np.random.randint(low=0, high=6, size=xx.shape)
levels = np.arange(0, 6)

fig, axes = plt.subplots(nrows=2, ncols=2)
axes[0, 0].contourf(xx, yy, z)
axes[0, 1].contour(xx, yy, z, colors='k')
axes[1, 0].scatter(xx, yy, marker='.', c=z)
axes[1, 1].pcolormesh(xx, yy, z, cmap=cmap)  ##I added here the custom colormap for reference

##expanding the arrays
N = 20
x1 = np.linspace(-1.0, 1.0, num=N*num)
xx1, yy1 = np.meshgrid(x1,x1)
z1 = np.repeat(np.repeat(z, N,axis=0),N).reshape(xx1.shape)

fig2, ax2 = plt.subplots()
ax2.contour(xx1, yy1, z1, cmap=cmap, levels = levels)

plt.show()

produces this kind of plot:

As you can see, the lines are still not quite straight and sometimes two lines next to each other can be seen. This is because the gradients between different plateaus are not equal. I ran another example using N=200, in which case the lines are much straighter:

Hope this helps.

Answer 2

Possibly you simply forgot to provide the levels you want to show. For N labels, one would need 7 levels, e.g. for labels [0 1 2 3 4 5] one would choose levels such that the labels are in the middle of the level interval, [-0.5 0.5 1.5 2.5 3.5 4.5 5.5].

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors

np.random.seed(1234)

x = np.linspace(-1.0, 1.0, num=5)
xx, yy = np.meshgrid(x, x)
z = np.random.randint(low=0, high=6, size=xx.shape)

levels = np.arange(0, z.max()+2)-0.5

fig, ax = plt.subplots()
im = ax.contourf(xx, yy, z, levels=levels)

fig.colorbar(im, ax=ax, ticks=np.unique(z))
ax.contour(xx, yy, z, colors='k',levels=levels)
ax.scatter(xx, yy, marker='.', c=z)

plt.show()

Note that the colors of the contourf plot are slightly different than those of the scatter. The reason is explained in the answer to the question: How does pyplot.contourf choose colors from a colormap?