python - How can I generate a colormap array from a simple array in matplotlib

Question

In some functions of matplotlib, we have to pass an color argument instead of a cmap argument, like bar3d.

So we have to generate a Colormap manually. If I have a dz array like this:

dz = [1,2,3,4,5]

What I want is:

ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color=cm.jet(dz), zsort='average')

However, It does not work, it seems Colormap instances can only convert normalized arrays.

>>> dz = [1,2,3,4,5]
>>> cm.jet(dz)
array([[ 0.        ,  0.        ,  0.51782531,  1.        ],
       [ 0.        ,  0.        ,  0.53565062,  1.        ],
       [ 0.        ,  0.        ,  0.55347594,  1.        ],
       [ 0.        ,  0.        ,  0.57130125,  1.        ],
       [ 0.        ,  0.        ,  0.58912656,  1.        ]])

Of course, this is not what I want.

I have to do things like this:

>>> cm.jet(plt.Normalize(min(dz),max(dz))(dz))
array([[ 0.        ,  0.        ,  0.5       ,  1.        ],
       [ 0.        ,  0.50392157,  1.        ,  1.        ],
       [ 0.49019608,  1.        ,  0.47754586,  1.        ],
       [ 1.        ,  0.58169935,  0.        ,  1.        ],
       [ 0.5       ,  0.        ,  0.        ,  1.        ]])

How ugly the code is!

In matplotlib's document it is said:

Typically Colormap instances are used to convert data values (floats) from the interval [0, 1] to the RGBA color that the respective Colormap represents. For scaling of data into the [0, 1] interval see matplotlib.colors.Normalize. It is worth noting that matplotlib.cm.ScalarMappable subclasses make heavy use of this data->normalize->map-to-color processing chain.

So why I can't use just cm.jet(dz)?

Here are the imports that I am using

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm

Answer 1

The answer to your question is given in the snipplet of the documentation that you copied into your question:

...from the interval [0, 1] to the RGBA color...

But if you find your code ugly you could try to make it nicer:

1. You don't have to specify the limits to the normalization manually (iff you intent to use min/max):

   norm = plt.Normalize()
   colors = plt.cm.jet(norm(dz))
   

2. If you find that ugly (I don't understand why, though), you could go on and do it manually):

   colors = plt.cm.jet(np.linspace(0,1,len(dz)))
   

   However this is solution is limited to equally spaced colors (which is what you want given the dz in your example.

3. And then you can also replicate the functionality of Normalize (since you seem to not like it):

   lower = dz.min()
   upper = dz.max()
   colors = plt.cm.jet((dz-lower)/(upper-lower))
   

4. Use a helper function:

   def get_colors(inp, colormap, vmin=None, vmax=None):
       norm = plt.Normalize(vmin, vmax)
       return colormap(norm(inp))
   

   Now you can use it like this:

   colors = get_colors(dz, plt.cm.jet)