def ap_monsoon_hm(run, title): #Produce identical plot to above, but averaged over all latitudes for a uniform mld run data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') mn_dic = month_dic(1) tickspace = [1, 19, 37, 55] labels = ['1st Jan', '1st Apr', '1st Jul', '1st Oct'] # Set figure parameters rcParams['figure.figsize'] = 5, 3.5 rcParams['font.size'] = 14 fig, ax1 = plt.subplots() try: f1 = (data.precipitation * 86400.).mean('lon').plot.contourf( ax=ax1, x='xofyear', y='lat', levels=np.arange(0., 21., 2.), add_labels=False, extend='max', cmap='Blues', add_colorbar=False) except: f1 = ((data.condensation_rain + data.convection_rain) * 86400.).mean('lon').plot.contourf(ax=ax1, x='xofyear', y='lat', levels=np.arange(0., 21., 2.), add_labels=False, extend='max', cmap='Blues', add_colorbar=False) ax1.set_ylim(-45, 45) ax1.grid(True, linestyle=':') ax1.set_yticks(np.arange(-30., 31., 30.)) ax1.set_title(title) ax1.set_ylabel('Latitude') ax1.set_xticks(tickspace) ax1.set_xticklabels(labels, rotation=25) plt.subplots_adjust(left=0.18, right=0.97, top=0.9, bottom=0.2, hspace=0.3, wspace=0.2) #cb1=fig.colorbar(f1, ax=ax1, use_gridspec=True, orientation = 'horizontal',fraction=0.05, pad=0.15, aspect=60, shrink=0.5) #cb1.set_label('Precipitation, mm/day') # Save as a pdf plt.savefig(plot_dir + 'precip_monsoons_' + run + '.pdf', format='pdf') plt.close() data.close()
def overturning_plot(data, ax_in, lonin=[-1., 361.], do_xlabels=False, month_labels=True, thresh=0., nh=False): lons = pick_lons(data, lonin) psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons) psi /= 1.e9 edge_loc, psi_max, psi_max_loc = get_edge_psi(data, lonin=lonin, lev=500., thresh=thresh, nh=nh) f1 = psi.sel(pfull=500.).plot.contourf(ax=ax_in, x='xofyear', y='lat', levels=np.arange(-500., 510., 50.), add_colorbar=False, add_labels=False, extend='both') edge_loc.plot(color='k', ax=ax_in) psi_max_loc.plot(color='k', ax=ax_in) ax_in.set_ylabel('Latitude') ax_in.set_ylim(-60, 60) ax_in.set_yticks(np.arange(-60., 61., 30.)) ax_in.grid(True, linestyle=':') if month_labels: mn_dic = month_dic(1) tickspace = list(range(13, 72, 18)) labels = [mn_dic[(k + 5) / 6] for k in tickspace] ax_in.set_xlim((1, 72)) ax_in.set_xticks(tickspace) if do_xlabels: ax_in.set_xlabel('') ax_in.set_xticklabels(labels, rotation=25) return f1
def isca_monsoon_hm(run): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') mn_dic = month_dic(1) tickspace = [1, 19, 37, 55] labels = ['1st Jan', '1st Apr', '1st Jul', '1st Oct'] # Set figure parameters rcParams['figure.figsize'] = 15, 6 rcParams['font.size'] = 14 # Start figure with 8 subplots fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8)) = plt.subplots(2, 4, sharex='col', sharey='row') axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8] def precip_plot(data, ax, lons, title, levels=np.arange(0., 21., 2.)): lons = pick_lons(data, lons) f1 = (data.precipitation * 86400.).sel( lon=lons).mean('lon').plot.contourf(ax=ax, x='xofyear', y='lat', levels=levels, add_labels=False, extend='max', cmap='Blues', add_colorbar=False) ax.set_ylim(-45, 45) ax.grid(True, linestyle=':') ax.set_yticks(np.arange(-30., 31., 30.)) ax.set_title(title) return f1 f1 = precip_plot(data, ax1, [0., 10.], '0-10 E') precip_plot(data, ax2, [80., 100.], '80-100 E') #precip_plot(data, ax3, [125.,145.], '125-145 E') precip_plot(data, ax3, [170., 180.], '170-180 E') precip_plot(data, ax4, [180., 190.], '180-190 E') #precip_plot(data, ax5, [215.,235.], '215-235 E') precip_plot(data, ax5, [260., 280.], '260-280 E') #precip_plot(data, ax7, [305.,325.], '305-325 E') precip_plot(data, ax6, [350., 0.], '350-0 E') ax1.set_ylabel('Latitude') ax5.set_ylabel('Latitude') for ax in axes: ax.set_xticks(tickspace) ax.set_xticklabels(labels, rotation=25) plt.subplots_adjust(left=0.07, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.2) cb1 = fig.colorbar(f1, ax=axes, use_gridspec=True, orientation='horizontal', fraction=0.05, pad=0.15, aspect=60, shrink=0.5) cb1.set_label('Precipitation, mm/day') # Save as a pdf plt.savefig(plot_dir + 'precip_monsoons_' + run + '_diffsplit.pdf', format='pdf') plt.close() data.close()
def precip_mse_plot(data, ax_in, lonin=[-1., 361.], do_xlabels=False, plot_type=None, precip_contour=8., p_cent=True, mse_max=True, month_labels=True, lat_bound=45.): lons = pick_lons(data, lonin) try: precip_plot = (data.precipitation * 86400.).sel(lon=lons).mean('lon') except: precip_plot = ((data.convection_rain + data.condensation_rain) * 86400.).sel(lon=lons).mean('lon') mse_plot = (data.temp * cp + data.sphum * L + data.height * g).mean('lon') / 1000. if plot_type == None: # Default case, plot precip with mse overlaid. Choice of whether or not to highlight a specific precip contour f1 = precip_plot.plot.contourf(ax=ax_in, x='xofyear', y='lat', levels=np.arange(2., 15., 2.), add_colorbar=False, add_labels=False, extend='max', cmap='Blues') if not precip_contour == None: precip_plot.plot.contour(ax=ax_in, x='xofyear', y='lat', levels=np.arange(precip_contour - 100., 200., 100.), add_labels=False, add_colorbar=False, colors='k', linewidth=2) cs = mse_plot.sel(pfull=850.).plot.contour(ax=ax_in, x='xofyear', y='lat', levels=np.arange( 200., 401., 10.), add_labels=False, colors='0.7', add_colorbar=False, linewidths=2) plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f') if p_cent: data = precip_centroid(data, lonin=lonin, lat_bound=lat_bound) data.p_cent.plot.line(color='w', ax=ax_in) ax_in.set_xlabel('') elif plot_type == 'precip': # No mse, plot precip and precip centroid f1 = precip_plot.plot.contourf(ax=ax_in, x='xofyear', y='lat', levels=np.arange(3., 16., 3.), add_colorbar=False, add_labels=False, extend='max', cmap='Blues', linewidth=2) if p_cent: data = precip_centroid(data, lonin=lonin, lat_bound=lat_bound) data.p_cent.plot.line(color='k', ax=ax_in, linewidth=2) ax_in.set_xlabel('') elif plot_type == 'mse': # Plot mse in colour, overplot max mse f1 = mse_plot.sel(pfull=850.).plot.contourf(ax=ax_in, x='xofyear', y='lat', levels=np.arange( 200., 401., 10.), add_labels=False, extend='both', add_colorbar=False) if mse_max: data_mse = peak_mse(data, lonin=lonin) data_mse.mse_max_loc.plot.line('k', ax=ax_in) ax_in.set_xlabel('') ax_in.set_ylabel('Latitude') ax_in.set_ylim(-60, 60) ax_in.set_yticks(np.arange(-60., 61., 30.)) ax_in.grid(True, linestyle=':') if month_labels: mn_dic = month_dic(1) tickspace = range(13, 72, 18) labels = [mn_dic[(k + 5) / 6] for k in tickspace] ax_in.set_xlim((1, 72)) ax_in.set_xticks(tickspace) if do_xlabels: ax_in.set_xlabel('') ax_in.set_xticklabels(labels, rotation=25) return f1
def vort_eq_hm(run, lev=150., lonin=[-1.,361.], planetary_only=False, month_labels=True, rot_fac=1., no_eddies=False, add_psi=True): '''Plot vorticity budget Hovmoller. RG 3/11/2017 Imputs: run = run_name lev = level to plot lonin = longitude range to average over planetary_only = only plot planetary vorticity terms month_labels = label x axis with month of year (no good for non earth year length) rot_fac = scale factor for Earth's rotation rate no_eddies = don't plot transient eddies too''' # Call the above function to get fields to plot. For a Hovmoller, run is by definition not steady state. Pass other inputs on. ds = vort_budg_terms(run, lonin=lonin, rot_fac=rot_fac, planetary_only=planetary_only, no_eddies=no_eddies) if add_psi: data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') if lonin[1]>lonin[0]: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]] else: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]] psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons) psi /= 1.e9 # Determine figure size based on number of panels if planetary_only or no_eddies: rcParams['figure.figsize'] = 15, 6.25 else: rcParams['figure.figsize'] = 15, 8 rcParams['font.size'] = 18 rcParams['text.usetex'] = True levels = np.arange(-1.5,1.6,0.25) print('starting plotting') # Set number of panels and declare which are the left column, bottom row. if planetary_only or no_eddies: f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row') left_column = [ax1,ax4] bottom_row = [ax4,ax5,ax6] all_plots = [ax1,ax2,ax3,ax4,ax5,ax6] else: f, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row') left_column = [ax1,ax4,ax7] bottom_row = [ax7,ax8,ax9] all_plots = [ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9] plt.set_cmap('RdBu_r') # Select the pressure level, if needed if run in ['ap_2','full_qflux']: ds = ds else: ds = ds.sel(pfull=lev) # Plot! f2 = ds.horiz_md.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax1, extend = 'both', add_labels=False) if add_psi: ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax1.set_title('Horizontal advection', fontsize=17) ax1.text(-15, 60, 'a)') ds.dvordy.plot.contourf(x='xofyear', y='lat', levels=np.arange(-1.e-10,1.05e-10,0.1e-10), ax=ax2, extend = 'both', add_labels=False) if planetary_only: ax2.set_title('$\partial f /\partial y$', fontsize=17) else: ax2.set_title('$\partial (\overline{\zeta} + f) /\partial y$', fontsize=17) ax2.text(-7, 60, 'b)') ds.v.plot.contourf(x='xofyear', y='lat', levels=np.arange(-12.,13,2.), ax=ax3, extend = 'both', add_labels=False) ax3.set_title('$\overline{v}$', fontsize=17) ax3.text(-7, 60, 'c)') ds.stretching.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax4, extend = 'both', add_labels=False) if add_psi: ax4.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax4.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax4.set_title('Vortex stretching', fontsize=17) ax4.text(-15, 60, 'd)') ds.div.plot.contourf(x='xofyear', y='lat', levels=np.arange(-0.6,0.7,0.1), ax=ax5, extend = 'both', add_labels=False) ax5.set_title('Divergence', fontsize=17) ax5.text(-7, 60, 'e)') ds.vor.plot.contourf(x='xofyear', y='lat', levels=np.arange(-12.,13.,2.), ax=ax6, extend = 'both', add_labels=False) if planetary_only: ax6.set_title('Planetary vorticity', fontsize=17) else: ax6.set_title('Absolute vorticity', fontsize=17) ax6.text(-7, 60, 'f)') # Plot transients if needed if not (planetary_only or no_eddies): ds.transient_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax7, extend = 'both', add_labels=False) if add_psi: ax7.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax7.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax7.set_title('Transient eddy vorticity tendency', fontsize=17) ax7.text(-15, 60, 'g)') ds.transient_h_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax8, extend = 'both', add_labels=False) ax8.set_title('Transient horizontal adv.', fontsize=17) ax8.text(-7, 60, 'h)') ds.transient_s_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax9, extend = 'both', add_labels=False) ax9.set_title('Transient vortex stretching', fontsize=17) ax9.text(-7, 60, 'i)') # Add grid, set y limits and ticks for ax in all_plots: ax.grid(True,linestyle=':') ax.set_ylim(-60,60) ax.set_yticks(np.arange(-60,61,30)) # Label left column for ax in left_column: ax.set_ylabel('Latitude') # If month labels are being used, load in the month dictionary and set up labelling, then label bottom row if month_labels: mn_dic = month_dic(1) tickspace = list(range(13,72,18)) labels = [mn_dic[(k+5)/6 ] for k in tickspace] for ax in bottom_row: ax.set_xticks(tickspace) ax.set_xticklabels(labels,rotation=25) plt.subplots_adjust(right=0.97, left=0.08, top=0.93, bottom=0.1, hspace=0.25, wspace=0.15) # Set plot name based on type of plot. if lonin == [-1.,361.]: lon_tag = '' else: lon_tag = '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) if planetary_only: f_tag = '_fonly' else: f_tag = '' figname = 'vort_breakdown_' + run + lon_tag + f_tag + '.pdf' plt.savefig(plot_dir + figname, format='pdf') plt.close()
def vort_eq_hm(run, lev=150., lonin=[-1.,361.], month_labels=True, rot_fac=1., add_psi=True): '''Plot vorticity budget Hovmoller. RG 3/11/2017 Imputs: run = run_name lev = level to plot lonin = longitude range to average over planetary_only = only plot planetary vorticity terms month_labels = label x axis with month of year (no good for non earth year length) rot_fac = scale factor for Earth's rotation rate no_eddies = don't plot transient eddies too''' # Call the above function to get fields to plot. For a Hovmoller, run is by definition not steady state. Pass other inputs on. ds = vort_budg_terms(run, lonin=lonin, rot_fac=rot_fac) if add_psi: data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') if lonin[1]>lonin[0]: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]] else: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]] psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons) psi /= 1.e9 # Determine figure size based on number of panels rcParams['figure.figsize'] = 15, 4 rcParams['font.size'] = 14 rcParams['text.usetex'] = True levels = np.arange(-1.5,1.6,0.25) print('starting plotting') # Set number of panels and declare which are the left column, bottom row. f, (ax1, ax2, ax3) = plt.subplots(1, 3, sharey='row') left_column = [ax1] all_plots = [ax1,ax2,ax3] plt.set_cmap('RdBu_r') # Select the pressure level, if needed if run in ['ap_2','full_qflux']: ds = ds else: ds = ds.sel(pfull=lev) # Plot! f2 = ds.div_uvor.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax1, extend = 'both', add_labels=False) if add_psi: ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax1.set_title('Relative vorticty part', fontsize=17) ax1.text(-15, 60, 'a)') ds.div_uf.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax2, extend = 'both', add_labels=False) if add_psi: ax2.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax2.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax2.set_title('Planetary vorticity part', fontsize=17) ax2.text(-7, 60, 'b)') ds.transient_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax3, extend = 'both', add_labels=False) if add_psi: ax3.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax3.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2) ax3.set_title('Transient eddy vorticity tendency', fontsize=17) ax3.text(-7, 60, 'c)') # Add grid, set y limits and ticks for ax in all_plots: ax.grid(True,linestyle=':') ax.set_ylim(-60,60) ax.set_yticks(np.arange(-60,61,30)) if month_labels: mn_dic = month_dic(1) tickspace = list(range(13,72,18)) labels = [mn_dic[(k+5)/6 ] for k in tickspace] ax.set_xticks(tickspace) ax.set_xticklabels(labels,rotation=25) # Label left column for ax in left_column: ax.set_ylabel('Latitude') plt.subplots_adjust(right=0.97, left=0.08, top=0.93, bottom=0.1, hspace=0.25, wspace=0.15) # Set plot name based on type of plot. if lonin == [-1.,361.]: lon_tag = '' else: lon_tag = '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) figname = 'alt_breakdown_' + run + lon_tag + '.pdf' plt.savefig(plot_dir + figname, format='pdf') plt.close()
def mom_budg_hm(run, lev=150, lonin=[-1., 361.]): #NB run is a dummy variable rcParams['figure.figsize'] = 12, 8 rcParams['font.size'] = 18 rcParams['text.usetex'] = True plot_dir = '/scratch/rg419/plots/paper_1_figs/revisions/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/era_mom_vars.nc') # Take pentad means data.coords['pentad'] = data.day_of_yr // 5 + 1. data = data.groupby('pentad').mean(('day_of_yr')) if lonin[1] > lonin[0]: lons = [ data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1] ] else: lons = [ data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1] ] #advective terms partition_advection(data, lons, lev=lev) #Coriolis omega = 7.2921150e-5 f = 2 * omega * np.sin(data.lat * np.pi / 180) fv = data.vcomp * f * 86400. fv_mean = fv.mean('lon') fv_local = (fv - fv_mean).sel(lon=lons).mean('lon') #Geopotential gradient dphidx = gr.ddx(data.phi) dphidx = -86400. * dphidx.sel(lon=lons).mean('lon') 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_cross1 = data.u_dudx_cross1 + data.v_dudy_cross1 + data.w_dudp_cross1 mom_cross2 = 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_crossu = data.u_dudx_cross1 + data.u_dudx_cross2 mom_crossv = data.v_dudy_cross1 + data.v_dudy_cross2 mom_crossw = data.w_dudp_cross1 + data.w_dudp_cross2 #data.w_dudp_cross1.plot.contourf(x='pentad', y='lat', extend='both',levels = np.arange(-20,21.1,2.)) #plt.yticks(np.arange(-60.,61.,30.)) #plt.ylim(-60,60) #plt.figure(2) #data.w_dudp_cross2.plot.contourf(x='pentad', y='lat', extend='both',levels = np.arange(-20,21.1,2.)) #plt.ylim(-60,60) #plt.yticks(np.arange(-60.,61.,30.)) #plt.show() #print mom_stat mom_sum = fv_local + fv_mean + dphidx + mom_mean + mom_trans + mom_stat + mom_cross levels = np.arange(-20, 21.1, 2.) mn_dic = month_dic(1) tickspace = list(range(13, 72, 18)) labels = [mn_dic[(k + 5) / 6] for k in tickspace] # Six subplots fig, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row') plt.set_cmap('RdBu_r') #First plot f1 = fv_mean.plot.contourf(ax=ax1, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax1.set_ylabel('Latitude') ax1.set_title('Zonal mean Coriolis', fontsize=17) ax1.set_ylim(-60, 60) ax1.grid(True, linestyle=':') ax1.set_yticks(np.arange(-60., 61., 30.)) ax1.text(-15, 60, 'a)') #Second plot mom_mean.plot.contourf(ax=ax2, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax2.grid(True, linestyle=':') ax2.set_title('Mean state advection', fontsize=17) ax2.set_ylim(-60, 60) ax2.text(-5, 60, 'b)') #Third plot mom_sum.plot.contourf(ax=ax3, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax3.grid(True, linestyle=':') ax3.set_ylim(-60, 60) ax3.set_title('Residual', fontsize=17) ax3.text(-5, 60, 'c)') #Fourth plot mom_trans.plot.contourf(ax=ax4, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax4.grid(True, linestyle=':') ax4.set_ylabel('Latitude') ax4.set_ylim(-60, 60) ax4.set_title('Transient eddy flux conv.', fontsize=17) ax4.set_yticks(np.arange(-60., 61., 30.)) ax4.text(-15, 60, 'd)') #Fifth plot mom_stat.plot.contourf(ax=ax5, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax5.grid(True, linestyle=':') ax5.set_title('Stat. eddy flux conv.', fontsize=17) ax5.set_ylim(-60, 60) ax5.text(-5, 60, 'e)') #Sixth plot mom_cross.plot.contourf(ax=ax6, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax6.grid(True, linestyle=':') ax6.set_ylim(-60, 60) ax6.set_title('Stat. eddy cross terms', fontsize=17) ax6.text(-5, 60, 'f)') #Fourth plot dphidx.plot.contourf(ax=ax7, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax7.grid(True, linestyle=':') ax7.set_ylabel('Latitude') ax7.set_xticks(tickspace) ax7.set_xticklabels(labels, rotation=25) ax7.set_ylim(-60, 60) ax7.set_title('Geopotential gradient', fontsize=17) ax7.set_yticks(np.arange(-60., 61., 30.)) ax7.text(-15, 60, 'g)') #Fifth plot fv_local.plot.contourf(ax=ax8, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax8.grid(True, linestyle=':') ax8.set_xticks(tickspace) ax8.set_xticklabels(labels, rotation=25) ax8.set_title('Stat. eddy Coriolis', fontsize=17) ax8.set_ylim(-60, 60) ax8.text(-5, 60, 'h)') #Sixth plot fv_ageo.plot.contourf(ax=ax9, x='pentad', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False) ax9.grid(True, linestyle=':') ax9.set_xticks(tickspace) ax9.set_xticklabels(labels, rotation=25) ax9.set_ylim(-60, 60) ax9.set_title('Ageostrophic stat. eddy Coriolis', fontsize=17) ax9.text(-5, 60, 'i)') plt.subplots_adjust(right=0.97, left=0.1, top=0.95, bottom=0., hspace=0.25, wspace=0.12) #Colorbar cb1 = fig.colorbar(f1, ax=[ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9], use_gridspec=True, orientation='horizontal', fraction=0.15, pad=0.15, aspect=30, shrink=0.5) #cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True,fraction=0.15, aspect=30) #cb1.set_label('$ms^{-1}day^{-1}$') if lonin == [-1., 361.]: figname = 'zon_mom_budg_era.pdf' else: figname = 'zon_mom_budg_era_' + str(int(lonin[0])) + '_' + str( int(lonin[1])) + '.pdf' plt.savefig(plot_dir + figname, format='pdf') plt.close()
def mom_budg_hm(run, lev=150, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.]): a = 6376.0e3 rcParams['figure.figsize'] = 12, 5.5 rcParams['font.size'] = 18 rcParams['text.usetex'] = True plot_dir = '/scratch/rg419/plots/itcz_theory_testing/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/'+run+'.nc') if lonin[1]>lonin[0]: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]] else: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]] #advective terms partition_advection(data, lons, lev=150) #Coriolis omega = 7.2921150e-5 f = 2 * omega * np.sin(data.lat *np.pi/180.) fu = -1.*data.ucomp.sel(pfull=lev) * f * 86400. fu = fu.sel(lon=lons).mean('lon') metric_term = (-86400. * data.ucomp.sel(pfull=lev)*data.ucomp.sel(pfull=lev)*np.tan(data.lat*np.pi/180.)/a).mean('lon') #totp = ((data.convection_rain + data.condensation_rain)*86400.).sel(lon=lons).mean('lon') #abs_vort = (data.vor + f).sel(lon=lons).mean('lon')*86400. #Geopotential gradient dphidy = gr.ddy(data.height.sel(pfull=lev), vector=False) dphidy = -86400. * 9.8 * dphidy.sel(lon=lons).mean('lon') mom_mean = data.v_dvdy_zav + data.w_dvdp_zav mom_trans = data.vv_trans_dy + data.vw_trans_dp mom_stat = data.v_dvdy_stat + data.w_dvdp_stat mom_sum = mom_mean + mom_trans + mom_stat + fu + dphidy + metric_term levels = np.arange(-300.,300.1,30.) mn_dic = month_dic(1) tickspace = list(range(13,72,18)) labels = [mn_dic[(k+5)/6 ] for k in tickspace] # Six subplots fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row') plt.set_cmap('RdBu_r') #First plot f1=fu.plot.contourf(ax=ax1, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) #totp.plot.contour(ax=ax1, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax1.set_ylabel('Latitude') ax1.set_title('Coriolis', fontsize=17) ax1.set_ylim(-60,60) ax1.grid(True,linestyle=':') ax1.set_yticks(np.arange(-60.,61.,30.)) ax1.text(-15, 60, 'a)') #Second plot mom_mean.plot.contourf(ax=ax2, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) #ax2.contour(data.xofyear, data.lat, abs_vort.sel(pfull=lev).T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2) #abs_vort.sel(pfull=lev).plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-2.,6.,2.), add_colorbar=False, add_labels=False, colors='k', linewidths=2) #totp.plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax2.grid(True,linestyle=':') ax2.set_title('Mean state advection', fontsize=17) ax2.set_ylim(-60,60) ax2.text(-5, 60, 'b)') #Third plot dphidy.plot.contourf(ax=ax3, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) #totp.plot.contour(ax=ax3, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax3.grid(True,linestyle=':') ax3.set_ylim(-60,60) ax3.set_title('Geopotential gradient', fontsize=17) ax3.text(-5, 60, 'c)') #Fourth plot mom_trans.plot.contourf(ax=ax4, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) #totp.plot.contour(ax=ax4, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax4.grid(True,linestyle=':') ax4.set_ylabel('Latitude') ax4.set_xticks(tickspace) ax4.set_xticklabels(labels,rotation=25) ax4.set_ylim(-60,60) ax4.set_title('Transient eddy flux conv.', fontsize=17) ax4.set_yticks(np.arange(-60.,61.,30.)) ax4.text(-15, 60, 'd)') #Fifth plot metric_term.plot.contourf(ax=ax5, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) #totp.plot.contour(ax=ax5, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax5.grid(True,linestyle=':') ax5.set_xticks(tickspace) ax5.set_xticklabels(labels,rotation=25) ax5.set_title('Metric term', fontsize=17) ax5.set_ylim(-60,60) ax5.text(-5, 60, 'e)') #Sixth plot mom_sum.plot.contourf(ax=ax6, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) #totp.plot.contour(ax=ax6, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax6.grid(True,linestyle=':') ax6.set_xticks(tickspace) ax6.set_xticklabels(labels,rotation=25) ax6.set_ylim(-60,60) ax6.set_title('Residual', fontsize=17) ax6.text(-5, 60, 'f)') plt.subplots_adjust(right=0.97, left=0.1, top=0.95, bottom=0., hspace=0.25, wspace=0.12) #Colorbar cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.15, aspect=30, shrink=0.5) #cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True,fraction=0.15, aspect=30) #cb1.set_label('$ms^{-1}day^{-1}$') if lonin == [-1.,361.]: figname = 'merid_mom_budg_' +run+ '.pdf' else: figname = 'merid_mom_budg_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf' plt.savefig(plot_dir + figname, format='pdf') plt.close()
def mom_budg_hm(run, lev=150, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.]): rcParams['figure.figsize'] = 12, 8 rcParams['font.size'] = 18 rcParams['text.usetex'] = True plot_dir = '/scratch/rg419/plots/paper_1_figs/revisions/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/'+run+'.nc') if lonin[1]>lonin[0]: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]] else: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]] #advective terms partition_advection(data, lons, lev=150) #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).sel(lon=lons).mean('lon') try: totp = ((data.precipitation)*86400.).sel(lon=lons).mean('lon') except: totp = ((data.convection_rain + data.condensation_rain)*86400.).sel(lon=lons).mean('lon') #abs_vort = (data.vor + f).sel(lon=lons).mean('lon')*86400. #Geopotential gradient dphidx = gr.ddx(data.height.sel(pfull=lev)) dphidx = -86400. * 9.8 * dphidx.sel(lon=lons).mean('lon') 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_cross1 = data.u_dudx_cross1 + data.v_dudy_cross1 + data.w_dudp_cross1 mom_cross2 = 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_crossu = data.u_dudx_cross1 + data.u_dudx_cross2 mom_crossv = data.v_dudy_cross1 + data.v_dudy_cross2 mom_crossw = data.w_dudp_cross1 + data.w_dudp_cross2 #print mom_stat mom_sum = fv_local + fv_mean + dphidx + mom_mean + mom_trans + mom_stat + mom_cross levels = np.arange(-20,21.1,2.) mn_dic = month_dic(1) tickspace = range(13,72,18) labels = [mn_dic[(k+5)/6 ] for k in tickspace] # Six subplots fig, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row') plt.set_cmap('RdBu_r') #First plot f1=fv_mean.plot.contourf(ax=ax1, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) totp.plot.contour(ax=ax1, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax1.set_ylabel('Latitude') ax1.set_title('Zonal mean Coriolis', fontsize=17) ax1.set_ylim(-60,60) ax1.grid(True,linestyle=':') ax1.set_yticks(np.arange(-60.,61.,30.)) ax1.text(-15, 60, 'a)') #Second plot mom_mean.plot.contourf(ax=ax2, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) #ax2.contour(data.xofyear, data.lat, abs_vort.sel(pfull=lev).T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2) #abs_vort.sel(pfull=lev).plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-2.,6.,2.), add_colorbar=False, add_labels=False, colors='k', linewidths=2) totp.plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax2.grid(True,linestyle=':') ax2.set_title('Mean state advection', fontsize=17) ax2.set_ylim(-60,60) ax2.text(-5, 60, 'b)') #Third plot mom_sum.plot.contourf(ax=ax3, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) totp.plot.contour(ax=ax3, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax3.grid(True,linestyle=':') ax3.set_ylim(-60,60) ax3.set_title('Residual', fontsize=17) ax3.text(-5, 60, 'c)') #Fourth plot mom_trans.plot.contourf(ax=ax4, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) totp.plot.contour(ax=ax4, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax4.grid(True,linestyle=':') ax4.set_ylabel('Latitude') ax4.set_ylim(-60,60) ax4.set_title('Transient eddy flux conv.', fontsize=17) ax4.set_yticks(np.arange(-60.,61.,30.)) ax4.text(-15, 60, 'd)') #Fifth plot mom_stat.plot.contourf(ax=ax5, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) totp.plot.contour(ax=ax5, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax5.grid(True,linestyle=':') ax5.set_title('Stat. eddy flux conv.', fontsize=17) ax5.set_ylim(-60,60) ax5.text(-5, 60, 'e)') #Sixth plot mom_cross.plot.contourf(ax=ax6, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) totp.plot.contour(ax=ax6, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax6.grid(True,linestyle=':') ax6.set_ylim(-60,60) ax6.set_title('Stat. eddy cross terms', fontsize=17) ax6.text(-5, 60, 'f)') #Fourth plot dphidx.plot.contourf(ax=ax7, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) totp.plot.contour(ax=ax7, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax7.grid(True,linestyle=':') ax7.set_ylabel('Latitude') ax7.set_xticks(tickspace) ax7.set_xticklabels(labels,rotation=25) ax7.set_ylim(-60,60) ax7.set_title('Geopotential gradient', fontsize=17) ax7.set_yticks(np.arange(-60.,61.,30.)) ax7.text(-15, 60, 'g)') #Fifth plot fv_local.plot.contourf(ax=ax8, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) totp.plot.contour(ax=ax8, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax8.grid(True,linestyle=':') ax8.set_xticks(tickspace) ax8.set_xticklabels(labels,rotation=25) ax8.set_title('Stat. eddy Coriolis', fontsize=17) ax8.set_ylim(-60,60) ax8.text(-5, 60, 'h)') #Sixth plot fv_ageo.plot.contourf(ax=ax9, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) totp.plot.contour(ax=ax9, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) ax9.grid(True,linestyle=':') ax9.set_xticks(tickspace) ax9.set_xticklabels(labels,rotation=25) ax9.set_ylim(-60,60) ax9.set_title('Ageostrophic stat. eddy Coriolis', fontsize=17) ax9.text(-5, 60, 'i)') plt.subplots_adjust(right=0.97, left=0.1, top=0.95, bottom=0., hspace=0.25, wspace=0.12) #Colorbar cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9], use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.15, aspect=30, shrink=0.5) #cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True,fraction=0.15, aspect=30) #cb1.set_label('$ms^{-1}day^{-1}$') if lonin == [-1.,361.]: figname = 'zon_mom_budg_' +run+ '_rev.pdf' else: figname = 'zon_mom_budg_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '_rev.pdf' plt.savefig(plot_dir + figname, format='pdf') plt.close()
plot_dir = '/scratch/rg419/plots/paper_1_figs/revisions/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/gpcp_detrendedpentads.nc') #lons = pick_lons(data, [60.,150.]) #'Asia' #lons = pick_lons(data, [15.,60.]) # South Africa lons = pick_lons(data, [100., 150.]) # East Asia lons = pick_lons(data, [70., 100.]) # India data['totp'] = data.precip_clim.sel(lon=lons).mean('lon') mn_dic = month_dic(1) tickspace = [13, 31, 49, 68] labels = ['Mar', 'Jun', 'Sep', 'Dec'] plevels = np.arange(2., 15., 2.) # Begin plotting: 3 subplots #First plot f1 = data.totp.plot.contourf(x='xofyear', y='lat', levels=plevels, add_labels=False, extend='max', cmap='Blues') #data.totp.plot.contour(x='xofyear', y='lat',levels=np.arange(-92.,109.,100.), add_labels = False, add_colorbar=False, colors='k')
def heat_budg_hm(run, lev=850, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.], plot_precip=True, do_theta=False): rcParams['figure.figsize'] = 12, 8 rcParams['font.size'] = 18 rcParams['text.usetex'] = True plot_dir = '/scratch/rg419/plots/heat_budg/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/'+run+'.nc') if lonin[1]>lonin[0]: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]] else: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]] #advective terms partition_advection(data, lons, lev=lev) try: precip_centroid(data) except: data['precipitation'] = data.convection_rain + data.condensation_rain precip_centroid(data) psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 heat_mean = data.u_dTdx_zav #+ data.w_dTdp_zav #+data.v_dTdy_zav # heat_cross = data.u_dTdx_cross1 + data.v_dTdy_cross1 + data.w_dTdp_cross1 + data.u_dTdx_cross2 + data.v_dTdy_cross2 + data.w_dTdp_cross2 heat_cross1 = data.u_dTdx_cross1 + data.v_dTdy_cross1 + data.w_dTdp_cross1 heat_cross2 = data.u_dTdx_cross2 + data.v_dTdy_cross2 + data.w_dTdp_cross2 heat_trans = data.uT_trans_dx + data.vT_trans_dy + data.wT_trans_dp heat_stat = data.u_dTdx_stat + data.v_dTdy_stat + data.w_dTdp_stat heat_crossu = data.u_dTdx_cross1 + data.u_dTdx_cross2 heat_crossv = data.v_dTdy_cross1 + data.v_dTdy_cross2 heat_crossw = data.w_dTdp_cross1 + data.w_dTdp_cross2 tdt_rad = data.tdt_rad.sel(pfull=lev).sel(lon=lons).mean('lon')*86400. tdt_conv = data.dt_tg_convection.sel(pfull=lev).sel(lon=lons).mean('lon')*86400. tdt_cond = data.dt_tg_condensation.sel(pfull=lev).sel(lon=lons).mean('lon')*86400. tdt_diff = data.dt_tg_diffusion.sel(pfull=lev).sel(lon=lons).mean('lon')*86400. if do_theta: convTtotheta=(1000./data.pfull)**(2./7.) diabatic = (tdt_rad + tdt_conv + tdt_cond + tdt_diff)*convTtotheta.sel(pfull=lev) heat_sum = heat_mean + heat_trans + heat_stat + heat_cross + diabatic asc_cool = heat_sum * 0. Tdt = gr.ddt(data.temp.sel(pfull=lev)).sel(lon=lons).mean('lon') *86400. else: diabatic = tdt_rad + tdt_conv + tdt_cond + tdt_diff rho = data.pfull*100./data.temp/287.04 asc_cool = (data.omega /(1004.64 * rho)).sel(pfull=lev).sel(lon=lons).mean('lon') *86400. heat_sum = heat_mean + heat_trans + heat_stat + heat_cross + diabatic + asc_cool Tdt = gr.ddt(data.temp.sel(pfull=lev)).sel(lon=lons).mean('lon') *86400. #levels = np.arange(-20,21.1,2.) #levels = np.arange(-10,10.1,1.) levels = np.arange(-1,1.1,0.1) mn_dic = month_dic(1) tickspace = list(range(13,72,18)) ticklabels = [mn_dic[(k+5)/6 ] for k in tickspace] # Nine subplots fig, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row') plt.set_cmap('RdBu_r') plot_vars = [heat_mean, heat_trans, heat_sum, asc_cool, diabatic, tdt_rad, tdt_diff, tdt_conv, tdt_cond] axes = [ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9] labels = ['a)','b)','c)','d)','e)','f)','g)','h)','i)'] for i in range(9): f1=plot_vars[i].plot.contourf(ax=axes[i], x='xofyear', y='lat', extend='both', add_labels=False, cmap='RdBu_r', levels = levels, add_colorbar=False) if plot_precip: psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(-500.,0.,100.), add_labels=False, colors='0.7', linewidths=2, linestyles='--') psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(0.,510.,100.), add_labels=False, colors='0.7', linewidths=2) psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(-1000.,1010.,1000.), add_labels=False, colors='0.5', linewidths=2) data.p_cent.plot.line(ax=axes[i],color='k', linewidth=2) axes[i].set_xlabel(' ') axes[i].set_ylabel(' ') #totp.plot.contour(ax=axes[i], x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) axes[i].set_ylim(-60,60) axes[i].grid(True,linestyle=':') axes[i].set_yticks(np.arange(-60.,61.,30.)) axes[i].set_xticks(tickspace) for i in [0,3,6]: axes[i].set_ylabel('Latitude') axes[i].text(-15, 60, labels[i]) for i in [1,2,4,5,7,8]: axes[i].text(-5, 60, labels[i]) for i in [6,7,8]: axes[i].set_xticklabels(ticklabels,rotation=25) # set titles ax1.set_title('Mean state advection', fontsize=17) ax2.set_title('Transient eddy flux conv.', fontsize=17) ax3.set_title('Residual', fontsize=17) ax4.set_title('w cooling', fontsize=17) ax5.set_title('Total diabatic heating', fontsize=17) ax6.set_title('Radiative heating', fontsize=17) ax7.set_title('Diffusive heating', fontsize=17) ax8.set_title('Convective heating', fontsize=17) ax9.set_title('Condensational heating', fontsize=17) plt.subplots_adjust(right=0.97, left=0.1, top=0.95, bottom=0., hspace=0.25, wspace=0.12) #Colorbar cb1=fig.colorbar(f1, ax=axes, use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.15, aspect=30, shrink=0.5) if lonin == [-1.,361.]: figname = 'zon_heat_budg_' +run+ '_u_850.pdf' else: figname = 'zon_heat_budg_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf' plt.savefig(plot_dir + figname, format='pdf') plt.close()
def mom_budg_hm(run, lev=150, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.], plot_precip=True, rot_fac=1.): rcParams['figure.figsize'] = 12, 6 rcParams['font.size'] = 18 rcParams['text.usetex'] = True plot_dir = '/scratch/rg419/plots/mom_budg_merid/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/'+run+'.nc') if lonin[1]>lonin[0]: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]] else: lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]] #advective terms partition_advection(data, lons, lev=150) #Coriolis omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat *np.pi/180) fu = -1. * data.ucomp.sel(pfull=lev) * f * 86400. fu_mean = fu.mean('lon') fu_local = (fu - fu_mean).sel(lon=lons).mean('lon') if plot_precip: try: totp = ((data.precipitation)*86400.).sel(lon=lons).mean('lon') except: totp = ((data.convection_rain + data.condensation_rain)*86400.).sel(lon=lons).mean('lon') #abs_vort = (data.vor + f).sel(lon=lons).mean('lon')*86400. #Geopotential gradient dphidy = gr.ddy(data.height.sel(pfull=lev), vector=False) dphidy = -86400. * 9.8 * dphidy.sel(lon=lons).mean('lon') fu_ageo = fu_local + fu_mean + dphidy dvdt = gr.ddt(data.vcomp.sel(pfull=lev)).sel(lon=lons).mean('lon')*86400. metric = (-86400.* data.ucomp_sq * np.tan(data.lat * np.pi/180.)/6376.0e3).sel(pfull=lev).sel(lon=lons).mean('lon') vert_term = (-86400.*data.vcomp_omega/6376.0e3).sel(pfull=lev).sel(lon=lons).mean('lon') mom_mean = data.u_dvdx_zav + data.v_dvdy_zav + data.w_dvdp_zav mom_cross = data.u_dvdx_cross1 + data.v_dvdy_cross1 + data.w_dvdp_cross1 + data.u_dvdx_cross2 + data.v_dvdy_cross2 + data.w_dvdp_cross2 mom_cross1 = data.u_dvdx_cross1 + data.v_dvdy_cross1 + data.w_dvdp_cross1 mom_cross2 = data.u_dvdx_cross2 + data.v_dvdy_cross2 + data.w_dvdp_cross2 mom_trans = data.uv_trans_dx + data.vv_trans_dy + data.vw_trans_dp mom_stat = data.u_dvdx_stat + data.v_dvdy_stat + data.w_dvdp_stat mom_crossu = data.u_dvdx_cross1 + data.u_dvdx_cross2 mom_crossv = data.v_dvdy_cross1 + data.v_dvdy_cross2 mom_crossw = data.w_dvdp_cross1 + data.w_dvdp_cross2 mom_sum = fu_local + fu_mean + dphidy + mom_mean + mom_trans + mom_stat + mom_cross + metric #levels = np.arange(-40,41.1,4.) levels = np.arange(-300.,300.1,30.) mn_dic = month_dic(1) tickspace = list(range(13,72,18)) ticklabels = [mn_dic[(k+5)/6 ] for k in tickspace] # Nine subplots fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row') plt.set_cmap('RdBu_r') plot_vars = [fu_mean, mom_mean, mom_sum, mom_trans, dphidy, metric] axes = [ax1,ax2,ax3,ax4,ax5,ax6] labels = ['a)','b)','c)','d)','e)','f)'] for i in range(len(plot_vars)): f1=plot_vars[i].plot.contourf(ax=axes[i], x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False) if plot_precip: totp.plot.contour(ax=axes[i], x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2) axes[i].set_ylim(-60,60) axes[i].grid(True,linestyle=':') axes[i].set_yticks(np.arange(-60.,61.,30.)) axes[i].set_xticks(tickspace) for i in [0,3]: axes[i].set_ylabel('Latitude') axes[i].text(-15, 60, labels[i]) for i in [1,2,4,5]: axes[i].text(-5, 60, labels[i]) for i in [3,4,5]: axes[i].set_xticklabels(ticklabels,rotation=25) # set titles ax1.set_title('Zonal mean Coriolis', fontsize=17) ax2.set_title('Mean state advection', fontsize=17) ax3.set_title('Residual', fontsize=17) ax4.set_title('Transient eddy flux conv.', fontsize=17) ax5.set_title('Geopotential gradient', fontsize=17) ax6.set_title('Metric term', fontsize=17) plt.subplots_adjust(right=0.97, left=0.1, top=0.95, bottom=0., hspace=0.25, wspace=0.12) #Colorbar cb1=fig.colorbar(f1, ax=axes, use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.15, aspect=30, shrink=0.5) if lonin == [-1.,361.]: figname = 'merid_mom_budg_' +run+ '.pdf' else: figname = 'merid_mom_budg_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf' plt.savefig(plot_dir + figname, format='pdf') plt.close()