def energy_eq(run, month, filename='plev_daily', lev=150.):
#Load in dataset
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc'
name = name_temp % month
#read data into xarray
data = xr.open_dataset(name, decode_times=False)
# Calculate planetary vorticity
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
uchi = data.ucomp - data.upsi
vchi = data.vcomp - data.vpsi
pe_to_ke = -1. * (uchi * gr.ddx(data.height * 9.8) +
vchi * gr.ddy(data.height * 9.8, vector=False))
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
conv_1 = vor * (data.upsi * vchi - uchi * data.vpsi)
conv_2 = -1. * omega * (data.upsi * gr.ddp(uchi) +
data.vpsi * gr.ddp(vchi))
conv_3 = -1. * (gr.ddx(uchi) + gr.ddy(vchi)) * (data.upsi**2. +
data.vpsi**2.) / 2.
ds = xr.Dataset(
{
'pe_to_ke': (['time', 'pfull', 'lat', 'lon'], pe_to_ke),
'conv_1': (['time', 'pfull', 'lat', 'lon'], conv_1),
'conv_2': (['time', 'pfull', 'lat', 'lon'], conv_2),
'conv_3': (['time', 'pfull', 'lat', 'lon'], conv_3),
'upsi': (['time', 'pfull', 'lat', 'lon'], data.upsi),
'vpsi': (['time', 'pfull', 'lat', 'lon'], data.vpsi),
'uchi': (['time', 'pfull', 'lat', 'lon'], uchi),
'vchi': (['time', 'pfull', 'lat', 'lon'], vchi)
},
coords={
'time': ('time', data.time),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat),
'lon': ('lon', data.lon)
})
ds.coords['xofyear'] = np.mod(ds.time - 1., 360.) // 5 + 1.
dsout = ds.groupby('xofyear').mean(('time'))
fileout = '/scratch/rg419/Data_moist/' + run + '/run%03d/energy_eq.nc'
fileout = fileout % month
dsout.to_netcdf(path=fileout)
print 'data written to ', fileout
return dsout
def partition_advection(data, lons, lev=150):
#First do uu terms
uu_trans_dx = -86400. * gr.ddx(
(data.ucomp_sq -
data.ucomp**2).sel(pfull=lev)) # <u'u'> = <uu> - <u><u>
u = data.ucomp.sel(pfull=lev) # u
u_dx = -86400. * gr.ddx(u) # dudx
u_dudx_zav = u.sel(lon=lons).mean('lon') * u_dx.sel(lon=lons).mean(
'lon') # [u][dudx]
u_dudx_stat = (u * u_dx).sel(
lon=lons).mean('lon') - u_dudx_zav # u*dudx* = [ududx] - [u][dudx]
data['uu_trans_dx'] = (('xofyear', 'lat'),
uu_trans_dx.sel(lon=lons).mean('lon'))
data['u_dudx_stat'] = (('xofyear', 'lat'), u_dudx_stat)
data['u_dudx_zav'] = (('xofyear', 'lat'), u_dudx_zav)
print 'uu terms done'
#Next do uv terms
uv_trans_dy = -86400. * gr.ddy(
(data.ucomp_vcomp - data.ucomp * data.vcomp).sel(pfull=lev), uv=True)
v = data.vcomp.sel(pfull=lev).load() # v
u_dy = -86400. * gr.ddy(u) # dudy
v_dudy_zav = v.sel(lon=lons).mean('lon') * u_dy.sel(lon=lons).mean(
'lon') # [v][dudy]
v_dudy_stat = (v * u_dy).sel(
lon=lons).mean('lon') - v_dudy_zav # v*dudy* = [vdudy] - [v][dudy]
data['uv_trans_dy'] = (('xofyear', 'lat'),
uv_trans_dy.sel(lon=lons).mean('lon'))
data['v_dudy_stat'] = (('xofyear', 'lat'), v_dudy_stat)
data['v_dudy_zav'] = (('xofyear', 'lat'), v_dudy_zav)
print 'uv terms done'
#Finally do uw terms
uw_trans_dp = -86400. * gr.ddp(
(data.ucomp_omega - data.ucomp * data.omega).sel(lon=lons).mean('lon'))
w = data.omega.sel(pfull=lev).load() # w
u_dp = -86400. * (gr.ddp(data.ucomp)).sel(pfull=lev) # dudp
w_dudp_zav = w.sel(lon=lons).mean('lon') * u_dp.sel(lon=lons).mean('lon')
w_dudp_stat = (w * u_dp).sel(lon=lons).mean('lon') - w_dudp_zav
data['uw_trans_dp'] = (('xofyear', 'lat'), uw_trans_dp.sel(pfull=lev))
data['w_dudp_stat'] = (('xofyear', 'lat'), w_dudp_stat)
data['w_dudp_zav'] = (('xofyear', 'lat'), w_dudp_zav)
print 'uw terms done'
def trop_height(temp):
# Use temperature to calculate tropopause level
# Take vertical gradient wrt height and find pressure level at which dT/dz goes and stays above -2 K/km
levs = temp.pfull[temp.pfull <= 500.]
pfull_big = np.arange(1000., 0., -10.)
f = spint.interp1d(temp.pfull,
temp.mean('lon'),
axis=1,
fill_value='extrapolate')
temp_big = f(pfull_big)
temp_big = xr.DataArray(temp_big,
coords=[temp.xofyear, pfull_big, temp.lat],
dims=['xofyear', 'pfull', 'lat'])
rho_big = temp_big.pfull * 100. / mc.rdgas / temp_big
dtdp_big = gr.ddp(temp_big)
levs_big = dtdp_big.pfull[dtdp_big.pfull <= 500.]
dtdz_big = (-1. * dtdp_big * rho_big * mc.grav * 1000.).sel(pfull=levs_big)
rho = temp.pfull * 100. / mc.rdgas / temp.mean('lon')
dtdp = gr.ddp(temp.mean('lon'))
dtdz = (-1. * dtdp * rho * mc.grav * 1000.) #.sel(pfull=levs)
f = spint.interp1d(dtdz.pfull, dtdz, axis=1, fill_value='extrapolate')
dtdz_big = f(levs_big)
# dtdz_mask = xr.DataArray( np.where(dtdz >= -2., 1., 0.), coords=[dtdz.xofyear, dtdz.pfull, dtdz.lat], dims=['xofyear', 'pfull', 'lat'])
dtdz_mask = np.where(dtdz >= -2., 1., 0.)
#print dtdz_mask[0,:-1,0]
#print dtdz_mask[0,1:,0]
#print (dtdz_mask[0,:-1,0] != dtdz_mask[0,1:,0])
#print np.where(dtdz_mask[0,:-1,0] != dtdz_mask[0,1:,0])
trop = np.zeros([len(dtdz.xofyear), len(dtdz.lat)])
for i in range(len(dtdz.xofyear)):
for j in range(len(dtdz.lat)):
#trop[i,j] = min(dtdz.pfull[np.where(dtdz[i,:,j] <= -2.)])
trop[i, j] = min(levs_big[np.where(dtdz_big[i, :, j] <= -2.)])
# trop[i,j] = dtdz.pfull[ np.where(dtdz_mask[i,:-1,j] != dtdz_mask[i,1:,j])[0] ]
plt.contourf(trop, levels=np.arange(90., 260., 10.))
plt.colorbar()
print trop[20, :]
print dtdz[0, :, 0]
print np.where(dtdz[0, :, 0] >= -2.)
plt.figure(2)
plt.contourf(dtdz_big[18, :, :], levels=np.arange(-12., 13., 1.))
plt.show()
def vort_eq(run, month, filename='plev_daily', period_fac=1.):
#Load in dataset
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/'+filename+'.nc'
name = name_temp % month
#read data into xarray
data = xr.open_dataset( name, decode_times=False)
# Calculate planetary vorticity
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat *np.pi/180)
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
#vor_rel = v_dx - u_dy
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
dvordp = gr.ddp(vor)
# Calculate horizontal material derivative
horiz_md = -1. * (data.ucomp * dvordx + data.vcomp * dvordy)
# Calculate vertical material derivative
vert_md = -1. * data.omega * dvordp
# Now do the terms involving gradients of windspeed
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching = -1. * vor * div
tilting = gr.ddp(data.ucomp) * gr.ddy(data.omega, vector=False) - gr.ddp(data.vcomp) * gr.ddx(data.omega)
total = horiz_md + vert_md + stretching + tilting
ds = xr.Dataset({'horiz_md': (['time', 'pfull', 'lat', 'lon'], horiz_md),
'vert_md': (['time', 'pfull', 'lat', 'lon'], vert_md),
'stretching': (['time', 'pfull', 'lat', 'lon'], stretching),
'tilting': (['time', 'pfull', 'lat', 'lon'], tilting),
'ucomp': (['time', 'pfull', 'lat', 'lon'], data.ucomp),
'vcomp': (['time', 'pfull', 'lat', 'lon'], data.vcomp),
'total': (['time', 'pfull', 'lat', 'lon'], total)},
coords={'time': ('time', data.time),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat),
'lon': ('lon', data.lon)})
dsout = ds#.mean(('time'))
fileout = '/scratch/rg419/Data_moist/' + run + '/run%03d/vort_eq.nc'
fileout = fileout % month
dsout.to_netcdf(path=fileout)
print 'data written to ', fileout
return dsout
def trop_height(temp, check=False):
# Use temperature to calculate tropopause level
# Take vertical gradient wrt height and find pressure level at which dT/dz goes and stays above -2 K/km
pfull_big = np.arange(500., 0., -10.)
rho = temp.pfull * 100. / mc.rdgas / temp.mean('lon')
dtdp = gr.ddp(temp.mean('lon'))
dtdz = (-1. * dtdp * rho * mc.grav * 1000.) #.sel(pfull=levs)
f = spint.interp1d(dtdz.pfull, dtdz, axis=1, fill_value='extrapolate')
dtdz_big = xr.DataArray(f(pfull_big),
coords=[temp.xofyear, pfull_big, temp.lat],
dims=['xofyear', 'pfull', 'lat'])
trop = np.zeros([len(dtdz.xofyear), len(dtdz.lat)])
for i in range(len(dtdz.xofyear)):
for j in range(len(dtdz.lat)):
trop[i, j] = min(pfull_big[np.where(dtdz_big[i, :, j] <= -2.)])
trop = xr.DataArray(trop,
coords=[temp.xofyear, temp.lat],
dims=['xofyear', 'lat'])
if check:
trop.plot.contourf(levels=np.arange(90., 270., 10.))
plt.figure(2)
dtdz_big.sel(xofyear=68).plot.contourf(levels=np.arange(-12., 13., 1.))
plt.show()
return trop
def partition_advection(data, lons, lev=150):
#Do vv terms
vv_trans_dy = -86400. * gr.ddy(
(data.vcomp_sq - data.vcomp * data.vcomp).sel(pfull=lev), uv=True)
v = data.vcomp.sel(pfull=lev).load() # v
v_dy = -86400. * gr.ddy(v) # dudy
v_dvdy_zav = v.sel(lon=lons).mean('lon') * v_dy.sel(lon=lons).mean(
'lon') # [v][dudy]
v_dvdy_stat = (v * v_dy).sel(
lon=lons).mean('lon') - v_dvdy_zav # v*dudy* = [vdudy] - [v][dudy]
data['vv_trans_dy'] = (('xofyear', 'lat'),
vv_trans_dy.sel(lon=lons).mean('lon'))
data['v_dvdy_stat'] = (('xofyear', 'lat'), v_dvdy_stat)
data['v_dvdy_zav'] = (('xofyear', 'lat'), v_dvdy_zav)
print 'vv terms done'
#Do vw terms
vw_trans_dp = -86400. * gr.ddp(
(data.vcomp_omega - data.vcomp * data.omega).sel(lon=lons).mean('lon'))
w = data.omega.sel(pfull=lev).load() # w
v_dp = -86400. * (gr.ddp(data.vcomp)).sel(pfull=lev) # dudp
w_dvdp_zav = w.sel(lon=lons).mean('lon') * v_dp.sel(lon=lons).mean('lon')
w_dvdp_stat = (w * v_dp).sel(lon=lons).mean('lon') - w_dvdp_zav
data['vw_trans_dp'] = (('xofyear', 'lat'), vw_trans_dp.sel(pfull=lev))
data['w_dvdp_stat'] = (('xofyear', 'lat'), w_dvdp_stat)
data['w_dvdp_zav'] = (('xofyear', 'lat'), w_dvdp_zav)
print 'vw terms done'
def mape(run):
# NB results are larger than in O'Gorman and Schneider papers, but are similar to those in Li et al 2007 - Lorenz Energy Cycle of the Global Atmosphere (GRL), and Lambert 1984 - A global available potential energy-kinetic energy budget (Atmos-Ocean), Saltzman 1970 - Large-Scale Atmospheric Energetics in the Wave Number Domain (rev Geophys), Peixoto and Oort 1974 - The annual distribution of atmospheric energy on a planetary scale.
# Get smaller values if confine average over a smaller storm track region and look at annual mean - might be what's going on in OG + S, but overall I think my code is fine. May get smaller values in my perpetual equinox runs for example
p00 = 1000.
dp = 50.
prefac = mc.cp_air * p00 * 100. / 2. / mc.grav
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
theta = data.temp * (p00 / data.pfull)**mc.kappa
lowest_trop = trop_height(data.temp).max()
#print 'Lowest tropopause level is ' + str(lowest_trop.values)
theta_zav = theta.mean('lon')
thetasq_zav = theta_zav**2.
dthdp = gr.ddp(theta_zav)
strack_n = theta.lat[theta.lat >= 20.]
strack_s = theta.lat[theta.lat <= -20.]
levs = data.pfull[(data.pfull >= lowest_trop) & (data.pfull <= 900.)]
theta_n = latmean(theta_zav, strack_n)
theta_s = latmean(theta_zav, strack_s)
thetasq_n = latmean(thetasq_zav, strack_n)
thetasq_s = latmean(thetasq_zav, strack_s)
dthdp_n = latmean(dthdp, strack_n)
dthdp_s = latmean(dthdp, strack_s)
gamma_n = -1. * mc.kappa / (data.pfull * 100.) / dthdp_n
gamma_s = -1. * mc.kappa / (data.pfull * 100.) / dthdp_s
theta_var_n = thetasq_n - theta_n**2.
theta_var_s = thetasq_s - theta_s**2.
mape_integrand_n = (data.pfull / p00)**mc.kappa * gamma_n * theta_var_n
mape_integrand_s = (data.pfull / p00)**mc.kappa * gamma_s * theta_var_s
#print mape_integrand_s[:,0]
mape_n = prefac * mape_integrand_n.sel(pfull=levs).sum('pfull') * dp / p00
mape_s = prefac * mape_integrand_s.sel(pfull=levs).sum('pfull') * dp / p00
#print mape_n/1.e6
#print mape_s/1.e6
mape = (mape_n + mape_s) / 2.
return mape_n, mape_s, mape
def lapse_rate(run, pentad, period_fac=1.0):
#Load in data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
g = 9.8
Ra = 287.04
mn_dic = month_dic(1)
tickspace = np.arange(13, 72, 18) * period_fac
labels = [mn_dic[(k + 5) / 6] for k in range(13, 72, 18)]
levels = np.arange(-1.5, 1.6, 0.25)
dlnTdp = gr.ddp(np.log(data.temp.mean('lon')))
dTdz = data.pfull * 100. * g / Ra * dlnTdp * 1000.
#dTdz.mean('xofyear').plot.contourf(x='lat', y='pfull', extend = 'both', add_labels=False, yincrease=False, levels=np.arange(-10.,10.5,2.))
#plt.xlabel('Latitude')
#plt.ylabel('Pressure, hPa')
#plt.grid(True,linestyle=':')
#plt.tight_layout()
#figname = 'dTdz_' + run + '.pdf'
#plt.savefig(plot_dir + figname, format='pdf')
#plt.close()
#ind = abs(dTdz.isel(pfull=range(10,20)) -2.).mean('xofyear').argmin('pfull')
ind = abs(dTdz.isel(pfull=range(10, 20)) -
2.).isel(xofyear=pentad).argmin('pfull')
trop_height = np.zeros([64, 1])
for i in range(0, 64):
#trop_height[i] = data.height.mean(('lon','xofyear'))[ind.values[i]+10,i]
trop_height[i] = data.height.mean(
('lon')).isel(xofyear=pentad)[ind.values[i] + 10, i]
print run, trop_height[32]
return trop_height
def mom_budg(run, lev=150):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
#First do uu terms
uu_trans_dx = -86400. * gr.ddx(
(data.ucomp_sq -
data.ucomp**2).sel(pfull=lev)) # <u'u'> = <uu> - <u><u>
u = data.ucomp.sel(pfull=lev) # u
u_dx = -86400. * gr.ddx(u) # dudx
u_ed = u - u.mean('lon')
u_dx_ed = u_dx - u_dx.mean('lon')
u_dudx_zav = u.mean('lon') * u_dx.mean(
'lon') # [u][dudx], where brackets denote mean over all longitudes
u_dudx_cross1 = u.mean('lon') * u_dx_ed # [u]dudx*
u_dudx_cross2 = u_ed * u_dx.mean('lon') # u*[dudx]
u_dudx_stat = u_ed * u_dx_ed # u*dudx*
data['uu_trans_dx'] = (('xofyear', 'lat', 'lon'), uu_trans_dx)
data['u_dudx_cross1'] = (('xofyear', 'lat', 'lon'), u_dudx_cross1)
data['u_dudx_cross2'] = (('xofyear', 'lat', 'lon'), u_dudx_cross2)
data['u_dudx_stat'] = (('xofyear', 'lat', 'lon'), u_dudx_stat)
data['u_dudx_zav'] = (('xofyear', 'lat'), u_dudx_zav)
#Next do uv terms
uv_trans_dy = -86400. * gr.ddy(
(data.ucomp_vcomp - data.ucomp * data.vcomp).sel(pfull=lev), uv=True)
v = data.vcomp.sel(pfull=lev).load() # v
u_dy = -86400. * gr.ddy(u) # dudy
v_ed = v - v.mean('lon')
u_dy_ed = u_dy - u_dy.mean('lon')
v_dudy_zav = v.mean('lon') * u_dy.mean('lon') # [v][dudy]
v_dudy_cross1 = v.mean('lon') * u_dy_ed # [v]dudy*
v_dudy_cross2 = v_ed * u_dy.mean('lon') # v*[dudy]
v_dudy_stat = v_ed * u_dy_ed # v*dudy*
data['uv_trans_dy'] = (('xofyear', 'lat', 'lon'), uv_trans_dy)
data['v_dudy_cross1'] = (('xofyear', 'lat', 'lon'), v_dudy_cross1)
data['v_dudy_cross2'] = (('xofyear', 'lat', 'lon'), v_dudy_cross2)
data['v_dudy_stat'] = (('xofyear', 'lat', 'lon'), v_dudy_stat)
data['v_dudy_zav'] = (('xofyear', 'lat'), v_dudy_zav)
#Finally do uw terms
uw_trans_dp = -86400. * gr.ddp(
(data.ucomp_omega - data.ucomp * data.omega))
w = data.omega.sel(pfull=lev).load() # w
u_dp = -86400. * (gr.ddp(data.ucomp)).sel(pfull=lev) # dudp
w_ed = w - w.mean('lon')
u_dp_ed = u_dp - u_dp.mean('lon')
w_dudp_zav = w.mean('lon') * u_dp.mean('lon') # [w][dudp]
w_dudp_cross1 = w.mean('lon') * u_dp_ed # [w]dudp*
w_dudp_cross2 = w_ed * u_dp.mean('lon') # w*[dudp]
w_dudp_stat = w_ed * u_dp_ed # w*dudp*
data['uw_trans_dp'] = (('xofyear', 'lat', 'lon'),
uw_trans_dp.sel(pfull=lev))
data['w_dudp_cross1'] = (('xofyear', 'lat', 'lon'), w_dudp_cross1)
data['w_dudp_cross2'] = (('xofyear', 'lat', 'lon'), w_dudp_cross2)
data['w_dudp_stat'] = (('xofyear', 'lat', 'lon'), w_dudp_stat)
data['w_dudp_zav'] = (('xofyear', 'lat'), w_dudp_zav)
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
fv = data.vcomp.sel(pfull=lev) * f * 86400.
fv_mean = fv.mean('lon')
fv_local = fv - fv_mean
totp = (data.convection_rain + data.condensation_rain) * 86400.
#Geopotential gradient
dphidx = gr.ddx(data.height.sel(pfull=lev))
dphidx = -86400. * 9.8 * dphidx
fv_ageo = fv_local + dphidx
mom_mean = data.u_dudx_zav + data.v_dudy_zav + data.w_dudp_zav
mom_cross = data.u_dudx_cross1 + data.v_dudy_cross1 + data.w_dudp_cross1 + data.u_dudx_cross2 + data.v_dudy_cross2 + data.w_dudp_cross2
mom_trans = data.uu_trans_dx + data.uv_trans_dy + data.uw_trans_dp
mom_stat = data.u_dudx_stat + data.v_dudy_stat + data.w_dudp_stat
mom_sum = fv_local + fv_mean + dphidx + mom_mean + mom_trans + mom_stat + mom_cross
data['mom_mean'] = (('xofyear', 'lat'), mom_mean)
data['mom_cross'] = (('xofyear', 'lat', 'lon'), mom_cross)
data['mom_trans'] = (('xofyear', 'lat', 'lon'), mom_trans)
data['mom_stat'] = (('xofyear', 'lat', 'lon'), mom_stat)
data['fv_local'] = (('xofyear', 'lat', 'lon'), fv_local)
data['fv_ageo'] = (('xofyear', 'lat', 'lon'), fv_ageo)
data['fv_mean'] = (('xofyear', 'lat'), fv_mean)
data['dphidx'] = (('xofyear', 'lat', 'lon'), dphidx)
data['mom_sum'] = (('xofyear', 'lat', 'lon'), mom_sum)
return data
def vort_eq(run, months, filename='plev_daily', lev=150.):
#Load in dataset
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc'
names = [name_temp % m for m in range(months[0], months[1])]
#read data into xarray
data = xr.open_mfdataset(
names,
decode_times=False, # no calendar so tell netcdf lib
# choose how data will be broken down into manageable chunks.
chunks={'time': 30})
print 'files opened'
# Calculate planetary vorticity
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
#vor_rel = v_dx - u_dy
print 'absolute vorticity calculated'
dvordx = gr.ddx(vor.sel(pfull=lev))
dvordy = gr.ddy(vor.sel(pfull=lev), vector=False)
dvordp = gr.ddp(vor)
# Calculate horizontal material derivative
horiz_md = -1. * (data.ucomp * dvordx + data.vcomp * dvordy)
# Calculate vertical material derivative
vert_md = -1. * data.omega.sel(pfull=lev) * dvordp.sel(pfull=lev)
print 'advective terms done'
# Now do the terms involving gradients of windspeed
div = gr.ddx(data.ucomp.sel(pfull=lev)) + gr.ddy(data.vcomp.sel(pfull=lev))
stretching = -1. * vor * div
tilting = (gr.ddp(data.ucomp)).sel(pfull=lev) * gr.ddy(
data.omega.sel(pfull=lev), vector=False) - (gr.ddp(
data.vcomp)).sel(pfull=lev) * gr.ddx(data.omega.sel(pfull=lev))
print 'stretching and tilting done'
total = horiz_md + vert_md + stretching + tilting
ds = xr.Dataset(
{
'horiz_md': (['time', 'lat', 'lon'], horiz_md),
'vert_md': (['time', 'lat', 'lon'], vert_md),
'stretching': (['time', 'lat', 'lon'], stretching),
'tilting': (['time', 'lat', 'lon'], tilting),
'total': (['time', 'lat', 'lon'], total)
},
coords={
'time': ('time', data.time),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat),
'lon': ('lon', data.lon)
})
ds.coords['xofyear'] = np.mod(ds.time - 1., 360.) // 5 + 1.
dsout = ds.groupby('xofyear').mean(('time'))
GFDL_DATA = os.environ['GFDL_DATA']
fileout = GFDL_DATA + 'climatologies/' + run + '_' + str(
int(lev)) + '_vort_eq.nc'
dsout.to_netcdf(path=fileout)
print 'data written to ', fileout
return dsout
def partition_advection(data, lons, lev=150):
#First do uu terms
uu_trans_dx = -86400. * gr.ddx(
(data.ucomp_sq -
data.ucomp.sel(pfull=lev)**2)) # <u'u'> = <uu> - <u><u>
u = data.ucomp.sel(pfull=lev) # u
u_dx = -86400. * gr.ddx(u) # dudx
u_ed = u - u.mean('lon')
u_dx_ed = u_dx - u_dx.mean('lon')
u_dudx_zav = u.mean('lon') * u_dx.mean(
'lon') # [u][dudx], where brackets denote mean over all longitudes
u_dudx_cross1 = (u.mean('lon') * u_dx_ed).sel(lon=lons).mean(
'lon') # [u]dudx*
u_dudx_cross2 = (u_ed * u_dx.mean('lon')).sel(lon=lons).mean(
'lon') # u*[dudx]
u_dudx_stat = (u_ed * u_dx_ed).sel(lon=lons).mean('lon') # u*dudx*
data['uu_trans_dx'] = (('pentad', 'lat'),
uu_trans_dx.sel(lon=lons).mean('lon'))
data['u_dudx_cross1'] = (('pentad', 'lat'), u_dudx_cross1)
data['u_dudx_cross2'] = (('pentad', 'lat'), u_dudx_cross2)
data['u_dudx_stat'] = (('pentad', 'lat'), u_dudx_stat)
data['u_dudx_zav'] = (('pentad', 'lat'), u_dudx_zav)
print 'uu terms done'
#Next do uv terms
uv_trans_dy = -86400. * gr.ddy(
(data.ucomp_vcomp - data.ucomp.sel(pfull=lev) * data.vcomp), uv=True)
v = data.vcomp.load() # v
u_dy = -86400. * gr.ddy(u) # dudy
v_ed = v - v.mean('lon')
u_dy_ed = u_dy - u_dy.mean('lon')
v_dudy_zav = v.mean('lon') * u_dy.mean('lon') # [v][dudy]
v_dudy_cross1 = (v.mean('lon') * u_dy_ed).sel(lon=lons).mean(
'lon') # [v]dudy*
v_dudy_cross2 = (v_ed * u_dy.mean('lon')).sel(lon=lons).mean(
'lon') # v*[dudy]
v_dudy_stat = (v_ed * u_dy_ed).sel(lon=lons).mean('lon') # v*dudy*
data['uv_trans_dy'] = (('pentad', 'lat'),
uv_trans_dy.sel(lon=lons).mean('lon'))
data['v_dudy_cross1'] = (('pentad', 'lat'), v_dudy_cross1)
data['v_dudy_cross2'] = (('pentad', 'lat'), v_dudy_cross2)
data['v_dudy_stat'] = (('pentad', 'lat'), v_dudy_stat)
data['v_dudy_zav'] = (('pentad', 'lat'), v_dudy_zav)
print 'uv terms done'
#Finally do uw terms
uw_trans_dp = -86400. * gr.ddp(
(data.ucomp_omega - data.ucomp * data.omega).sel(lon=lons).mean('lon'))
w = data.omega.sel(pfull=lev).load() # w
u_dp = -86400. * (gr.ddp(data.ucomp)).sel(pfull=lev) # dudp
w_ed = w - w.mean('lon')
u_dp_ed = u_dp - u_dp.mean('lon')
w_dudp_zav = w.mean('lon') * u_dp.mean('lon') # [w][dudp]
w_dudp_cross1 = (w.mean('lon') * u_dp_ed).sel(lon=lons).mean(
'lon') # [w]dudp*
w_dudp_cross2 = (w_ed * u_dp.mean('lon')).sel(lon=lons).mean(
'lon') # w*[dudp]
w_dudp_stat = (w_ed * u_dp_ed).sel(lon=lons).mean('lon') # w*dudp*
data['uw_trans_dp'] = (('pentad', 'lat'), uw_trans_dp.sel(pfull=lev))
data['w_dudp_cross1'] = (('pentad', 'lat'), w_dudp_cross1)
data['w_dudp_cross2'] = (('pentad', 'lat'), w_dudp_cross2)
data['w_dudp_stat'] = (('pentad', 'lat'), w_dudp_stat)
data['w_dudp_zav'] = (('pentad', 'lat'), w_dudp_zav)
print 'uw terms done'
def partition_mse_advection(data, lons, lev=850):
cp = 287.04 / 2 * 7
L = 2.500e6
data['mse_u'] = cp * data.ucomp_temp + L * data.sphum_u
data['mse_v'] = cp * data.vcomp_temp + L * data.sphum_v
data['mse_w'] = cp * data.omega_temp + L * data.sphum_w
data['mse'] = cp * data.temp + L * data.sphum
#First do uu terms
umse_trans_dx = -1. * gr.ddx(
(data.mse_u -
data.ucomp * data.mse).sel(pfull=lev)) # <u'q'> = <uq> - <u><q>
u = data.ucomp.sel(pfull=lev).load()
mse = data.mse.sel(pfull=lev) # mse
mse_dx = -1. * gr.ddx(mse) # dmsedx
u_dmsedx_zav = u.sel(lon=lons).mean('lon') * mse_dx.sel(lon=lons).mean(
'lon') # [u][dmsedx]
u_dmsedx_stat = (u * mse_dx).sel(lon=lons).mean(
'lon') - u_dmsedx_zav # u*dudx* = [udmsedx] - [u][dmsedx]
data['umse_trans_dx'] = (('xofyear', 'lat'),
umse_trans_dx.sel(lon=lons).mean('lon'))
data['u_dmsedx_stat'] = (('xofyear', 'lat'), u_dmsedx_stat)
data['u_dmsedx_zav'] = (('xofyear', 'lat'), u_dmsedx_zav)
print 'umse terms done'
#Next do uv terms
vmse_trans_dy = -1. * gr.ddy(
(data.mse_v - data.mse * data.vcomp).sel(pfull=lev), uv=True)
v = data.vcomp.sel(pfull=lev).load() # v
mse_dy = -1. * gr.ddy(mse, vector=False) # dmsedy
v_dmsedy_zav = v.sel(lon=lons).mean('lon') * mse_dy.sel(lon=lons).mean(
'lon') # [v][dudy]
v_dmsedy_stat = (v * mse_dy).sel(
lon=lons).mean('lon') - v_dmsedy_zav # v*dudy* = [vdudy] - [v][dudy]
data['vmse_trans_dy'] = (('xofyear', 'lat'),
vmse_trans_dy.sel(lon=lons).mean('lon'))
data['v_dmsedy_stat'] = (('xofyear', 'lat'), v_dmsedy_stat)
data['v_dmsedy_zav'] = (('xofyear', 'lat'), v_dmsedy_zav)
print 'vmse terms done'
#Finally do uw terms
wmse_trans_dp = -1. * gr.ddp(
(data.mse_w - data.mse * data.omega).sel(lon=lons).mean('lon'))
w = data.omega.sel(pfull=lev).load() # w
mse_dp = -1. * (gr.ddp(data.mse)).sel(pfull=lev) # dudp
w_dmsedp_zav = w.sel(lon=lons).mean('lon') * mse_dp.sel(
lon=lons).mean('lon')
w_dmsedp_stat = (w * mse_dp).sel(lon=lons).mean('lon') - w_dmsedp_zav
data['wmse_trans_dp'] = (('xofyear', 'lat'), wmse_trans_dp.sel(pfull=lev))
data['w_dmsedp_stat'] = (('xofyear', 'lat'), w_dmsedp_stat)
data['w_dmsedp_zav'] = (('xofyear', 'lat'), w_dmsedp_zav)
print 'wmse terms done'
