使用 matplotlib 上的子圖繪制帶有圖例的多條曲線 animation

Question

我使用 matplotlib 生成帶有子圖的 animation，但對於某些子圖，我想包含多條帶圖例的曲線。 我沒能讓它發揮作用。 我的代碼如下所示：

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 個子圖的 animation。 我希望獎勵曲線與相應的圖例位於同一子圖中。 當然，負載和獎勵不是隨機的，但我添加了它以准確顯示它們是什么：numpy arrays。

為了生成圖例，我使用了以下代碼：

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

原答案

添加以下行怎么樣（例如，在您 plot 行的循環之后）：

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

您需要將線條實例作為句柄傳遞給圖例 function 以包含來自其他子圖的線條。

編輯

以下代碼對我有用。 請注意，我稍微修改了您的示例以使其運行。 您可能還應該查看圖例指南，了解有關如何放置和修改圖例的更多詳細信息。 我在您的繪圖循環下方插入了圖例代碼。 有兩種選擇，要么為所有線條繪制一個圖例，要么為每條線條繪制一個圖例。

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