Colorbar for matplotlib plot_surface command

You can use a proxy mappable object since your surface array is not mapped. The mappable simply converts the values of any array to RGB colors defined by a colormap.

In your case you want to do this with the z1 array:

import matplotlib.cm as cm
m = cm.ScalarMappable(cmap=cm.jet)
m.set_array(z1)
plt.colorbar(m)