Cómo hacer la leyenda del diagrama de dispersión en Matplotlib

1. Añadir una leyenda a la gráfica de dispersión 2D en Matplotlib
2. Agregar una leyenda al gráfico de dispersión 3D en Matplotlib

La leyenda es simplemente la descripción de varios elementos de una figura. Podemos generar una leyenda de gráfica de dispersión utilizando la función matplotlib.pyplot.legend.

Añadir una leyenda a la gráfica de dispersión 2D en Matplotlib

import numpy as np
import matplotlib.pyplot as plt

x = [1, 2, 3, 4, 5]

y1 = [i ** 2 for i in x]
y2 = [2 * i + 1 for i in x]

plt.scatter(x, y1, marker="x", color="r", label="x**2")
plt.scatter(x, y2, marker="o", color="b", label="2*x+1")
plt.legend()
plt.show()

Resultado:

Tenemos dos diagramas de dispersión separados en la figura: uno representado por la x y otro por la o. Asignamos la label a cada gráfica de dispersión usada como etiqueta mientras se genera la leyenda. Luego, creamos la leyenda en la figura usando la función legend() y finalmente mostramos la figura completa usando el método show().

También podemos pasar directamente una tupla de etiquetas como argumento al método legend() y cambiar la posición de la leyenda utilizando el parámetro loc del método legend().

import numpy as np
import matplotlib.pyplot as plt

x = [1, 2, 3, 4, 5]

y1 = [i ** 2 for i in x]
y2 = [2 * i + 1 for i in x]

plt.scatter(x, y1, marker="x", color="r")
plt.scatter(x, y2, marker="o", color="b")
plt.legend(("x**2", "2*x+1"), loc="center left")
plt.show()

Resultado:

Este proceso crea una figura con dos diagramas de dispersión y una leyenda colocada en el centro a la izquierda de la caja fronteriza de los ejes.

Agregar una leyenda al gráfico de dispersión 3D en Matplotlib

import matplotlib.pyplot as plt

x = [1, 2, 3, 4, 5]
y = [2, 1, 4, 5, 6]

z1 = [i + j for (i, j) in zip(x, y)]
z2 = [3 * i - j for (i, j) in zip(x, y)]

axes = plt.subplot(111, projection="3d")
axes.plot(x, y, z1, "x", label="x+y")
axes.plot(x, y, z2, "o", label="3*x-y")

plt.legend(loc="upper left")
plt.show()

Resultado:

Para crear una leyenda para el gráfico de dispersión 3D, usamos el método plot() en lugar del método scatter(); es porque el método legend() no soporta Patch3DCollection devuelto por el método scatter() de la instancia Axes3D.