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.)