如何在 matplotlib 中绘制动画矩阵

Question

我需要在视觉上一步一步做一些数值计算算法，如下图：(gif)

Font

如何使用 matplotlib 制作此动画？有什么方法可以直观地呈现这些过渡吗？作为矩阵的转换，求和，转置，使用循环并呈现转换等。 我的目标不是使用图形，而是使用相同的矩阵表示。这是为了便于理解算法。

Answer 1

由于可以使用imshow 轻松绘制矩阵，因此可以使用imshow 绘图创建此类表格，并根据当前动画步骤调整数据。

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

#####################
# Array preparation
#####################

#input array
a = np.random.randint(50,150, size=(5,5))
# kernel
kernel = np.array([[ 0,-1, 0], [-1, 5,-1], [ 0,-1, 0]])

# visualization array (2 bigger in each direction)
va = np.zeros((a.shape[0]+2, a.shape[1]+2), dtype=int)
va[1:-1,1:-1] = a

#output array
res = np.zeros_like(a)

#colorarray
va_color = np.zeros((a.shape[0]+2, a.shape[1]+2)) 
va_color[1:-1,1:-1] = 0.5

#####################
# Create inital plot
#####################
fig = plt.figure(figsize=(8,4))

def add_axes_inches(fig, rect):
    w,h = fig.get_size_inches()
    return fig.add_axes([rect[0]/w, rect[1]/h, rect[2]/w, rect[3]/h])

axwidth = 3.
cellsize = axwidth/va.shape[1]
axheight = cellsize*va.shape[0]

ax_va  = add_axes_inches(fig, [cellsize, cellsize, axwidth, axheight])
ax_kernel  = add_axes_inches(fig, [cellsize*2+axwidth,
                                   (2+res.shape[0])*cellsize-kernel.shape[0]*cellsize,
                                   kernel.shape[1]*cellsize,  
                                   kernel.shape[0]*cellsize])
ax_res = add_axes_inches(fig, [cellsize*3+axwidth+kernel.shape[1]*cellsize,
                               2*cellsize, 
                               res.shape[1]*cellsize,  
                               res.shape[0]*cellsize])
ax_kernel.set_title("Kernel", size=12)

im_va = ax_va.imshow(va_color, vmin=0., vmax=1.3, cmap="Blues")
for i in range(va.shape[0]):
    for j in range(va.shape[1]):
        ax_va.text(j,i, va[i,j], va="center", ha="center")

ax_kernel.imshow(np.zeros_like(kernel), vmin=-1, vmax=1, cmap="Pastel1")
for i in range(kernel.shape[0]):
    for j in range(kernel.shape[1]):
        ax_kernel.text(j,i, kernel[i,j], va="center", ha="center")


im_res = ax_res.imshow(res, vmin=0, vmax=1.3, cmap="Greens")
res_texts = []
for i in range(res.shape[0]):
    row = []
    for j in range(res.shape[1]):
        row.append(ax_res.text(j,i, "", va="center", ha="center"))
    res_texts.append(row)    


for ax  in [ax_va, ax_kernel, ax_res]:
    ax.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
    ax.yaxis.set_major_locator(mticker.IndexLocator(1,0))
    ax.xaxis.set_major_locator(mticker.IndexLocator(1,0))
    ax.grid(color="k")

###############
# Animation
###############
def init():
    for row in res_texts:
        for text in row:
            text.set_text("")

def animate(ij):
    i,j=ij
    o = kernel.shape[1]//2
    # calculate result
    res_ij = (kernel*va[1+i-o:1+i+o+1, 1+j-o:1+j+o+1]).sum()
    res_texts[i][j].set_text(res_ij)
    # make colors
    c = va_color.copy()
    c[1+i-o:1+i+o+1, 1+j-o:1+j+o+1] = 1.
    im_va.set_array(c)

    r = res.copy()
    r[i,j] = 1
    im_res.set_array(r)

i,j = np.indices(res.shape)
ani = matplotlib.animation.FuncAnimation(fig, animate, init_func=init, 
                                         frames=zip(i.flat, j.flat), interval=400)
ani.save("algo.gif", writer="imagemagick")
plt.show()

Answer 2

此示例在 Jupyter 笔记本中设置内联动画。我想可能还有一种方法可以导出为 gif，但到目前为止我还没有研究过。

不管怎样，首先要做的就是摆好桌子。我从Export a Pandas dataframe as a table image 大量借用render_mpl_table 代码。

这个问题的（改编的）版本是：

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML
import six

width = 8
data = pd.DataFrame([[0]*width,
                     [0, *np.random.randint(95,105,size=width-2), 0],
                     [0, *np.random.randint(95,105,size=width-2), 0],
                     [0, *np.random.randint(95,105,size=width-2), 0]])

def render_mpl_table(data, col_width=3.0, row_height=0.625, font_size=14,
                     row_color="w", edge_color="black", bbox=[0, 0, 1, 1],
                     ax=None, col_labels=data.columns, 
                     highlight_color="mediumpurple",
                     highlights=[], **kwargs):
    if ax is None:
        size = (np.array(data.shape[::-1]) + np.array([0, 1])) *
                np.array([col_width, row_height])
        fig, ax = plt.subplots(figsize=size)
        ax.axis('off')

    mpl_table = ax.table(cellText=data.values, bbox=bbox, colLabels=col_labels,
                         **kwargs)

    mpl_table.auto_set_font_size(False)
    mpl_table.set_fontsize(font_size)

    for k, cell in six.iteritems(mpl_table._cells):
        cell.set_edgecolor(edge_color)
        if k in highlights:
            cell.set_facecolor(highlight_color)
        elif data.iat[k] > 0:
            cell.set_facecolor("lightblue")
        else:
            cell.set_facecolor(row_color)
    return fig, ax, mpl_table

fig, ax, mpl_table = render_mpl_table(data, col_width=2.0, col_labels=None,
                                  highlights=[(0,2),(0,3),(1,2),(1,3)])

在这种情况下，以不同颜色突出显示的单元格由指定行和列的元组数组给出。

对于动画，我们需要设置一个绘制不同高光的表格的函数：

def update_table(i, *args, **kwargs):
    r = i//(width-1)
    c = i%(width-1)
    highlights=[(r,c),(r,c+1),(r+1,c),(r+1,c+1)]
    for k, cell in six.iteritems(mpl_table._cells):
        cell.set_edgecolor("black")
        if k in highlights:
            cell.set_facecolor("mediumpurple")
        elif data.iat[k] > 0:
            cell.set_facecolor("lightblue")
        else:
            cell.set_facecolor("white")
    return (mpl_table,)

这会强制更新表格中所有单元格的颜色。 highlights 数组是根据当前帧计算的。在此示例中，表格的宽度和高度是硬编码的，但根据输入数据的形状进行更改应该不会太难。

我们基于现有的 fig 和更新函数创建动画：

a = animation.FuncAnimation(fig, update_table, (width-1)*3,
                               interval=750, blit=True)

最后我们在笔记本中显示它：

HTML(a.to_jshtml())

我把这些放在了github上的一个notebook里，见https://github.com/gurudave/so_examples/blob/master/mpl_animation.ipynb

希望这足以让您朝着正确的方向前进！