Как построить анимированную матрицу в matplotlib

Question

Мне нужно визуально сделать шаг за шагом некоторые алгоритмы численного расчета, как показано на рисунке ниже: (gif)

шрифт

Как я могу сделать эту анимацию с помощью 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, но я пока не рассматривал это.

В любом случае, первое, что нужно сделать, это накрыть стол. Я сильно позаимствовал у Экспорт кадра данных Pandas в виде изображения таблицы для кода 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 вычисляется на основе текущего кадра. Ширина и высота таблицы в этом примере жестко запрограммированы, но это не должно быть слишком сложно изменить в зависимости от формы ваших входных данных.

Мы создаем анимацию на основе существующего фига и функции обновления:

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

И, наконец, мы показываем это в нашей записной книжке:

HTML(a.to_jshtml())

Я собрал это в блокноте на github, см. https://github.com/gurudave/so_examples/blob/master/mpl_animation.ipynb

Надеюсь, этого достаточно, чтобы вы пошли в правильном направлении!