<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Two-Layer QG Model Example
Here is a quick overview of how to use the two-layer model. See the
Step1:  Initialize and Run the Model
Here we set up a model which will run for 10 years and start averaging
after 5 years. There are lots of parameters that can be specified as
keyword arguments but we are just using the defaults.
Step2:  Convert Model Outpt to an xarray Dataset
Model variables, coordinates, attributes, and metadata can be stored conveniently as an xarray Dataset. (Notice that this feature requires xarray to be installed on your machine. See here for installation instructions
Step3:  Visualize Output
Let's assign a new data variable, q_upper, as the upper layer PV anomaly. We access the PV values in the Dataset as m_ds.q, which has two levels and a corresponding background PV gradient, m_ds.Qy.
Step4:  Plot Diagnostics
The model automatically accumulates averages of certain diagnostics. We can 
find out what diagnostics are available by calling
Step5:  To look at the wavenumber energy spectrum, we plot the KEspec diagnostic.
(Note that summing along the l-axis, as in this example, does not give us
a true isotropic wavenumber spectrum.)
Step6:  We can also plot the spectral fluxes of energy and enstrophy.
<ASSISTANT_TASK:>
Python Code:
import numpy as np
from matplotlib import pyplot as plt
%matplotlib inline
import pyqg
from pyqg import diagnostic_tools as tools
Explanation: Two-Layer QG Model Example
Here is a quick overview of how to use the two-layer model. See the
:py:class:pyqg.QGModel api documentation for further details.
First import numpy, matplotlib, and pyqg:
End of explanation
year = 24*60*60*360.
m = pyqg.QGModel(tmax=10*year, twrite=10000, tavestart=5*year)
m.run()
Explanation: Initialize and Run the Model
Here we set up a model which will run for 10 years and start averaging
after 5 years. There are lots of parameters that can be specified as
keyword arguments but we are just using the defaults.
End of explanation
m_ds = m.to_dataset().isel(time=-1)
m_ds
Explanation: Convert Model Outpt to an xarray Dataset
Model variables, coordinates, attributes, and metadata can be stored conveniently as an xarray Dataset. (Notice that this feature requires xarray to be installed on your machine. See here for installation instructions: http://xarray.pydata.org/en/stable/getting-started-guide/installing.html#instructions)
End of explanation
m_ds['q_upper'] = m_ds.q.isel(lev=0) + m_ds.Qy.isel(lev=0)*m_ds.y
m_ds['q_upper'].attrs = {'long_name': 'upper layer PV anomaly'}
m_ds.q_upper.plot.contourf(levels=18, cmap='RdBu_r');
Explanation: Visualize Output
Let's assign a new data variable, q_upper, as the upper layer PV anomaly. We access the PV values in the Dataset as m_ds.q, which has two levels and a corresponding background PV gradient, m_ds.Qy.
End of explanation
m.describe_diagnostics()
Explanation: Plot Diagnostics
The model automatically accumulates averages of certain diagnostics. We can 
find out what diagnostics are available by calling
End of explanation
kr, kespec_upper = tools.calc_ispec(m, m_ds.KEspec.isel(lev=0).data)
_, kespec_lower = tools.calc_ispec(m, m_ds.KEspec.isel(lev=1).data)
plt.loglog(kr, kespec_upper, 'b.-', label='upper layer')
plt.loglog(kr, kespec_lower, 'g.-', label='lower layer')
plt.legend(loc='lower left')
plt.ylim([1e-14,1e-8])
plt.xlabel(r'k (m$^{-1}$)'); plt.grid()
plt.title('Kinetic Energy Spectrum');
Explanation: To look at the wavenumber energy spectrum, we plot the KEspec diagnostic.
(Note that summing along the l-axis, as in this example, does not give us
a true isotropic wavenumber spectrum.)
End of explanation
kr, APEgenspec = tools.calc_ispec(m, m_ds.APEgenspec.data)
_, APEflux     = tools.calc_ispec(m, m_ds.APEflux.data)
_, KEflux      = tools.calc_ispec(m, m_ds.KEflux.data)
_, KEfrictionspec = tools.calc_ispec(m, m_ds.KEfrictionspec.data)
_, Dissspec    = tools.calc_ispec(m, m_ds.Dissspec.data)
ebud = [ APEgenspec,
         APEflux,
         KEflux,
         KEfrictionspec,
         Dissspec]
ebud.append(-np.vstack(ebud).sum(axis=0))
ebud_labels = ['APE gen','APE flux','KE flux','Bottom drag','Diss.','Resid.']
[plt.semilogx(kr, term) for term in ebud]
plt.legend(ebud_labels, loc='upper right')
plt.xlabel(r'k (m$^{-1}$)'); plt.grid()
plt.title('Spectral Energy Transfer');
_, ENSflux    = tools.calc_ispec(m, m_ds.ENSflux.data.squeeze())
_, ENSgenspec = tools.calc_ispec(m, m_ds.ENSgenspec.data.squeeze())
_, ENSfrictionspec = tools.calc_ispec(m, m_ds.ENSfrictionspec.data.squeeze())
_, ENSDissspec = tools.calc_ispec(m, m_ds.ENSDissspec.data.squeeze())
ebud = [ ENSgenspec,
         ENSflux,
         ENSDissspec,
         ENSfrictionspec]
ebud.append(-np.vstack(ebud).sum(axis=0))
ebud_labels = ['ENS gen','ENS flux div.','Dissipation','Friction','Resid.']
[plt.semilogx(kr, term) for term in ebud]
plt.legend(ebud_labels, loc='best')
plt.xlabel(r'k (m$^{-1}$)'); plt.grid()
plt.title('Spectral Enstrophy Transfer');
Explanation: We can also plot the spectral fluxes of energy and enstrophy.
End of explanation
<END_TASK>