Search code examples
pythonmatplotlibseaborn

Is it possible to keep the different color palettes of two lineplots on the same axes in a manually changed legend?


I have two lineplots plotted on the same axes. Each plot represents a different group and consists of two lines representing one of two conditions. Additonally each group is given a different colour palette (Blues for Group 1 and Reds for Group 2).

The automatically created legend keeps track of the colour palette defined for each plot but misses the information about the group.

How can I manually change the legend while keeping the colours of both plots the same.

Or is there even another method to include the group dimension to the legend in matplotlib?

I have the following dataframe:

data = {'group': [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
        'block': [1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3],
        'cond': ['c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2'],
        'value': [1, 2, 1.2, 2.1, 1.1, 2.2, 3, 4, 3.3, 4.1, 3.1, 4]}

Output:

      group  block cond  value
0       1      1   c1    1.0
1       1      1   c2    2.0
2       1      2   c1    1.2
3       1      2   c2    2.1
4       1      3   c1    1.1
5       1      3   c2    2.2
6       2      1   c1    3.0
7       2      1   c2    4.0
8       2      2   c1    3.3
9       2      2   c2    4.1
10      2      3   c1    3.1
11      2      3   c2    4.0
sns.lineplot(data=df.loc[df['group'] == 1], x='block', y='value', hue='cond',
             palette='Blues', marker='o')

sns.lineplot(data=df.loc[df['group'] == 2], x='block', y='value', hue='cond',
             palette='Reds', marker='o')
plt.show()

This is what I get if I normally plot both lines in a graph

This  is what I get if i normally plot both lines in a graph

If i now try to manually rename the conditions in the legend by adding the group the lines belong to using the following line:

plt.legend(title='Legend', labels=['Group 1: C1', 'Group 1: C2', 'Group 2: C1', 'Group 2: C2'])

The colour information of one of the plots is getting lost in the legend

This is what happens


Solution

  • You will need to get the legend handles from the lineplot plots as well before you set the legend and modify the text. Please see if the updated code is what you are looking for... Hope this helps.

    data = {'group': [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
            'block': [1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3],
            'cond': ['c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2'],
            'value': [1, 2, 1.2, 2.1, 1.1, 2.2, 3, 4, 3.3, 4.1, 3.1, 4]}
    df=pd.DataFrame(data)
    sns.lineplot(data=df.loc[df['group'] == 1], x='block', y='value', hue='cond',
                 palette='Blues', marker='o')
    sns.lineplot(data=df.loc[df['group'] == 2], x='block', y='value', hue='cond',
                 palette='Reds', marker='o')
    
    ## Get the legend handles and add them to plt.legend, so that the right handles are addigned
    legend_handles, _= plt.gca().get_legend_handles_labels()
    plt.legend(handles = legend_handles, title='Legend', loc='right',
               labels=['Group 1: C1', 'Group 1: C2', 'Group 2: C1', 'Group 2: C2'])
    plt.show()
    

    Output plot

    enter image description here