Configure matplotlib colorbar to match 3D surface values

Question

I'm trying to display the result of a Fast Fourier Transform as a 3D-surface from a 2D-matrix (shape: 2048x1024) with a colorbar on the side that should match with my values. The display works fine but the colorbar's "color" and the graph's color don't match. How should I configure the colorbar and the graph to match?

I have tried to set vmin and vmax values according to my FFT results but then the color of the graph is totally wrong. I have also tried the function clim(vmin, vmax) or the method .set_clim(vmin, vmax) but the problem is still the same.

Here is my code: rp is a 2048x1024 matrix.

import numpy as np
from pylab import *
from mpl_toolkits.mplot3d import Axes3D


figure(2, figsize=(9.6, 7.2))
ax1 = gca(projection='3d')
X = np.arange(0, 1024)
Y = np.arange(0, 2048)
X, Y = np.meshgrid(X, Y)
Z = 20 * np.log10(abs(rp))
# Plot the surface.
surf1 = ax1.plot_surface(X, Y, Z, cmap='jet', antialiased=False, vmin=np.min(Z), vmax=np.max(Z))
# Add a color bar which maps values to colors.
colorbar(surf1)
title('3D frequency profile')

show()

Here is what I get without configuring vmin and vmax values for the colorbar. The max color (red) doesn't match with the maximun pikes and the colorbar max and min value don't match with my min and max values from the graph (NB: min=-24 and max=145).

first_figure

And here is what I get when I set vmin and vmax. All the graph is roughly the same shade whereas I should have to two pikes in red.

second_figure

Any ideas?

Answer 1

I have finally found the reason of my problem. The problem was the stride of the colormap for the graph (it's kind of the resolution of it, if I understand it well - cstride and rstride documentation). More precisely, it uses a surface of cstride * rstride to average the energy of the signal on the surface and makes the color of map correspond to this mean. My pikes were really thin that is why they were not in the good shade. I had to specify the parameters cstride and rstride with a smaller value. As their default value was 10.

So now I have:

figure(1, figsize=(12, 6))
X = np.arange(0, 1024)
Y = np.arange(0, 2048)
X, Y = np.meshgrid(X, Y)
Z = 20 * np.log10(abs(rp))
mappable = cm.ScalarMappable()
mappable.set_array(Z)
suptitle('3D-plot and 2D-plot of frequency profile')
ax1 = subplot(121, projection='3d')
ax1.plot_surface(X, Y, Z, cmap=mappable.cmap, norm=mappable.norm, linewidth=0,
                 antialiased=False, cstride=10, rstride=1)
ax2 = subplot(122)
ax2.imshow(Z, cmap=mappable.cmap, norm=mappable.norm,
           extent=(np.min(X), np.max(X), np.min(Y), np.max(Y)),
           interpolation=None)
colorbar(mappable)

show()

NB: I could have left the default value for cstride as I only need to change the row stride. But it could be helpful for someone else.

And here is the result I get and it corresponds to what I want.

NB2: Be careful by reducing the stride, the bigger the input data, the slower the script.