def animate_shallow_water(sol):
    """Create a matplotlib animation of the result."""
    import matplotlib.pyplot as plt
    from matplotlib import animation

    quiver_stride = (
        slice(1, -1, ny_global // MAX_QUIVERS[0]),
        slice(1, -1, nx_global // MAX_QUIVERS[1]),
    )

    # set up figure
    fig = plt.figure(figsize=(6, 4))
    ax = plt.gca()

    # surface plot of height anomaly
    cs = ax.pcolormesh(
        0.5 * (x_global[:-1] + x_global[1:]) / 1e3,
        0.5 * (y_global[:-1] + y_global[1:]) / 1e3,
        sol[0].h[1:-1, 1:-1] - DEPTH,
        vmin=-PLOT_ETA_RANGE,
        vmax=PLOT_ETA_RANGE,
        cmap="RdBu_r",
    )

    # quiver plot of velocity
    cq = ax.quiver(
        xx_global[quiver_stride] / 1e3,
        yy_global[quiver_stride] / 1e3,
        sol[0].u[quiver_stride],
        sol[0].v[quiver_stride],
        clip_on=True,
    )

    # time indicator
    t = ax.text(
        s="",
        x=0.05,
        y=0.95,
        ha="left",
        va="top",
        backgroundcolor=(1, 1, 1, 0.8),
        transform=ax.transAxes,
    )

    ax.set(
        aspect="equal",
        xlim=(x_global[1] / 1e3, x_global[-2] / 1e3),
        ylim=(y_global[1] / 1e3, y_global[-2] / 1e3),
        xlabel="$x$ (km)",
        ylabel="$y$ (km)",
    )

    plt.colorbar(
        cs,
        orientation="horizontal",
        label="Surface height anomaly (m)",
        pad=0.2,
        shrink=0.8,
    )
    fig.tight_layout()

    def animate(i):
        state = sol[i]

        eta = state.h - DEPTH
        cs.set_array(eta[1:-1, 1:-1].flatten())
        cq.set_UVC(state.u[quiver_stride], state.v[quiver_stride])

        current_time = PLOT_EVERY * dt * i
        t.set_text(f"t = {current_time / DAY_IN_SECONDS:.2f} days")
        return (cs, cq, t)

    anim = animation.FuncAnimation(
        fig, animate, frames=len(sol), interval=50, blit=True, repeat_delay=3_000
    )
    return anim
Create a matplotlib animation of the result.
examples/shallow_water.py