Plot contours from discrete data in matplotlib

Question

How do I make a contourf plot where the areas are supposed to be discrete (integer array instead of float)?

The values should discretely mapped to color indices. Instead matplotlib just scales the result across the whole set of colors.

Example:

import numpy as np
from matplotlib import pyplot as plt

axes = (np.linspace(-2, 2, 100), np.linspace(-2, 2, 100))
xx, yy = np.meshgrid(*axes, indexing="xy")
fig, ax = plt.subplots()
z = np.abs(xx * yy).astype(int)  # values 0, 1, 2, 3, 4
z[z==0] = 4
ax.contourf(xx, yy, z, cmap="Set1")

Answer 1

Now I got it :) Thanks @jared, pcolormesh was the right function, but I have to explicitly map the colors as the plotted variable:

import numpy as np
from matplotlib import pyplot as plt

axes = (np.linspace(-2, 2, 100), np.linspace(-2, 2, 100))
xx, yy = np.meshgrid(*axes, indexing="xy")
fig, ax = plt.subplots()
z = np.abs(xx * yy).astype(int)  # values 0, 1, 2, 3, 4
z[z==0] = 4

cmap = plt.get_cmap("Set1")   
z_color = cmap(z)  # shape (100, 100, 4) with `z` as index
ax.pcolormesh(xx, yy, z_color)