Matplotlib: múltiples leyendas para la trama de contorno para múltiples variables de contorno - python, matplotlib, legend, contour

Necesito hacer trazados de contornos múltiples de variosVariables en la misma página. Puedo hacer esto con MATLAB (ver más abajo para el código de MATLAB). No puedo obtener matplotlib para mostrar varias leyendas. Cualquier ayuda sería muy apreciada.

Código Python:

import numpy as np
from matplotlib import cm as cm
from matplotlib import pyplot as plt

delta = 0.25
x = np.arange(-3.0, 3.0, delta)
y = np.arange(-2.0, 2.0, delta)
X, Y = np.meshgrid(x, y)
Z1 = X*np.exp(-X**2-Y**2)
Z2 = Y*np.exp(-X**2-Y**2)

plt.figure()
CS = plt.contour(X, Y, Z1, colors="k")
plt.clabel(CS, inline=1, fontsize=10)
CS = plt.contour(X, Y, Z2, colors="r")
plt.clabel(CS, inline=1, fontsize=10)
plt.legend(["case 1", "case 2"])

plt.show()

Código MATLAB:

[X,Y] = meshgrid(-2:.2:2,-2:.2:3);
Z1 = X.*exp(-X.^2-Y.^2);
Z2 = Y.*exp(-X.^2-Y.^2);

[C,h] = contour(X,Y,Z1, "color", "k");
set(h,"ShowText","on","TextStep",get(h,"LevelStep")*2);
hold on

[C,h] = contour(X,Y,Z2, "color", "r");
set(h,"ShowText","on","TextStep",get(h,"LevelStep")*2);

fn = {"case 1", "case 2"};

legend(fn,"Location","NorthWest");

Respuestas

Sería de ayuda si mostrara su salida deseada de Matlab. Por ejemplo, ¿realmente quieres múltiples leyendas? ¿O realmente quieres decir 1 leyenda con varios elementos?

Dado que las parcelas de contorno (can) tienen un estilo diferentepara cada nivel, no es obvio cómo querrías trazar eso en una leyenda. Pero para comenzar, puede acceder a cada línea individual examinando la matriz CS.collections.

Así por ejemplo:

delta = 0.25
x = np.arange(-3.0, 3.0, delta)
y = np.arange(-2.0, 2.0, delta)
X, Y = np.meshgrid(x, y)
Z1 = X*np.exp(-X**2-Y**2)
Z2 = Y*np.exp(-X**2-Y**2)

fig, ax = plt.subplots()

CS1 = ax.contour(X, Y, Z1, colors="k")
ax.clabel(CS1, inline=1, fontsize=10)

CS2 = ax.contour(X, Y, Z2, colors="r")
ax.clabel(CS2, inline=1, fontsize=10)

lines = [ CS1.collections[0], CS1.collections[-1], CS2.collections[0], CS2.collections[-1]]
labels = ["CS1_neg","CS1_pos","CS2_neg","CS2_pos"]

plt.legend(lines, labels)

Resultados en:

Tal vez algo como plt.legend(CS2.legend_elements()[0], CS2.legend_elements()[1]), también puede ser útil para ti.