在 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"])

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

编辑

以下代码适用于我。请注意，我稍微修改了您的示例以使其成为运行。您可能还应该查看 legend guide 以了解有关如何放置和修改图例的更多详细信息。我在您的绘图循环下方插入了图例代码。有两种选择，要么为所有行绘制一个图例，要么为每行绘制一个图例。

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()

在 matplotlib 动画上使用子图绘制带有图例的多条曲线

Plotting multiple curves with a legend using subplot on matplotlib animation

python

animation

matplotlib

legend

subplot