def set_vars(data, plottype=None, lonin=[-1.,361.], thresh=0.): # Load up variables to plot, depending on plot type edge_loc, psi_max, psi_max_loc = get_edge_psi(data, thresh=thresh, lev=lev, lonin=lonin, nh=True, sanity_check=True) if plottype == 'mse': data_mse = peak_mse(data, lonin=lonin) vars_out = (data_mse.mse_max_loc, psi_max) elif plottype == 'pcent': data = precip_centroid(data,lonin=lonin) vars_out = (data.p_cent, psi_max) elif plottype == None: vars_out = (edge_loc, psi_max) return vars_out
def set_vars(data, type=None, lonin=[-1., 361.], thresh=0.): edge_loc, psi_max, psi_max_loc = get_edge_psi(data, thresh=thresh, lev=lev, lonin=lonin) if type == 'mse': data_mse = peak_mse(data, lonin=lonin) vars_out = (data_mse.mse_max_loc, psi_max) elif type == 'pcent': data = precip_centroid(data, lonin=lonin) vars_out = (data.p_cent, psi_max) elif type == None: vars_out = (edge_loc, psi_max) return vars_out
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