def psi_plot(run, ax, title=''): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) mse = (mc.cp_air * data.temp + mc.L * data.sphum + mc.grav * data.height).mean('lon').sel(pfull=850.) / 1000. dmsedy = gr.ddy(mse, vector=False) * 100000. dtsurfdy = gr.ddy(data.t_surf.mean('lon'), vector=False) * 100000. psi /= 1.e9 psi = np.abs(psi).max('pfull') #mse.plot.contour(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(200.,401.,10.), colors='w') dmsedy.plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-3., 3.1, 0.25), add_colorbar=False) #dtsurfdy.plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-1.5,1.6,0.1), add_colorbar=False) psi.plot.contour(ax=ax, x='xofyear', y='lat', extend='both', add_labels=False, levels=np.arange(0, 101., 100.), add_colorbar=False, alpha=0.4) ax.set_title(title, fontsize=17) ax.set_ylim(-60, 60) ax.grid(True, linestyle=':') ax.set_yticks(np.arange(-60., 61., 30.)) ax.set_xticks(np.arange(0., 72., 18.))
def partition_advection(data, lons, lev=150): #Do vv terms vv_trans_dy = -86400. * gr.ddy( (data.vcomp_sq - data.vcomp * data.vcomp).sel(pfull=lev) , uv=True) v = data.vcomp.sel(pfull=lev).load() # v v_dy = -86400. * gr.ddy( v ) # dudy v_dvdy_zav = v.sel(lon=lons).mean('lon') * v_dy.sel(lon=lons).mean('lon') # [v][dudy] v_dvdy_stat = (v * v_dy).sel(lon=lons).mean('lon') - v_dvdy_zav # v*dudy* = [vdudy] - [v][dudy] data['vv_trans_dy'] = (('xofyear','lat'), vv_trans_dy.sel(lon=lons).mean('lon')) data['v_dvdy_stat'] = (('xofyear','lat'), v_dvdy_stat) data['v_dvdy_zav'] = (('xofyear','lat'), v_dvdy_zav ) print('vv terms done') #Do vw terms vw_trans_dp = -86400. * gr.ddp( (data.vcomp_omega - data.vcomp * data.omega).sel(lon=lons).mean('lon') ) w = data.omega.sel(pfull=lev).load() # w v_dp = -86400. * (gr.ddp(data.vcomp)).sel(pfull=lev) # dudp w_dvdp_zav = w.sel(lon=lons).mean('lon') * v_dp.sel(lon=lons).mean('lon') w_dvdp_stat = (w * v_dp).sel(lon=lons).mean('lon') - w_dvdp_zav data['vw_trans_dp'] = (('xofyear','lat'), vw_trans_dp.sel(pfull=lev)) data['w_dvdp_stat'] = (('xofyear','lat'), w_dvdp_stat) data['w_dvdp_zav'] = (('xofyear','lat'), w_dvdp_zav ) print('vw terms done')
def fig_9(run, ax, pentad=40, lev=850.): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') heating = (data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad).mean('lon') * 86400. rho = data.pfull * 100. / mc.rdgas / data.temp dtdy = gr.ddy(data.temp.mean('lon'), vector=False) dtdp = gr.ddp(data.temp.mean('lon')) dtdt = gr.ddt(data.temp.mean('lon')) * 86400. vt_eddy = data.vcomp_temp.mean( 'lon') - data.vcomp.mean('lon') * data.temp.mean('lon') wt_eddy = data.omega_temp.mean( 'lon') - data.omega.mean('lon') * data.temp.mean('lon') vdtdy_mean = -1. * data.vcomp.mean('lon') * dtdy * 86400. wdtdp_mean = -1. * data.omega.mean('lon') * dtdp * 86400. expansion_term = (data.omega / rho / mc.cp_air).mean('lon') * 86400. div_vt_eddy = -1. * gr.ddy(vt_eddy, vector=True) * 86400. div_wt_eddy = -1. * gr.ddp(wt_eddy) * 86400. lats = [ data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -40. and data.lat[i] <= 40. ] heating.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='k') dtdt.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') #vdtdy_mean.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C1') #(wdtdp_mean + heating).sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C2') expansion_term.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C1') (vdtdy_mean + wdtdp_mean + expansion_term).sel(pfull=lev, xofyear=pentad, lat=lats).plot( ax=ax, color='k', linestyle='--') #vdtdy_mean.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C2', linestyle='--') #wdtdp_mean.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C3', linestyle='--') #(vdtdy_mean + wdtdp_mean + expansion_term).sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--') #(-dthetadt + vdthetady_mean + wdthetadp_mean + heating_theta + div_vt_eddy + div_wt_eddy).sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') (div_vt_eddy + div_wt_eddy).sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.') ax.set_title('') ax.set_xlabel('') #ax.set_ylim(-0.25,0.25) ax.set_ylim(-5., 5.) ax.set_xlim(-30., 30.) ax.grid(True, linestyle=':') ax.set_ylabel('Heating rate, K/day')
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 get_budget_terms(energy, uenergy, venergy, wenergy, diab): u_e_eddy = uenergy - energy * data.ucomp v_e_eddy = venergy - energy * data.vcomp w_e_eddy = wenergy - energy * data.omega eddy_conv = -1.*(gr.ddx(u_e_eddy) + gr.ddy(v_e_eddy) + gr.ddp(w_e_eddy)) horiz_adv = -1.*(data.ucomp * gr.ddx(energy) + data.vcomp * gr.ddy(energy, vector=False)) denergydt = gr.ddt(energy)*20. total = (diab + horiz_adv + vert_adv + eddy_conv)*10. return diab, horiz_adv, vert_adv, eddy_conv, denergydt, total
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 fig_9(run, ax, pentad=40): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') convTtotheta=(1000./data.pfull)**mc.kappa heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad rho = data.pfull*100./mc.rdgas/data.temp heating_theta = (heating*convTtotheta).mean('lon') theta = data.temp*convTtotheta dthetady = gr.ddy(theta.mean('lon'), vector=False) dthetadp = gr.ddp(theta.mean('lon')) dthetadt = gr.ddt(theta.mean('lon')) vcomp_theta = data.vcomp_temp*convTtotheta vcomp_theta_eddy = vcomp_theta.mean('lon') - data.vcomp.mean('lon')*theta.mean('lon') vdthetady_mean = data.vcomp.mean('lon') * dthetady wdthetadp_mean = data.omega.mean('lon') * dthetadp def column_int(var_in): var_int = mc.cp_air * var_in.sum('pfull')*5000./mc.grav return var_int vdtdy_mean_int = -1. * column_int(vdthetady_mean) wdtdp_mean_int = -1. * column_int(wdthetadp_mean) vt_eddy_int = -1. * column_int(vcomp_theta_eddy) div_vt_eddy_int = gr.ddy(vt_eddy_int, vector=True) w_850 = data.omega.sel(pfull=850.).mean('lon') w_max, w_max_lat = get_latmax(-1.*w_850) heating_theta_int = column_int(heating_theta) dthetadt_int = column_int(dthetadt) lats = [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -40. and data.lat[i] <= 40.] dthetadt_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2') heating_theta_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--') #wdtdp_mean_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C1') #vdtdy_mean_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--') (vdtdy_mean_int + wdtdp_mean_int).sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.') #(-dthetadt_int + vdtdy_mean_int + wdtdp_mean_int + heating_theta_int + div_vt_eddy_int).sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') #(vdtdy_mean_int + wdtdp_mean_int + heating_theta_int + div_vt_eddy_int).sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') div_vt_eddy_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') ax.plot([w_max_lat.sel(xofyear=pentad),w_max_lat.sel(xofyear=pentad)], [-500.,500.], color='0.7', linestyle=':') ax.set_title('') ax.set_xlabel('') ax.set_ylim(-500.,500.) ax.set_xlim(-35.,35.) ax.grid(True,linestyle=':') ax.set_ylabel('Heating rate, W/m$^2$')
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 gross_stability(run, moist=False, i=1): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') convTtotheta = (1000. / data.pfull)**mc.kappa theta = data.temp * convTtotheta theta_equiv = (data.temp + mc.L / mc.cp_air * data.sphum / (1 - data.sphum)) * convTtotheta dthetady = gr.ddy(theta.mean('lon'), vector=False) dthetadp = gr.ddp(theta.mean('lon')) dthetady_equiv = gr.ddy(theta_equiv.mean('lon'), vector=False) dthetadp_equiv = gr.ddp(theta_equiv.mean('lon')) vdthetady_mean = data.vcomp.mean('lon') * dthetady wdthetadp_mean = data.omega.mean('lon') * dthetadp vdthetady_mean_equiv = data.vcomp.mean('lon') * dthetady_equiv wdthetadp_mean_equiv = data.omega.mean('lon') * dthetadp_equiv def column_int(var_in): var_int = mc.cp_air * var_in.sum('pfull') * 5000. / mc.grav return var_int div_vt_mean_int = -1. * column_int(wdthetadp_mean + vdthetady_mean) div_vt_mean_int_equiv = -1. * column_int(wdthetadp_mean_equiv + vdthetady_mean_equiv) vt_mean_int = column_int(data.vcomp.mean('lon') * theta.mean('lon')) vt_mean_int_equiv = column_int( data.vcomp.mean('lon') * theta_equiv.mean('lon')) #vt_mean_int.plot.contourf(x='xofyear', y='lat') psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.) #psi /= 1.e9 psi = np.abs(psi).max('pfull') #plt.figure(2) #psi.plot.contourf(x='xofyear', y='lat') gross_stab = (2. * np.pi * mc.a * np.cos(data.lat * np.pi / 180.) * np.abs(vt_mean_int)) / psi gross_moist_stab = (2. * np.pi * mc.a * np.cos(data.lat * np.pi / 180.) * np.abs(vt_mean_int_equiv)) / psi plt.figure(i) gross_moist_stab.plot.contourf(x='xofyear', y='lat', levels=np.arange(0., 2.e5, 2.e4))
def incompressibility(run, lev=850.): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') rho = data.pfull * 100. / data.temp / mc.rdgas rho = data.pfull * 100. / mc.rdgas / data.temp / ( 1 + (mc.rvgas - mc.rdgas) / mc.rdgas * data.sphum) dudx = gr.ddx(data.ucomp * rho) dvdy = gr.ddy(data.vcomp * rho) dwdp = gr.ddp(data.omega * rho) divU = (dudx + dvdy + dwdp).mean('lon').sel(pfull=lev) * 86400. rcParams['figure.figsize'] = 5.5, 3. rcParams['font.size'] = 14 fig, ax = plt.subplots(1, sharex=True) divU.plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False) ax.set_xlabel('Pentad') ax.set_ylabel('Latitude') plt.subplots_adjust(left=0.15, right=0.95, top=0.97, bottom=0.2) plt.savefig(plot_dir + 'incompressibility_rho_' + run + '.pdf', format='pdf') plt.close()
def fig_8(run, ax, pentad=45, rotfac=1., dayfac=1.): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad convTtotheta=(1000./data.pfull)**mc.kappa theta = data.temp*convTtotheta heating_theta = heating*convTtotheta dthetady = gr.ddy(theta.mean('lon'), vector=False) dthetadp = gr.ddp(theta.mean('lon')) dthetadt = gr.ddt(theta.mean('lon')) vdthetady_mean = data.vcomp.mean('lon') * dthetady wdthetadp_mean = data.omega.mean('lon') * dthetadp adv_heating = -1. *(vdthetady_mean + wdthetadp_mean) sinphi = np.sin(data.lat * np.pi/180.) cosphi = np.cos(data.lat * np.pi/180.) w_850 = data.omega.sel(pfull=850.).mean('lon') w_max, w_max_lat = get_latmax(-1.*w_850) theta_850 = theta.sel(pfull=850.).mean('lon') theta_max, theta_max_lat = get_latmax(theta_850) sinphimax = np.sin(theta_max_lat * np.pi/180.) theta_m = theta_max - 300.*(mc.omega * rotfac)**2.*mc.a**2./(2.*mc.grav*14000.) * (sinphi**2.-sinphimax**2.)**2./cosphi**2. def adj_theta(theta_in, adj_by, dayfac=dayfac): if 'lon' in adj_by.coords: return theta_in + adj_by.sel(pfull=850.).mean('lon') * 86400. * dayfac else: return theta_in + adj_by.sel(pfull=850.) * 86400. * dayfac theta_rc = adj_theta(theta_850, heating_theta) theta_adv = adj_theta(theta_850, adv_heating) theta_net = adj_theta(theta_850, dthetadt, dayfac=20.) lats = [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -60. and data.lat[i] <= 60.] theta_m.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7') theta_rc.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--') theta_adv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.') theta_850.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C0') theta_net.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2') ax.plot([w_max_lat.sel(xofyear=pentad),w_max_lat.sel(xofyear=pentad)], [295,315], color='0.7', linestyle=':') ax.set_title('') ax.set_xlabel('') ax.set_ylim(295.,315.) ax.set_xlim(-35.,35.) ax.grid(True,linestyle=':') ax.set_ylabel('$\Theta$, K')
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 vor_plot(run, ax, tf, linestyle='-', rot_fac=1., lev=150.): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') # 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. div = (gr.ddx(data.ucomp) + gr.ddy(data.vcomp)).sel(pfull=lev) vor = (make_sym(vor, asym=True)).mean('lon') div = (make_sym(div)).mean('lon') * 8640. m = ((mc.omega * mc.a**2. * np.cos(data.lat * np.pi / 180.)**2. + data.ucomp.mean('lon') * mc.a * np.cos(data.lat * np.pi / 180.)).sel( pfull=lev)) / (mc.omega * mc.a**2.) #dvordt = gr.ddt(vor.mean('lon'))*86400. #stretching = (-86400. * vor * div).mean('lon') #adv = -86400. * (data.vcomp.sel(pfull=lev) * gr.ddy(vor, vector=False)).mean('lon') #dvordy = gr.ddy(vor, vector=False).mean('lon')*8640. #dvordy[tf[0]:tf[1],:].mean('xofyear').plot(ax=ax, color='r', linestyle=linestyle) #m[:,tf[0]:tf[1]].mean('xofyear').plot(ax=ax, color='r', linestyle=linestyle) vor[tf[0]:tf[1], :].mean('xofyear').plot(ax=ax, color='k', linestyle=linestyle) #stretching[tf[0]:tf[1],:].mean('xofyear').plot(ax=ax, color='b', linestyle=linestyle) #adv[tf[0]:tf[1],:].mean('xofyear').plot(ax=ax, color='r', linestyle=linestyle) #(stretching + adv)[tf[0]:tf[1],:].mean('xofyear').plot(ax=ax, color='r', linestyle=linestyle) ax.set_title('') ax.set_xlabel('') #ax.set_ylim(0.5,1.) ax.set_ylim(-10, 10) ax.grid(True, linestyle=':') ax.set_ylabel('Absolute vorticity, day$^{-1}$')
def mom_budg_fun(run, lev=150): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') #First do uu terms uu_trans_dx = -86400. * gr.ddx( (data.ucomp_sq - data.ucomp**2).sel(pfull=lev)) # <u'u'> = <uu> - <u><u> u = data.ucomp.sel(pfull=lev) # u u_dx = -86400. * gr.ddx(u) # dudx u_ed = u - u.mean('lon') u_dx_ed = u_dx - u_dx.mean('lon') u_dudx_zav = u.mean('lon') * u_dx.mean( 'lon') # [u][dudx], where brackets denote mean over all longitudes u_dudx_cross1 = u.mean('lon') * u_dx_ed # [u]dudx* u_dudx_cross2 = u_ed * u_dx.mean('lon') # u*[dudx] u_dudx_stat = u_ed * u_dx_ed # u*dudx* data['uu_trans_dx'] = (('xofyear', 'lat', 'lon'), uu_trans_dx) data['u_dudx_cross1'] = (('xofyear', 'lat', 'lon'), u_dudx_cross1) data['u_dudx_cross2'] = (('xofyear', 'lat', 'lon'), u_dudx_cross2) data['u_dudx_stat'] = (('xofyear', 'lat', 'lon'), u_dudx_stat) data['u_dudx_zav'] = (('xofyear', 'lat'), u_dudx_zav) #Next do uv terms uv_trans_dy = -86400. * gr.ddy( (data.ucomp_vcomp - data.ucomp * data.vcomp).sel(pfull=lev), uv=True) v = data.vcomp.sel(pfull=lev).load() # v u_dy = -86400. * gr.ddy(u) # dudy v_ed = v - v.mean('lon') u_dy_ed = u_dy - u_dy.mean('lon') v_dudy_zav = v.mean('lon') * u_dy.mean('lon') # [v][dudy] v_dudy_cross1 = v.mean('lon') * u_dy_ed # [v]dudy* v_dudy_cross2 = v_ed * u_dy.mean('lon') # v*[dudy] v_dudy_stat = v_ed * u_dy_ed # v*dudy* data['uv_trans_dy'] = (('xofyear', 'lat', 'lon'), uv_trans_dy) data['v_dudy_cross1'] = (('xofyear', 'lat', 'lon'), v_dudy_cross1) data['v_dudy_cross2'] = (('xofyear', 'lat', 'lon'), v_dudy_cross2) data['v_dudy_stat'] = (('xofyear', 'lat', 'lon'), v_dudy_stat) data['v_dudy_zav'] = (('xofyear', 'lat'), v_dudy_zav) #Finally do uw terms uw_trans_dp = -86400. * gr.ddp( (data.ucomp_omega - data.ucomp * data.omega)) w = data.omega.sel(pfull=lev).load() # w u_dp = -86400. * (gr.ddp(data.ucomp)).sel(pfull=lev) # dudp w_ed = w - w.mean('lon') u_dp_ed = u_dp - u_dp.mean('lon') w_dudp_zav = w.mean('lon') * u_dp.mean('lon') # [w][dudp] w_dudp_cross1 = w.mean('lon') * u_dp_ed # [w]dudp* w_dudp_cross2 = w_ed * u_dp.mean('lon') # w*[dudp] w_dudp_stat = w_ed * u_dp_ed # w*dudp* data['uw_trans_dp'] = (('xofyear', 'lat', 'lon'), uw_trans_dp.sel(pfull=lev)) data['w_dudp_cross1'] = (('xofyear', 'lat', 'lon'), w_dudp_cross1) data['w_dudp_cross2'] = (('xofyear', 'lat', 'lon'), w_dudp_cross2) data['w_dudp_stat'] = (('xofyear', 'lat', 'lon'), w_dudp_stat) data['w_dudp_zav'] = (('xofyear', 'lat'), w_dudp_zav) #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 #Geopotential gradient dphidx = gr.ddx(data.height.sel(pfull=lev)) dphidx = -86400. * 9.8 * dphidx 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_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_sum = fv_local + fv_mean + dphidx + mom_mean + mom_trans + mom_stat + mom_cross data['mom_mean'] = (('xofyear', 'lat'), mom_mean) data['mom_cross'] = (('xofyear', 'lat', 'lon'), mom_cross) data['mom_trans'] = (('xofyear', 'lat', 'lon'), mom_trans) data['mom_stat'] = (('xofyear', 'lat', 'lon'), mom_stat) data['fv_local'] = (('xofyear', 'lat', 'lon'), fv_local) data['fv_ageo'] = (('xofyear', 'lat', 'lon'), fv_ageo) data['fv_mean'] = (('xofyear', 'lat'), fv_mean) data['dphidx'] = (('xofyear', 'lat', 'lon'), dphidx) data['mom_sum'] = (('xofyear', 'lat', 'lon'), mom_sum) return data
def stretching_plot(run, ax_h, ax_s, ax_v, pentad, tibet=True): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc') land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc' land = xr.open_dataset(land_file) lons = [ data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= 60. and data.lon[i] < 150. ] lats = [ data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -30. and data.lat[i] < 60. ] omega = 7.2921150e-5 f = 2 * omega * np.sin(data.lat * np.pi / 180) v_dx = gr.ddx(data.vcomp.sel(pfull=150)) # dvdx u_dy = gr.ddy(data.ucomp.sel(pfull=150)) # dudy vor = v_dx - u_dy + f div = gr.ddx(data.ucomp.sel(pfull=150)) + gr.ddy(data.vcomp.sel(pfull=150)) dvordx = gr.ddx(vor) dvordy = gr.ddy(vor, vector=False) horiz_md_mean = -86400.**2. * (data.ucomp.sel(pfull=150) * dvordx + data.vcomp.sel(pfull=150) * dvordy) stretching_mean = -86400.**2. * vor * div levels = np.arange(-1.5, 1.6, 0.5) fh = horiz_md_mean[pentad, :, :].plot.contourf(x='lon', y='lat', levels=levels, ax=ax_h, add_labels=False, extend='both', add_colorbar=False) land.land_mask.plot.contour(x='lon', y='lat', levels=np.arange(0., 2., 1.), ax=ax_h, colors='k', add_colorbar=False, add_labels=False) if tibet: land.zsurf.plot.contourf(x='lon', y='lat', ax=ax_h, levels=np.arange(2500., 100001., 97500.), add_labels=False, extend='neither', add_colorbar=False, alpha=0.5, cmap='Greys_r') fs = stretching_mean[pentad, :, :].plot.contourf(x='lon', y='lat', levels=np.arange( -1.5, 1.6, 0.5), ax=ax_s, add_labels=False, extend='both', add_colorbar=False) land.land_mask.plot.contour(x='lon', y='lat', levels=np.arange(0., 2., 1.), ax=ax_s, colors='k', add_colorbar=False, add_labels=False) if tibet: land.zsurf.plot.contourf(x='lon', y='lat', ax=ax_s, levels=np.arange(2500., 100001., 97500.), add_labels=False, extend='neither', add_colorbar=False, alpha=0.5, cmap='Greys_r') fv = (vor[pentad, :, :] * 86400.).plot.contourf(x='lon', y='lat', levels=np.arange( -12., 13., 2.), ax=ax_v, add_labels=False, extend='both', add_colorbar=False) land.land_mask.plot.contour(x='lon', y='lat', levels=np.arange(0., 2., 1.), ax=ax_v, colors='k', add_colorbar=False, add_labels=False) if tibet: land.zsurf.plot.contourf(x='lon', y='lat', ax=ax_v, levels=np.arange(2500., 100001., 97500.), add_labels=False, extend='neither', add_colorbar=False, alpha=0.5, cmap='Greys_r') for ax in [ax_h, ax_s, ax_v]: ax.set_xlim(60, 150) ax.set_ylim(-30, 60) ax.set_xticks(np.arange(60., 151., 30.)) ax.set_yticks(np.arange(-30., 61., 30.)) ax.grid(True, linestyle=':') ax_h.set_xticklabels('') ax_s.set_xticklabels('') ax_v.set_xticklabels(['60', '90', '120', '150']) return fh, fs, fv
def plot_vort_dev(video=False, threed=True): data_w = xr.open_dataset('/disca/share/rg419/jra_omega_pentad_clim.nc') data_u = xr.open_dataset('/disca/share/rg419/jra_ucomp_pentad_clim.nc') data_v_temp = xr.open_dataset( '/disca/share/rg419/jra_vcomp_pentad_clim.nc') # v has different time coord to u, presumably due to how Stephen has downloaded/averaged. I think the two are equivalent, so just substitute the time dimension into v data_v = xr.DataArray(data_v_temp.var34.values, coords={ 'pentad': data_u.pentad, 'lat': data_u.lat, 'lon': data_u.lon }, dims=('pentad', 'lat', 'lon')) print('files opened') data_u = data_u.var33 data_w = data_w.var39 zon_adv = data_u.sel(lev=20000.) * gr.ddx(data_u.sel(lev=20000.)) merid_adv = data_v * gr.ddy(data_u.sel(lev=20000.)) vert_adv = data_w * (gr.ddp(data_u, pname='lev')).sel(lev=20000.) * 100. sinphi = np.sin(data_u.lat * np.pi / 180.) f = 2. * mc.omega * sinphi if threed: rossby = (zon_adv + merid_adv + vert_adv) / (f * data_v) else: rossby = merid_adv / (f * data_v) # Start figure with 1 subplots rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 14 for i in range(72): fig, ax1 = plt.subplots() title = 'Pentad ' + str(int(data_u.pentad[i])) f1 = rossby.sel(pentad=i + 1).plot.contourf(x='lon', y='lat', ax=ax1, add_labels=False, add_colorbar=False, extend='both', zorder=1, levels=np.arange( 0., 1.1, 0.1)) # data_p.precipitation.sel(pentad=i+1).plot.contour(x='lon', y='lat', ax=ax1, add_labels=False, extend='both', zorder=1, levels=np.arange(2.,21.,2.), colors='k') ax1.grid(True, linestyle=':') ax1.set_ylim(-60., 60.) ax1.set_yticks(np.arange(-60., 61., 30.)) ax1.set_xticks(np.arange(0., 361., 90.)) ax1.set_title(title) land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc' land = xr.open_dataset(land_mask) land.lsm[0, :, :].plot.contour(ax=ax1, x='longitude', y='latitude', levels=np.arange(-1., 2., 1.), add_labels=False, colors='k') ax1.set_ylabel('Latitude') ax1.set_xlabel('Longitude') plt.subplots_adjust(left=0.1, right=0.97, top=0.93, bottom=0.05, hspace=0.25, 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) vidstr = '' if video: vidstr = 'video/' if threed: plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/jra/' + vidstr + '/rossby3d/' else: plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/jra/' + vidstr + '/rossby/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) if video: plt.savefig(plot_dir + 'rossby_and_precip_' + str(int(data_u.pentad[i])) + '.png', format='png') else: plt.savefig(plot_dir + 'rossby_and_precip_' + str(int(data_u.pentad[i])) + '.pdf', format='pdf') plt.close()
print('v loaded') data_t = xr.open_mfdataset( '/disca/share/reanalysis_links/NCEP_NCAR/air.mon.mean.nc') data_t = data_t['air'].load().loc['1948-01':'2016-12'] + 273.15 print('t loaded') data_q = xr.open_mfdataset( '/disca/share/reanalysis_links/NCEP_NCAR/shum.mon.mean.nc') data_q = data_q['shum'].load().loc['1948-01':'2016-12'] / 1000. print('q loaded') data_z = xr.open_mfdataset( '/disca/share/reanalysis_links/NCEP_NCAR/hgt.mon.mean.nc') data_z = data_z['hgt'].load().loc['1948-01':'2016-12'] * 9.81 print('z loaded') dvdx = gr.ddx(data_v, a=6371.0e3) dudy = gr.ddy(data_u, a=6371.0e3) data_vo = (dvdx.values - dudy.values) * 86400. data_vo = xr.DataArray(data_vo, dims=data_u.dims, coords=data_u.coords) lh = 2.5e6 cp = 1005. data_mse = (data_q * lh + data_t * cp + data_z) / 1000. print('Data loaded ok') land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc' land = xr.open_dataset(land_mask) def plot_winter_climate(data, title, levels_clim=np.arange(-50., 51., 5.), levels=np.arange(-5., 5.1, 0.5),
def vort_budg_terms(run, lonin=[-1.,361.], rot_fac=1.): '''Evaluate pentad mean vorticity budget and differences from daily snapshot budget RG 3/11/2017 Imputs: run = run_name lonin = longitude range to average over do_ss = run is steady state rot_fac = scale factor for Earth's rotation rate planetary_only = only plot planetary vorticity terms no_eddies = don't plot transient eddies too ll = keep longitude dimension for a lat-lon plot''' #Load in vorticity budget term means data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/vort_eq_'+run+'.nc') if run in ['ap_2', 'full_qflux']: data_vort = data data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/vort_eq_uv'+run+'.nc') print('vorticity budget data loaded') 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]] # Calculate vertical component of pentad mean vorticity parts: f and dv/dx - du/dy omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat *np.pi/180) v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy # Evaluate relative vorticity vor = v_dx - u_dy # Take divergencc of u*vor dvorudx = gr.ddx(vor*data.ucomp) dvorvdy = gr.ddy(vor*data.vcomp) div_uvor = -86400.**2 * (dvorudx + dvorvdy) dfudx = gr.ddx(data.ucomp * f) dfvdy = gr.ddy(data.vcomp * f) div_uf = -86400.**2 * (dfudx + dfvdy) # If run is ap_2 or full_qflux, then vorticity budget and velocities were saved separately. Load the budget up now if run in ['ap_2','full_qflux']: data = data_vort # Another glitch with ap_2 and full_qflux - these are only on 150 hPa. coord_dict = {'xofyear': ('xofyear', data.pentad), 'lat': ('lat', data.lat)} dim_list = ['xofyear', 'lat'] # For a Hovmoller, keep time, pressure, and lat else: coord_dict = {'xofyear': ('xofyear', data.pentad), 'pfull': ('pfull', data.pfull), 'lat': ('lat', data.lat)} dim_list = ['xofyear', 'pfull', 'lat'] # Specify output dictionary to be written output_dict = {'div_uvor': (dim_list, div_uvor.sel(lon=lons).mean('lon')), 'div_uf': (dim_list, div_uf.sel(lon=lons).mean('lon'))} if run in ['ap_2','full_qflux']: # Single level transients from ap_2, full_qflux mistakes transient_hm = ((data.stretching.sel(pfull=150).values + data.horiz_md.sel(pfull=150).values) * 86400.**2. - (div_uvor + div_uf).sel(pfull=150)) else: # Average out longitude for other plot types transient_hm = ((data.stretching + data.horiz_md)*86400.**2. - div_uvor - div_uf).sel(lon=lons).mean('lon') output_dict.update({'transient_hm': (dim_list, transient_hm)}) # Create a dataset of the terms to be plotted and return ds = xr.Dataset(output_dict, coords=coord_dict) return ds
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 heat_budg_hm_fluxform(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') uT_dx = -86400. * gr.ddx(data.ucomp_temp.sel(pfull=lev)).mean('lon') vT_dy = -86400. * gr.ddy(data.vcomp_temp.sel(pfull=lev)).mean('lon') wT_dp = -86400. * gr.ddp(data.omega_temp).sel(pfull=lev).mean('lon') uT_eddy = data.ucomp_temp - data.ucomp * data.temp vT_eddy = data.vcomp_temp - data.vcomp * data.temp wT_eddy = data.omega_temp - data.omega * data.temp #uT_dx = -86400. * (data.ucomp * gr.ddx(data.temp)).sel(pfull=lev).mean('lon') #vT_dy = -86400. * (data.vcomp * gr.ddy(data.temp, vector=False)).sel(pfull=lev).mean('lon') #wT_dp = -86400. * (data.omega * gr.ddp(data.temp)).sel(pfull=lev).mean('lon') uT_dx_eddy = -86400. * gr.ddx(uT_eddy.sel(pfull=lev)).mean('lon') vT_dy_eddy = -86400. * gr.ddy(vT_eddy.sel(pfull=lev)).mean('lon') wT_dp_eddy = -86400. * gr.ddp(wT_eddy).sel(pfull=lev).mean('lon') tdt_rad = data.tdt_rad.sel(pfull=lev).mean('lon')*86400. tdt_conv = data.dt_tg_convection.sel(pfull=lev).mean('lon')*86400. tdt_cond = data.dt_tg_condensation.sel(pfull=lev).mean('lon')*86400. tdt_diff = data.dt_tg_diffusion.sel(pfull=lev).mean('lon')*86400. diabatic = tdt_rad + tdt_conv + tdt_cond + tdt_diff eddies = uT_dx_eddy + vT_dy_eddy + wT_dp_eddy rho = data.pfull*100./data.temp/mc.rdgas asc_cool = (data.omega /(mc.cp_air * rho)).sel(pfull=lev).mean('lon') *86400. heat_sum = uT_dx + vT_dy + wT_dp + diabatic + asc_cool #+ eddies Tdt = gr.ddt(data.temp.sel(pfull=lev)).mean('lon') *86400. horiz_adv = uT_dx + vT_dy vertical_term = wT_dp + asc_cool levels = np.arange(-100,101.1,10.) #levels = np.arange(-20,21.,2.) #levels = np.arange(-3,3.1,0.2) #levels = np.arange(-1,1.1,0.1) # Nine subplots fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row') plt.set_cmap('RdBu_r') plot_vars = [horiz_adv, vertical_term, diabatic, heat_sum, Tdt, eddies] axes = [ax1,ax2,ax3,ax4,ax5,ax6] labels = ['a)','b)','c)','d)','e)','f)'] for i in range(6): 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) axes[i].set_ylim(-60,60) axes[i].grid(True,linestyle=':') axes[i].set_yticks(np.arange(-60.,61.,30.)) 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('Horizontal advection', fontsize=17) ax2.set_title('Vertical advection and expansion', fontsize=17) ax3.set_title('Diabatic', fontsize=17) ax4.set_title('Residual', fontsize=17) ax5.set_title('Temperature tendency', fontsize=17) ax6.set_title('Eddy convergence', 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) figname = 'heat_budg_fluxform_' +run+ '.pdf' plt.savefig(plot_dir + figname, format='pdf') plt.close()
def bs_mombudg(run, ax, pentad=45, lev=150., dayfac=3., rotfac=1.): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') damping = data.dt_ug_diffusion # + data.tdt_diss_rdamp dudy = gr.ddy(data.ucomp.mean('lon')) duvdy = gr.ddy(data.ucomp_vcomp.mean('lon'), uv=True) dudp = gr.ddp(data.ucomp.mean('lon')) duwdp = gr.ddp(data.ucomp_omega.mean('lon')) dudt = gr.ddt(data.ucomp.mean('lon')) vdudy_mean = data.vcomp.mean('lon') * dudy wdudp_mean = data.omega.mean('lon') * dudp adv_tend = -1. * (vdudy_mean + wdudp_mean) vert_adv = -1. * wdudp_mean horiz_adv = -1. * vdudy_mean eddy_tend = -1. * duvdy - 1. * duwdp - adv_tend vert_eddy = -1. * duwdp - vert_adv horiz_eddy = -1. * duvdy - horiz_adv sinphi = np.sin(data.lat * np.pi / 180.) cosphi = np.cos(data.lat * np.pi / 180.) coriolis = 2. * rotfac * mc.omega * sinphi * data.vcomp def get_latmax(data_in): data_max = data_in.max('lat') data_max_lat = data_in.lat.values[data_in.argmax('lat').values] data_max_lat = xr.DataArray(data_max_lat, coords=[data_in.xofyear], dims=['xofyear']) return data_max, data_max_lat w_850 = data.omega.sel(pfull=850.).mean('lon') w_max, w_max_lat = get_latmax(-1. * w_850) convTtotheta = (1000. / data.pfull)**mc.kappa theta_equiv_850 = ((data.temp + mc.L / mc.cp_air * data.sphum / (1 - data.sphum)) * convTtotheta).sel(pfull=850.).mean('lon') theta_max, theta_max_lat = get_latmax(theta_equiv_850) cosphimax = np.cos(theta_max_lat * np.pi / 180.) u_150 = data.ucomp.sel(pfull=lev).mean('lon') u_m = rotfac * mc.omega * mc.a * (cosphimax**2. - cosphi**2.) / cosphi def adj_u(u_in, adj_by, dayfac=dayfac): if 'lon' in adj_by.coords: return u_in + adj_by.sel(pfull=lev).mean('lon') * 86400. * dayfac else: return u_in + adj_by.sel(pfull=lev) * 86400. * dayfac u_damp = adj_u(u_150, damping) u_cor = adj_u(u_150, coriolis) u_adv = adj_u(u_150, adv_tend) u_vadv = adj_u(u_150, horiz_adv) u_hadv = adj_u(u_150, vert_adv) u_eddy = adj_u(u_150, eddy_tend) u_veddy = adj_u(u_150, horiz_eddy) u_heddy = adj_u(u_150, vert_eddy) resid = coriolis + damping + adv_tend + eddy_tend u_resid = adj_u(u_150, resid, dayfac=10. * dayfac) u_net = adj_u(u_150, dudt, dayfac=10. * dayfac) lats = [ data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -60. and data.lat[i] <= 60. ] u_m.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7') u_cor.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--') u_adv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.') u_eddy.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') u_150.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C0') u_net.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2') ax.plot([w_max_lat.sel(xofyear=pentad), w_max_lat.sel(xofyear=pentad)], [-75., 75.], color='0.7', linestyle=':') ax.set_title('') ax.set_xlabel('') ax.set_xlim(-35., 35.) ax.set_ylim(-75., 75.) ax.grid(True, linestyle=':') ax.set_ylabel('u, m/s')
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 plot_vort_dev(run, land_mask=None, lev=200, qscale=150., windtype='full', ref_arrow=10., video=False): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') sinphi = np.sin(data.lat * np.pi / 180.) zeta = (2. * mc.omega * sinphi - 1. * gr.ddy(data.ucomp)) * 86400. # Take zonal anomaly data_zanom = data - data.mean('lon') # Get rotational and divergent components of the flow w = VectorWind(data.ucomp.sel(pfull=lev), data.vcomp.sel(pfull=lev)) streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() uchi_zanom = (uchi - uchi.mean('lon')).sortby('lat') vchi_zanom = (vchi - vchi.mean('lon')).sortby('lat') upsi_zanom = (upsi - upsi.mean('lon')).sortby('lat') vpsi_zanom = (vpsi - vpsi.mean('lon')).sortby('lat') # Start figure with 1 subplots rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 14 for i in range(72): fig, ax1 = plt.subplots() title = 'Pentad ' + str(int(data.xofyear[i])) f1 = zeta.sel(xofyear=i + 1, pfull=lev).plot.contourf(x='lon', y='lat', ax=ax1, add_labels=False, add_colorbar=False, extend='both', zorder=1, levels=np.arange(-10., 10., 2.)) if windtype == 'div': b = ax1.quiver(data.lon[::6], data.lat[::3], uchi_zanom[i, ::3, ::6], vchi_zanom[i, ::3, ::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) ax1.quiverkey(b, 0., -0.5, ref_arrow, str(ref_arrow) + ' m/s', fontproperties={ 'weight': 'bold', 'size': 10 }, color='k', labelcolor='k', labelsep=0.03, zorder=10) elif windtype == 'rot': b = ax1.quiver(data.lon[::6], data.lat[::3], upsi_zanom[i, ::3, ::6], vpsi_zanom[i, ::3, ::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) ax1.quiverkey(b, 0., -0.5, ref_arrow, str(ref_arrow) + ' m/s', fontproperties={ 'weight': 'bold', 'size': 10 }, color='k', labelcolor='k', labelsep=0.03, zorder=10) elif windtype == 'full': b = ax1.quiver(data.lon[::6], data.lat[::3], data_zanom.ucomp.sel(pfull=lev)[i, ::3, ::6], data_zanom.vcomp.sel(pfull=lev)[i, ::3, ::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) ax1.quiverkey(b, 0., -0.5, ref_arrow, str(ref_arrow) + ' m/s', fontproperties={ 'weight': 'bold', 'size': 10 }, color='k', labelcolor='k', labelsep=0.03, zorder=10) else: windtype = 'none' ax1.grid(True, linestyle=':') ax1.set_ylim(-60., 60.) ax1.set_yticks(np.arange(-60., 61., 30.)) ax1.set_xticks(np.arange(0., 361., 90.)) ax1.set_title(title) if not land_mask == None: land = xr.open_dataset(land_mask) land.land_mask.plot.contour(x='lon', y='lat', ax=ax1, levels=np.arange(-1., 2., 1.), add_labels=False, colors='k') ax1.set_ylabel('Latitude') ax1.set_xlabel('Longitude') plt.subplots_adjust(left=0.1, right=0.97, top=0.93, bottom=0.05, hspace=0.25, 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) levstr = '' windtypestr = '' msestr = '' vidstr = '' if lev != 850: levstr = '_' + str(lev) if windtype != 'full': windtypestr = '_' + windtype if video: vidstr = 'video/' plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/' + run + '/' + vidstr + windtype + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) if video: plt.savefig(plot_dir + 'wind_and_vort_zanom_' + str(int(data.xofyear[i])) + levstr + windtypestr + '.png', format='png') else: plt.savefig(plot_dir + 'wind_and_vort_zanom_' + str(int(data.xofyear[i])) + levstr + windtypestr + '.pdf', format='pdf') plt.close()
def fig_9_moist(run, ax, pentad=40): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') convTtotheta = (1000. / data.pfull)**mc.kappa heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad hum_tend = data.dt_qg_condensation + data.dt_qg_convection + data.dt_qg_diffusion heating_theta_equiv = ((heating + mc.L / mc.cp_air * hum_tend / (1 - data.sphum)**2.) * convTtotheta).mean('lon') theta_equiv = (data.temp + mc.L / mc.cp_air * data.sphum / (1 - data.sphum)) * convTtotheta dthetady_equiv = gr.ddy(theta_equiv.mean('lon'), vector=False) dthetadp_equiv = gr.ddp(theta_equiv.mean('lon')) dthetadt_equiv = gr.ddt(theta_equiv.mean('lon')) vcomp_theta_equiv = (data.vcomp_temp + mc.L / mc.cp_air * data.sphum_v) * convTtotheta vcomp_theta_eddy_equiv = vcomp_theta_equiv.mean( 'lon') - data.vcomp.mean('lon') * theta_equiv.mean('lon') vdthetady_mean_equiv = data.vcomp.mean('lon') * dthetady_equiv wdthetadp_mean_equiv = data.omega.mean('lon') * dthetadp_equiv def column_int(var_in): var_int = mc.cp_air * var_in.sum('pfull') * 5000. / mc.grav return var_int vdtdy_mean_int_equiv = -1. * column_int(vdthetady_mean_equiv) wdtdp_mean_int_equiv = -1. * column_int(wdthetadp_mean_equiv) vt_eddy_int_equiv = -1. * column_int(vcomp_theta_eddy_equiv) div_vt_eddy_int_equiv = gr.ddy(vt_eddy_int_equiv, vector=True) heating_theta_int_equiv = column_int(heating_theta_equiv) dthetadt_int_equiv = column_int(dthetadt_equiv) lats = [ data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -60. and data.lat[i] <= 60. ] w_850 = data.omega.sel(pfull=850.).mean('lon') w_max, w_max_lat = get_latmax(-1. * w_850) dthetadt_int_equiv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2') heating_theta_int_equiv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--') (vdtdy_mean_int_equiv + wdtdp_mean_int_equiv).sel( xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.') #(vdtdy_mean_int_equiv + wdtdp_mean_int_equiv + heating_theta_int_equiv + div_vt_eddy_int_equiv).sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') div_vt_eddy_int_equiv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') ax.plot([w_max_lat.sel(xofyear=pentad), w_max_lat.sel(xofyear=pentad)], [-250., 250.], color='0.7', linestyle=':') ax.set_title('') ax.set_xlabel('') ax.set_ylim(-250., 250.) ax.set_xlim(-35., 35.) ax.grid(True, linestyle=':') ax.set_ylabel('Heating rate, W/m$^2$')
def vort_budg_terms(run, lonin=[-1.,361.], do_ss=False, rot_fac=1., planetary_only=False, no_eddies=False, ll=False): '''Evaluate pentad mean vorticity budget and differences from daily snapshot budget RG 3/11/2017 Imputs: run = run_name lonin = longitude range to average over do_ss = run is steady state rot_fac = scale factor for Earth's rotation rate planetary_only = only plot planetary vorticity terms no_eddies = don't plot transient eddies too ll = keep longitude dimension for a lat-lon plot''' #Load in vorticity budget term means data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_'+run+'.nc') if run in ['ap_2', 'full_qflux']: data_vort = data data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_uv'+run+'.nc') print('vorticity budget data loaded') 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]] # Calculate vertical component of pentad mean absolute vorticity = f + dv/dx - du/dy omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat *np.pi/180) v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy # If only want to look at parts relating to planetary vorticity, substract relative part (this way turns f into a 4d array) if planetary_only: vor = v_dx - u_dy + f - v_dx + u_dy else: vor = v_dx - u_dy + f # Take gradients of vorticity dvordx = gr.ddx(vor) dvordy = gr.ddy(vor, vector=False) # Horizontal material derivative horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy) # Calculate divergence and stretching term div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp) stretching_mean = -86400.**2. * vor * div if ll: #Keep lon dimension if want a steady state lat-lon plot horiz_md_av = horiz_md_mean stretching_av = stretching_mean dvordy_av = dvordy vor_av = vor*86400. v_av = data.vcomp div_av = div*86400. else: # Take zonal mean of all over specified longitude range horiz_md_av = horiz_md_mean.sel(lon=lons).mean('lon') stretching_av = stretching_mean.sel(lon=lons).mean('lon') dvordy_av = dvordy.sel(lon=lons).mean('lon') vor_av = vor.sel(lon=lons).mean('lon')*86400. v_av = data.vcomp.sel(lon=lons).mean('lon') div_av = div.sel(lon=lons).mean('lon')*86400. # If run is ap_2 or full_qflux, then vorticity budget and velocities were saved separately. Load the budget up now if run in ['ap_2','full_qflux']: data = data_vort # For a steady state lat-pressure run, only need pfull and lat if do_ss: coord_dict = {'pfull': ('pfull', data.pfull), 'lat': ('lat', data.lat)} dim_list = ['pfull', 'lat'] # For a steady state lat-lon run, also need lon. Keep pfull to allow a choice of plotting level elif ll: if 'pentad' in horiz_md_av.coords: coord_dict = {'xofyear': ('xofyear', data.pentad), 'pfull': ('pfull', data.pfull), 'lat': ('lat', data.lat), 'lon': ('lon', data.lon)} dim_list = ['xofyear', 'pfull', 'lat', 'lon'] else: coord_dict = {'pfull': ('pfull', data.pfull), 'lat': ('lat', data.lat), 'lon': ('lon', data.lon)} dim_list = ['pfull', 'lat', 'lon'] # Another glitch with ap_2 and full_qflux - these are only on 150 hPa. elif run in ['ap_2','full_qflux']: coord_dict = {'xofyear': ('xofyear', data.pentad), 'lat': ('lat', data.lat)} dim_list = ['xofyear', 'lat'] # For a Hovmoller, keep time, pressure, and lat else: coord_dict = {'xofyear': ('xofyear', data.pentad), 'pfull': ('pfull', data.pfull), 'lat': ('lat', data.lat)} dim_list = ['xofyear', 'pfull', 'lat'] # Specify output dictionary to be written output_dict = {'horiz_md': (dim_list, horiz_md_av), 'stretching': (dim_list, stretching_av), 'dvordy': (dim_list, dvordy_av), 'v': (dim_list, v_av), 'vor': (dim_list, vor_av), 'div': (dim_list, div_av)} if not (planetary_only or no_eddies): if run in ['ap_2','full_qflux']: # Single level transients from ap_2, full_qflux mistakes transient_s_hm = (data.stretching.sel(pfull=150).values * 86400.**2. - stretching_mean).sel(lon=lons).mean('lon') transient_h_hm = (data.horiz_md.sel(pfull=150).values * 86400.**2. - horiz_md_mean).sel(lon=lons).mean('lon') elif ll: # Lat lon transients transient_s_hm = data.stretching.values * 86400.**2. - stretching_mean transient_h_hm = data.horiz_md.values * 86400.**2. - horiz_md_mean else: # Average out longitude for other plot types transient_s_hm = (data.stretching.values * 86400.**2. - stretching_mean).sel(lon=lons).mean('lon') transient_h_hm = (data.horiz_md.values * 86400.**2. - horiz_md_mean).sel(lon=lons).mean('lon') transient_hm = transient_s_hm + transient_h_hm output_dict.update({'transient_hm': (dim_list, transient_hm), 'transient_s_hm': (dim_list, transient_s_hm), 'transient_h_hm': (dim_list, transient_h_hm)}) # Create a dataset of the terms to be plotted and return ds = xr.Dataset(output_dict, coords=coord_dict) return ds
def fig_8_moist(run, ax, pentad=45, rotfac=1., dayfac=5.): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad hum_tend = data.dt_qg_condensation + data.dt_qg_convection + data.dt_qg_diffusion convTtotheta = (1000. / data.pfull)**mc.kappa theta = data.temp * convTtotheta heating_theta = heating * convTtotheta heating_equiv_theta = heating_theta + mc.L / mc.cp_air * hum_tend / ( 1 - data.sphum)**2. * convTtotheta sinphi = np.sin(data.lat * np.pi / 180.) cosphi = np.cos(data.lat * np.pi / 180.) theta_850 = theta.sel(pfull=850.).mean('lon') theta_equiv = (data.temp + mc.L / mc.cp_air * data.sphum / (1 - data.sphum)) * convTtotheta theta_equiv_850 = theta_equiv.sel(pfull=850.).mean('lon') lats = [ data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -10. and data.lat[i] <= 10. ] Tt = (data.temp.sel(pfull=150.).mean(('lon')) * cosphi).sel(lat=lats).sum('lat') / cosphi.sel(lat=lats).sum('lat') Ts = (data.t_surf.mean(('lon')) * cosphi).sel(lat=lats).sum('lat') / cosphi.sel(lat=lats).sum('lat') dthetady_equiv = gr.ddy(theta_equiv.mean('lon'), vector=False) dthetadp_equiv = gr.ddp(theta_equiv.mean('lon')) dthetadt_equiv = gr.ddt(theta_equiv.mean('lon')) vdthetady_mean = data.vcomp.mean('lon') * dthetady_equiv wdthetadp_mean = data.omega.mean('lon') * dthetadp_equiv adv_heating_equiv = -1. * (vdthetady_mean + wdthetadp_mean) w_850 = data.omega.sel(pfull=850.).mean('lon') w_max, w_max_lat = get_latmax(-1. * w_850) theta_max, theta_max_lat = get_latmax(theta_equiv_850) cosphimax = np.cos(theta_max_lat * np.pi / 180.) chi = (rotfac * mc.omega)**2 * mc.a**2. / mc.cp_air / (Ts - Tt) theta_m_equiv = theta_max * np.exp( -1. * chi * (cosphimax**2. - cosphi**2.)**2. / cosphi**2.) def adj_theta(theta_in, adj_by, dayfac=dayfac): if 'lon' in adj_by.coords: return theta_in + adj_by.sel( pfull=850.).mean('lon') * 86400. * dayfac else: return theta_in + adj_by.sel(pfull=850.) * 86400. * dayfac theta_equiv_rc = adj_theta(theta_equiv_850, heating_equiv_theta, dayfac=dayfac) theta_equiv_adv = adj_theta(theta_equiv_850, adv_heating_equiv, dayfac=dayfac) theta_equiv_net = adj_theta(theta_equiv_850, dthetadt_equiv, dayfac=dayfac * 3.) lats = [ data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -60. and data.lat[i] <= 60. ] theta_m_equiv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7') theta_equiv_rc.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--') theta_equiv_adv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.') theta_equiv_850.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C0') theta_equiv_net.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2') ax.plot([w_max_lat.sel(xofyear=pentad), w_max_lat.sel(xofyear=pentad)], [300., 380.], color='0.7', linestyle=':') ax.set_title('') ax.set_xlabel('') ax.set_ylim(300., 380.) ax.set_xlim(-35., 35.) ax.grid(True, linestyle=':') ax.set_ylabel('$\Theta$, K')
def vort_eq(run, month, filename='plev_daily', period_fac=1., rot_fac=1., do_ss=False, day_fac=1.): '''Evaluate terms in the vorticity budget using daily data from simulations. RG 2/11/17 Inputs: run = run name month = month to evaluate filename = name of file to look for, default is pressure interpolated daily data period_fac = ratio of the orbital period to Earth's rot_fac = ratio of the rotation rate to Earth's do_ss = evaluate for a steady state simulation, do not average daily data, average all saved data together ''' #Load in dataset try: name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc' name = name_temp % month #read data into xarr data = xr.open_dataset(name, decode_times=False) except: name_temp = '/scratch/rg419/Data_moist/' + run + '/run%04d/' + filename + '.nc' name = name_temp % month #read data into xarr data = xr.open_dataset(name, decode_times=False) # Calculate planetary vorticity omega = 7.2921150e-5 * rot_fac f = 2 * omega * np.sin(data.lat * np.pi / 180) # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy v_dx = gr.ddx(data.vcomp) # dvdx u_dy = gr.ddy(data.ucomp) # dudy vor = v_dx - u_dy + f # Take gradients of vorticity dvordx = gr.ddx(vor) dvordy = gr.ddy(vor, vector=False) dvordp = gr.ddp(vor) # Calculate horizontal material derivative horiz_md = -1. * (data.ucomp * dvordx + data.vcomp * dvordy) # Calculate vertical material derivative vert_md = -1. * data.omega * dvordp # Now do the terms involving gradients of windspeed div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp) stretching = -1. * vor * div tilting = gr.ddp(data.ucomp) * gr.ddy(data.omega, vector=False) - gr.ddp( data.vcomp) * gr.ddx(data.omega) total = horiz_md + vert_md + stretching + tilting ds = xr.Dataset( { 'horiz_md': (['time', 'pfull', 'lat', 'lon'], horiz_md), 'vert_md': (['time', 'pfull', 'lat', 'lon'], vert_md), 'stretching': (['time', 'pfull', 'lat', 'lon'], stretching), 'tilting': (['time', 'pfull', 'lat', 'lon'], tilting), 'ucomp': (['time', 'pfull', 'lat', 'lon'], data.ucomp), 'vcomp': (['time', 'pfull', 'lat', 'lon'], data.vcomp), 'total': (['time', 'pfull', 'lat', 'lon'], total) }, coords={ 'time': ('time', data.time), 'pfull': ('pfull', data.pfull), 'lat': ('lat', data.lat), 'lon': ('lon', data.lon) }) # If a steady state run, leave time dimension in to average out later. Otherwise take pentad means for climatology if do_ss: dsout = ds else: ds.coords['xofyear'] = np.mod(ds.time / day_fac - 1., 360. * period_fac) // 5 + 1. dsout = ds.groupby('xofyear').mean(('time')) try: fileout = '/scratch/rg419/Data_moist/' + run + '/run%03d/vort_eq.nc' fileout = fileout % month dsout.to_netcdf(path=fileout) except: fileout = '/scratch/rg419/Data_moist/' + run + '/run%04d/vort_eq.nc' fileout = fileout % month dsout.to_netcdf(path=fileout) print 'data written to ', fileout return dsout
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 fig_8(run, ax, pentad=45, lev=200., dayfac=3., rotfac=1.): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') sinphi = np.sin(data.lat * np.pi/180.) cosphi = np.cos(data.lat * np.pi/180.) zeta = 2.* rotfac * mc.omega*sinphi -1.* gr.ddy(data.ucomp.mean('lon')) #* 86400. nu = gr.ddp(data.ucomp.mean('lon')) #* 86400. dzetady = gr.ddy(zeta, vector=False) dzetadp = gr.ddp(zeta) dwdy = gr.ddy(data.omega, vector=False) dvdy = gr.ddy(data.vcomp) dwdp = gr.ddp(data.omega) dzetadt = gr.ddt(zeta) vdzetady_mean = data.vcomp.mean('lon') * dzetady wdzetadp_mean = data.omega.mean('lon') * dzetadp adv_tend = -1. *(vdzetady_mean + wdzetadp_mean) vert_adv = -1.*wdzetadp_mean horiz_adv = -1.*vdzetady_mean stretching = -1.* zeta * dvdy tilting = nu * dwdy eddy_tend = dzetadt - adv_tend - stretching - tilting w_850 = data.omega.sel(pfull=850.).mean('lon') w_max = w_850.min('lat') w_max_lat = data.lat.values[w_850.argmin('lat').values] w_max_lat = xr.DataArray(w_max_lat, coords=[data.xofyear], dims=['xofyear']) cosphimax = np.cos(w_max_lat * np.pi/180.) u_m = rotfac * mc.omega * mc.a * (cosphimax**2. - cosphi**2.)/cosphi zeta_m = 2.*mc.omega*sinphi -1.* gr.ddy(u_m) def adj_zeta(zeta_in, adj_by, dayfac=dayfac): if 'lon' in adj_by.coords: return zeta_in + adj_by.sel(pfull=lev).mean('lon') * 86400.**2. * dayfac else: return zeta_in + adj_by.sel(pfull=lev) * 86400.**2. * dayfac zeta_150 = zeta.sel(pfull=150.)*86400. zeta_adv = adj_zeta(zeta_150, adv_tend); zeta_vadv = adj_zeta(zeta_150, horiz_adv); zeta_hadv = adj_zeta(zeta_150, vert_adv) zeta_stretching = adj_zeta(zeta_150, stretching) zeta_tilting = adj_zeta(zeta_150, tilting) zeta_eddy = adj_zeta(zeta_150, eddy_tend); zeta_net = adj_zeta(zeta_150, dzetadt, dayfac=10.*dayfac) lats = [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -60. and data.lat[i] <= 60.] zeta_m.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7', linestyle='-') zeta_stretching.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--') zeta_tilting.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7', linestyle='--') zeta_adv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.') #u_hadv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C1', linestyle='-.') #u_vadv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2', linestyle='-.') zeta_eddy.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':') zeta_150.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C0') zeta_net.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2') ax.plot([w_max_lat.sel(xofyear=pentad),w_max_lat.sel(xofyear=pentad)], [-7.5,7.5], color='0.7', linestyle=':') ax.set_title('') ax.set_xlabel('') ax.set_xlim(-35.,35.) ax.set_ylim(-7.5,7.5) ax.grid(True,linestyle=':') ax.set_ylabel('zeta, /s')
def h_w_mass_flux_monthly(run, lev=500., dp=5000.): data = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/era_v_clim_alllevs.nc') data_u = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/era_u_clim_alllevs.nc') plot_dir = '/scratch/rg419/plots/overturning_monthly/era/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data.coords['month'] = (data.xofyear - 1) // 6 + 1 data = data.groupby('month').mean(('xofyear')) # Create a VectorWind instance to handle the computation w = VectorWind(data.ucomp.sel(pfull=np.arange(50., 950., 50.)), data.vcomp.sel(pfull=np.arange(50., 950., 50.))) # Compute variables streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() coslat = np.cos(data.lat * np.pi / 180) # Evaluate mass fluxes for the zonal and meridional components (Walker and Hadley) following Schwendike et al. 2014 mass_flux_zon = (gr.ddx(uchi)).cumsum('pfull') * dp * coslat / mc.grav mass_flux_merid = (gr.ddy(vchi)).cumsum('pfull') * dp * coslat / mc.grav # Set figure parameters rcParams['figure.figsize'] = 15, 11 rcParams['font.size'] = 14 # Start figure with 12 subplots fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8), (ax9, ax10, ax11, ax12)) = plt.subplots(3, 4, sharex='col', sharey='row') axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12] i = 0 for ax in axes: f1 = mass_flux_merid.sel(pfull=lev)[i, :, :].plot.contourf( ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, levels=np.arange(-0.0065, 0.0066, 0.001), extend='both') mass_flux_zon.sel(pfull=lev)[i, :, :].plot.contour(ax=ax, x='lon', y='lat', add_labels=False, colors='k', levels=np.arange( 0.0005, 0.0066, 0.001)) mass_flux_zon.sel(pfull=lev)[i, :, :].plot.contour( ax=ax, x='lon', y='lat', add_labels=False, colors='0.5', levels=np.arange(-0.0065, -0.00049, 0.001)) i = i + 1 ax.set_ylim(-60, 60) ax.set_xticks(np.arange(0, 361, 90)) ax.set_yticks(np.arange(-60, 61, 30)) ax.grid(True, linestyle=':') plt.subplots_adjust(left=0.05, right=0.97, top=0.95, bottom=0.1, hspace=0.2, wspace=0.2) cb1 = fig.colorbar(f1, ax=axes, use_gridspec=True, orientation='horizontal', fraction=0.05, pad=0.1, aspect=30, shrink=0.5) cb1.set_label( 'Vertical mass flux associated with meridional circulation, kgm$^{-2}$s$^{-1}$' ) figname = plot_dir + 'h_w_' + run + '.pdf' plt.savefig(figname, format='pdf') plt.close()