Нормализация цветов в matplotlib

Я пытаюсь построить поверхность с помощью matplotlib, используя приведенный ниже код:

from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import axes3d, Axes3D
import pylab as p

vima=0.5

fig = plt.figure()
ax = fig.gca(projection='3d')
X = np.arange(0, 16.67, vima)
Y = np.arange(0, 12.5, vima)
X, Y = np.meshgrid(X, Y)

Z = np.sqrt(((1.2*Y+0.6*X)**2+(0.2*Y+1.6*X)**2)/(0.64*Y**2+0.36*X**2))

surf = ax.plot_surface(X, Y, Z,rstride=1, cstride=1, alpha=1,cmap=cm.jet,  linewidth=0)
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

Если вы запустите его, вы увидите синюю поверхность, но я хочу использовать весь цветовой диапазон струи ... Я знаю есть класс matplotlib.colors.Normalize, но я не знаю, как им пользоваться. Не могли бы вы добавить необходимый код для этого?