How can I add an R^2 value to the legend of a seaborn barplot?

user_15293617 0
Viewed 332

I have a seaborn barplot and a regression line graphed on top of it that looks like this. As you can see, I have a legend that is automatically created with seaborn.barplot() and I am attempting to add the R^2 score with this:

g = sns.barplot(x='City/Town', y="Value", hue="Metric", data=df, ax=ax1)
h, l = g.get_legend_handles_labels()
g.legend(h + [lin_reg.score(X, Y)], l + ['R^2 score'], title="Legend")

It doesn't throw an error, in fact I know it's working because it changes the title to "Legend", but it also doesn't add the R^2.

seaborn matplotlib python
1 Answers

The legend() function expects a handle in the first argument, and I don't think you can use text as one of those. You can read more help page for matplotlib legend

One quick solution I can think of, is to make a blank rectangle for the line intended for R^2, below is an example using iris as an example:

import seaborn as sns
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from sklearn.linear_model import LinearRegression 

df = sns.load_dataset("iris")
lin_reg = LinearRegression().fit(df[['petal_length']], df['sepal_length'])
r2 = lin_reg.score(df[['petal_length']], df['sepal_length'])

blank = Rectangle((0, 0), 1, 1, fc="w", fill=False, edgecolor='none', linewidth=0)

fig, ax = plt.subplots(figsize=(10,5))
sns.scatterplot(x='sepal_width', y="sepal_length", hue="species", data=df, ax=ax)

h, l = ax.get_legend_handles_labels()
ax.legend(h + [blank], l + [f'R^2 score = {r2:.3f}'], title="Legend")

enter image description here

I noticed you have a barplot with a few categories, so I am not sure how you can calculate R^2 from that. In any case, with the code above you should be able to add the R^2

user_122584590
Related