如何在matplotlib中绘制动画矩阵

Question

I need to do step by step some numerical calculation algorithms visually, as in the figure below: (gif)

Font

How can I do this animation with matplotlib? Is there any way to visually present these transitions? As transformation of matrices, sum, transposition, using a loop and it presenting the transitions etc. My goal is not to use graphics but the same matrix representation. This is to facilitate the understanding of the algorithms.

Answer 1

Since matrices can be plotted easily with imshow , one could create such table with an imshow plot and adjust the data according to the current animation step.

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

This example sets up the animation inline in a Jupyter notebook. I suppose there's probably also a way to export as a gif, but I haven't looked into that so far.

Anyway, first thing to do is set up the table. I borrowed heavily from Export a Pandas dataframe as a table image for the render_mpl_table code.

The (adapted) version for this problem is:

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

In this case, the cells to highlight in a different color are given by an array of tuples that specify the row and column.

For the animation, we need to set up a function that draws the table with different highlights:

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

This forcibly updates the colors for all cells in the table. The highlights array is computed based on the current frame. The width and height of the table are kind of hard-coded in this example, but that shouldn't be super hard to change based on the shape of your input data.

We create an animation based on the existing fig and update function:

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

And lastly we show it inline in our notebook:

HTML(a.to_jshtml())

I put this together in a notebook on github, see https://github.com/gurudave/so_examples/blob/master/mpl_animation.ipynb

Hope that's enough to get you going in the right direction!