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,128,len(years))), [('pentad', range(1,73) ), ('lon', rundata.lon), ('year', years)])
v = 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(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[:,2,37,:].groupby('pentad').mean(('time'))
v[:,:,i] = rundata.vcomp[:,2,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'))
pd_dic = pentad_dic(1)
tickspace = range(20,61,5)
labels = [pd_dic[k] for k in tickspace]
#plot
plt.figure(0)
tot_p_clim.plot.contourf(x='lon', y='pentad',levels=np.arange(3.,31.,3.), add_label = False)
plt.yticks(tickspace,labels,rotation=25)
plt.quiver(u_clim.lon[::3], u_clim.pentad[::2], u_clim[::2,::3], v_clim[::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('Precip and 850hPa wind vectors at 15N, mm/day')
plt.savefig('/scratch/rg419/plots/mom_budg_work/precip_hm3.png')
plt.clf()
def bs_plot(inp_fol, years, lonrange=['all']):
data = load_bs_vars(inp_fol, years, lonrange)
pd_dic = pentad_dic(1)
tickspace = range(10, 80, 10)
labels = [pd_dic[k] for k in tickspace]
data.totp.plot.contourf(x='pentad', y='lat', levels=np.arange(9., 21., 3.))
plt.ylim((-40, 40))
plt.xlim((1, 73))
plt.xlabel('Pentad')
plt.xticks(tickspace, labels, rotation=25)
plt.ylabel('Latitude')
plt.title(inp_fol)
if lonrange[0] == 'all':
plt.savefig('/scratch/rg419/plots/' + inp_fol +
'/rainfall_plot_bs.png')
else:
plt.savefig('/scratch/rg419/plots/' + inp_fol + '/rainfall_plot_bs_' +
str(int(lonrange[0])) + '_' + str(int(lonrange[1])) +
'.png')
plt.clf()
def itcz_plot(inp_fols, years, lonrange=['all']):
for inp_fol in inp_fols:
itcz_lat = itcz_lat_fn(inp_fol, years, lonrange)
itcz_lat.plot.line() #'x',markersize=3)
pd_dic = pentad_dic(1)
tickspace = range(10, 80, 10)
labels = [pd_dic[k] for k in tickspace]
plt.ylim((-60, 60))
plt.xlim((1, 73))
plt.xlabel('Pentad')
plt.ylabel('Peak SST latitude')
plt.xticks(tickspace, labels, rotation=25)
plt.legend(inp_fols, loc='lower right')
if lonrange[0] == 'all':
plt.savefig('/scratch/rg419/plots/itcz_lat.png')
else:
plt.savefig('/scratch/rg419/plots/itcz_lat_' + str(int(lonrange[0])) +
'_' + str(int(lonrange[1])) + '.png')
plt.clf()
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()
#data_p = time_means('q_flux_test', [121,481], filename='plev_pentad', timeav='pentad')
data = xr.open_dataset('q_flux_test.nc', decode_times=False)
#data['totp'] = (('xofyear','lat','lon'), (data_p.convection_rain + data_p.condensation_rain)*86400.)
#data.totp.load()
#data_p.close()
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', data.lat),
('lon', data.lon)])
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
mn_dic = month_dic(1)
pd_dic = pentad_dic(1)
print 'ready to start!'
def u_gradients(data):
data['dudx'] = (('xofyear', 'pfull', 'lat', 'lon'),
cfd(data.ucomp.values, data.lon * np.pi / 180., 3))
data['dudx'] = (('xofyear', 'pfull', 'lat', 'lon'), data.dudx / coslat / a)
data['dudy'] = (('xofyear', 'pfull', 'lat', 'lon'),
cfd((data.ucomp * coslat).values, data.lat * np.pi / 180.,
2))
data['dudy'] = (('xofyear', 'pfull', 'lat', 'lon'), data.dudy / coslat / a)
def plot_pd_mom_fn(inp_fol):
data_file = '/scratch/rg419/plots/monsoon_analysis/' + inp_fol + '/mom_stat_data.nc'
data = xr.open_dataset(data_file)
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', data.lat),
('lon', data.lon)])
pd_dic = pentad_dic(1)
#mom_sum = (data.dphidx*9.8 + data.fv - data.mom_trans - data.mom_stat - data.mom_mean)*10000.
#mom_sum[60,9,:,:].plot.contourf(x='lon', y='lat',levels=np.arange(-2.,2.1,0.2))
#plt.show()
def plot_mom_var(var, levels, qscale):
if var == 'mom_sum':
plot_data = data.dphidx * 9.8 + data.fv - data.mom_trans - data.mom_stat - data.mom_mean
else:
plot_data = data.data_vars[var]
if var == 'dphidx':
plot_data = plot_data * 9.8
if var == 'fv':
plot_data = plot_data - data.mom_mean
plot_data = plot_data * 10000.
#print plot_data.values.shape
for mnth in range(0, 12):
if var == 'fv':
g = plot_data[:, mnth * 6:mnth * 6 + 6,
9, :].plot.contourf(x='lon',
y='lat',
levels=levels,
col='xofyear',
col_wrap=3,
extend='both')
else:
g = plot_data[mnth * 6:mnth * 6 + 6,
9, :, :].plot.contourf(x='lon',
y='lat',
levels=levels,
col='xofyear',
col_wrap=3,
extend='both')
plt.ylim(-30, 60)
plt.xlim(60, 180)
for i, ax in enumerate(g.axes.flat):
land_plot.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='k',
add_colorbar=False,
add_labels=False,
ax=ax)
#ax.quiver(data.lon[::3], data.lat[::1], data.u_av[mnth*6 +i,::1,::3], data.v_av[mnth*6+i,::1,::3], scale=qscale,headwidth=5)
ax.set_title(pd_dic[mnth * 6 + i + 1])
plt.savefig('/scratch/rg419/plots/mom_budg_stat/' + inp_fol + '/' +
var + str(mnth + 1) + '.png')
plt.close()
vars = ['fv', 'dphidx', 'mom_trans', 'mom_stat', 'mom_sum']
for var in vars:
plot_mom_var(var, np.arange(-2., 2.1, 0.2), 500.)