That's Geophysical Fluid Dynamics.

Regardless of what we’re doing, we’ll need to do some importing: NumPy (for numerical work), Matplotlib (for plotting), and pyqg (for the obvious):

import numpy as np
from matplotlib import pyplot as plt
import pyqg

We’ll start with two layer QG!

Two Layer QG

Let’s set up a model that runs for 10 years and starts averaging after 5 years. It will record the solution in increments of 10000 seconds:

year = 24*60*60*360.
m = pyqg.QGModel(tmax=10*year, twrite=10000, tavestart=5*year)
m.run()

And it’s as simple as that! The object m stores prognostic variables as attributes. So to access the PV, we would simply use m.q. The background PV gradient is stored as m.Qy.

The following code plots the PV distribution in the upper layer:

q_upper = m.q[0] + m.Qy[0]*m.y
plt.contourf(m.x, m.y, q_upper, 12, cmap='RdBu_r')
plt.xlabel('x'); plt.ylabel('y'); plt.title('Upper Layer PV')
plt.colorbar();

To get a diagnostic variable, one can use the get_diagnostic function and pass the name of the diagnostic as an argument. The following plots the kinetic energy spectrum in the upper and lower layers:

kespec_u = m.get_diagnostic('KEspec')[0].sum(axis=0)
kespec_l = m.get_diagnostic('KEspec')[1].sum(axis=0)
plt.loglog( m.kk, kespec_u, '.-' )
plt.loglog( m.kk, kespec_l, '.-' )
plt.legend(['upper layer','lower layer'], loc='lower left')
plt.ylim([1e-9,1e-3]); plt.xlim([m.kk.min(), m.kk.max()])
plt.xlabel(r'k (m$^{-1}$)'); plt.grid()
plt.title('Kinetic Energy Spectrum');

One can also plot the spectral energy fluxes:

ape_gen_spec = m.get_diagnostic('APEgenspec')
ape_flux = m.get_diagnostic('APEflux')
ke_flux = m.get_diagnostic('KEflux')
dissip = -m.rek * m.del2 * m.get_diagnostic('KEspec')[1] * m.M**2

ebud = [ape_gen_spec.sum(axis=0), ape_flux.sum(axis=0), ke_flux.sum(axis=0), dissip.sum(axis=0)]
resid = -np.vstack(ebud).sum(axis=0)
ebud.append(resid)

ebud_labels = ['APE gen','APE flux','KE flux','Diss.','Resid.']
[plt.semilogx(m.kk, term) for term in ebud]
plt.legend(ebud_labels, loc='upper right')
plt.xlim([m.kk.min(), m.kk.max()])
plt.xlabel(r'k (m$^{-1}$)'); plt.grid()
plt.title('Spectral Energy Transfers');

Next we’ll look at the baroclinic instability of a 3-layer flow.