CMAQ Animation

animate_cmaq.py

import time
import glob
import cmaqsatproc as csp
import matplotlib.pyplot as plt
import matplotlib.colors as mc
import matplotlib.animation as animation
import numpy as np

# Get a list of paths
pat = '/work/ROMO/2021_MP/cmaq54mp/2021hb_MP_cb6r5ae7_21k/12US1/output/CCTM_CGRID_2021hb_MP_cb6r5ae7_21k.12US1_35_2021010?.nc'
paths = sorted(glob.glob(pat))

# Options
outpath = 'example.gif'
varkey = 'O3'
layidx = 0
norm = plt.Normalize(vmin=0, vmax=0.08)
cmap = 'viridis'
# Or a complex norm/cmap definition
# edges = np.array([0.03, 0.04, 0.05, 0.06, 0.07, 0.08])
# mids = (edges[1:] + edges[:-1]) / 2
# colors = plt.colormaps[cmap]((mids - mids.min()) / (mids.max() - mids.min()))
# cmap, norm = mc.from_levels_and_colors(edges, colors)
# cmap.set_under('k')
# cmap.set_over('r')

# Open a figure with a colorbar axis
gskw = dict(left=0.1, right=0.85)
fig, ax = plt.subplots(gridspec_kw=gskw)
cax = fig.add_axes([.87, .1, .025, .8])
ax.set(facecolor='gainsboro')
# Add tile plots (all to the same figure)
containers = []
for path in paths:
    f = csp.open_ioapi(path)
    for Z in f[varkey][:, layidx]:
        # xarray plot adds useful axes labels and title
        qm = Z.plot(norm=norm, cmap=cmap, ax=ax, add_colorbar=False)
        t = fig.text(0.5, 0.92, ax.title.get_text(), size=16, horizontalalignment='center')
        containers.append([qm, t])

ax.set(title='')
# Add colorbar to figure
label = f'{Z.long_name.strip()} [{Z.units.strip()}]'
fig.colorbar(qm, cax=cax, label=label, extend='both')
# Add state map
f.csp.cno.drawstates(ax=ax)

# Convert tile plots to an animation
ani = animation.ArtistAnimation(fig=fig, artists=containers, interval=400)
# Save as a GIF with 400ms intervals
ani.save(filename=outpath, writer="pillow")