Set discrete colorbar in matplotlib in layered cross-plot

I have a 5-layered model with constant parameters. How to make the colorbar discrete in matplotlib?

import matplotlib.pyplot as plt
import numpy as np
AI_t = np.zeros((3236, 6))
tborders = np.array([800, 1371, 2371, 2871, 3235])
AI_ = np.array([ 400.,  490.,  400.,  800., 1320.])
for n in range(5):
    if n==0:
        AI_t[:tborders[n]] = AI_[n]
    else:
        AI_t[tborders[n-1]:tborders[n]+1] = AI_[n]
R_t = np.zeros((AI_t.shape))
for n in range(1, R_t.shape[0]):
    R_t[n] = (AI_t[n]-AI_t[n-1])/(AI_t[n]+AI_t[n-1])
Xt, Zt = np.meshgrid(np.arange(0, 1200, 200), np.arange(0, 3.236, 0.001))
fig, ax = plt.subplots()
im = ax.contourf(Xt, -Zt, AI_t, cmap='viridis', extend='both')
fig.colorbar(im, orientation='vertical')
plt.show()

So, here in the colorbar should be only 4 unique values for AI_t = [400., 490., 800., 1320.]