Adjust legend Seaborn jointplot with several labels for one single dataframe

Question

Suppose I have a pandas dataframe with this structure:

Column 1   Column2 Column3
x1         y1       0
x2         y2       0
x3         y3       1
x4         y4       1
x5         y5       1
          ....
x_n-1      y_n-1    5
x_n        y_n      5

I want to create a jointplot where I assign different colors based on the values of Column3. The command I use is

h = sns.jointplot(x="Column1", y="Column2", data=data, hue="Column3")

So I have all my points with 6 different colors. The legend that comes out from the previous command has labels "0", "1", ... "5", which are not explanatory. Instead of them, I would like to have "label0", "label1", and so on.

I tried to use the following command:

h.ax_joint.legend([data.loc[data['Column3'] == 0], data.loc[data['Column3'] == 1], data.loc[data['Column3'] == 2], data.loc[data['Column3'] == 3], data.loc[data['Column3'] == 4], data.loc[data['Column3'] == 5]], ['label0', 'label1', 'label2', 'label3', 'label4', 'label5'])

But executing it I have the following message:

A proxy artist may be used instead. See: https://matplotlib.org/users/legend_guide.html#creating-artists-specifically-for-adding-to-the-legend-aka-proxy-artists

And of course it doesn't plot any legend anymore. I have been looking in the suggested documentation, but I couldn't figure out how to improve this. Does somebody have an idea? Thanks in advance!

Answer 1

The simplest way, and the most true to Seaborn's spirit, would be to (temporarily) rename the labels of the hue column:

import seaborn as sns
import pandas as pd
import numpy as np

data = pd.DataFrame({"Column1": np.random.randn(36) * 10,
                     "Column2": np.arange(36) % 6 + np.random.randn(36) / 4,
                     "Column3": np.arange(36) % 6})
labels = ['label0', 'label1', 'label2', 'label3', 'label4', 'label5']
g = sns.jointplot(data=data.replace({"Column3": {i: label for i, label in enumerate(labels)}}),
                  x="Column1", y="Column2", hue="Column3", palette="turbo")
g.ax_joint.invert_yaxis()

Another option would be to create the legend a second time and provide new labels. The second legend will replace the default one. This can be useful if you also want to change other properties, such as the legend's location or remove its title:

g = sns.jointplot(x="Column1", y="Column2", data=data, hue="Column3", palette="turbo")
handles, labels = g.ax_joint.get_legend_handles_labels()
g.ax_joint.legend(handles=handles, labels=['label0', 'label1', 'label2', 'label3', 'label4', 'label5'],
                  title="Column3")

PS: Here is an example with sns.jointplot(..., kind="kde"):

import seaborn as sns
import pandas as pd
import numpy as np

data = pd.DataFrame({"Column1": np.random.randn(36) * 10,
                     "Column2": np.arange(36) % 6 + np.random.randn(36) / 4,
                     "Column3": np.arange(36) % 6})
labels = ['label0', 'label1', 'label2', 'label3', 'label4', 'label5']
g = sns.jointplot(data=data.replace({"Column3": {i: label for i, label in enumerate(labels)}}),
                  x="Column1", y="Column2", hue="Column3", palette="turbo", kind="kde")
g.ax_joint.invert_yaxis()