python - 动画Matplotlib Sankey图

我正在尝试制作matplotlib Sankey图表的动画，其中流的大小发生了变化。我有一个示例可以部分工作，但是它不断添加更多Sankey图，如计数器所示； len（sankey.diagrams）。这意味着它将无法正常工作，带有多个箭头或持续很长时间。如何正确设置带有多个箭头的matplotlib Sankey图的动画？

from matplotlib import pyplot as plt
from matplotlib import animation
from matplotlib.sankey import Sankey
import numpy as np

time = np.linspace(0, np.pi * 2, 100)
flow_in = 4 * np.ones_like(time)
flow_out_a = 2 + np.sin(time)
flow_out_b = flow_in - flow_out_a

fig = plt.figure()
ax = fig.add_subplot(111)
sankey = Sankey(ax=ax, scale=1 / max(flow_in), format=u'%0.3g')
diags_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)

def init():
    diags_text.set_text('')
    sankey.diagrams=[]
    a = flow_in[0]
    b = -flow_out_a[0]
    c = -flow_out_b[0]
    sankey.add(flows=[a, b, c],
               orientations=[0, 0, -1],
               alpha=0.2, lw=2.0, trunklength=1, pathlengths=[1, 1, 1])
    diag = sankey.finish()
    return(diag, diags_text)

def anim(idx):
    for diag in sankey.diagrams:
        diag.patch.set_visible(False)
        diag.text.set_visible(False)
        for txt in diag.texts:
            txt.set_visible(False)
    a = flow_in[idx]
    b = -flow_out_a[idx]
    c = -flow_out_b[idx]
    sankey.add(flows=[a, b, c],
               orientations=[0, 0, -1],
               alpha=0.2, lw=2.0, trunklength=1, pathlengths=[1, 1, 1])

    diag = sankey.finish()
    diags_text.set_text('len(sankey.diagrams) = {l}'.format(l=len(sankey.diagrams)))
    return(diag, diags_text)

frames, = time.shape

anim = animation.FuncAnimation(fig, anim, init_func=init,
                               frames=frames, interval=20, blit=False)

最佳答案

可以在anim函数中添加线sankey.diagrams = []。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib import sankey
import subprocess

N = 10
time = np.linspace(0, np.pi * 2, N)
flow_in = 4 * np.ones_like(time)
flow_out_a = 2 + np.sin(time)
flow_out_b = flow_in - flow_out_a

fig, ax = plt.subplots()
sk = sankey.Sankey(ax=ax, scale=1 / max(flow_in), format=u'%0.3g')
diags_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)

def init():
    a = flow_in[0]
    b = -flow_out_a[0]
    c = -flow_out_b[0]
    sk.add(flows=[a, b, c],
           orientations=[0, 0, -1],
           alpha=0.2, lw=2.0, trunklength=1, pathlengths=[1, 1, 1])
    diags = sk.finish()
    return []

def anim(idx):
    for diag in sk.diagrams:
        # remove called instead of set_visible(False) to keep
        # `len(ax.get_children())` constant
        diag.patch.remove()
        diag.text.remove()
        for txt in diag.texts:
            txt.remove()
    sk.diagrams = []

    a = flow_in[idx]
    b = -flow_out_a[idx]
    c = -flow_out_b[idx]

    sk.add(flows=[a, b, c],
           orientations=[0, 0, -1],
           alpha=0.2, lw=2.0, trunklength=1, pathlengths=[1, 1, 1])
    diags = sk.finish()
    diags_text.set_text('len(sk.diagrams) = {l} | len(ax.get_children()) = {gc}'
                        .format(l=len(sk.diagrams), gc=len(ax.get_children())))
    ax.set_xlim(-2, 3)
    ax.set_ylim(-3, 2)
    fig.savefig('frame_{:03d}.png'.format(idx))
    return []

ax.set_xlim(-2, 3)
ax.set_ylim(-3, 2)

anim = animation.FuncAnimation(fig, anim, init_func=init, repeat=False,
                               frames=N, interval=20, blit=True)

# anim.save('mymovie.mp4', writer='mencoder')
plt.show()
subprocess.call(['convert', '-delay', '10', '-loop', '0', '*.png', 'out.gif'])