def sf_spinup(run, months_list, filenames=['plev_pentad']): # Function to open files for a specfied month range and filename. def open_files(run, months, filename): name_temp = '/disca/share/rg419/Data_moist/' + run + '/run%04d/' + filename + '.nc' names = [name_temp % m for m in range(months[0], months[1])] data = xr.open_mfdataset(names, decode_times=False, chunks={'time': 30}) # Reduce dataset so that only vcomp is retained data = xr.Dataset( { 'vcomp': data.vcomp, 'precipitation': data.precipitation }, coords=data.vcomp.coords) #data.coords['month'] = data.time//30 + 1 #data = data.groupby('month').mean(('time')) return data arrays = [] i = 0 for filename in filenames: data = open_files(run, months_list[i], filename) arrays.append(data) i = i + 1 data = xr.concat(arrays, dim='time') precip_centroid(data) adj_time = np.min( data.p_cent.time.where(data.p_cent >= 15., drop=True).values) return data.p_cent, adj_time
def p_cent_rate(data, days=False): """ Inputs: data - xarray dataset climatology including precipitation as either precipitation or as convection_rain and condensation_rain days - instructs if data is in daily or pentad means Returns: dpcentdt - rate of movement of the precipitation centroid dpcentdt2 - rate of acceleration of the precipitation centroid """ # Get total precip try: data['precipitation'] = data.condensation_rain + data.convection_rain except: data['precipitation'] = data.precipitation data['precipitation'] = make_sym(data.precipitation) # Locate precipitation centroid precip_centroid(data) # If units in are days rather than pentads (e.g. for shorter years) convert to pentads if days: dpcentdt = gr.ddt(data.p_cent, secperunit=86400.) * 86400. else: dpcentdt = gr.ddt(data.p_cent) * 86400. dpcentdt2 = gr.ddt(dpcentdt, secperunit=86400.) * 86400. return dpcentdt, dpcentdt2
def rate_at_eq(runs, do_make_sym=True, days=None): dpdt_eq = [] if days == None: days = [False] * len(runs) j = 0 for run in runs: data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') if do_make_sym: # If symmetric data is wanted, average over both hemispheres (NB currently only set up for a climatology 2/05/18) data['precipitation'] = make_sym(data.precipitation) # Locate precipitation centroid precip_centroid(data) # Get rate of movement of precip centroid if days[j]: dpcentdt = gr.ddt(data.p_cent, secperunit=86400.) * 86400. else: dpcentdt = gr.ddt(data.p_cent) * 86400. #dpcentdt_max = dpcentdt_pmask.where(dpcentdt_pmask==dpcentdt_pmask.max('xofyear'),drop=True) # Find the maximum rate p_cent = np.abs(data.p_cent.where(dpcentdt >= 0.)) dpdt_eq_j = dpcentdt.where(p_cent == p_cent.min('xofyear'), drop=True) dpdt_eq.append(dpdt_eq_j.values[0]) j = j + 1 return np.asarray(dpdt_eq)
def precip_centroid_rate(run, ax_in, ylim_in=1.): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') precip_centroid(data) dpcentdt = gr.ddt(data.p_cent) * 86400. dpcentdt_ppos = dpcentdt.where(data.p_cent>=0.) dpcentdt_max = dpcentdt_ppos.where(dpcentdt_ppos==dpcentdt_ppos.max(),drop=True) pcent_dtmax = data.p_cent.sel(xofyear=dpcentdt_max.xofyear) print((dpcentdt_max.values, pcent_dtmax.values)) ax_twin = ax_in.twinx() data.p_cent.plot(ax=ax_twin, color='b') dpcentdt.plot(ax=ax_in, color='k') ax_in.set_xlabel('') ax_in.set_ylabel('') ax_twin.set_ylabel('') ax_twin.set_ylim([-30,30]) ax_in.set_ylim([-1.*ylim_in, ylim_in]) ax_in.set_title(run, fontsize=14) #ax_in.set_yticks([-8.,-4.,0.,4.,8.]) ax_twin.spines['right'].set_color('blue') plt.tight_layout() return ax_twin
def abs_vort_dt_at_pcent(run, rot_fac=1., lev=150.,lat_bound=45., res=0.01, interp=True): '''Calculate the (normalised) vorticity tendency at the precipitation centroid''' data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') # Interpolate in time if wanted if interp: times_new = np.arange(1., 72.2, 0.2) precip_new = time_interp(data.precipitation, times_new) u_new = time_interp(data.ucomp.sel(pfull=lev), times_new) v_new = time_interp(data.vcomp.sel(pfull=lev), times_new) else: precip_new = data.precipitation u_new = data.ucomp.sel(pfull=lev) v_new = data.vcomp.sel(pfull=lev) # Find precipitation centroid precip_new = make_sym(precip_new) data = xr.Dataset({'precipitation': precip_new}, coords=precip_new.coords) precip_centroid(data, lat_bound=lat_bound, res=res) # Calculate vorticity v_dx = gr.ddx(v_new) # dvdx u_dy = gr.ddy(u_new) # dudy omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat *np.pi/180) vor = (v_dx - u_dy + f)*86400. div = gr.ddx(u_new) + gr.ddy(v_new) stretching_mean = (-86400. * vor * div).mean('lon') vor = make_sym(vor, asym=True) stretching_mean = make_sym(stretching_mean, asym=True) # Take time derivative of absolute vorticity if interp: dvordt = gr.ddt(vor.mean('lon'))*86400.*5. else: dvordt = gr.ddt(vor.mean('lon'))*86400. # Also normalise this by the value of vorticity dvordtvor = dvordt/vor.mean('lon') #dvordtvor = stretching_mean/vor.mean('lon') # Interpolate vorticity in latitude to match precipitation centroid lats lats = [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -lat_bound and data.lat[i] <= lat_bound] lats_new = np.arange(-lat_bound, lat_bound+res, res) stretching_mean = lat_interp(stretching_mean.sel(lat=lats), lats_new) dvordt = lat_interp(dvordt.sel(lat=lats), lats_new) dvordtvor = lat_interp(dvordtvor.sel(lat=lats), lats_new) # Get and return dvordt and dvordtvor at the precipitation centroid, as well as the precipitation centroid itself st_pcent = [float(stretching_mean[i,:].sel(lat=data.p_cent.values[i]).values) for i in range(len(data.xofyear))] st_pcent = xr.DataArray(np.asarray(st_pcent), coords=[stretching_mean.xofyear.values], dims=['xofyear']) dvordt_pcent = [float(dvordt[i,:].sel(lat=data.p_cent.values[i]).values) for i in range(len(data.xofyear))] dvordt_pcent = xr.DataArray(np.asarray(dvordt_pcent), coords=[dvordt.xofyear.values], dims=['xofyear']) dvordtvor_pcent = [float(dvordtvor[i,:].sel(lat=data.p_cent.values[i]).values) for i in range(len(data.xofyear))] dvordtvor_pcent = xr.DataArray(np.asarray(dvordtvor_pcent), coords=[dvordtvor.xofyear.values], dims=['xofyear']) return dvordt_pcent, dvordtvor_pcent, data.p_cent, st_pcent
def pcent_vs_cell_edge(data, symbol='xk'): precip_centroid(data) x = np.linspace(0., 360., 72.) sinewave = -1.*np.sin(x*np.pi/180.) coswave = 3.*np.cos(3.*x*np.pi/180.) + 20.*np.cos(3. + x*np.pi/180.) plt.plot(sinewave, data.p_cent, symbol) plt.plot(sinewave, coswave)
def div_at_pcent(run, rot_fac=1., lev=150., lat_bound=45., res=0.01, interp=True): '''Calculate the (normalised) vorticity tendency at the precipitation centroid''' data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') # Interpolate in time if wanted if interp: times_new = np.arange(1., 72.2, 0.2) precip_new = time_interp(data.precipitation, times_new) u_new = time_interp(data.ucomp.sel(pfull=lev), times_new) v_new = time_interp(data.vcomp.sel(pfull=lev), times_new) else: precip_new = data.precipitation u_new = data.ucomp.sel(pfull=lev) v_new = data.vcomp.sel(pfull=lev) # Find precipitation centroid precip_new = make_sym(precip_new) data = xr.Dataset({'precipitation': precip_new}, coords=precip_new.coords) precip_centroid(data, lat_bound=lat_bound, res=res) # Calculate divergence v_dx = gr.ddx(v_new) # dvdx u_dy = gr.ddy(u_new) # dudy div = (gr.ddx(u_new) + gr.ddy(v_new)).mean('lon') * 86400. omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat * np.pi / 180) vor = (-1. * gr.ddy(u_new) * 0. + f).mean('lon') * 86400. div = make_sym(div) vor = make_sym(vor, asym=True) div = div * vor div = vor # Interpolate divergence in latitude to match precipitation centroid lats lats = [ data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -lat_bound and data.lat[i] <= lat_bound ] lats_new = np.arange(-lat_bound, lat_bound + res, res) div = lat_interp(div.sel(lat=lats), lats_new) # Get and return dvordt and dvordtvor at the precipitation centroid, as well as the precipitation centroid itself div_pcent = [ float(div[i, :].sel(lat=data.p_cent.values[i]).values) for i in range(len(data.xofyear)) ] div_pcent = xr.DataArray(np.asarray(div_pcent), coords=[div.xofyear.values], dims=['xofyear']) return div_pcent, data.p_cent
def abs_vort_dt_at_pcent(run, rot_fac=1., lev=150., lat_bound=45.): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') data['precipitation'] = make_sym(data.precipitation) precip_centroid(data, lat_bound=lat_bound) v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat * np.pi / 180) vor = (v_dx - u_dy + f).sel(pfull=lev) * 86400. vor = make_sym(vor, asym=True) # Take time derivative of absolute vorticity dvordt = gr.ddt(vor.mean('lon')) * 86400. # Select latitudes over which to evaluate precip centroid lats = [ data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -lat_bound and data.lat[i] <= lat_bound ] dvordt_lats = dvordt.sel(lat=lats).values #f_lats = f.sel(lat=lats).values # Interpolate vorticity tendency and f in latitude lats_new = np.arange(-lat_bound, lat_bound + 0.1, 0.1) fun = spint.interp1d(lats, dvordt_lats, axis=-1, fill_value='extrapolate', kind='quadratic') dvordt_new = fun(lats_new) dvordt = xr.DataArray(dvordt_new, coords=[data.xofyear.values, lats_new], dims=['xofyear', 'lat']) #fun = spint.interp1d(lats, f_lats, axis=-1, fill_value='extrapolate', kind='quadratic') #f_new = fun(lats_new) #f = xr.DataArray(f_new, coords=[lats_new], dims=['lat']) dvordt_pcent = [ float(dvordt[i, :].sel(lat=data.p_cent.values[i]).values) for i in range(len(data.xofyear)) ] #f_pcent = [float(f.sel(lat=data.p_cent.values[i]).values*86400.) for i in range(len(data.xofyear))] dvordt_pcent = xr.DataArray(np.asarray(dvordt_pcent), coords=[data.xofyear.values], dims=['xofyear']) return dvordt_pcent
def pcent_egrad_scatter(run, ax_in, lonin=[-1.,361.]): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') precip_centroid(data) dsdy = 100000.* subcloud_entropy_gradient(data, lonin=lonin) ax_in.plot(dsdy[:,31:33].mean('lat'), data.p_cent, 'xk') ax_in.set_xlabel('') ax_in.set_ylabel('') ax_in.set_ylim([-30,30]) ax_in.set_xlim([-8,8])
def pcent_grad_scatter(run, ax_in, ylim_in=1.): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') precip_centroid(data) dpcentdt = gr.ddt(data.p_cent) * 86400. ax_in.plot(data.p_cent, dpcentdt, 'xk') ax_in.set_xlabel('') ax_in.set_ylabel('') ax_in.set_xlim([-30,30]) ax_in.set_ylim([-1.*ylim_in, ylim_in]) ax_in.set_title(run, fontsize=14)
def abs_vort_mean(run, rot_fac=1., lev=150.): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') precip_centroid(data) v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat * np.pi / 180) vor = (v_dx - u_dy + f).sel(pfull=lev) # Take gradients of vorticity dvordx = gr.ddx(vor) dvordy = gr.ddy(vor, vector=False) # Horizontal material derivative horiz_md_mean = (-86400.**2. * (data.ucomp.sel(pfull=lev) * dvordx + data.vcomp.sel(pfull=lev) * dvordy)).mean('lon') # Calculate divergence and stretching term div = gr.ddx(data.ucomp.sel(pfull=lev)) + gr.ddy(data.vcomp.sel(pfull=lev)) stretching_mean = (-86400.**2. * vor * div).mean('lon') vor = vor.mean('lon') vor_mag = np.abs(vor) stretching_mag = stretching_mean stretching_mag[:, 0:32] = stretching_mag[:, 0:32] * -1. lat_itcz, lat_nh, lat_sh = get_streamfun_zeros(data, lev=lev) def take_cell_average(var, lats_min, lats_max): cell_av = np.zeros(len(data.xofyear.values)) for i in range(len(data.xofyear.values)): lats = [ data.lat[j] for j in range(len(data.lat)) if data.lat[j] > lats_min[i] and data.lat[j] < lats_max[i] ] cell_av[i] = var[i].sel(lat=lats).mean('lat') return cell_av vor_av_sh = take_cell_average(vor_mag, lat_sh, lat_itcz) stretching_av_sh = take_cell_average(stretching_mag, lat_sh, lat_itcz) plt.plot(stretching_av_sh, 'k') plt.figure(2) plt.plot(data.p_cent, stretching_av_sh, 'kx') plt.show()
def p_cent_rate(data, days=False): # Locate precipitation centroid precip_centroid(data) # If units in are days rather than pentads (e.g. for shorter years) convert to pentads if days: data['xofyear'] = data['xofyear'] / 5. # Get rate of movement of precip centroid dpcentdt = gr.ddt(data.p_cent) * 86400. dpcentdt2 = gr.ddt(dpcentdt, secperunit=86400.) * 86400. return dpcentdt, dpcentdt2
def p_cent_rate_max(runs, days=None): # Get the maximum rate of change of the precipitation centroid latitude, and the latitude at which this occurs. max_rate = [] max_rate_lat = [] if days==None: days=[False]*len(runs) j=0 for run in runs: # Open dataset data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') # Get total precip try: data['precipitation'] = data.condensation_rain + data.convection_rain except: data['precipitation'] = data.precipitation data['precipitation'] = make_sym(data.precipitation) # Locate precipitation centroid precip_centroid(data) # Get rate of movement of precip centroid if days[j]: dpcentdt = gr.ddt(data.p_cent, secperunit = 86400.) * 86400. else: dpcentdt = gr.ddt(data.p_cent) * 86400. dpcentdt_ppos = dpcentdt.where(data.p_cent>=0.) # Find precip centroid rate where precip centroid is in the northern hemisphere dpcentdt_max = dpcentdt_ppos.where(dpcentdt_ppos==dpcentdt_ppos.max(),drop=True) # Find the maximum of the above if len(dpcentdt_max) > 1: dpcentdt_max = dpcentdt_max[0] pcent_dtmax = data.p_cent.sel(xofyear=dpcentdt_max.xofyear) # Find the location of the preciptiation when the rate is maximum print(dpcentdt_max.values, pcent_dtmax.values) # Print both max_rate.append(dpcentdt_max) max_rate_lat.append(pcent_dtmax) j=j+1 max_rate = np.asarray(max_rate) max_rate_lat = np.asarray(max_rate_lat) return max_rate, max_rate_lat
def plot_bs_fig_9_moist(run, pentads=[35, 40, 45, 50], rotfac=1.): plot_dir = '/scratch/rg419/plots/paper_2_figs/revisions/heatbudg_moist/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) rcParams['figure.figsize'] = 5.5, 7. rcParams['font.size'] = 14 data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') data['precipitation'] = make_sym(data.precipitation) precip_centroid(data) # Locate precipitation centroid dpcentdt = gr.ddt(data.p_cent) * 86400. p_cent_pos = np.abs(data.p_cent.where(dpcentdt >= 0.)) eq_time = p_cent_pos.xofyear[p_cent_pos.argmin('xofyear').values] peak_rate_time = dpcentdt.xofyear[dpcentdt.argmax('xofyear').values] max_lat_time = data.p_cent.xofyear[data.p_cent.argmax('xofyear').values] edge_loc, psi_max, psi_max_loc = get_edge_psi(data, lev=850., thresh=0., intdown=True) fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, sharex=True) axes = [ax1, ax2, ax3, ax4] pentads = [ eq_time.values, peak_rate_time.values, np.floor((peak_rate_time.values + max_lat_time.values) / 2.), max_lat_time.values ] for i in range(4): fig_9_moist(run, ax=axes[i], pentad=pentads[i]) edge = edge_loc.sel(xofyear=pentads[i]) axes[i].plot([edge, edge], [-250., 250.], 'k:') axes[i].text(20, 240., 'Pentad ' + str(int(pentads[i]))) ax4.set_xlabel('Latitude') #ax1.text(-55, 315., 'a)') #ax2.text(-55, 315., 'b)') #ax3.text(-55, 315., 'c)') #ax4.text(-55, 315., 'd)') plt.subplots_adjust(left=0.15, right=0.95, top=0.97, bottom=0.1) plt.savefig(plot_dir + 'sb08_fig9_moist_' + run + '.pdf', format='pdf') plt.close()
def bootstrap_method(data, days=False): max_rate = [] max_rate_lat = [] max_lat = [] eq_rate = [] for boot in range(1000): print(boot) sample = bootstrap_sample(data.year_no.size) precip_sample = data.sel(year_no=sample).mean('year_no') precip_centroid(precip_sample) # Next step is to unpick pcent_rate_max function so that you can just hand it the data for the sample mean # Get the numbers for 1000 or so samples, find 5 and 95 percentile if days: dpcentdt = gr.ddt(precip_sample.p_cent, secperunit=86400.) * 86400. else: dpcentdt = gr.ddt(precip_sample.p_cent) * 86400. dpcentdt_ppos = dpcentdt.where( precip_sample.p_cent >= 0. ) # Find precip centroid rate where precip centroid is in the northern hemisphere dpcentdt_max = dpcentdt_ppos.where( dpcentdt_ppos == dpcentdt_ppos.max('xofyear'), drop=True) dpcentdt_max = dpcentdt_max[0] pcent_max = precip_sample.p_cent.max('xofyear') pcent_dtmax = precip_sample.p_cent.sel( xofyear=dpcentdt_max.xofyear ) # Find the location of the preciptiation when the rate is maximum p_cent_abs = np.abs(precip_sample.p_cent.where(dpcentdt >= 0.)) dpdt_eq = (dpcentdt.where(p_cent_abs == p_cent_abs.min('xofyear'), drop=True)).values[0] #dpdt_eq = dpcentdt.where(precip_sample.p_cent == precip_sample.p_cent.min('xofyear'), drop=True).max() max_rate.append(float(dpcentdt_max.values)) max_rate_lat.append(float(pcent_dtmax.values)) max_lat.append(float(pcent_max.values)) eq_rate.append(float(dpdt_eq)) return max_rate, max_rate_lat, max_lat, eq_rate
def load_pcent_find_zero(data, lat_bound=45, lonin=[-1., 361.]): precip_centroid(data) pcent_sign = np.ones(data.p_cent.values.shape) pcent_sign[data.p_cent.values <= 0.] = 0. pcent_diff = np.diff(pcent_sign) loc = np.argmax(pcent_diff) len1 = np.max(pcent_sign.shape) - loc pcent_repos = np.zeros(data.p_cent.values.shape) pcent_repos[:len1] = data.p_cent[loc:] pcent_repos[len1:] = data.p_cent[:loc] pcent_repos = xr.DataArray(pcent_repos, coords=[data.xofyear.values], dims=['xofyear']) return pcent_repos
def p_cent_rate(data, days=False): """ Inputs: data - xarray dataset climatology including precipitation as either precipitation or as convection_rain and condensation_rain days - instructs if data is in daily or pentad means Returns: dpcentdt - rate of movement of the precipitation centroid dpcentdt2 - rate of acceleration of the precipitation centroid """ # Get total precip try: data['precipitation'] = data.condensation_rain + data.convection_rain except: data['precipitation'] = data.precipitation precip_temp = np.zeros(data.precipitation.values.shape) n = len(data.xofyear.values) // 2 for i in range(0, n): precip_temp[i, :, :] = (data.precipitation[i, :, :].values + data.precipitation[i + n, ::-1, :].values) / 2. precip_temp[i + n, :, :] = precip_temp[i, ::-1, :] precip_temp = xr.DataArray( precip_temp, coords=[data.xofyear.values, data.lat, data.lon], dims=['xofyear', 'lat', 'lon']) data['precipitation'] = precip_temp # Locate precipitation centroid precip_centroid(data) # If units in are days rather than pentads (e.g. for shorter years) convert to pentads if days: dpcentdt = gr.ddt(data.p_cent, secperunit=86400.) * 86400. else: dpcentdt = gr.ddt(data.p_cent) * 86400. dpcentdt2 = gr.ddt(dpcentdt, secperunit=86400.) * 86400. return dpcentdt, dpcentdt2
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 div_plot(run, ax, rot_fac=1., lev=150.): '''Plot dvordt or 1/vor * dvordt''' #Load data data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') #Calculate psi to overplot psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 psi = make_sym(psi, asym=True) # Make precip symmetric and find the precip centroid data['precipitation'] = make_sym(data.precipitation) precip_centroid(data) # Calculate divergence u_dx = gr.ddx(data.ucomp) # dudx v_dy = gr.ddy(data.vcomp) # dvdy div = (u_dx + v_dy).sel(pfull=lev) * 86400. div = make_sym(div) div = div.mean('lon') f1 = div.plot.contourf(ax=ax, x='xofyear', y='lat', levels=np.arange(-1.2, 1.3, 0.2), add_colorbar=False, add_labels=False) psi.sel(pfull=500).plot.contour(ax=ax, 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=ax, 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=ax, x='xofyear', y='lat', levels=np.arange(-1000., 1010., 1000.), add_labels=False, colors='0.5', linewidths=2) data.p_cent.plot.line(ax=ax, color='k', linewidth=2) ax.set_ylim([-60, 60]) ax.set_ylabel('Latitude') ax.set_xticks([12, 24, 36, 48, 60, 72]) ax.set_yticks([-60, -30, 0, 30, 60]) ax.set_xlabel('Pentad') ax.grid(True, linestyle=':') box = ax.get_position() axColor = plt.axes( [box.x0 + box.width * 0.92, box.y0 + box.y0 * 0.12, 0.015, box.height]) cb1 = fig.colorbar(f1, cax=axColor, orientation='vertical')
def abs_vort_dt_plot(run, ax, rot_fac=1., lev=150., dvordt_flag=False): '''Plot dvordt or 1/vor * dvordt''' #Load data data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') #Calculate psi to overplot psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 psi = make_sym(psi, asym=True) # Make precip symmetric and find the precip centroid data['precipitation'] = make_sym(data.precipitation) precip_centroid(data) # Calculate vorticity v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat *np.pi/180) vor = (v_dx - u_dy + f).sel(pfull=lev)*86400. vor = make_sym(vor, asym=True) # Take time derivative of absolute vorticity dvordt = gr.ddt(vor.mean('lon'))*86400. # Also normalise this by the value of vorticity dvordtvor = dvordt/vor.mean('lon') # Plot! plot_dir = '/scratch/rg419/plots/paper_2_figs/abs_vort_dt/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) rcParams['figure.figsize'] = 5.5, 4.3 rcParams['font.size'] = 14 #fig = plt.figure() #ax1 = fig.add_subplot(111) if dvordt_flag: f1 = dvordt.plot.contourf(ax=ax, x='xofyear', y='lat', levels=np.arange(-0.06,0.07,0.01), add_colorbar=False, add_labels=False) else: f1 = dvordtvor.plot.contourf(ax=ax, x='xofyear', y='lat', levels=np.arange(-0.05,0.055,0.005), add_colorbar=False, add_labels=False) #psi.sel(pfull=500).plot.contour(ax=ax, 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=ax, 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=ax, x='xofyear', y='lat', levels=np.arange(-1000.,1010.,1000.), add_labels=False, colors='0.5', linewidths=2) data.p_cent.plot.line(ax=ax, color='k', linewidth=2) ax.set_ylim([-60,60]) ax.set_ylabel('Latitude') ax.set_xticks([12,24,36,48,60,72]) ax.set_yticks([-60,-30,0,30,60]) ax.set_xlabel('Pentad') ax.grid(True,linestyle=':') #originalSize = get(gca, 'Position') #set(f1, 'Position', originalSize) #plt.subplots_adjust(left=0.15, right=0.95, top=0.95, bottom=0.05) box = ax.get_position() #ax.set_position([box.x0*1.05, box.y0, box.width, box.height]) axColor = plt.axes([box.x0 + box.width*0.92, box.y0+box.y0*0.12, 0.015, box.height]) #cb1=fig.colorbar(f1, ax=ax, use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.2, aspect=40) #cb1=plt.colorbar(f1, ax=axColor, use_gridspec=True, orientation = 'vertical',fraction=0.15, pad=0.05, aspect=30) cb1=fig.colorbar(f1, cax=axColor, orientation = 'vertical')
def surface_plot(run, lonin=[-1., 361.], diff_run=None, scale_fac=1., do_make_sym=True): print('what?') rcParams['figure.figsize'] = 15, 10 rcParams['font.size'] = 14 plot_dir = '/scratch/rg419/plots/surface_fluxes/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') #Calculate psi to overplot psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 psi = make_sym(psi, asym=True) # Make precip symmetric and find the precip centroid data['precipitation'] = make_sym(data.precipitation) precip_centroid(data) data['flux_lw_up'] = data.t_surf**4. * mc.stefan data['dTdt'] = gr.ddt(data.t_surf) * 86400. * scale_fac if do_make_sym: for field in [ 't_surf', 'dTdt', 'flux_sw', 'flux_lw', 'flux_lw_up', 'flux_t', 'flux_lhe' ]: data[field] = make_sym(data[field]) lons = pick_lons(data, lonin) if not diff_run == None: diff_data = xr.open_dataset( '/disca/share/rg419/Data_moist/climatologies/' + diff_run + '.nc') data = (data - diff_data).sel(lon=lons).mean('lon') levels_t = np.arange(-10., 11., 1.) levels_dt = np.arange(-0.5, 0.51, 0.05) levels_flux = np.arange(-80., 81., 5.) else: data = data.sel(lon=lons).mean('lon') levels_t = np.arange(270., 306., 2.5) levels_dt = np.arange(-0.3, 0.31, 0.05) levels_flux = np.arange(-300., 305., 20.) f, ((ax1, ax2), (ax3, ax4), (ax5, ax6)) = plt.subplots(3, 2, sharex='col', sharey='row') left_column = [ax1, ax3, ax5] bottom_row = [ax5, ax6] all_plots = [ax1, ax2, ax3, ax4, ax5, ax6] #check = data.flux_sw + (data.flux_lw - flux_lw_up) - data.flux_t - data.flux_lhe #check = data.flux_sw - data.flux_lhe data.t_surf.plot.contourf(x='xofyear', y='lat', levels=levels_t, ax=ax1, extend='both', add_labels=False) ax1.set_title('SST') data.dTdt.plot.contourf(x='xofyear', y='lat', ax=ax2, levels=levels_dt, extend='both', add_labels=False, cmap='RdBu_r') ax2.set_title('d(SST)/dt') data.flux_sw.plot.contourf(x='xofyear', y='lat', levels=levels_flux, ax=ax3, extend='both', add_labels=False, cmap='RdBu_r') ax3.set_title('Net downward SW flux') (data.flux_lw - data.flux_lw_up).plot.contourf(x='xofyear', y='lat', levels=levels_flux, ax=ax4, extend='both', add_labels=False, cmap='RdBu_r') ax4.set_title('Net downward LW flux') (-1. * data.flux_t).plot.contourf(x='xofyear', y='lat', levels=levels_flux, ax=ax5, extend='both', add_labels=False, cmap='RdBu_r') #check.plot.contourf(x='xofyear', y='lat', levels=np.arange(-300.,305.,20.), ax=ax5, extend = 'both', add_labels=False, cmap='RdBu_r') ax5.set_title('Downward sensible heat flux') (-1. * data.flux_lhe).plot.contourf(x='xofyear', y='lat', levels=levels_flux, ax=ax6, extend='both', add_labels=False, cmap='RdBu_r') ax6.set_title('Downward latent heat flux') i = 0 labels = ['a)', 'b)', 'c)', 'd)', 'e)', 'f)'] for ax in all_plots: ax.grid(True, linestyle=':') ax.set_xticks([12, 24, 36, 48, 60, 72]) ax.set_ylim([-60, 60]) ax.set_yticks([-60, -30, 0, 30, 60]) ax.text(-8, 60., labels[i]) psi.sel(pfull=500).plot.contour(ax=ax, 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=ax, 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=ax, x='xofyear', y='lat', levels=np.arange(-1000., 1010., 1000.), add_labels=False, colors='0.5', linewidths=2) data.p_cent.plot.line(ax=ax, color='k', linewidth=2) ax.set_ylabel('') ax.set_xlabel('') i = i + 1 for ax in left_column: ax.set_ylabel('Latitude') for ax in bottom_row: ax.set_xlabel('Pentad') plt.subplots_adjust(left=0.07, right=0.97, top=0.97, bottom=0.05, hspace=0.2, wspace=0.1) if lonin == [-1., 361.]: if not diff_run == None: plt.savefig(plot_dir + 'surface_fluxes_' + run + '_' + diff_run + '.pdf', format='pdf') else: plt.savefig(plot_dir + 'surface_fluxes_' + run + '.pdf', format='pdf') else: figname = plot_dir + 'surface_fluxes_' + run + '_' + str(int( lonin[0])) + '_' + str(int(lonin[1])) + '.pdf' plt.savefig(figname, format='pdf') plt.close()
def vdtdy_plot(run, ax, lev=850.): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 try: data['precipitation'] = make_sym(data.precipitation) except: data['precipitation'] = data.convection_rain + data.condensation_rain data['precipitation'] = make_sym(data.precipitation) precip_centroid(data) convTtotheta = (1000. / data.pfull)**(2. / 7.) theta = data.temp * convTtotheta rho = lev * 100. / Rd / data.temp.sel( pfull=lev) / (1 + (Rv - Rd) / Rd * data.sphum.sel(pfull=lev)) expansion_term = (1. / rho / cp * data.omega.sel(pfull=lev)).mean('lon') * 86400. #expansion_term = (1./rho/cp ).mean('lon') * 86400. v = data.vcomp.sel(pfull=lev) # v T_dy = -86400. * gr.ddy(data.temp.sel(pfull=lev), vector=False) # dTdy vdTdy = v.mean('lon') * T_dy.mean('lon') # [v][dTdy] w = data.omega.sel(pfull=lev) # w T_dp = -86400. * (gr.ddp(data.temp)).sel(pfull=lev) # dTdp wdTdp = w.mean('lon') * T_dp.mean('lon') # [w][dTdp] dTdt = gr.ddt(data.temp).sel(pfull=lev).mean('lon') * 86400. #dvtdy = -1.*(gr.ddy(data.vcomp * data.temp)).mean('lon').sel(pfull=lev)*86400. #dwtdp = -1.*(gr.ddp(data.omega * data.temp)).mean('lon').sel(pfull=lev)*86400. #dvtdy_colav = -1.*gr.ddy((data.vcomp * theta).mean('lon').sel(pfull=np.arange(50.,501.,50.)).mean('pfull')*5000./9.8)*1004.64 #dvHdy_colav = -1.*gr.ddy((data.vcomp * (data.temp*1004.64 + data.sphum*2.500e6)).mean('lon').mean('pfull')*5000./9.8) #f1 = (dvtdy_colav).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-30.,33.,3.), extend='both', add_colorbar=False) #f1 = (vdTdy).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-5.,5.5,0.5), extend='both', add_colorbar=False) f1 = (vdTdy + wdTdp).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-3., 3.1, 0.5), extend='both', add_colorbar=False) #f1 = (dTdt).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-0.2,0.21,0.02), extend='both', add_colorbar=False) #f1 = (vdTdy + wdTdp + expansion_term).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, extend='both', add_colorbar=False) #f1 = (-1.*expansion_term-T_dp.mean('lon')).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, extend='both', add_colorbar=False) psi.sel(pfull=500).plot.contour(ax=ax, 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=ax, 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=ax, x='xofyear', y='lat', levels=np.arange(-1000., 1010., 1000.), add_labels=False, colors='0.5', linewidths=2) data.p_cent.plot.line(ax=ax, color='k', linewidth=2) ax.set_xlabel('') ax.set_ylim([-60, 60]) ax.set_ylabel('Latitude') ax.set_xticks([12, 24, 36, 48, 60, 72]) ax.set_yticks([-60, -30, 0, 30, 60]) ax.grid(True, linestyle=':') return f1
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 abs_vort_dt_plot(run, rot_fac=1., lev=150.): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 data['precipitation'] = make_sym(data.precipitation) precip_centroid(data) v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat * np.pi / 180) vor = (v_dx - u_dy + f).sel(pfull=lev) * 86400. vor = make_sym(vor, asym=True) psi = make_sym(psi, asym=True) # Take time derivative of absolute vorticity dvordt = gr.ddt(vor.mean('lon')) * 86400. plot_dir = '/scratch/rg419/plots/paper_2_figs/abs_vort_dt/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) rcParams['figure.figsize'] = 5.5, 4.3 rcParams['font.size'] = 16 fig = plt.figure() ax1 = fig.add_subplot(111) f1 = dvordt.plot.contourf(ax=ax1, x='xofyear', y='lat', levels=np.arange(-0.06, 0.07, 0.01), add_colorbar=False, add_labels=False) psi.sel(pfull=500).plot.contour(ax=ax1, 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=ax1, 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=ax1, x='xofyear', y='lat', levels=np.arange(-1000., 1010., 1000.), add_labels=False, colors='0.5', linewidths=2) data.p_cent.plot.line(color='k', linewidth=2) ax1.set_ylim([-60, 60]) ax1.set_ylabel('Latitude') ax1.set_xticks([12, 24, 36, 48, 60, 72]) ax1.set_yticks([-60, -30, 0, 30, 60]) ax1.set_xlabel('Pentad') ax1.grid(True, linestyle=':') plt.subplots_adjust(left=0.15, right=0.95, top=0.95, bottom=0.05) cb1 = fig.colorbar(f1, ax=ax1, use_gridspec=True, orientation='horizontal', fraction=0.15, pad=0.2, aspect=40) cb1.set_label('Absolute vorticity tendency, day$^{-2}$') plt.savefig(plot_dir + 'abs_vort_dt_' + run + '.pdf', format='pdf') plt.close()
def ang_mom_grad_plot(run, rot_fac=1., lev=150.): '''Plot dvordt or 1/vor * dvordt''' #Load data data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') #Calculate psi to overplot psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 psi = make_sym(psi, asym=True) # Make precip symmetric and find the precip centroid data['precipitation'] = make_sym(data.precipitation) precip_centroid(data) data['ucomp'] = make_sym(data.ucomp) data['vcomp'] = make_sym(data.vcomp, asym=True) data['omega'] = make_sym(data.omega) # Calculate vorticity v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy u_dp = gr.ddp(data.ucomp) # dudp omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat * np.pi / 180) vor = (v_dx - u_dy + f).sel(pfull=lev).mean('lon') * 86400. vertical_term = (-1. * data.omega / data.vcomp * u_dp).sel(pfull=lev).mean('lon') * 86400. ang_mom_grad = vor + vertical_term # Plot! plot_dir = '/scratch/rg419/plots/paper_2_figs/ang_mom_grad/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) rcParams['figure.figsize'] = 5.5, 4.3 rcParams['font.size'] = 14 fig = plt.figure() ax1 = fig.add_subplot(111) f1 = ang_mom_grad.plot.contourf(ax=ax1, x='xofyear', y='lat', levels=np.arange(-6., 6.1, 1.), add_colorbar=False, add_labels=False, extend='both') psi.sel(pfull=lev).plot.contour(ax=ax1, x='xofyear', y='lat', levels=np.arange(-500., 0., 100.), add_labels=False, colors='0.7', linewidths=2, linestyles='--') psi.sel(pfull=lev).plot.contour(ax=ax1, x='xofyear', y='lat', levels=np.arange(0., 510., 100.), add_labels=False, colors='0.7', linewidths=2) psi.sel(pfull=lev).plot.contour(ax=ax1, x='xofyear', y='lat', levels=np.arange(-1000., 1010., 1000.), add_labels=False, colors='0.5', linewidths=2) data.p_cent.plot.line(ax=ax1, color='k', linewidth=2) ax1.set_ylim([-60, 60]) ax1.set_ylabel('Latitude') ax1.set_xticks([12, 24, 36, 48, 60, 72]) ax1.set_yticks([-60, -30, 0, 30, 60]) ax1.set_xlabel('Pentad') ax1.grid(True, linestyle=':') plt.subplots_adjust(left=0.17, right=0.9, top=0.97, bottom=0.07, hspace=0.3) cb1 = fig.colorbar(f1, ax=ax1, use_gridspec=True, orientation='horizontal', fraction=0.15, pad=0.2, aspect=40) cb1.set_label('Angular momentum gradient') plt.savefig(plot_dir + 'ang_mom_grad_' + run + '.pdf', format='pdf') plt.close()
lhe_temp[i, :, :] = (data.flux_lhe[i, :, :].values + data.flux_lhe[i + n, ::-1, :].values) / 2. lhe_temp[i + n, :, :] = lhe_temp[i, ::-1, :] precip_temp = xr.DataArray(precip_temp, coords=[data.xofyear.values, data.lat, data.lon], dims=['xofyear', 'lat', 'lon']) lhe_temp = xr.DataArray(lhe_temp, coords=[data.xofyear.values, data.lat, data.lon], dims=['xofyear', 'lat', 'lon']) data['precipitation'] = precip_temp data['flux_lhe'] = lhe_temp # Locate precipitation centroid precip_centroid(data) data = data.mean('lon') levels_flux = np.arange(-300., 305., 20.) fig = plt.figure() ax1 = plt.subplot(111) f1 = (-1. * data.flux_lhe).plot.contourf(x='xofyear', y='lat', ax=ax1, levels=levels_flux, extend='both', add_labels=False, cmap='RdBu_r', add_colorbar=False)
def p_cent_rate_max(runs, do_make_sym=True, months=None, days=None, period_fac=1.): # Get the maximum rate of change of the precipitation centroid latitude, and the latitude at which this occurs. # Ensure runs is a list not a string if not isinstance(runs, list): runs = [runs] if not isinstance(period_fac, list): period_fac = [period_fac] * len(runs) # Set up empty lists max_rate = [] max_rate_lat = [] max_lat = [] if not do_make_sym: # If we're not averaging both hemispheres, get southern peak magnitude and location too max_rate_s = [] max_rate_lat_s = [] max_lat_s = [] if days == None: days = [False] * len(runs) j = 0 for run in runs: if months == None: # If no months provided, look for a climatology # Open dataset data = xr.open_dataset( '/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') # Get total precip try: data[ 'precipitation'] = data.condensation_rain + data.convection_rain except: data['precipitation'] = data.precipitation else: # If months are provided, load the data for these, and reshape to years/day of year try: # First make sure months is a list of lists len(months[0]) except: months = [months] data = precip_load_and_reshape(run, months=months[j], period_fac=period_fac[j]) if do_make_sym: # If symmetric data is wanted, average over both hemispheres (NB currently only set up for a climatology 2/05/18) data['precipitation'] = make_sym(data.precipitation) # Locate precipitation centroid precip_centroid(data) # Get rate of movement of precip centroid if days[j]: dpcentdt = gr.ddt(data.p_cent, secperunit=86400.) * 86400. else: dpcentdt = gr.ddt(data.p_cent) * 86400. def get_max( dpcentdt_pmask, sh=False ): # Function to find the maximum ratea and it's latitude and the maximum latitude. dpcentdt_max = dpcentdt_pmask.where( dpcentdt_pmask == dpcentdt_pmask.max('xofyear'), drop=True) # Find the maximum rate numdims = len( dpcentdt_max.shape ) # Check if this is just 1 year or if it's many years if numdims != 1: # If it's many years, evaluate the metrics for each year dpdt_max_temp = [] #First set up empty lists pcent_dtmax = [] pcent_max = [] for y in range(dpcentdt_max.shape[0]): # Loop over years if sh: #If it's the southern hemisphere, look for the min latitude of pcent, and make it positive #data.p_cent.sel(year_no=y).plot() #plt.show() pcent_max.append( -1. * data.p_cent.sel(year_no=y).min('xofyear')) else: # Otherwise look for the maximum latitude pcent_max.append( data.p_cent.sel(year_no=y).max('xofyear')) # Find the maximum rate for the given year, remove any nans that have been added. If there are multiple values, use the smallest dpdt_max_temp.append( dpcentdt_max[y, :].dropna('xofyear')[0]) # Find the location of the max rate for that year pcent_dtmax.append( data.p_cent.sel( year_no=y, xofyear=dpdt_max_temp[y].xofyear.values)) # Convert all results to datarrays pcent_dtmax = xr.DataArray(np.asarray(pcent_dtmax), coords=[data.year_no.values], dims=['year_no']) dpcentdt_max = xr.DataArray(np.asarray(dpdt_max_temp), coords=[data.year_no.values], dims=['year_no']) pcent_max = xr.DataArray(np.asarray(pcent_max), coords=[data.year_no.values], dims=['year_no']) else: # If it's only 1D if len( dpcentdt_max ) > 1: # If there are multiple values for the maximum take the smallest dpcentdt_max = dpcentdt_max[0].expand_dims('xofyear', axis=0) if sh: # If it's the southern hemisphere find the min lat of pcent, and make it positive pcent_max = -1. * data.p_cent.min('xofyear') else: # Otherwise find the maximum latitude pcent_max = data.p_cent.max('xofyear') pcent_max = np.expand_dims(pcent_max, axis=1) pcent_dtmax = data.p_cent.sel( xofyear=dpcentdt_max.xofyear ) # Find the location of the preciptiation when the rate is maximum return dpcentdt_max, pcent_dtmax, pcent_max dpcentdt_ppos = dpcentdt.where( data.p_cent >= 0. ) # Find precip centroid rate where precip centroid is in the northern hemisphere dpcentdt_max, pcent_dtmax, pcent_max = get_max( dpcentdt_ppos) # Get the maximum etc max_rate.append(dpcentdt_max) # Add values to lists for different runs max_rate_lat.append(pcent_dtmax) max_lat.append(pcent_max) if not do_make_sym: dpcentdt_pneg = -1. * dpcentdt.where( data.p_cent <= 0. ) # Find precip centroid rate where precip centroid is in the southern hemisphere, and make poleward positive dpcentdt_max_s, pcent_dtmax_s, pcent_max_s = get_max( dpcentdt_pneg, sh=True) # Get the maximum etc max_rate_s.append( dpcentdt_max_s) # Add values to lists for different runs max_rate_lat_s.append(-1. * pcent_dtmax_s) max_lat_s.append(pcent_max_s) j = j + 1 # Add 1 to counter to get values for next run # Convert all output to xarrays if do_make_sym: if months == None: coords_yrno = ['clim'] else: coords_yrno = data.year_no.values max_rate = xr.DataArray(np.asarray(max_rate), coords=[runs, coords_yrno], dims=['run', 'year_no']) max_rate_lat = xr.DataArray(np.asarray(max_rate_lat), coords=[runs, coords_yrno], dims=['run', 'year_no']) max_lat = xr.DataArray(np.asarray(max_lat), coords=[runs, coords_yrno], dims=['run', 'year_no']) else: if months == None: coords_yrno = ['clim'] else: coords_yrno = data.year_no.values max_rate = (xr.DataArray(np.asarray([max_rate, max_rate_s]), coords=[['n', 's'], runs, coords_yrno], dims=['hemisphere', 'run', 'year_no' ])).transpose('run', 'year_no', 'hemisphere') max_rate_lat = (xr.DataArray(np.asarray([max_rate_lat, max_rate_lat_s]), coords=[['n', 's'], runs, coords_yrno], dims=['hemisphere', 'run', 'year_no'])).transpose( 'run', 'year_no', 'hemisphere') max_lat = (xr.DataArray(np.asarray([max_lat, max_lat_s]), coords=[['n', 's'], runs, coords_yrno], dims=['hemisphere', 'run', 'year_no' ])).transpose('run', 'year_no', 'hemisphere') return max_rate, max_rate_lat, max_lat
def evap_parts(run, lonin=[-1., 361.], do_make_sym=True): rcParams['figure.figsize'] = 6, 9 rcParams['font.size'] = 14 plot_dir = '/scratch/rg419/plots/surface_fluxes/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') #Calculate psi to overplot psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 psi = make_sym(psi, asym=True) # Make precip symmetric and find the precip centroid data['precipitation'] = make_sym(data.precipitation) precip_centroid(data) lons = pick_lons(data, lonin) # Density of lowest level air #rho_a = 95000./Rd/data.temp.sel(pfull=950.) / (1 + (Rv - Rd)/Rd*data.sphum.sel(pfull=950.)) # Look at an average over a zonal region v_a = data.wind_speed.sel(lon=lons).mean('lon') q_a = data.q_atm.sel(lon=lons).mean('lon') q_s = data.q_surf.sel(lon=lons).mean('lon') qdiff = (q_a - q_s) * 1000. rho_a = data.rho.sel(lon=lons).mean('lon') #rho_a.plot.contourf() #plt.colorbar() #plt.show() #(v_a*qdiff).plot.contourf(x='xofyear', y='lat', levels=np.arange(-0.05,0.05,0.005)) #plt.show() #q_a.plot.contourf(x='xofyear', y='lat', levels=np.arange(0,0.02,0.001)) #plt.figure(2) #q_s.plot.contourf(x='xofyear', y='lat', levels=np.arange(0,0.02,0.001)) #plt.show() if do_make_sym: qdiff = make_sym(qdiff) rho_a = make_sym(rho_a) v_a = make_sym(v_a) f, (ax1, ax2, ax3) = plt.subplots(3, sharex=True) qdiff.plot.contourf(x='xofyear', y='lat', ax=ax1, extend='both', add_labels=False, levels=np.arange(-10., 10.1, 1.)) ax1.set_title('q$_{a}$ - q$_{s}$') rho_a.plot.contourf(x='xofyear', y='lat', ax=ax2, extend='both', add_labels=False, levels=np.arange(1.12, 1.26, 0.01)) ax2.set_title(r'$\rho_{a}$') v_a.plot.contourf(x='xofyear', y='lat', levels=np.arange(0., 20., 2.), ax=ax3, extend='both', add_labels=False) ax3.set_title('v$_{a}$') labels = ['a)', 'b)', 'c)'] i = 0 for ax in [ax1, ax2, ax3]: #psi.sel(pfull=500).plot.contour(ax=ax, 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=ax, 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=ax, x='xofyear', y='lat', levels=np.arange(-1000.,1010.,1000.), add_labels=False, colors='0.5', linewidths=2) #data.p_cent.plot.line(ax=ax, color='k', linewidth=2) #ax.set_xlabel('') ax.set_ylim([-60, 60]) ax.set_yticks([-60, -30, 0, 30, 60]) ax.set_ylabel('Latitude') ax.text(-10, 60., labels[i]) ax.grid(True, linestyle=':') i = i + 1 ax3.set_xticks([12, 24, 36, 48, 60, 72]) ax3.set_xlabel('Pentad') plt.subplots_adjust(left=0.12, right=0.99, top=0.95, bottom=0.06) plt.savefig(plot_dir + 'evap_parts_revisions_' + run + '.pdf', format='pdf') plt.close()
def rossby_plot(run, ax, rot_fac=1., lev=200., type='vor_only', plottype='rossby'): '''Plot dvordt or 1/vor * dvordt''' #Load data data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') #Calculate psi to overplot psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) psi /= 1.e9 psi = make_sym(psi, asym=True) # Make precip symmetric and find the precip centroid data['precipitation'] = make_sym(data.precipitation) data['ucomp'] = make_sym(data.ucomp) data['vcomp'] = make_sym(data.vcomp, asym=True) data['omega'] = make_sym(data.omega) precip_centroid(data) # Calculate vorticity v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy u_dp = gr.ddp(data.ucomp) # dudp omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat * np.pi / 180) vor = (v_dx - u_dy).sel(pfull=lev) metric = (data.ucomp / mc.a * np.tan(data.lat * np.pi / 180)).sel(pfull=lev) vertical_term = (-1. * data.omega / data.vcomp * u_dp).sel(pfull=lev) #vor = make_sym(vor, asym=True) #metric = make_sym(metric, asym=True) #vertical_term = make_sym(vertical_term, asym=True) if type == 'vor_only': rossby = (-1. * vor / f).mean('lon') elif type == 'metric': rossby = (-1. * (metric + vor) / f).mean('lon') elif type == 'vertical': rossby = (-1. * (vor + vertical_term) / f).mean('lon') elif type == 'full': rossby = (-1. * (metric + vor + vertical_term) / f).mean('lon') levels = np.arange(-0.9, 1.0, 0.3) if plottype == 'drodt': rossby = gr.ddt(rossby) * 84600. levels = np.arange(-0.1, 0.1, 0.01) f1 = rossby.plot.contourf(ax=ax, x='xofyear', y='lat', levels=levels, add_colorbar=False, add_labels=False, extend='both') psi.sel(pfull=lev).plot.contour(ax=ax, x='xofyear', y='lat', levels=np.arange(-500., 0., 100.), add_labels=False, colors='0.7', linewidths=2, linestyles='--') psi.sel(pfull=lev).plot.contour(ax=ax, x='xofyear', y='lat', levels=np.arange(0., 510., 100.), add_labels=False, colors='0.7', linewidths=2) psi.sel(pfull=lev).plot.contour(ax=ax, x='xofyear', y='lat', levels=np.arange(-1000., 1010., 1000.), add_labels=False, colors='0.5', linewidths=2) data.p_cent.plot.line(ax=ax, color='k', linewidth=2) ax.set_ylim([-30, 30]) ax.set_ylabel('Latitude') ax.set_xticks([12, 24, 36, 48, 60, 72]) ax.set_yticks([-30, -15, 0, 15, 30]) ax.set_xlabel('Pentad') ax.grid(True, linestyle=':') #plt.subplots_adjust(left=0.15, right=0.95, top=0.95, bottom=0.05) cb1 = fig.colorbar(f1, ax=ax, use_gridspec=True, orientation='horizontal', fraction=0.15, pad=0.2, aspect=40) cb1.set_label('Rossby number')