Search code examples
pythonmatplotlibplotdata-visualizationfigure

how to get different line colors depending on one variable for different plots in one single figure in python?


Let's say I have one figure with a certain number of plots, which resembles like this one: enter image description here

where the colors of the single plots are decided automatically by matplotlib. The code to obtain this is very simple:

for i in range(len(some_list)):
    x, y = some_function(dataset, some_list[i])
    plt.plot(x, y) 

Now suppose that all these lines depend on a third variable z. I would like to include this information plotting the given lines with a color that gives information about the magnitude of z, possibly using a colormap and a colorbar on the right side of the figure. What would you suggest me to do? I exclude to use a legend since in my figures I have many more lines that the ones I am showing. All information I can find is about how to draw one single line with different colors, but this is not what I am looking for. I thank you in advance!


Solution

  • Here it is some code that, in my opinion, you can easily adapt to your problem

    import numpy as np
    import matplotlib.pyplot as plt
    from random import randint
    
    # generate some data
    N, vmin, vmax = 12, 0, 20
    rd = lambda: randint(vmin, vmax)
    segments_z = [((rd(),rd()),(rd(),rd()),rd()) for _ in range(N)]
    
    # prepare for the colorization of the lines,
    # first the normalization function and the colomap we want to use
    norm = plt.Normalize(vmin, vmax)
    cm = plt.cm.rainbow
    # most important, plt.plot doesn't prepare the ScalarMappable
    # that's required to draw the colorbar, so we'll do it instead
    sm = plt.cm.ScalarMappable(cmap=cm, norm=norm)
    
    # plot the segments, the segment color depends on z
    for p1, p2, z in segments_z:
        x, y = zip(p1,p2)
        plt.plot(x, y, color=cm(norm(z)))
    
    # draw the colorbar, note that we pass explicitly the ScalarMappable
    plt.colorbar(sm)
    
    # I'm done, I'll show the results,
    # you probably want to add labels to the axes and the colorbar.
    plt.show()
    

    enter image description here