Matplotlib Legend with Axis

Question

I'm having a hard time with this specific plot when it comes to adding a legend on it. I have reviewed similar questions but I don't see any where the challenge is how to label three elements such as bars in this case on a legend appended to the axis object.

How would I add a legend so that the first bar would show as "Training Errors", the second as "Val Errors" and third bar as "Test Errors"?

plt.figure(figsize=(20,10))
ax = plt.subplot(111)

x1 = [i-0.2 for i in range(len(train_errors))]
x2 = [i for i in range(len(train_errors))]
x3 = [i+0.2 for i in range(len(train_errors))]

ax.bar(x1, train_errors, width=0.2, color='b', align='center')
ax.bar(x2, val_errors, width=0.2, color='g', align='center')
ax.bar(x3, test_errors, width=0.2, color='r', align='center')
ax.set_xticklabels(X)
ax.xaxis.set_major_locator(ticker.FixedLocator([i-0.05 for i in x2]))

ax.legend((bar), ('label1'))

ax.set_xlabel('Models')
ax.set_ylabel('RMSE')
ax.set_title('Regression Models Comparison')
plt.show()

Thanks!

Answer 1

A legend entry for a bar plot, similar to many other artists, is created by specifiying the label argument.

ax.bar(...., label="my label")
ax.legend()

This is also shown in the first example in the documentation.

Complete example:

import matplotlib.pyplot as plt

fig, ax = plt.subplots()

train_errors = [1,2,3,4,5]
val_errors = [2,1,4,2,3]
test_errors = [5,4,3,2,1]
X = list("ABCDE")

x1 = [i-0.2 for i in range(len(train_errors))]
x2 = [i for i in range(len(train_errors))]
x3 = [i+0.2 for i in range(len(train_errors))]

ax.bar(x1, train_errors, width=0.2, color='b', label="Train Errors")
ax.bar(x2, val_errors, width=0.2, color='g', label="Val Errors")
ax.bar(x3, test_errors, width=0.2, color='r', label="Test Errors")

ax.set_xticks(x2)
ax.set_xticklabels(X)

ax.legend()

ax.set_xlabel('Models')
ax.set_ylabel('RMSE')
ax.set_title('Regression Models Comparison')
plt.show()