Matplotlib: Múltiplas legendas para gráfico de contorno para múltiplas variáveis ​​de contorno - python, matplotlib, legenda, contorno

Eu preciso fazer vários contornos gráficos de váriosvariáveis ​​na mesma página. Eu posso fazer isso com o MATLAB (veja abaixo o código do MATLAB). Não consigo obter o matplotlib para mostrar várias legendas. Qualquer ajuda seria muito 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");

Respostas:

Ajudaria se você mostrasse sua saída desejada do Matlab. Por exemplo, você realmente quer várias legendas? Ou você realmente quer dizer 1 legenda com vários itens?

Como os gráficos de contorno (podem) têm um estilo diferentepara cada nível, não é óbvio como você gostaria de traçar isso em uma legenda. Mas, para começar, você pode acessar cada linha individual examinando o array CS.collections.

Então, por exemplo:

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)

Resulta em:

Talvez algo como plt.legend(CS2.legend_elements()[0], CS2.legend_elements()[1]), também pode ser útil para você.