source: dynamico_lmdz/aquaplanet/CMIP5a/RUN_AQUAPLANET_DYNAMICO_LMDZ5/figs_fluxes.py @ 4100

import numpy as np
import netCDF4 as cdf
# select non-interactive backend, cf http://stackoverflow.com/questions/4931376/generating-matplotlib-graphs-without-a-running-x-server
import matplotlib
matplotlib.use('Agg') 
import matplotlib.pyplot as plt

def getdims(nc, *names): return [len(nc.dimensions[name]) for name in names]
def getvars(nc, *names): return [nc.variables[name] for name in names]

#--------------------------- MAIN ----------------------------

infile = 'diagflux.nc'
nc = cdf.Dataset(infile, "r")
print nc
llm, nlat = getdims(nc, 'lev','lat_dom_out')
Lv, Cpd, kappa, g, perim = 2.5e6, 1004.5, 0.2857143, 9.80616, 4e7

(lat,p,qflux,q,epotflux,epot,
 ekinflux,ekin,cpTflux,cpT,mass,mflux) = getvars(nc,
                                                'lat_dom_out','p','qflux_lat','qmass_t','epotflux_lat','epot_t','ekinflux_lat','ekin_t',
                                                'enthalpyflux_lat','enthalpy_t','mass_t','massflux_lat')

lat,p,epotflux,epot,ekinflux,ekin,cpTflux,cpT,mass,mflux = [ x[:] for x in lat,p,epotflux,epot,ekinflux,ekin,cpTflux,cpT,mass,mflux ] # read data from file
qflux,q = qflux[0,:,:], q[0,:,:]

perim = 4e7*np.cos(lat*np.pi/180.) # geometric factor
stream = mflux.cumsum(axis=0)
for i in range(lat.size):
    stream[:,i]=stream[:,i]*perim[i]
p=p.mean(axis=1)

print qflux.shape, q.shape, stream.shape
# cpT est en (J/m^2)
# cpTflux est en (J/m^2)*(m/s)

cpTflux = cpTflux - (cpT/mass)*mflux # flux "turbulents"
qflux = qflux - (q/mass)*mflux # flux "turbulents"

qflux,cpTflux,mass = [ x.sum(axis=0) for x in qflux,cpTflux,mass ] # somme sur la verticale
# cpTflux est en (J/m^2)*(m/s)=W/m

print cpTflux.shape, p.shape

cpTflux = cpTflux*perim
qflux   = qflux*perim
# cpTflux est en W

plt.figure()
plt.plot(lat,cpTflux)
plt.title('Meridional enthalpy flux (W)')
plt.savefig('enthalpy_flux.png')
plt.close()

plt.figure()
plt.plot(lat,Lv*qflux)
plt.title('Meridional latent heat flux (W)')
plt.savefig('latent_heat_flux.png')
plt.close()

plt.figure()
plt.contourf(lat,p,stream,20)
plt.ylim(1e5, 0)
plt.title('Meridional stream function (kg/s)')
plt.colorbar()
plt.savefig('stream.png')
plt.close()