How to add a legend entry for a curve in 'plt.contour'?

Question

Based on this question, How to plot a function ax^2+bxy+cy^2+d=0 in Python?. I cite code from @Michael Szczesny’s answer:

import matplotlib.pyplot as plt

def f(x,y):
    return 0.01*x**2+0.008*x*y-0.013*y**2+0.15*x+0.003*y+1.0097

x = np.arange(-10.0,10.0,0.1)
y = np.arange(-10.0,10.0,0.1)
X, Y = np.meshgrid(x,y)
plt.contour(x, y, f(X, Y), [0]);

If I try to add a label for curve f(X,Y):

import matplotlib.pyplot as plt

def f(x,y):
    return 0.01*x**2+0.008*x*y-0.013*y**2+0.15*x+0.003*y+1.0097

x = np.arange(-10.0,10.0,0.1)
y = np.arange(-10.0,10.0,0.1)
X, Y = np.meshgrid(x,y)
plt.contour(x, y, f(X, Y), [0], label='class 1')
plt.legend()
plt.show()

Error: UserWarning: The following kwargs were not used by contour:

Asked By: Hermi

||

Source

Answer 1

Matplotlib’s contour seems to be designed with the idea of drawing a colorbar instead of a legend. However, in this case, with only one level, a legend can be forced assigning a label to individual contours:

import matplotlib.pyplot as plt
import numpy as np

def f(x, y):
    return 0.01 * x ** 2 + 0.008 * x * y - 0.013 * y ** 2 + 0.15 * x + 0.003 * y + 1.0097

x = np.arange(-10.0, 10.0, 0.1)
y = np.arange(-10.0, 10.0, 0.1)
X, Y = np.meshgrid(x, y)
cnt = plt.contour(x, y, f(X, Y), [0])
cnt.collections[0].set_label('class 1')
plt.legend()
plt.show()

Here is an example of the same approach with multiple levels:

levels = [0, 1, 2, 3]
cnt = plt.contour(x, y, f(X, Y), levels)
for level, contour in zip(levels, cnt.collections):
    contour.set_label('level {level}')
plt.legend(bbox_to_anchor=[1.01, 1], loc='upper left')
plt.tight_layout()
plt.show()

Answered By: JohanC

Answer 2

The answer from 2021 no longer works on Matplotlib 3.5.2. It will not render correctly:

Instead, we should use legend_elements with custom labels:

import matplotlib.pyplot as plt
import numpy as np

def f(x, y):
    return 0.01 * x ** 2 + 0.008 * x * y - 0.013 * y ** 2 + 0.15 * x + 0.003 * y + 1.0097

x = np.arange(-10.0, 10.0, 0.1)
y = np.arange(-10.0, 10.0, 0.1)
X, Y = np.meshgrid(x, y)
cnt = plt.contour(x, y, f(X, Y), [0])
artists, labels = cnt.legend_elements()
plt.legend(artists, ['class 1'])
plt.show()

# multiple levels
levels = [0, 1, 2, 3]
cnt = plt.contour(x, y, f(X, Y), levels)
artists, labels = cnt.legend_elements()
custom_labels = []
for level, contour in zip(levels, cnt.collections):
    custom_labels.append(f'level {level}')
plt.legend(artists, custom_labels, bbox_to_anchor=[1.01, 1], loc='upper left')
plt.tight_layout()
plt.show()

It will produce the correct result:

Answered By: jiegec

How to add a legend entry for a curve in 'plt.contour'?

Question:

Answers: