Search code examples
pythonmatplotlibcolorbar

Drawing a colorbar aside a line plot, using Matplotlib


I'm trying to add a color bar in a graph, but I don't understand how it works. The problem is that I make my own colorcode by:

x = np.arange(11)

ys = [i+x+(i*x)**2 for i in range(11)]

colors = cm.rainbow(np.linspace(0, 1, len(ys)))

and colors[i] will give me a new color. Then I use (homemade) functions to select the relevant data and plot them accordingly. This would look something like this:

function(x,y,concentration,temperature,1,37,colors[0])
function(x,y,concentration,temperature,2,37,colors[1])
# etc

Now I want to add the colors in a color bar, with labels I can change. How do I do this?

I have seen several examples where you plot all the data as one array, with automated color bars, but here I plot the data one by one (by using functions to select the relevant data).

EDIT:

function(x,y,concentration,temperature,1,37,colors[0]) looks like this (simplified):

def function(x,y,c,T,condition1,condition2,colors):

  import matplotlib.pyplot as plt

  i=0

  for element in c:
   if element == condition1:
      if T[i]==condition2:
          plt.plot(x,y,color=colors,linewidth=2)

  i=i+1

return

Solution

  • Drawing a colorbar aside a line plot

    Please map my solution (I used simply 11 sines of different amplitudes) to your problem (as I told you, it is difficult to understand from what you wrote in your Q).

    import matplotlib
    import numpy as np
    from matplotlib import pyplot as plt
    
    # an array of parameters, each of our curves depend on a specific
    # value of parameters
    parameters = np.linspace(0,10,11)
    
    # norm is a class which, when called, can normalize data into the
    # [0.0, 1.0] interval.
    norm = matplotlib.colors.Normalize(
        vmin=np.min(parameters),
        vmax=np.max(parameters))
    
    # choose a colormap
    c_m = matplotlib.cm.cool
    
    # create a ScalarMappable and initialize a data structure
    s_m = matplotlib.cm.ScalarMappable(cmap=c_m, norm=norm)
    s_m.set_array([])
    
    # plotting 11 sines of varying amplitudes, the colors are chosen
    # calling the ScalarMappable that was initialised with c_m and norm
    x = np.linspace(0,np.pi,31)
    for parameter in parameters:
        plt.plot(x,
                 parameter*np.sin(x),
                 color=s_m.to_rgba(parameter))
    
    # having plotted the 11 curves we plot the colorbar, using again our
    # ScalarMappable
    plt.colorbar(s_m)
    
    # That's all, folks
    plt.show()
    

    Example

    Curves+colormap

    Acknowledgements

    A similar problem, about a scatter plot

    Update — April 14, 2021

    1. With recent versions of Matplotlib, the statement s_m.set_array([]) is not required any more. On the other hand, it does no harm.
    2. When plotting, in place of color=s_m.to_rgba(parameter) one may want to use the (slightly) more obvious color=c_m(norm(parameter)).