def heatbudg_2d_pd_fn(run_fol, years, trange, heatbudg=True, humbudg=False):
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
mngrp = rundata.ucomp.groupby('pentad').mean(('time'))
trl = trange[1] - trange[0]
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
#Initialise arrays to load into
u = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
v = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
w = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
if heatbudg == True:
t = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
if humbudg == True:
q = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
if heatbudg == False and humbudg == False:
print 'No budget specified'
return
for year in years:
print year
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
u[:, :, :, :, year - years[0]] = rundata.ucomp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
v[:, :, :, :, year - years[0]] = rundata.vcomp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
w[:, :, :, :, year - years[0]] = rundata.omega.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
if heatbudg == True:
t[:, :, :, :,
year - years[0]] = rundata.temp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
if humbudg == True:
q[:, :, :, :,
year - years[0]] = rundata.sphum.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
#Take averages: u, v, uv, w, phi, damping terms
u_av = u.mean(('year'))
v_av = v.mean(('year'))
w_av = w.mean(('year'))
#Evaluate eddy quantities
u_ed = u - u_av
v_ed = v - v_av
w_ed = w - w_av
if heatbudg == True:
t_av = t.mean(('year'))
t_ed = t - t_av
ut_ed = (u_ed * t_ed).mean(('year'))
vt_ed = (v_ed * t_ed).mean(('year')) * coslat
wt_ed = (w_ed * t_ed).mean(('year'))
#Evaluate gradients needed
dtdx_av = xr.DataArray(
cfd(t_av.values, t_av.lon * np.pi / 180, 3),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
udtdx_av = u_av * dtdx_av / coslat / a
dtdy_av = xr.DataArray(
cfd(t_av.values, t_av.lat * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
vdtdy_av = v_av * dtdy_av / a
dtdp_av = xr.DataArray(
cfd(t_av.values, t_av.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
wdtdp_av = w_av * dtdp_av
duteddx_av = xr.DataArray(
cfd(ut_ed.values, ut_ed.lon * np.pi / 180, 3),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
duteddx_av = duteddx_av / coslat / a
dvteddy_av = xr.DataArray(
cfd(vt_ed.values, vt_ed.lat * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
dvteddy_av = dvteddy_av / coslat / a
dwteddp_av = xr.DataArray(
cfd(wt_ed.values, wt_ed.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
#evaluate sums of the terms
heat_eddy = -(duteddx_av + dvteddy_av + dwteddp_av)
heat_mean = -(udtdx_av + vdtdy_av + wdtdp_av)
heat_sum = heat_eddy + heat_mean
if humbudg == True:
q_av = q.mean(('year'))
q_ed = q - q_av
uq_ed = (u_ed * q_ed).mean(('year'))
vq_ed = (v_ed * q_ed).mean(('year')) * coslat
wq_ed = (w_ed * q_ed).mean(('year'))
#Evaluate gradients needed
dqdx_av = xr.DataArray(
cfd(q_av.values, q_av.lon * np.pi / 180, 3),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
udqdx_av = u_av * dqdx_av / coslat / a
dqdy_av = xr.DataArray(
cfd(q_av.values, q_av.lat * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
vdqdy_av = v_av * dqdy_av / a
dqdp_av = xr.DataArray(
cfd(q_av.values, q_av.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
wdqdp_av = w_av * dqdp_av
duqeddx_av = xr.DataArray(
cfd(uq_ed.values, uq_ed.lon * np.pi / 180, 3),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
duqeddx_av = duqeddx_av / coslat / a
dvqeddy_av = xr.DataArray(
cfd(vq_ed.values, vq_ed.lat * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
dvqeddy_av = dvqeddy_av / coslat / a
dwqeddp_av = xr.DataArray(
cfd(wq_ed.values, wq_ed.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
#evaluate sums of the terms
hum_eddy = -(duqeddx_av + dvqeddy_av + dwqeddp_av)
hum_mean = -(udqdx_av + vdqdy_av + wdqdp_av)
hum_sum = hum_eddy + hum_mean
if heatbudg == True and humbudg == True:
data_out = xr.Dataset({
'hum_eddy': hum_eddy,
'hum_mean': hum_mean,
'hum_sum': hum_sum,
'heat_eddy': heat_eddy,
'heat_mean': heat_mean,
'heat_sum': heat_sum,
'u_av': u_av,
'v_av': v_av
})
elif heatbudg == True:
data_out = xr.Dataset({
'heat_eddy': heat_eddy,
'heat_mean': heat_mean,
'heat_sum': heat_sum,
'u_av': u_av,
'v_av': v_av
})
elif humbudg == True:
data_out = xr.Dataset({
'hum_eddy': hum_eddy,
'hum_mean': hum_mean,
'hum_sum': hum_sum,
'u_av': u_av,
'v_av': v_av
})
return data_out
def mom_adv_terms_fn(run_fol,years, trange):
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
mngrp = rundata.ucomp.groupby('pentad').mean(('time'))
trl = trange[1]-trange[0]
#Define constants
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi/180)
f = 2 * omega * np.sin(rundata.lat *np.pi/180)
#Initialise arrays to load into
u = xr.DataArray(np.zeros((trl,17,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', rundata.pfull), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
v = xr.DataArray(np.zeros((trl,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
w = xr.DataArray(np.zeros((trl,17,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', rundata.pfull), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
uu = xr.DataArray(np.zeros((trl,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
uv = xr.DataArray(np.zeros((trl,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
uw = xr.DataArray(np.zeros((trl,17,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', rundata.pfull), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
for year in years:
print year
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
u[:,:,:,:,year-years[0]] = rundata.ucomp.groupby('pentad').mean(('time'))[trange[0]:trange[1],:,:,:]
v[:,:,:,year-years[0]] = rundata.vcomp.groupby('pentad').mean(('time'))[trange[0]:trange[1],9,:,:]
w[:,:,:,:,year-years[0]] = rundata.omega.groupby('pentad').mean(('time'))[trange[0]:trange[1],:,:,:]
uu[:,:,:,year-years[0]] = rundata.ucomp_sq.groupby('pentad').mean(('time'))[trange[0]:trange[1],9,:,:]
uv[:,:,:,year-years[0]] = rundata.ucomp_vcomp.groupby('pentad').mean(('time'))[trange[0]:trange[1],9,:,:]
uw[:,:,:,:,year-years[0]] = rundata.ucomp_omega.groupby('pentad').mean(('time'))[trange[0]:trange[1],:,:,:]
#Take averages: u, v, uv, w, phi, damping terms
u_av = u.mean(('year'))
v_av = v.mean(('year'))
w_av = w.mean(('year'))
uu_av = uu.mean(('year'))
uv_av = uv.mean(('year'))
uw_av = uw.mean(('year'))
#Evaluate mean eddy products
uued_av = uu_av - u_av[:,9,:,:]*u_av[:,9,:,:]
uved_av = (uv_av - u_av[:,9,:,:]*v_av)*coslat*coslat
uwed_av = uw_av - u_av*w_av
#Evaluate gradients needed
dudx_av = xr.DataArray( cfd(u_av[:,9,:,:].values,u_av.lon*np.pi/180,2), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', u_av.lat), ('lon', u_av.lon )])
ududx_av = u_av[:,9,:,:]*dudx_av/coslat/a
dudy_av = xr.DataArray( cfd((u_av[:,9,:,:]*coslat).values,u_av.lat*np.pi/180,1), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', u_av.lat), ('lon', u_av.lon )])
vdudy_av = v_av*dudy_av/coslat/a
dudp_av = xr.DataArray( cfd(u_av.values,u_av.pfull*100,1), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', u_av.pfull ), ('lat', u_av.lat), ('lon', u_av.lon )])
wdudp_av = w_av[:,9,:,:]*dudp_av[:,9,:,:]
duueddx_av = xr.DataArray( cfd(uued_av.values,uued_av.lon*np.pi/180,2), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', u_av.lat), ('lon', uued_av.lon )])
duueddx_av = duueddx_av/coslat/a
duveddy_av = xr.DataArray( cfd(uved_av.values ,uved_av.lat*np.pi/180,1), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', uved_av.lat), ('lon', uved_av.lon )])
duveddy_av = duveddy_av/coslat/coslat/a
duweddp_av = xr.DataArray( cfd(uwed_av.values,uwed_av.pfull*100,1), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', uwed_av.pfull ), ('lat', u_av.lat), ('lon', uwed_av.lon )])
duweddp_av = duweddp_av[:,9,:,:]
data_out = xr.Dataset({'ududx_av': ududx_av, 'vdudy_av': vdudy_av, 'wdudp_av': wdudp_av,
'duueddx_av': duueddx_av, 'duveddy_av': duveddy_av, 'duweddp_av': duweddp_av, 'u_av': u_av[:,9,:,:], 'v_av': v_av })
return data_out
def precip_hm(run_fol,years):
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) -71
mngrp = rundata.ucomp.groupby('pentad').mean(('time'))
tot_p = xr.DataArray(np.zeros((72,128,len(years))), [('pentad', range(1,73) ), ('lon', rundata.lon), ('year', years)])
u = xr.DataArray(np.zeros((72,17,128,len(years))), [('pentad', range(1,73) ), ('pfull', rundata.pfull), ('lon', rundata.lon), ('year', years)])
v = xr.DataArray(np.zeros((72,17,128,len(years))), [('pentad', range(1,73) ), ('pfull', rundata.pfull), ('lon', rundata.lon), ('year', years)])
phi = xr.DataArray(np.zeros((72,17,128,len(years))), [('pentad', range(1,73) ), ('pfull', rundata.pfull), ('lon', rundata.lon), ('year', years)])
i=0
for year in years:
print year
rundata = load_year_xr(run_fol, year)
rundata.coords['pentad'] = (rundata.time // 5) -71
#take time mean
conv_p = rundata.convection_rain[:,37,:].groupby('pentad').mean(('time'))
cond_p = rundata.condensation_rain[:,37,:].groupby('pentad').mean(('time'))
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) -71
u[:,:,:,i] = rundata.ucomp[:,:,37,:].groupby('pentad').mean(('time'))
v[:,:,:,i] = rundata.vcomp[:,:,37,:].groupby('pentad').mean(('time'))
phi[:,:,:,i] = rundata.height[:,:,37,:].groupby('pentad').mean(('time'))
tot_p[:,:,i] = (conv_p + cond_p)*86400
i=i+1
tot_p_clim = tot_p.mean(('year'))
u_clim = u.mean(('year'))
v_clim = v.mean(('year'))
phi_clim = phi.mean(('year'))
phi_ed_clim = phi_clim - phi_clim.mean(('lon'))
pd_dic = pentad_dic(1)
tickspace = range(20,61,5)
labels = [pd_dic[k] for k in tickspace]
#plot
plt.figure(0)
phi_ed_clim[:,2,:].plot.contourf(x='lon', y='pentad', levels=np.arange(-40.,40.,2.), add_label = False)
tot_p_clim.plot.contour(x='lon', y='pentad',levels=[-1000.,6.,1000.], add_label = False, add_colorbar=False, colors='k')
plt.yticks(tickspace,labels,rotation=25)
plt.quiver(u_clim.lon[::3], u_clim.pentad[::2], u_clim[::2,2,::3], v_clim[::2,2,::3], headwidth=3)#, scale=500.)
plt.plot([float(rundata.lonb[21]),float(rundata.lonb[21])],[0.,60.],'k')
plt.plot([float(rundata.lonb[26]),float(rundata.lonb[26])],[0.,60.],'k')
plt.plot([float(rundata.lonb[30]),float(rundata.lonb[30])],[0.,60.],'k')
plt.plot([float(rundata.lonb[35]),float(rundata.lonb[35])],[0.,60.],'k')
plt.plot([float(rundata.lonb[39]),float(rundata.lonb[39])],[0.,60.],'k')
plt.ylim(20,60)
plt.xlim(60,180)
plt.xlabel('Longitude')
plt.ylabel('Pentad')
plt.title('850hPa geopotential anomalies, m')
plt.savefig('/scratch/rg419/plots/mom_budg_work/geopot_hm_850.png')
plt.clf()
plt.figure(0)
phi_ed_clim[:,9,:].plot.contourf(x='lon', y='pentad', levels=np.arange(-40.,40.,2.), add_label = False)
tot_p_clim.plot.contour(x='lon', y='pentad',levels=[-1000.,6.,1000.], add_label = False, add_colorbar=False, colors='k')
plt.yticks(tickspace,labels,rotation=25)
plt.quiver(u_clim.lon[::3], u_clim.pentad[::2], u_clim[::2,9,::3], v_clim[::2,9,::3], headwidth=3)#, scale=500.)
plt.plot([float(rundata.lonb[21]),float(rundata.lonb[21])],[0.,60.],'k')
plt.plot([float(rundata.lonb[26]),float(rundata.lonb[26])],[0.,60.],'k')
plt.plot([float(rundata.lonb[30]),float(rundata.lonb[30])],[0.,60.],'k')
plt.plot([float(rundata.lonb[35]),float(rundata.lonb[35])],[0.,60.],'k')
plt.plot([float(rundata.lonb[39]),float(rundata.lonb[39])],[0.,60.],'k')
plt.ylim(20,60)
plt.xlim(60,180)
plt.xlabel('Longitude')
plt.ylabel('Pentad')
plt.title('200hPa geopotential anomalies, m')
plt.savefig('/scratch/rg419/plots/mom_budg_work/geopot_hm_200.png')
plt.clf()
def load_bs_vars(run_fol, years, lonrange=['all']):
#set run name
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=False)
rundata.coords['pentad'] = (rundata.time // 5) - 71
mngrp = rundata.ucomp.groupby('pentad').mean(('time'))
#Initialise arrays to load into
cnvp = xr.DataArray(np.zeros((72, 64, len(years))),
[('pentad', range(1, 73)), ('lat', rundata.lat),
('year', years)])
cndp = xr.DataArray(np.zeros((72, 64, len(years))),
[('pentad', range(1, 73)), ('lat', rundata.lat),
('year', years)])
t_surf = xr.DataArray(np.zeros((72, 64, len(years))),
[('pentad', range(1, 73)), ('lat', rundata.lat),
('year', years)])
for year in years:
print year
rundata = load_year_xr(run_fol, year, pinterp=False)
rundata.coords['pentad'] = (rundata.time // 5) - 71
if lonrange[0] == 'all':
cnvp[:, :, year -
years[0]] = rundata.convection_rain.groupby('pentad').mean(
('time', 'lon'))
cndp[:, :, year -
years[0]] = rundata.condensation_rain.groupby('pentad').mean(
('time', 'lon'))
t_surf[:, :,
year - years[0]] = rundata.t_surf.groupby('pentad').mean(
('time', 'lon'))
else:
lonmin = np.min(
[k for k, j in enumerate(rundata.lon) if j >= lonrange[0]])
#print lonmin, rundata.lon[lonmin]
lonmax = np.max(
[k for k, j in enumerate(rundata.lon) if j <= lonrange[1]])
#print lonmax, rundata.lon[lonmax]
cnvp[:, :, year -
years[0]] = rundata.convection_rain[:, :,
lonmin:lonmax].groupby(
'pentad').mean(
('time', 'lon'))
cndp[:, :, year -
years[0]] = rundata.condensation_rain[:, :,
lonmin:lonmax].groupby(
'pentad').mean(
('time', 'lon'))
t_surf[:, :,
year - years[0]] = rundata.t_surf[:, :,
lonmin:lonmax].groupby(
'pentad').mean(
('time', 'lon'))
data_out = xr.Dataset({
'totp': (cnvp.mean(('year')) + cndp.mean(('year'))) * 86400.,
't_surf':
t_surf.mean(('year'))
})
return data_out
def plot_weather(inp_fol, years):
year = years[0]
rundata = load_year_xr(inp_fol, year)
tot_p = xr.DataArray(np.zeros((72, 64, 128, len(years))),
[('pentad', range(1, 73)), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
tsurf = xr.DataArray(np.zeros((72, 64, 128, len(years))),
[('pentad', range(1, 73)), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
#w = xr.DataArray(np.zeros((72,64,128,len(years))), [('pentad', range(1,73) ), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years)])
if not 'aqua' in inp_fol:
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/exp/topo_10m/input/land.nc'
land = xr.open_dataset(land_file)
land_plot = xr.DataArray(land.land_mask.values, [('lat', rundata.lat),
('lon', rundata.lon)])
i = 0
for year in years:
print year
rundata = load_year_xr(inp_fol, year)
rundata.coords['pentad'] = (rundata.time // 5) - 71
conv_p = rundata.convection_rain.groupby('pentad').mean(('time'))
cond_p = rundata.condensation_rain.groupby('pentad').mean(('time'))
tot_p[:, :, :, i] = (conv_p + cond_p) * 86400
tsurf[:, :, :, i] = rundata.t_surf.groupby('pentad').mean(('time'))
#rundata = load_year_xr(inp_fol, year, pinterp=True)
#rundata.coords['pentad'] = (rundata.time // 5) -71
#w[:,:,:,i] = rundata.omega[:,3,:,:].groupby('pentad').mean(('time'))
i = i + 1
tot_p_clim = tot_p.mean(('year'))
tsurf_clim = tsurf.mean(('year'))
#w_clim = w.mean(('year'))
pd_dic = pentad_dic(1)
for p in range(0, 72):
print p
tot_p_clim[p, 15:64, 21:66].plot.contourf(x='lon',
y='lat',
add_labels=False,
levels=range(0, 31, 3))
cs = tsurf_clim[p, 15:64,
21:66].plot.contour(x='lon',
y='lat',
colors='w',
add_labels=False,
add_colorbar=False,
levels=range(250, 321, 5))
if not 'aqua' in inp_fol:
land_plot.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='w',
add_colorbar=False,
add_labels=False)
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
plt.ylim(0, 90)
plt.xlim(60, 180)
plt.title(pd_dic[p + 1])
plt.tight_layout()
plt.savefig('/scratch/rg419/plots/monsoon_analysis/' + inp_fol +
'/rain_and_t' + str(p) + '.png')
plt.close()
def mombudg_pd_fn(run_fol, years):
#inp_fol = experiment name
#nproc = no of processors
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
f = 2 * omega * np.sin(rundata.lat * np.pi / 180)
#Initialise arrays to load into
u = xr.DataArray(np.zeros((72, 17, 64, len(years))),
[('pentad', range(1, 73)), ('pfull', rundata.pfull),
('lat', rundata.lat), ('year', years)])
v = xr.DataArray(np.zeros((72, 17, 64, len(years))),
[('pentad', range(1, 73)), ('pfull', rundata.pfull),
('lat', rundata.lat), ('year', years)])
w = xr.DataArray(np.zeros((72, 17, 64, len(years))),
[('pentad', range(1, 73)), ('pfull', rundata.pfull),
('lat', rundata.lat), ('year', years)])
uu = xr.DataArray(np.zeros((72, 17, 64, len(years))),
[('pentad', range(1, 73)), ('pfull', rundata.pfull),
('lat', rundata.lat), ('year', years)])
uv = xr.DataArray(np.zeros((72, 17, 64, len(years))),
[('pentad', range(1, 73)), ('pfull', rundata.pfull),
('lat', rundata.lat), ('year', years)])
uw = xr.DataArray(np.zeros((72, 17, 64, len(years))),
[('pentad', range(1, 73)), ('pfull', rundata.pfull),
('lat', rundata.lat), ('year', years)])
phi = xr.DataArray(np.zeros((72, 17, 64, len(years))),
[('pentad', range(1, 73)), ('pfull', rundata.pfull),
('lat', rundata.lat), ('year', years)])
ddamp = xr.DataArray(np.zeros((72, 17, 64, len(years))),
[('pentad', range(1, 73)), ('pfull', rundata.pfull),
('lat', rundata.lat), ('year', years)])
for year in years:
print year
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
u[:, :, :, year - years[0]] = rundata.ucomp.groupby('pentad').mean(
('time', 'lon'))
v[:, :, :, year - years[0]] = rundata.vcomp.groupby('pentad').mean(
('time', 'lon'))
w[:, :, :, year - years[0]] = rundata.omega.groupby('pentad').mean(
('time', 'lon'))
uu[:, :, :, year - years[0]] = rundata.ucomp_sq.groupby('pentad').mean(
('time', 'lon'))
uv[:, :, :,
year - years[0]] = rundata.ucomp_vcomp.groupby('pentad').mean(
('time', 'lon'))
uw[:, :, :,
year - years[0]] = rundata.ucomp_omega.groupby('pentad').mean(
('time', 'lon'))
phi[:, :, :, year - years[0]] = rundata.height.groupby('pentad').mean(
('time', 'lon'))
ddamp[:, :, :, year -
years[0]] = rundata.dt_ug_diffusion.groupby('pentad').mean(
('time', 'lon'))
#Take averages: u, v, uv, w, phi, damping terms
u_av = u.mean(('year'))
v_av = v.mean(('year'))
w_av = w.mean(('year'))
uu_av = uu.mean(('year'))
uv_av = uv.mean(('year'))
uw_av = uw.mean(('year'))
phi_av = phi.mean(('year')) * 9.8
ddamp_av = ddamp.mean(('year'))
fv_av = v_av * f
#Evaluate mean eddy products
uued_av = uu_av - u_av * u_av
uved_av = (uv_av - u_av * v_av) * coslat * coslat
uwed_av = uw_av - u_av * w_av
#Evaluate gradients needed
dudy_av = xr.DataArray(
cfd((u_av * coslat).values, u_av.lat * np.pi / 180, 2),
[('pentad', range(1, 73)), ('pfull', u_av.pfull), ('lat', u_av.lat)])
vdudy_av = v_av * dudy_av / coslat / a
dudp_av = xr.DataArray(cfd(u_av.values, u_av.pfull * 100, 1),
[('pentad', range(1, 73)), ('pfull', u_av.pfull),
('lat', u_av.lat)])
wdudp_av = w_av * dudp_av
duveddy_av = xr.DataArray(
cfd(uved_av.values, uved_av.lat * np.pi / 180, 2),
[('pentad', range(1, 73)), ('pfull', uued_av.pfull),
('lat', uued_av.lat)])
duveddy_av = duveddy_av / coslat / coslat / a
duweddp_av = xr.DataArray(cfd(uwed_av.values, uwed_av.pfull * 100, 1),
[('pentad', range(1, 73)),
('pfull', uued_av.pfull), ('lat', uued_av.lat)])
#evaluate sums of the terms
mom_eddy = -(duveddy_av + duweddp_av)
mom_mean = -(vdudy_av + wdudp_av)
mom_sum = fv_av + ddamp_av + mom_eddy + mom_mean
data_out = xr.Dataset({
'fv_av': fv_av,
'ddamp_av': ddamp_av,
'mom_eddy': mom_eddy,
'mom_mean': mom_mean,
'mom_sum': mom_sum,
'u_av': u_av,
'v_av': v_av
})
return data_out
def mombudg_lev_pd_fn(run_fol, years, trange):
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
mngrp = rundata.ucomp.groupby('pentad').mean(('time'))
trl = trange[1] - trange[0]
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
f = 2 * omega * np.sin(rundata.lat * np.pi / 180)
#Initialise arrays to load into
u = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
v = xr.DataArray(np.zeros((trl, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('lat', rundata.lat), ('lon', rundata.lon),
('year', years)])
w = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
uu = xr.DataArray(np.zeros((trl, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('lat', rundata.lat), ('lon', rundata.lon),
('year', years)])
uv = xr.DataArray(np.zeros((trl, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('lat', rundata.lat), ('lon', rundata.lon),
('year', years)])
uw = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
phi = xr.DataArray(np.zeros((trl, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('lat', rundata.lat), ('lon', rundata.lon),
('year', years)])
ddamp = xr.DataArray(np.zeros((trl, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('lat', rundata.lat), ('lon', rundata.lon),
('year', years)])
for year in years:
print year
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
u[:, :, :, :, year - years[0]] = rundata.ucomp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
v[:, :, :, year - years[0]] = rundata.vcomp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], 9, :, :]
w[:, :, :, :, year - years[0]] = rundata.omega.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
uu[:, :, :, year - years[0]] = rundata.ucomp_sq.groupby('pentad').mean(
('time'))[trange[0]:trange[1], 9, :, :]
uv[:, :, :,
year - years[0]] = rundata.ucomp_vcomp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], 9, :, :]
uw[:, :, :, :,
year - years[0]] = rundata.ucomp_omega.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
phi[:, :, :, year - years[0]] = rundata.height.groupby('pentad').mean(
('time'))[trange[0]:trange[1], 9, :, :]
ddamp[:, :, :, year -
years[0]] = rundata.dt_ug_diffusion.groupby('pentad').mean(
('time'))[trange[0]:trange[1], 9, :, :]
#Take averages: u, v, uv, w, phi, damping terms
u_av = u.mean(('year'))
v_av = v.mean(('year'))
w_av = w.mean(('year'))
uu_av = uu.mean(('year'))
uv_av = uv.mean(('year'))
uw_av = uw.mean(('year'))
phi_av = phi.mean(('year')) * 9.8
ddamp_av = ddamp.mean(('year'))
fv_av = v_av * f
#Evaluate mean eddy products
uued_av = uu_av - u_av[:, 9, :, :] * u_av[:, 9, :, :]
uved_av = (uv_av - u_av[:, 9, :, :] * v_av) * coslat * coslat
uwed_av = uw_av - u_av * w_av
#Evaluate gradients needed
dudx_av = xr.DataArray(
cfd(u_av[:, 9, :, :].values, u_av.lon * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)), ('lat', u_av.lat),
('lon', u_av.lon)])
ududx_av = u_av[:, 9, :, :] * dudx_av / coslat / a
dudy_av = xr.DataArray(
cfd((u_av[:, 9, :, :] * coslat).values, u_av.lat * np.pi / 180, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)), ('lat', u_av.lat),
('lon', u_av.lon)])
vdudy_av = v_av * dudy_av / coslat / a
dudp_av = xr.DataArray(cfd(u_av.values, u_av.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
wdudp_av = w_av[:, 9, :, :] * dudp_av[:, 9, :, :]
duueddx_av = xr.DataArray(
cfd(uued_av.values, uued_av.lon * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)), ('lat', u_av.lat),
('lon', uued_av.lon)])
duueddx_av = duueddx_av / coslat / a
duveddy_av = xr.DataArray(
cfd(uved_av.values, uved_av.lat * np.pi / 180, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)), ('lat', uved_av.lat),
('lon', uved_av.lon)])
duveddy_av = duveddy_av / coslat / coslat / a
duweddp_av = xr.DataArray(cfd(uwed_av.values, uwed_av.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', uwed_av.pfull), ('lat', u_av.lat),
('lon', uwed_av.lon)])
duweddp_av = duweddp_av[:, 9, :, :]
dphidx_av = xr.DataArray(cfd(phi_av.values, phi_av.lon * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('lat', u_av.lat), ('lon', phi_av.lon)])
dphidx_av = -1 * dphidx_av / coslat / a
#evaluate sums of the terms
mom_eddy = -(duueddx_av + duveddy_av + duweddp_av)
mom_mean = -(ududx_av + vdudy_av + wdudp_av)
mom_sum = fv_av + ddamp_av + dphidx_av + mom_eddy + mom_mean
data_out = xr.Dataset({
'dphidx_av': dphidx_av,
'fv_av': fv_av,
'ddamp_av': ddamp_av,
'mom_eddy': mom_eddy,
'mom_mean': mom_mean,
'mom_sum': mom_sum,
'u_av': u_av[:, 9, :, :],
'v_av': v_av
})
return data_out
# Plot the streamfunction and jets for the 2m aquaplanet experiment before and after onset, hoping to get a similar result to B+S
#onset pentad:
#topo: 41.25
#flat:41.8
#2m: 37.48
#10m:46.81
from physics import mass_streamfunction
from data_handling import load_year_xr, pentad_dic
import matplotlib.pyplot as plt
import xarray as xr
import numpy as np
run_fol = 'aquaplanet_10m'
rundata = load_year_xr('aquaplanet_10m', 21)
years = range(21, 41)
pd_dic = pentad_dic(1)
u = xr.DataArray(np.zeros((72, 40, 64, len(years))), [('pentad', range(1, 73)),
('pfull', rundata.pfull),
('lat', rundata.lat),
('year', years)])
psi = xr.DataArray(np.zeros((64, 72, 40, len(years))),
[('lat', rundata.lat), ('pentad', range(1, 73)),
('pfull', rundata.pfull), ('year', years)])
i = 0
for year in years:
print year
rundata = load_year_xr(run_fol, year)
def mombudg_2d_an_fn(inp_fol, nproc, years, era=True):
#inp_fol = experiment name
#nproc = no of processors
#set run name
run_fol = inp_fol + '/np' + nproc
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=True, era=era)
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
f = 2 * omega * np.sin(rundata.lat * np.pi / 180)
if era == True:
plevs = 17
else:
plevs = 40
#Initialise arrays to load into
u = xr.DataArray(np.zeros((plevs, 64, 128, len(years))),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
v = xr.DataArray(np.zeros((plevs, 64, 128, len(years))),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
w = xr.DataArray(np.zeros((plevs, 64, 128, len(years))),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
uu = xr.DataArray(np.zeros((plevs, 64, 128, len(years))),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
uv = xr.DataArray(np.zeros((plevs, 64, 128, len(years))),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
uw = xr.DataArray(np.zeros((plevs, 64, 128, len(years))),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
phi = xr.DataArray(np.zeros((plevs, 64, 128, len(years))),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
ddamp = xr.DataArray(np.zeros((plevs, 64, 128, len(years))),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
for year in years:
print year
rundata = load_year_xr(run_fol, year, pinterp=True, era=era)
rundata.coords['pentad'] = (rundata.time // 5) + 1
u[:, :, :, year - years[0]] = rundata.ucomp.mean(('time'))
v[:, :, :, year - years[0]] = rundata.vcomp.mean(('time'))
w[:, :, :, year - years[0]] = rundata.omega.mean(('time'))
uu[:, :, :, year - years[0]] = rundata.ucomp_sq.mean(('time'))
uv[:, :, :, year - years[0]] = rundata.ucomp_vcomp.mean(('time'))
uw[:, :, :, year - years[0]] = rundata.ucomp_omega.mean(('time'))
phi[:, :, :, year - years[0]] = rundata.height.mean(('time'))
ddamp[:, :, :, year - years[0]] = rundata.dt_ug_diffusion.mean(
('time'))
#Take averages: u, v, uv, w, phi, damping terms
u_av = u.mean(('year'))
v_av = v.mean(('year'))
w_av = w.mean(('year'))
uu_av = uu.mean(('year'))
uv_av = uv.mean(('year'))
uw_av = uw.mean(('year'))
phi_av = phi.mean(('year')) * 9.8
ddamp_av = ddamp.mean(('year'))
fv_av = v_av * f
#Evaluate mean eddy products
uued_av = uu_av - u_av * u_av
uved_av = (uv_av - u_av * v_av) * coslat * coslat
uwed_av = uw_av - u_av * w_av
#Evaluate gradients needed
dudx_av = xr.DataArray(cfd(u_av.values, u_av.lon * np.pi / 180, 2),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
ududx_av = u_av * dudx_av / coslat / a
dudy_av = xr.DataArray(
cfd((u_av * coslat).values, u_av.lat * np.pi / 180, 1),
[('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
vdudy_av = v_av * dudy_av / coslat / a
dudp_av = xr.DataArray(cfd(u_av.values, u_av.pfull * 100, 0),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
wdudp_av = w_av * dudp_av
duueddx_av = xr.DataArray(
cfd(uued_av.values, uued_av.lon * np.pi / 180, 2),
[('pfull', uued_av.pfull), ('lat', uued_av.lat), ('lon', uued_av.lon)])
duueddx_av = duueddx_av / coslat / a
duveddy_av = xr.DataArray(
cfd(uved_av.values, uved_av.lat * np.pi / 180, 1),
[('pfull', uued_av.pfull), ('lat', uued_av.lat), ('lon', uued_av.lon)])
duveddy_av = duveddy_av / coslat / coslat / a
duweddp_av = xr.DataArray(cfd(uwed_av.values, uwed_av.pfull * 100, 0),
[('pfull', uued_av.pfull), ('lat', uued_av.lat),
('lon', uued_av.lon)])
dphidx_av = xr.DataArray(cfd(phi_av.values, phi_av.lon * np.pi / 180, 2),
[('pfull', phi_av.pfull), ('lat', phi_av.lat),
('lon', phi_av.lon)])
dphidx_av = -1 * dphidx_av / coslat / a
#evaluate sums of the terms
mom_eddy = -(duueddx_av + duveddy_av + duweddp_av)
mom_mean = -(ududx_av + vdudy_av + wdudp_av)
mom_sum = fv_av + ddamp_av + dphidx_av + mom_eddy + mom_mean
data_out = xr.Dataset({
'dphidx_av': dphidx_av,
'fv_av': fv_av,
'ddamp_av': ddamp_av,
'mom_eddy': mom_eddy,
'mom_mean': mom_mean,
'mom_sum': mom_sum,
'u_av': u_av,
'v_av': v_av
})
return data_out
def plot_pd_mom_fn(inp_fol, years):
year = years[0]
rundata = load_year_xr(inp_fol, year)
tot_p = xr.DataArray(np.zeros((72, 128, len(years))),
[('pentad', range(1, 73)), ('lon', rundata.lon),
('year', years)])
i = 0
for year in years:
print year
rundata = load_year_xr(inp_fol, year)
rundata.coords['pentad'] = (rundata.time // 5) - 71
conv_p = rundata.convection_rain[:, 37, :].groupby('pentad').mean(
('time'))
cond_p = rundata.condensation_rain[:, 37, :].groupby('pentad').mean(
('time'))
tot_p[:, :, i] = (conv_p + cond_p) * 86400
i = i + 1
tot_p_clim = tot_p.mean(('year'))
data, wspd = mombudg_lon_pd_fn(inp_fol, years, [19, 61])
mn_dic = month_dic(1)
tickspace = range(25, 60, 6)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
def plot_mom_var(var, levels):
var_dic = {
'fv_av': 'fv',
'mom_eddy': 'Eddy advective terms',
'mom_mean': 'Mean state advective terms',
'mom_sum': 'Residual',
'dphidx_av': 'Geopotential gradient'
}
if var == 'fv_mn_imb':
plot_data = data.data_vars['fv_av'] + data.data_vars['mom_mean']
else:
plot_data = data.data_vars[var]
plot_data = plot_data * 10000.
g = plot_data.plot.contourf(x='lon',
y='pentad',
levels=levels,
add_labels=False,
add_colorbar=False,
extend='both')
cb1 = plt.colorbar(g)
cb1.set_label('$\displaystyle10^{-4}m/s^2$')
plt.quiver(data.lon[::3],
data.pentad[::2],
wspd.u_av[::2, ::3],
wspd.v_av[::2, ::3],
headwidth=3) #, scale=500.)
#tickspace = range(25,60,10)
#labels = [pd_dic[k] for k in tickspace]
plt.ylim(20, 60)
plt.xlim(60, 180)
if not 'aqua' in inp_fol:
plt.plot([float(rundata.lonb[21]),
float(rundata.lonb[21])], [0., 60.], 'k')
plt.plot([float(rundata.lonb[26]),
float(rundata.lonb[26])], [0., 60.], 'k')
plt.plot([float(rundata.lonb[30]),
float(rundata.lonb[30])], [0., 60.], 'k')
plt.plot([float(rundata.lonb[35]),
float(rundata.lonb[35])], [0., 60.], 'k')
plt.plot([float(rundata.lonb[39]),
float(rundata.lonb[39])], [0., 60.], 'k')
plt.fill_between(
[float(rundata.lonb[26]),
float(rundata.lonb[30])],
0.,
60.,
facecolor='gray',
alpha=0.5)
plt.fill_between(
[float(rundata.lonb[35]),
float(rundata.lonb[39])],
0.,
60.,
facecolor='gray',
alpha=0.5)
else:
plt.plot([float(rundata.lonb[21]),
float(rundata.lonb[21])], [0., 60.], 'k--')
plt.plot([float(rundata.lonb[26]),
float(rundata.lonb[26])], [0., 60.], 'k--')
plt.plot([float(rundata.lonb[30]),
float(rundata.lonb[30])], [0., 60.], 'k--')
plt.plot([float(rundata.lonb[35]),
float(rundata.lonb[35])], [0., 60.], 'k--')
plt.plot([float(rundata.lonb[39]),
float(rundata.lonb[39])], [0., 60.], 'k--')
tot_p_clim.plot.contour(x='lon',
y='pentad',
colors='k',
levels=[-1000., 6., 1000.],
add_labels=False,
add_colorbar=False)
#plt.contour(tot_p_clim.lon, tot_p_clim.pentad, tot_p_clim.values, colors='k', levels=[-1000.,6.,1000.], linewidth=10.)
plt.xlabel('Longitude')
plt.yticks(tickspace, labels, rotation=25)
#plt.ylabel('Pentad')
#plt.title(var_dic[var] + ', $\displaystyle10^{-4}m/s^2$')
plt.tight_layout()
plt.savefig('/scratch/rg419/plots/mom_budg_work/' + inp_fol + '/' +
var + '_30yr.png')
plt.close()
vars = [
'fv_mn_imb', 'fv_av', 'mom_eddy', 'mom_mean', 'mom_sum', 'dphidx_av'
]
for var in vars:
plot_mom_var(var, np.arange(-2., 2.1, 0.2))