シミュレーションの可視化コードがうまくいかない

前提・実現したいこと

シミュレーションするプログラミングについて勉強しています。
その時に、このサイトにあるサンプルコードを使用させていただきました。
セルオートマトン法を用いた森林火災シミュレーションのアニメーションを導き出したいです。
ご回答宜しくお願い致します。

発生している問題・エラーメッセージ

実行時にエラーメッセージはなく、アニメーション作成・GIF画像の部分が実行されていないようで、
step１の状態しか表示されません。
Google Colaboratoryを使用しています。

該当のソースコード

import numpy as np

np.random.seed(42)  # 乱数のシードを固定する。


# 状態
GROUND = 0  # 地面
TREE = 1  # 木
BURNING = 2  # 燃えている木
BURNED = 3  # 燃え尽きた木

# パラメータ
size = 10, 10  # 地面の大きさ
prob = 0.7  # 地面に木が生えてる確率

# (1) マップを生成する。
field = np.random.choice([GROUND, TREE], size, p=[1 - prob, prob])

# (2) マップの一箇所をランダムに選択し、2とする。
rows = np.arange(size[0])  # 行方向のインデックス一覧
cols = np.arange(size[1])  # 列方向のインデックス一覧
field[np.random.choice(rows), np.random.choice(cols)] = BURNING

history = [field]  # 履歴

while True:
    # (3) 各反復の処理
    new_field = field.copy()

    # (3.1) 現在の状態が BURNING のセルは、BURNED にする。
    new_field[field == BURNING] = BURNED

    # (3.2) 「現在の状態が TREE」かつ「上下左右のいずれかのセルが BURNING」のセルは、
    #       BURNING にする。
    bottom = np.take(field, rows - 1, axis=0, mode="clip")
    top = np.take(field, rows + 1, axis=0, mode="clip")
    left = np.take(field, cols - 1, axis=1, mode="clip")
    right = np.take(field, cols + 1, axis=1, mode="clip")

    indices_to_burn = (field == TREE) & (
        (bottom == BURNING) | (top == BURNING) | (left == BURNING) | (right == BURNING)
    )
    new_field[indices_to_burn] = BURNING

    # マップを更新する。
    field = new_field
    history.append(field)

    # (4) 燃えている木がない場合は終了する。
    if np.all(field != BURNING):
        break

# (5) 最終的に燃えてしまう土地の面積は全体の何割かを求める。
burned_ratio = np.count_nonzero(field == BURNED) / field.size
print(f"{burned_ratio:.2%}")  # 55.00%

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib import colors

# カラーマップを作成する。
cmap = colors.ListedColormap(["gray", "green", "red", "black"])
norm = colors.BoundaryNorm(np.arange(5), cmap.N)

fig = plt.figure()


def plot(i):
    plt.cla()
    plt.imshow(history[i], cmap=cmap, norm=norm)
    plt.title(f"step={i}")


# アニメーションを作成する。
anim = FuncAnimation(fig, plot, frames=len(history), interval=1000, repeat=True)

# gif 画像として保存する。
anim.save("animation.gif", writer="pillow")

試したこと

def plot(i):
plt.cla()
plt.imshow(history[i], cmap=cmap, norm=norm)
plt.title(f"step={i}")
の変数iの部分を数字に変更すると、その時の状態のシミュレーションを反映させるために
変数ⅰを０～１５まで代入して実行するコードを自分なりにやってみましたが、
すべてエラーになりました。