How to plot legend for a colormap?

[Felipe]

Hi guys,

I wrote a function that plot rectangles filled with colors defined by a colormap. Here's a part of the function:

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib import cm

def tabularS(X,dx,p,q,prism):

   vm = 3*max(p[:,3])
   unit = dx/len(X) 

  norm = plt.Normalize()
  colors = plt.cm.rainbow(norm(p[:,0]))

   for i in range(prism):
       xl = p[i,2] - unit
       cor = colors[i]

       rect = Rectangle((xl,p[i,3]),unit,vm-p[i,3],color=cor)

       ax.add_patch(rect)
       plt.xlim(X[0], X[-1])
       plt.ylim([0, vm])
       ax = plt.gca()
       ax.invert_yaxis()

   plt.title('Prisms')
   plt.xlabel('Position(m)')
   plt.ylabel('Depth(m)')
   plt.show()
   return

The function works and I can plot a figure composed of rectangles filled with the colors defined by the colormap.
Now I need to plot a legend to this colormap, but I don't know how to do it.
I appreciate the help.

[Felipe]

Hi all,

I wrote a more complete function which works well too, but I still have no idea, about how to display a color bar based on the colors that fills the rectangles. Here's the updated code:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib import cm 
from colour import Color

def tabularS(X,dx,p,q,prism):
   vm = 3*max(p[:,3])
   unit = dx/len(X) # half size for minimun value of MDT
   lenp=[]

   cont1 = 0

# Ordenating the p vector
   posit = np.argsort(p[:,0])
   p = p[posit,:]
   norm = plt.Normalize()
   colors = plt.cm.rainbow(norm(p[:,0]))

   ax3 = fig.add_subplot(313)

# Rectangle width
   for i in range(1,prism):
       aux = p[i,0]/p[0,0]
       lenp.append(unit*aux)

# Plot and fill
   for i in range(prism):
       xl = (p[i,2]-unit)

       cor = colors[cont1]
       cont1+=1
       rect = Rectangle((xl,p[i,3]),unit,vm-p[i,3],color=cor)
       ax3.add_patch(rect)

       plt.xlim(X[0], X[-1])
       plt.ylim([0, vm])
       plt.gca().invert_yaxis()
       plt.title('Position')


       if p[i,1] <= 0:
           rect = Rectangle((xl,p[i,3]),unit,vm-p[i,3],linestyle = '-',fill = False, hatch = '//')
           ax3.add_patch(rect)

           plt.xlim(X[0], X[-1])
           plt.ylim([0, vm])
           plt.gca().invert_yaxis()
           plt.title('Position')

       rect = Rectangle((xl,q[i,3]), unit, vm-q[i,3], linestyle = '--', fill = False)
       ax3.add_patch(rect)

       plt.xlim(X[0], X[-1])
       plt.ylim([0, vm])
       plt.gca().invert_yaxis()
       #plt.colorbar(colors,ax3)
       plt.title('Prisms Position')
       plt.ylabel('Depth (m)')
       plt.xlabel('Distance (m)')

   plt.show()
   return

If it is possible to help me with this problem, I will be very grateful.