Matplotlib Animation, Either Runs Slow and Replots Colorbar, or Missing Plot

I am attempting to animate data that I have previously saved by loading the numpy array for each frame.

My code is as follows:

fig, ax = plt.subplots(1, 1, constrained_layout = False)

def routine(omega):
    omega_plot = ax.pcolormesh(x_grid, y_grid, omega, cmap = 'coolwarm')
    
    ax.set_aspect('equal')
    ax.set_xlim(0, 2 * np.pi)
    ax.set_ylim(0, 2 * np.pi) 

    divider = make_axes_locatable(omega_plot.axes)
    cax = divider.append_axes("right", size = "5%", pad = 0.2)
    fig.colorbar(omega_plot, cax = cax, orientation = 'vertical')
    cax.set_ylabel(r"$\mathrm{mag}\left[ \vec{u} \right]$")

    ax.set_title(r"$\omega(t)")
    # some plot routine parts excluded

    return ax

def init():
    omega = np.load("pcolor Animation/temp/omega_ana, n = 0.npy")
    return routine(omega)

def update(n):
    print(n)
    omega = np.load("pcolor Animation/temp/omega_ana, n = {0}.npy".format(n))
    return routine(omega)

ANI = FuncAnimation(fig, update, frames = range(0, 2000, 10), init_func = init)
ANI.save("pcolor Animation/Taylor-Greene Vortex Animation.mp4", fps = 180, dpi = 200)

When I run it, I get the following: Now, this is actually correct. Taylor-Greene flow exponentially decays but otherwise does not change, but it is also obvious that the colorbar is being replotted each time, leading to a 'drift'. Additionally, the code slows down with each frame, because it is plotting on top of the previous ones.

If we add plt.clf() or fig.clear() or any of the other suggestions that various Stack Overflow answers suggest about matplotlib.animations to the top of routine(omega), we get the following: This time, our colorbar is correct, and the animation runs much more quickly, but the data itself is missing.

What am I missing? I've done my due diligence and have tried to fix this problem with the suggestions on various other similar questions, and have been having this problem for a few days. I'd appreciate any help, and thank you in advance.

To make the animation efficiently, we want to update as little as possible. So create the artists up front, and then just update the array used in the QuadMesh artist (which is returned from pcolormesh). Also updating the norm that the QuadMesh uses will make the colorbar's range change.

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.colors import Normalize
from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable
import numpy as np


fig, ax = plt.subplots()
omega_plot = ax.pcolormesh(np.zeros((10, 10)), cmap='Blues', vmin=0, vmax=100)
ax.set_aspect('equal')

divider = make_axes_locatable(omega_plot.axes)
cax = divider.append_axes("right", size = "5%", pad = 0.2)
fig.colorbar(omega_plot, cax = cax, orientation = 'vertical')
cax.set_ylabel(r"$\mathrm{mag}\left[ \vec{u} \right]$")

ax.set_title(r"$\omega(t)$")


def update(n):
    omega = np.arange(100).reshape(10, 10) * 0.99**n
    omega_plot.set_array(omega)
    
    # Change the valid range of the data (remove this line for the colorbar to be constant).
    omega_plot.set_norm(Normalize(vmin=omega.min(), vmax=omega.max()))


ANI = FuncAnimation(fig, update, frames = range(0, 200, 10))
ANI.save("test.gif", fps = 10)

The above code produces

If we remove the last line of the update function, we instead get