在matplotlib动画上用图例绘制多条曲线

Question

我用matplotlib生成一个带有子图的动画，但是对于一些子图，我想包含多条有图例的曲线。我一直没能使它发挥作用。我的代码如下所示：

load = np.rand(1000,24)
index = np.arange(24)
reward1 = np.rand(1000)
reward2 = np.rand(1000)
reward3 = np.rand(1000)
Z = np.arange(1000)
reward1 = np.vstack((Z,reward1))
reward2 = np.vstack((Z,reward2))
reward3 = np.vstack((Z,reward3))
fig, axes = plt.subplots(2,2,figsize=(15,9))
fig.tight_layout()
lines = []
for nd, ax in enumerate(axes.flatten()):
    if nd == 0:
        l, = ax.plot(load[0,:],index)
    if nd == 1:
        l, = ax.plot(reward1[0], reward1[0])
    if nd == 2:
        l, = ax.plot(reward2[0], reward2[0])
    if nd == 3:
        l, = ax.plot(reward3[0], reward3[0])
    lines.append(l)

def run(it):
    for nd, line in enumerate(lines):
        if nd == 0:
            line.set_data(index,load[it,:])
        if nd == 1:
            line.set_data(reward1[..., :it])
        if nd == 2:
            line.set_data(reward2[..., :it])
        if nd == 3:
            line.set_data(reward3[..., :it])

ani = animation.FuncAnimation(fig, run, frames=1000, interval=30, blit=True)
ani.save('figure.mp4')
plt.show()

这将创建一个包含4个子情节的动画。我希望奖励曲线和相应的传说在同一个子情节上。当然，负载和奖励并不是随机的，但我补充说，为了准确地显示它们是什么: numpy数组。

为了生成传奇，我使用了以下代码：

plt.figure()
na1, = plt.plot(reward1)
na2, = plt.plot(reward2)
na3, = plt.plot(reward3)
plt.legend((na1,na2,na3),('Reward 1','Reward2','Reward 3'))
plt.show()

我试图集成两段代码，但没有成功。那么，有什么简单的方法来做我想做的事吗？

Answer 1

原始答案

添加以下行怎么样(例如，在绘制行的循环之后)：

axes.flatten()[0].legend(handles=lines, labels=["index", "reward1", "reward2", "reward3"])

您需要将您的行实例作为句柄传递给图例函数，以包含来自其他子图的行。

编辑

下面的代码适用于我。注意，我对您的示例做了一些修改，使其运行。您可能还应该查看图例指南，了解有关如何放置和修改图例的更多细节。我在你的绘图循环下面插入了图例代码。有两种选择，要么为所有行绘制一个图例，要么为每行绘制一个图例。

import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np

# Dummy data
load = np.random.randn(1000, 24)
index = np.arange(24)
reward1 = np.random.randn(1000)
reward2 = np.random.randn(1000)
reward3 = np.random.randn(1000)
Z = np.arange(1000)
reward1 = np.vstack((Z, reward1))
reward2 = np.vstack((Z, reward2))
reward3 = np.vstack((Z, reward3))

fig, axes = plt.subplots(2, 2, figsize=(15, 9))
fig.tight_layout()

axes = axes.flatten()

lines = []
for nd, ax in enumerate(axes):
    if nd == 0:
        l, = ax.plot(index, load[0, :], label="index")
    if nd == 1:
        l, = ax.plot(reward1[0], reward1[1], label="reward1")
    if nd == 2:
        l, = ax.plot(reward2[0], reward2[1], label="reward2")
    if nd == 3:
        l, = ax.plot(reward3[0], reward3[1], label="reward3")
    lines.append(l)

# Legend
# Have one legend on axes[0] for all lines
axes[0].legend(handles=lines, loc=1)

# Legend alternative
# Have one legend per axes for one line each
# for ax in axes:
#     ax.legend()

def init():
    return lines

def run(it):
    for nd, line in enumerate(lines):
        if nd == 0:
            line.set_data(index, load[it, :])
        if nd == 1:
            line.set_data(reward1[0, :it], reward1[1, :it])
        if nd == 2:
            line.set_data(reward2[0, :it], reward2[1, :it])
        if nd == 3:
            line.set_data(reward3[0, :it], reward3[1, :it])
    return lines

ani = animation.FuncAnimation(
    fig, run, frames=np.arange(1000), interval=30, blit=True,
    init_func=init
)
ani.save('figure.mp4')
plt.show()