Contour plots and colors

Hello,

I am a big consumer of contour plots and it is great that matplotlib now features them. I didn't find the 'colors' argument to contour() intuitive to use though, and I wonder whether contour() should accept a colormap instance, much as imshow, so we can also display a colorbar.

Browsing through colors.py and cm.py it didn't appear clearly to me how all of that works. There is some info in the mailing list archives but I still didn't feel comfortable enough with such aspects of matplotlib to go ahead and modify contour().

Instead I wrote this simple class which returns a range of colors around the spectrum. There are as many colors as specified, and i use the same number as the number of levels in my contour plot. Perhaps this could be the default colors in contour()?

I don't mean to be reinventing the wheel; if there is a simpler way to do this with colormap instances, i'd love to know how.

···

===============
class mycolors:

     def __init__( self, nlevels ):
         jet6 = ( (0,0,1), (0,1,1), (0,1,0), (1,1,0), (1,0,0), (1,0,1) )
         self._jet6 = jet6

         if nlevels <= 6:
             jet = jet6[:nlevels]
         else:
             spectrum = linspace( 0, nlevels-1, 6 )
             for i in range( 6 ):
                 spectrum[i] = round( spectrum[i] )

             # Initialize colors to black
             jet = []
             for i in range( nlevels ):
                 jet.append( (0,0,0) )

             # Insert basic colors
             for i in range( 6 ):
                 jet[ int( spectrum[i] ) ] = jet6[i]

             # Insert spectrum in each bin
             for i in range( 5 ):
                 inthisbin = int( spectrum[i+1] - spectrum[i] - 1 )
                 eps = 1.0/(inthisbin + 1)
                 tones = linspace( eps, 1 - eps, inthisbin )
                 thistone = [ jet6[i+1][0] - jet6[i][0],
                             jet6[i+1][1] - jet6[i][1],
                             jet6[i+1][2] - jet6[i][2] ]
                 for j in range( inthisbin ):
                     thiscolor = [ jet6[i][0] + thistone[0] * tones[j],
                                 jet6[i][1] + thistone[1] * tones[j],
                                 jet6[i][2] + thistone[2] * tones[j] ]
                     jet[ int( spectrum[i] ) + j + 1 ] = tuple( thiscolor )

         self.jet = jet
         self.nlevels = nlevels

     def get_colors( self ):
         return self.jet

     def get_levels( self ):
         return self.nlevels

Here is an example script where i use the same number of colors as levels:

from pylab import *

def rosenbrock(x,y): return 10*(y-x**2)**2 + (x-1)**2

x = y = arange( -2, 2, 0.1 )
X, Y = meshgrid( x, y )
Z = rosenbrock( X, Y )
nlevels = 30
cols = mycolors( nlevels )
contour( Z,
          x = X,
          y = Y,
          levels = nlevels,
          colors = cols.get_colors(),
          origin = 'lower' )
show()

If this is useful to anyone, feel free to use it.

Dominique