def walker_hm(regions=[[0,10], [10,20], [20,30], [30,40]]): 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' ) # Take pentad means data.coords['pentad'] = data.day_of_yr //5 + 1. data = data.groupby('pentad').mean(('day_of_yr')) data_u.coords['pentad'] = data_u.day_of_yr //5 + 1. data_u = data_u.groupby('pentad').mean(('day_of_yr')) plot_dir = '/scratch/rg419/plots/overturning_monthly/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) # Create a VectorWind instance to handle the computation w = VectorWind(data_u.u.sel(pfull=np.arange(50.,950.,50.)), data.v.sel(pfull=np.arange(50.,950.,50.))) # Compute variables streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() # Set figure parameters rcParams['figure.figsize'] = 10, 7 rcParams['font.size'] = 14 # Start figure with 4 subplots fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex='col', sharey='row') axes = [ax1,ax2,ax3,ax4] i=0 for ax in axes: psi_w = walker_cell(uchi, latin=regions[i], dp_in=-50.) psi_w /= 1.e9 i=i+1 f1=psi_w.sel(pfull=500).plot.contourf(ax=ax, x='lon', y='pentad', add_labels=False, add_colorbar=False, levels=np.arange(-300.,301.,50.), extend='both') ax1.set_title('0-10 N') ax2.set_title('10-20 N') ax3.set_title('20-30 N') ax4.set_title('30-40 N') for ax in [ax1,ax2,ax3,ax4]: ax.grid(True,linestyle=':') ax.set_yticks(np.arange(0.,73., 18.)) ax.set_xticks([0.,90.,180.,270.,360.]) ax3.set_xlabel('Longitude') ax4.set_xlabel('Longitude') ax1.set_ylabel('Pentad') ax3.set_ylabel('Pentad') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) 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('Zonal overturning streamfunction') # Save as a pdf plt.savefig(plot_dir + 'walker_cell_hm_era.pdf', format='pdf') plt.close() data.close()
def walker_cell_monthly(run, latin=None): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') plot_dir = '/scratch/rg419/plots/overturning_monthly/' + run + '/' 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, and compute variables w = VectorWind(data.ucomp.sel(pfull=np.arange(50., 950., 50.)), data.vcomp.sel(pfull=np.arange(50., 950., 50.))) uchi, vchi, upsi, vpsi = w.helmholtz() psi_w = walker_cell(uchi, latin=latin, dp_in=-50.) psi_w /= 1.e9 # Set figure parameters rcParams['figure.figsize'] = 12, 7 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) axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12] i = 0 for ax in axes: psi_w.sel(month=i + 1).plot.contour(ax=ax, x='lon', y='pfull', yincrease=False, levels=np.arange(0., 301, 50.), colors='k', add_labels=False) psi_w.sel(month=i + 1).plot.contour(ax=ax, x='lon', y='pfull', yincrease=False, levels=np.arange(-300., 0., 50.), colors='k', linestyles='dashed', add_labels=False) i = i + 1 ax.set_xticks(np.arange(0., 361., 120.)) ax.set_yticks(np.arange(0., 1001., 250.)) ax.grid(True, linestyle=':') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) plt.savefig(plot_dir + 'walker_' + run + '.pdf', format='pdf') plt.close()
def walker_strength(data, lonin=[90.,180.], psi_min=True, sanity_check=False): ''' data - a climatology from Isca with dimensions including xofyear lonin - longitudes over which to look for min/max values psi_min - if true look for minimum of walker streamfunction, otherwise maximum sanity_check - produce a plot on screen to check the max/min has been identified correctly ''' lons = np.arange(lonin[0], lonin[1], 0.1) # create hi res longitude array to interpolate to # Create a VectorWind instance, calculate uchi, and use walker_cell to get the streamfunction from this. w = VectorWind(data.ucomp.sel(pfull=np.arange(50.,950.,50.)), data.vcomp.sel(pfull=np.arange(50.,950.,50.))) uchi, vchi, upsi, vpsi = w.helmholtz() psi_w = walker_cell(uchi, latin=[-20,20], dp_in=-50.) # Divide by 1e9 and select 500 hPa level psi_w /= 1.e9 psi_w = psi_w.sel(pfull=500.) # interpolate to hi res lons f = spint.interp1d(psi_w.lon, psi_w, axis=1, fill_value='extrapolate') psi_w = f(lons) psi_w = xr.DataArray(psi_w, coords=[uchi.xofyear, lons], dims=['xofyear', 'lon']) # mask to include only pos/neg values psi_mask = np.ones(psi_w.values.shape) if psi_min: psi_mask[psi_w > 0.] = np.nan else: psi_mask[psi_w < 0.] = np.nan psi_masked = psi_w * psi_mask # At some times there will be no pos/neg value. Find times where there are lons with non-NaN values times_defd = [] for i in range(0, len(psi_masked.xofyear)): if np.any(np.isfinite(psi_masked[i,:])): times_defd.append(np.float(psi_masked.xofyear[i])) # reduce masked array times to include only non-Nan values psi_red = psi_masked.sel(xofyear=times_defd) # Find min/max and its location if psi_min: walker_mag = psi_red.min('lon') walker_mag_loc = psi_red.lon.values[psi_red.argmin('lon').values] else: walker_mag = psi_red.max('lon') walker_mag_loc = psi_red.lon.values[psi_red.argmax('lon').values] walker_mag_loc = xr.DataArray(walker_mag_loc, coords=[psi_red.xofyear], dims=['xofyear']) # plot the streamfunction and the location of the maximum as a sanity check if sanity_check: psi_w.plot.contourf(x='xofyear',y='lon') walker_mag_loc.plot() #plt.figure(2) #plt.plot(walker_mag_loc,walker_mag) plt.show() return walker_mag
def cross_prod(run, months, filename='plev_pentad', timeav='month', period_fac=1., land_mask=False, level=9): plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/cross_prod/' + str( level) + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) mn_dic = month_dic(1) data = time_means(run, months, filename=filename, timeav=timeav, period_fac=period_fac) uwnd = data.ucomp[:, level, :, :] vwnd = data.vcomp[:, level, :, :] w = VectorWind(uwnd, vwnd) uchi, vchi, upsi, vpsi = w.helmholtz() cross_prod = (upsi * vchi - vpsi * uchi) for i in range(0, 12): ax = cross_prod[i, :, :].plot.contourf(x='lon', y='lat', levels=np.arange( -200., 201., 10.), extend='both', add_colorbar=False, add_label=False) cb1 = plt.colorbar(ax) cb1.set_label('Vpsi x Vchi') if land_mask: land = xr.open_dataset( '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc') land_plot = xr.DataArray(land.land_mask.values, [('lat', data.lat), ('lon', data.lon)]) land_plot.plot.contour(x='lon', y='lat', levels=np.arange(0., 2., 1.), colors='0.75', add_colorbar=False, add_labels=False) plt.ylabel('Latitude') plt.xlabel('Longitude') plt.ylim(-10, 45) plt.xlim(25, 150) plt.tight_layout() plt.savefig(plot_dir + str(i + 1) + '_' + str(mn_dic[i + 1]) + '.png') plt.close()
def ke_partition(run, months, filename='atmos_pentad', timeav='pentad', period_fac=1.): data = time_means(run, months, filename=filename, timeav=timeav, period_fac=period_fac) totp = (data.convection_rain + data.condensation_rain) * 86400. cell_ar = cell_area(42, '/scratch/rg419/GFDL_model/GFDLmoistModel/') cell_ar = xr.DataArray(cell_ar, [('lat', data.lat), ('lon', data.lon)]) uwnd = data.ucomp vwnd = data.vcomp w = VectorWind(uwnd, vwnd) uchi, vchi, upsi, vpsi = w.helmholtz() lats = [ i for i in range(len(data.lat)) if data.lat[i] >= 5. and data.lat[i] < 30. ] lons = [ i for i in range(len(data.lon)) if data.lon[i] >= 60. and data.lon[i] < 150. ] ke_chi = 0.5 * (uchi * uchi + vchi * vchi) * cell_ar ke_chi_av = ke_chi[:, :, lats, lons].sum('lat').sum('lon') / cell_ar[ lats, lons].sum('lat').sum('lon') ke_psi = 0.5 * (upsi * upsi + vpsi * vpsi) * cell_ar ke_psi_av = ke_psi[:, :, lats, lons].sum('lat').sum('lon') / cell_ar[ lats, lons].sum('lat').sum('lon') totp_av = (totp[:, lats, lons] * cell_ar[lats, lons]).sum('lat').sum( 'lon') / cell_ar[lats, lons].sum('lat').sum('lon') ke_chi_av[:, 36].plot() ke_psi_av[:, 36].plot() plt.legend(['KE_chi', 'KE_psi']) plt.xlabel('Pentad') plt.ylabel('Kinetic Energy, m2/s2') plt.savefig(plot_dir + 'KE_' + run + '.png') plt.close() totp_av.plot() plt.xlabel('Pentad') plt.ylabel('Precipitation, mm/day') plt.savefig(plot_dir + 'precip_' + run + '.png') plt.close()
def mass_streamfunction(data, a=6376.0e3, g=9.8, lons=[-1000], dp_in=0., intdown=True, use_v_locally=False): """Calculate the mass streamfunction for the atmosphere. Based on a vertical integral of the meridional wind. Ref: Physics of Climate, Peixoto & Oort, 1992. p158. `a` is the radius of the planet (default Isca value 6376km). `g` is surface gravity (default Earth 9.8m/s^2). lons allows a local area to be used by specifying boundaries as e.g. [70,100] dp_in - if no phalf and if using regularly spaced pressure levels, use this increment for integral. Units hPa. intdown - choose integratation direction (i.e. from surface to TOA, or TOA to surface). Returns an xarray DataArray of mass streamfunction. """ if lons[0] == -1000: #Use large negative value to use all data if no lons provided vbar = data.vcomp.mean('lon') elif use_v_locally: vbar = data.vcomp.sel(lon=lons).mean('lon').sortby('pfull', ascending=False) else: from windspharm.xarray import VectorWind # 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 uchi, vbar, upsi, vpsi = w.helmholtz() vbar = vbar.sel(lon=lons).mean('lon').sortby('pfull', ascending=False) c = 2 * np.pi * a * np.cos(vbar.lat * np.pi / 180) / g # take a diff of half levels, and assign to pfull coordinates if dp_in == 0.: dp = xr.DataArray(data.phalf.diff('phalf').values * 100., coords=[('pfull', data.pfull)]) return c * np.cumsum(vbar * dp, axis=vbar.dims.index('pfull')) else: dp = dp_in * 100. if intdown: psi = c * dp * np.flip(vbar, axis=vbar.dims.index('pfull')).cumsum( 'pfull') #[:,:,::-1] else: psi = -1. * c * dp * vbar.cumsum('pfull') #[:,:,::-1] #psi.plot.contourf(x='lat', y='pfull', yincrease=False) #plt.show() return psi #c*dp*vbar[:,::-1,:].cumsum('pfull')#[:,:,::-1]
def overturning_monthly(run, lonin=[-1., 361.]): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') plot_dir = '/scratch/rg419/plots/overturning_monthly/' + run + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data.coords['month'] = (data.xofyear - 1) // 6 + 1 data = data.groupby('month').mean(('xofyear')) #data['vcomp'] = data.vcomp.fillna(0.) #data['ucomp'] = data.ucomp.fillna(0.) # 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.))) #w = VectorWind(data.ucomp, data.vcomp) # Compute variables streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() #print(vchi.pfull) #print(data.pfull) #data.vcomp.mean('lon')[0,:,:].plot.contourf(x='lat', y='pfull', yincrease=False, add_labels=False) #plt.figure(2) #vchi.mean('lon')[0,:,:].plot.contourf(x='lat', y='pfull', yincrease=False, add_labels=False) #plt.show() ds_chi = xr.Dataset({'vcomp': (vchi)}, coords={ 'month': ('month', vchi.month), 'pfull': ('pfull', vchi.pfull), 'lat': ('lat', vchi.lat), 'lon': ('lon', vchi.lon) }) ds_psi = xr.Dataset({'vcomp': (vpsi)}, coords={ 'month': ('month', vchi.month), 'pfull': ('pfull', vchi.pfull), 'lat': ('lat', vchi.lat), 'lon': ('lon', vchi.lon) }) def get_lons(lonin, data): 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] ] return lons lons = get_lons(lonin, data) psi = mass_streamfunction(data, lons=lons, dp_in=50., use_v_locally=True) psi /= 1.e9 psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=50.) psi_chi /= 1.e9 # Set figure parameters rcParams['figure.figsize'] = 10, 7 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) axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12] i = 0 for ax in axes: psi[:, i, :].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0., 601, 100.), colors='k', add_labels=False) psi[:, i, :].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600., 0., 100.), colors='k', linestyles='dashed', add_labels=False) f1 = data.ucomp.sel(month=i + 1).sel(lon=lons).mean('lon').plot.contourf( ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-50., 50.1, 5.), extend='both', add_labels=False, add_colorbar=False) m = mc.omega * mc.a**2. * np.cos( psi.lat * np.pi / 180.)**2. + data.ucomp.sel( lon=lons).mean('lon') * mc.a * np.cos(psi.lat * np.pi / 180.) m_levs = mc.omega * mc.a**2. * np.cos( np.arange(-60., 1., 5.) * np.pi / 180.)**2. m.sel(month=i + 1).plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=m_levs, colors='0.7', add_labels=False) i = i + 1 ax.set_xlim(-35, 35) ax.set_xticks(np.arange(-30, 31, 15)) ax.grid(True, linestyle=':') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) 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('Zonal wind speed, m/s') if lonin == [-1., 361.]: plt.savefig(plot_dir + 'psi_u_' + run + '.pdf', format='pdf') else: figname = plot_dir + 'psi_u_' + run + '_' + str(int( lonin[0])) + '_' + str(int(lonin[1])) + '.pdf' plt.savefig(figname, format='pdf') plt.close() # Start figure with 12 subplots fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8), (ax9, ax10, ax11, ax12)) = plt.subplots(3, 4) axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12] i = 0 for ax in axes: psi_chi[:, i, :].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0., 601, 100.), colors='k', add_labels=False) psi_chi[:, i, :].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600., 0., 100.), colors='k', linestyles='dashed', add_labels=False) f1 = uchi.sel(month=i + 1).sel(lon=lons).mean('lon').plot.contourf( ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-3., 3.1, 0.5), extend='both', add_labels=False, add_colorbar=False) i = i + 1 ax.set_xlim(-35, 35) ax.set_xticks(np.arange(-30, 31, 15)) ax.grid(True, linestyle=':') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) 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('Zonal wind speed, m/s') if lonin == [-1., 361.]: plt.savefig(plot_dir + 'psi_chi_u_' + run + '.pdf', format='pdf') else: figname = plot_dir + 'psi_chi_u_' + run + '_' + str(int( lonin[0])) + '_' + str(int(lonin[1])) + '.pdf' plt.savefig(figname, format='pdf') plt.close()
def overturning_hm(run, regions=[[350, 10], [80, 100], [170, 190], [260, 280]]): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') plot_dir = '/scratch/rg419/plots/overturning_monthly/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) # 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() ds_chi = xr.Dataset({'vcomp': (vchi)}, coords={ 'xofyear': ('xofyear', vchi.xofyear), 'pfull': ('pfull', vchi.pfull), 'lat': ('lat', vchi.lat), 'lon': ('lon', vchi.lon) }) def get_lons(lonin, data): 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] ] return lons # Set figure parameters rcParams['figure.figsize'] = 10, 7 rcParams['font.size'] = 14 # Start figure with 4 subplots fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex='col', sharey='row') axes = [ax1, ax2, ax3, ax4] i = 0 for ax in axes: lons = get_lons(regions[i], data) psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=50.) psi_chi /= 1.e9 i = i + 1 f1 = psi_chi.sel(pfull=500).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, add_colorbar=False, levels=np.arange( -500., 501., 100.), extend='both') ax1.set_title('West coast') ax2.set_title('Land') ax3.set_title('East coast') ax4.set_title('Ocean') for ax in [ax1, ax2, ax3, ax4]: ax.grid(True, linestyle=':') ax.set_ylim(-60, 60) ax.set_yticks(np.arange(-60., 61., 30.)) ax.set_xticks([0, 18, 36, 54, 72]) ax3.set_xlabel('Pentad') ax4.set_xlabel('Pentad') ax1.set_ylabel('Latitude') ax3.set_ylabel('Latitude') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) 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('Overturning streamfunction') # Save as a pdf plt.savefig(plot_dir + 'regional_overturning_hm_' + run + '.pdf', format='pdf') plt.close() data.close()
def h_w_mass_flux_hm(run, dp=5000., lonin=[170, 190]): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') plot_dir = '/scratch/rg419/plots/zonal_asym_runs/' + run + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) # 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 lats = [ uchi.lat[i] for i in range(len(uchi.lat)) if uchi.lat[i] >= 0 and uchi.lat[i] <= 30 ] 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] ] coslat = np.cos(uchi.lat * np.pi / 180.) mass_flux_zon_areaweight = mass_flux_zon * coslat mass_flux_zon_mean = (mass_flux_zon_areaweight.sel(lat=lats).sum('lat') / coslat.sel(lat=lats).sum('lat')).sel(pfull=500) mass_flux_merid_mean = mass_flux_merid.sel(lon=lons).mean('lon').sel( pfull=500.) # Set figure parameters rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 14 fig, (ax1, ax2) = plt.subplots(1, 2) f1 = mass_flux_zon_mean.plot.contourf(ax=ax1, x='lon', y='xofyear', add_labels=False, add_colorbar=False, levels=np.arange( -0.0065, 0.0066, 0.001), extend='both') f1 = mass_flux_merid_mean.plot.contourf(ax=ax2, x='xofyear', y='lat', add_labels=False, add_colorbar=False, levels=np.arange( -0.0065, 0.0066, 0.001), extend='both') ax2.set_ylim(-60, 60) ax2.set_yticks(np.arange(-60, 61, 30)) ax2.set_xticks(np.arange(0, 73, 12)) ax1.set_yticks(np.arange(0, 73, 12)) ax1.set_xticks(np.arange(0, 361, 90)) ax1.grid(True, linestyle=':') ax2.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=[ax1, ax2], use_gridspec=True, orientation='horizontal', fraction=0.05, pad=0.1, aspect=30, shrink=0.5) cb1.set_label('Vertical mass flux, kgm$^{-2}$s$^{-1}$') figname = plot_dir + 'h_w_hm_' + run + '.pdf' plt.savefig(figname, format='pdf') plt.close()
def plot_gill_dev(lev=85000, qscale=100., windtype='full', ref_arrow=5, mse=False, video=False): if mse: data_t = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/jra_temp_daily_850.nc', chunks={'time': 30}) data_t = data_t['var11'].load().loc['1979-01':'2016-12'] data_q = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/jra_sphum_daily_850.nc', chunks={'time': 30}) data_q = data_q['var51'].load().loc['1979-01':'2016-12'] data_h = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/jra_height_daily_850.nc', chunks={'time': 30}) data_h = data_h['var7'].load().loc['1979-01':'2016-12'] data_t = pentad_mean_climatology(data_t, np.arange(1979, 2017)) data_q = pentad_mean_climatology(data_q, np.arange(1979, 2017)) data_h = pentad_mean_climatology(data_h, np.arange(1979, 2017)) data_mse = (mc.cp_air * data_t + mc.L * data_q + 9.81 * data_h) / 1000. else: data_precip = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/CMAP_precip.pentad.mean.nc', chunks={'time': 30}) data_precip = data_precip.load() data_precip.coords['pentad'] = (('time'), np.tile(np.arange(1, 74), 38)) data_precip = data_precip.groupby('pentad').mean('time') if windtype is not 'none': data_u = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/jra_ucomp_daily_850.nc', chunks={'time': 30}) data_u = data_u['var33'].load().loc['1979-01':'2016-12'] data_v_temp = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/jra_vcomp_daily_850.nc', chunks={'time': 30}) data_v_temp = data_v_temp['var34'].load().loc['1979-01':'2016-12'] # 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.sel(lev=85000.).values, coords={ 'time': data_u.time, 'lat': data_u.lat, 'lon': data_u.lon }, dims=('time', 'lat', 'lon')) data_u = pentad_mean_climatology(data_u, np.arange(1979, 2017)) data_v = pentad_mean_climatology(data_v, np.arange(1979, 2017)) data_u = data_u - data_u.mean('lon') data_v = data_v - data_v.mean('lon') if windtype is not 'full': # Get rotational and divergent components of the flow w = VectorWind(data_u, data_v) streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() #print(uchi.lat) #print(data_u.lat) 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') data_slp = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/jra_slp_daily.nc', chunks={'time': 30}) data_slp = data_slp['var2'].load().loc['1979-01':'2016-12'] data_slp = data_slp.load() data_slp = pentad_mean_climatology(data_slp / 100., np.arange(1979, 2017)) data_slp = data_slp - data_slp.mean('lon') print('files opened') rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 14 for i in range(73): fig, ax1 = plt.subplots() title = 'Pentad ' + str(int(data_u.pentad[i])) if mse: f1 = data_mse.sel(pentad=i + 1, lev=85000.).plot.contourf(x='lon', y='lat', ax=ax1, add_labels=False, add_colorbar=False, extend='both', cmap='Blues', zorder=1, levels=np.arange( 290., 341., 5.)) else: f1 = data_precip.precip[i, :, :].plot.contourf(x='lon', y='lat', ax=ax1, levels=np.arange( 2., 15., 2.), add_labels=False, add_colorbar=False, extend='both', cmap='Blues', zorder=1) ax1.contour(data_slp.lon, data_slp.lat, data_slp[i, :, :], levels=np.arange(0., 16., 3.), colors='0.4', alpha=0.5, zorder=2) ax1.contour(data_slp.lon, data_slp.lat, data_slp[i, :, :], levels=np.arange(-15., 0., 3.), colors='0.4', alpha=0.5, linestyle='--', zorder=2) if windtype == 'div': b = ax1.quiver(uchi_zanom.lon[::6], uchi_zanom.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, 5., 65., 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(upsi_zanom.lon[::6], upsi_zanom.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, 5., 65., 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_u.lon[::6], data_u.lat[::3], data_u[i, ::3, ::6], data_v[i, ::3, ::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) ax1.quiverkey(b, 5., 65., 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) 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) windtypestr = '' msestr = '' vidstr = '' if windtype != 'full': windtypestr = '_' + windtype if mse: msestr = '_mse' if video: vidstr = 'video/' plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/jra/' + vidstr + windtype + msestr + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) if video: plt.savefig(plot_dir + 'wind_and_slp_zanom_' + str(int(data_u.pentad[i])) + windtypestr + msestr + '.png', format='png') else: plt.savefig(plot_dir + 'wind_and_slp_zanom_' + str(int(data_u.pentad[i])) + windtypestr + msestr + '.pdf', format='pdf') plt.close()
def add_vars_netcdf(filename, sf=True, vp=True, ups=True, vps=True): #first load up the dataset as an xarray and calculate the streamfunction data = xr.open_dataset(filename,decode_times=False) uwnd = data.ucomp vwnd = data.vcomp # Create a VectorWind instance to handle the computation w = VectorWind(uwnd, vwnd) # Compute variables streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() # Reverse latitude axes streamfun_out = streamfun.data[:,:,::-1,:] vel_pot_out = vel_pot.data[:,:,::-1,:] uchi_out = uchi.data[:,:,::-1,:] upsi_out = upsi.data[:,:,::-1,:] vchi_out = vchi.data[:,:,::-1,:] vpsi_out = vpsi.data[:,:,::-1,:] # Open file using netcdf dsin= Dataset(filename, 'a', format='NETCDF3_CLASSIC') #Add variables, updating if they already exist if sf: try: dsin['streamfun'][:] = streamfun_out except: sf_out = dsin.createVariable('streamfun', 'f4', ('time', 'pfull', 'lat', 'lon',)) sf_out.setncatts({k: dsin.variables['ps'].getncattr(k) for k in dsin.variables['ps'].ncattrs()}) sf_out.setncattr('long_name', 'streamfunction') sf_out.setncattr('units', 'm**2 s**-1') sf_out[:] = streamfun_out if vp: try: dsin['vel_pot'][:] = vel_pot_out except: vp_out = dsin.createVariable('vel_pot', 'f4', ('time', 'pfull', 'lat', 'lon',)) vp_out.setncatts({k: dsin.variables['ps'].getncattr(k) for k in dsin.variables['ps'].ncattrs()}) vp_out.setncattr('long_name', 'velocity potential') vp_out.setncattr('units', 'm**2 s**-1') vp_out[:] = vel_pot_out if ups: try: dsin['upsi'][:] = upsi_out except: sf_out = dsin.createVariable('upsi', 'f4', ('time', 'pfull', 'lat', 'lon',)) sf_out.setncatts({k: dsin.variables['ps'].getncattr(k) for k in dsin.variables['ps'].ncattrs()}) sf_out.setncattr('long_name', 'nondivergent zonal wind') sf_out.setncattr('units', 'm/sec') sf_out[:] = upsi_out if vps: try: dsin['vpsi'][:] = vpsi_out except: sf_out = dsin.createVariable('vpsi', 'f4', ('time', 'pfull', 'lat', 'lon',)) sf_out.setncatts({k: dsin.variables['ps'].getncattr(k) for k in dsin.variables['ps'].ncattrs()}) sf_out.setncattr('long_name', 'nondivergent meridional wind') sf_out.setncattr('units', 'm/sec') sf_out[:] = vpsi_out dsin.close()
def ke_components(run, months, filename='plev_pentad', timeav='month', period_fac=1., land_mask=False, level=9): plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/ke/' + str( level) + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) mn_dic = month_dic(1) data = time_means(run, months, filename=filename, timeav=timeav, period_fac=period_fac) uwnd = data.ucomp[:, level, :, :] vwnd = data.vcomp[:, level, :, :] w = VectorWind(uwnd, vwnd) uchi, vchi, upsi, vpsi = w.helmholtz() ke_chi = 0.5 * (uchi * uchi + vchi * vchi) ke_psi = 0.5 * (upsi * upsi + vpsi * vpsi) for i in range(0, 12): ax = ke_psi[i, :, :].plot.contourf(x='lon', y='lat', levels=np.arange(0., 250., 10.), extend='both', add_colorbar=False, add_label=False) cb1 = plt.colorbar(ax) cb1.set_label('Nondivergent KE') if land_mask: land = xr.open_dataset( '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc') land_plot = xr.DataArray(land.land_mask.values, [('lat', data.lat), ('lon', data.lon)]) land_plot.plot.contour(x='lon', y='lat', levels=np.arange(0., 2., 1.), colors='0.75', add_colorbar=False, add_labels=False) cs = ke_chi[i, :, :].plot.contour(x='lon', y='lat', levels=np.arange(0., 36., 6.), extend='both', add_colorbar=False, colors='k', add_label=False) plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f') plt.ylabel('Latitude') plt.xlabel('Longitude') plt.ylim(-10, 45) plt.xlim(25, 150) plt.tight_layout() plt.savefig(plot_dir + str(i + 1) + '_' + str(mn_dic[i + 1]) + '.png') plt.close()
def h_w_mass_flux_monthly(lev=50000., dp=5000.): # Load in data data_u = xr.open_dataset( '/disca/share/reanalysis_links/jra_55/1958_2016/ucomp_monthly/atmos_monthly_together.nc' ) data_v = xr.open_dataset( '/disca/share/reanalysis_links/jra_55/1958_2016/vcomp_monthly/atmos_monthly_together.nc' ) land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc' land = xr.open_dataset(land_mask) data_u = data_u.sel(time=slice('1979', '2016')) data_v = data_v.sel(time=slice('1979', '2016')) # Make climatologies u_clim = data_u.groupby('time.month').mean('time') v_clim = data_v.groupby('time.month').mean('time') print('means taken') plot_dir = '/scratch/rg419/plots/monsoon_review_figs/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) # Create a VectorWind instance to handle the computation w = VectorWind(u_clim.var33, v_clim.var34) # Compute variables streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() coslat = np.cos(data_u.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.sel(lev=np.arange(5000., 100000., 5000.))) ).cumsum('lev') * dp * coslat / mc.grav mass_flux_merid = (gr.ddy(vchi.sel(lev=np.arange(5000., 100000., 5000.))) ).cumsum('lev') * dp * coslat / mc.grav # Set figure parameters rcParams['figure.figsize'] = 15, 9 rcParams['font.size'] = 16 # Start figure with 4 subplots fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex='col', sharey='row') axes = [ax1, ax2, ax3, ax4] f1 = mass_flux_merid.sel(lev=lev, month=[5, 6, 7, 8, 9]).mean('month').plot.contourf( ax=ax1, x='lon', y='lat', add_labels=False, add_colorbar=False, levels=np.arange(-0.006, 0.007, 0.001), extend='both') f1 = mass_flux_zon.sel(lev=lev, month=[5, 6, 7, 8, 9]).mean('month').plot.contourf( ax=ax2, x='lon', y='lat', add_labels=False, add_colorbar=False, levels=np.arange(-0.006, 0.007, 0.001), extend='both') f1 = mass_flux_merid.sel(lev=lev, month=[11, 12, 1, 2, 3]).mean('month').plot.contourf( ax=ax3, x='lon', y='lat', add_labels=False, add_colorbar=False, levels=np.arange(-0.006, 0.007, 0.001), extend='both') f1 = mass_flux_zon.sel(lev=lev, month=[11, 12, 1, 2, 3]).mean('month').plot.contourf( ax=ax4, x='lon', y='lat', add_labels=False, add_colorbar=False, levels=np.arange(-0.006, 0.007, 0.001), extend='both') i = 0 for ax in axes: land.lsm[0, :, :].plot.contour(ax=ax, x='longitude', y='latitude', levels=np.arange(-1., 2., 1.), add_labels=False, colors='k', linewidths=2.) ax.fill_between([0, 360], [-30, -30], [-10, -10], alpha=0.1, color='k') ax.fill_between([0, 360], [10, 10], [30, 30], alpha=0.1, color='k') 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=':') ax.set_xlim(0, 360) ax1.text(-40, 60, 'a)') ax2.text(-30, 60, 'b)') ax3.text(-40, 60, 'c)') ax4.text(-30, 60, 'd)') plt.subplots_adjust(left=0.05, right=0.97, top=0.95, bottom=0.05, hspace=0.1, wspace=0.1) 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, kgm$^{-2}$s$^{-1}$') figname = plot_dir + 'hadley_walker_jra.pdf' plt.savefig(figname, format='pdf') plt.close()
def plot_gill_dev(lev=85000, qscale=100., windtype='', ref_arrow=5): plot_dir = '/scratch/rg419/plots/zonal_asym_runs/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data_u = xr.open_dataset( '/disca/share/reanalysis_links/jra_55/1958_2016/ucomp_monthly/atmos_monthly_together.nc' ) data_v = xr.open_dataset( '/disca/share/reanalysis_links/jra_55/1958_2016/vcomp_monthly/atmos_monthly_together.nc' ) data_slp = xr.open_dataset( '/disca/share/reanalysis_links/jra_55/1958_2016/slp_monthly/atmos_monthly_together.nc' ) data_precip = xr.open_dataset('/disca/share/rg419/CMAP_precip.mon.mean.nc') data_u = data_u['var33'].load().loc['1979-01':'2016-12'].groupby( 'time.month').mean('time') data_v = data_v['var34'].load().loc['1979-01':'2016-12'].groupby( 'time.month').mean('time') data_slp = data_slp['var2'].load().loc['1979-01':'2016-12'].groupby( 'time.month').mean('time') / 100. data_precip = data_precip['precip'].load( ).loc['1979-01':'2016-12'].groupby('time.month').mean('time') data_u = data_u - data_u.mean('lon') data_v = data_v - data_v.mean('lon') data_slp = data_slp - data_slp.mean('lon') # Get rotational and divergent components of the flow w = VectorWind(data_u.sel(lev=lev), data_v.sel(lev=lev)) streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() #print(uchi.lat) #print(data_u.lat) 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 12 subplots rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 14 fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row') axes = [ax1, ax2, ax3, ax4, ax5, ax6] titles = ['March', 'April', 'May', 'June', 'July', 'August'] blist = [] for i in range(6): f1 = data_precip[i + 2, :, :].plot.contourf(x='lon', y='lat', ax=axes[i], levels=np.arange( 2., 15., 2.), add_labels=False, add_colorbar=False, extend='both', cmap='Blues', zorder=1) #data_slp[i+3,:,:].plot.contour(x='lon', y='lat', ax=axes[i], levels = np.arange(-15.,16.,3.), add_labels=False, add_colorbar=False, extend='both', colors='0.5', alpha=0.5) axes[i].contour(data_slp.lon, data_slp.lat, data_slp[i + 2, :, :], levels=np.arange(0., 16., 3.), colors='0.4', alpha=0.5, zorder=2) axes[i].contour(data_slp.lon, data_slp.lat, data_slp[i + 2, :, :], levels=np.arange(-15., 0., 3.), colors='0.4', alpha=0.5, linestyle='--', zorder=2) if windtype == 'div': b = axes[i].quiver(uchi_zanom.lon[::6], uchi_zanom.lat[::3], uchi_zanom[i + 2, ::3, ::6], vchi_zanom[i + 2, ::3, ::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) elif windtype == 'rot': b = axes[i].quiver(upsi_zanom.lon[::6], upsi_zanom.lat[::3], upsi_zanom[i + 2, ::3, ::6], vpsi_zanom[i + 2, ::3, ::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) else: b = axes[i].quiver(data_u.lon[::6], data_u.lat[::3], data_u.sel(lev=lev)[i + 2, ::3, ::6], data_v.sel(lev=lev)[i + 2, ::3, ::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) blist.append(b) axes[i].grid(True, linestyle=':') axes[i].set_ylim(-60., 60.) axes[i].set_yticks(np.arange(-60., 61., 30.)) axes[i].set_xticks(np.arange(0., 361., 90.)) axes[i].set_title(titles[i]) #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=axes[i], x='longitude', y='latitude', levels=np.arange(-1.,2.,1.), add_labels=False, colors='k') for ax in [ax1, ax4]: ax.set_ylabel('Latitude') for ax in [ax4, ax5, ax6]: ax.set_xlabel('Longitude') ax4.quiverkey(blist[3], 0., -0.5, ref_arrow, str(ref_arrow) + ' m/s', fontproperties={ 'weight': 'bold', 'size': 10 }, color='k', labelcolor='k', labelsep=0.03, zorder=10) 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=axes, use_gridspec=True, orientation='horizontal', fraction=0.05, pad=0.15, aspect=60, shrink=0.5) levstr = '' windtypestr = '' if lev != 85000: levstr = '_' + str(lev) if windtype != '': windtypestr = '_' + windtype plt.savefig(plot_dir + 'wind_and_slp_zanom_obs' + levstr + windtypestr + '.pdf', format='pdf') plt.close()
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()
def ke_hm(run, months, filename='plev_pentad', timeav='pentad', period_fac=1., lonin=[-1., 361.], level=10): plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/hm/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) mn_dic = month_dic(1) tickspace = range(7, 72, 12) labels = [mn_dic[(k + 5) / 6] for k in tickspace] data = time_means(run, months, filename=filename, timeav=timeav, period_fac=period_fac) if lonin[1] > lonin[0]: lons = [ i for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1] ] else: lons = [ i for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1] ] uwnd = data.ucomp[:, level, :, :].load() vwnd = data.vcomp[:, level, :, :].load() w = VectorWind(uwnd, vwnd) uchi, vchi, upsi, vpsi = w.helmholtz() ke_chi = 0.5 * (uchi * uchi + vchi * vchi)[:, :, lons].mean('lon') ke_psi = 0.5 * (upsi * upsi + vpsi * vpsi)[:, :, lons].mean('lon') ax = ke_psi.plot.contourf(x='xofyear', y='lat', levels=np.arange(0., 250., 10.), extend='both', add_colorbar=False, add_label=False) plt.grid(True, linestyle=':') cb1 = plt.colorbar(ax) cb1.set_label('Nondivergent KE') cs = uwnd[:, :, lons].mean('lon').plot.contour(x='xofyear', y='lat', levels=np.arange( -1000., 1001., 1000.), extend='both', colors='w', add_colorbar=False, add_label=False) cs = ke_chi.plot.contour(x='xofyear', y='lat', levels=np.arange(0., 36., 6.), extend='both', add_colorbar=False, colors='k', add_label=False) plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f') plt.xlabel('') plt.xticks(tickspace, labels, rotation=25) plt.yticks(range(-60, 61, 30)) plt.ylabel('Latitude') plt.tight_layout() plt.savefig(plot_dir + 'ke.png') plt.close()
else: dp = dp_in * 100. if intdown: psi = c * dp * np.flip(ubar, axis=ubar.dims.index('pfull')).cumsum( 'pfull') #[:,:,::-1] else: psi = -1. * c * dp * ubar.cumsum('pfull') #[:,:,::-1] #psi.plot.contourf(x='lat', y='pfull', yincrease=False) #plt.show() return psi #c*dp*vbar[:,::-1,:].cumsum('pfull')#[:,:,::-1] if __name__ == '__main__': # example calculating Walker cell for a GFDL dataset from windspharm.xarray import VectorWind d = xarray.open_dataset( '/disca/share/rg419/Data_moist/climatologies/3q_shallow.nc') w = VectorWind(d.ucomp.sel(pfull=np.arange(50., 950., 50.)), d.vcomp.sel(pfull=np.arange(50., 950., 50.))) # Compute variables uchi, vchi, upsi, vpsi = w.helmholtz() walker = walker_cell(uchi, dp_in=-50., intdown=True) walker /= 1.e9 walker.sel(xofyear=1).plot.contourf(x='lon', y='pfull', yincrease=False, levels=np.arange(-300., 301., 50.)) plt.show()
def overturning_monthly(run, lonin=[-1.,361.]): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') plot_dir = '/scratch/rg419/plots/overturning_monthly_ageo/' + run + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data.coords['month'] = (data.xofyear - 1) //6 + 1 data = data.groupby('month').mean(('xofyear')) #Coriolis omega = 7.2921150e-5 f = 2 * omega * np.sin(data.lat *np.pi/180) dphidx = gr.ddx(data.height) dphidx = 9.8 * dphidx v_ageo = data.vcomp - dphidx/f v_geo = dphidx/f dphidy = gr.ddy(data.height, vector=False) dphidy = 9.8 * dphidy u_ageo = data.ucomp - dphidy/f u_geo = dphidx/f # 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.))) w = VectorWind(u_ageo.sel(pfull=np.arange(50.,950.,50.)), v_ageo.sel(pfull=np.arange(50.,950.,50.))) # Compute variables streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() ds_chi = xr.Dataset({'vcomp': (vchi)}, coords={'month': ('month', vchi.month), 'pfull': ('pfull', vchi.pfull), 'lat': ('lat', vchi.lat), 'lon': ('lon', vchi.lon)}) w = VectorWind(u_geo.sel(pfull=np.arange(50.,950.,50.)), v_geo.sel(pfull=np.arange(50.,950.,50.))) # Compute variables streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() ds_chi_geo = xr.Dataset({'vcomp': (vchi)}, coords={'month': ('month', vchi.month), 'pfull': ('pfull', vchi.pfull), 'lat': ('lat', vchi.lat), 'lon': ('lon', vchi.lon)}) def get_lons(lonin, data): 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]] return lons lons = get_lons(lonin,data) psi = mass_streamfunction(data, lons=lons, dp_in=50.) psi /= 1.e9 psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=-50., intdown=False) psi_chi /= 1.e9 psi_chi_geo = mass_streamfunction(ds_chi_geo, lons=lons, dp_in=-50., intdown=False) psi_chi_geo /= 1.e9 # Set figure parameters rcParams['figure.figsize'] = 10, 7 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) axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12] i=0 for ax in axes: psi_chi[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0.,601,100.), colors='k', add_labels=False) psi_chi[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600.,0.,100.), colors='k', linestyles='dashed', add_labels=False) i=i+1 ax.set_xlim(-35,35) ax.set_xticks(np.arange(-30,31,15)) ax.grid(True,linestyle=':') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) if lonin == [-1.,361.]: plt.savefig(plot_dir + 'psi_chi_' + run + '.pdf', format='pdf') else: figname = plot_dir + 'psi_chi_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf' plt.savefig(figname, format='pdf') plt.close() # Start figure with 12 subplots fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8), (ax9, ax10, ax11, ax12)) = plt.subplots(3, 4) axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12] i=0 for ax in axes: psi_chi_geo[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0.,601,100.), colors='k', add_labels=False) psi_chi_geo[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600.,0.,100.), colors='k', linestyles='dashed', add_labels=False) i=i+1 ax.set_xlim(-35,35) ax.set_xticks(np.arange(-30,31,15)) ax.grid(True,linestyle=':') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) if lonin == [-1.,361.]: plt.savefig(plot_dir + 'psi_chi_geo_' + run + '.pdf', format='pdf') else: figname = plot_dir + 'psi_chi_geo_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf' plt.savefig(figname, format='pdf') plt.close()
def plot_gill_dev(run, land_mask=None, lev=850, qscale=100., windtype='full', ref_arrow=5, mse=False, video=False): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') data['mse'] = (mc.cp_air * data.temp + mc.L * data.sphum + mc.grav * data.height)/1000. data['precipitation'] = (data.precipitation*86400.) # Take zonal anomaly data_zanom = data - data.mean('lon') # Get rotational and divergent components of the flow if windtype == 'div' or windtype=='rot': 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])) if mse: f1 = data.mse.sel(xofyear=i+1, pfull=850.).plot.contourf(x='lon', y='lat', ax=ax1, add_labels=False, add_colorbar=False, extend='both', cmap='Blues', zorder=1, levels = np.arange(290.,341.,5.)) else: f1 = data.precipitation[i,:,:].plot.contourf(x='lon', y='lat', ax=ax1, levels = np.arange(2.,21.,2.), add_labels=False, add_colorbar=False, extend='both', cmap='Blues', zorder=1) #data_zanom.slp[i+4,:,:].plot.contour(x='lon', y='lat', ax=axes[i], levels = np.arange(-15.,16.,3.), add_labels=False, colors='0.5', alpha=0.5) ax1.contour(data_zanom.lon, data_zanom.lat, data_zanom.slp[i,:,:], levels = np.arange(0.,16.,3.), colors='0.4', alpha=0.5, zorder=2) ax1.contour(data_zanom.lon, data_zanom.lat, data_zanom.slp[i,:,:], levels = np.arange(-15.,0.,3.), colors='0.4', alpha=0.5, linestyle='--', zorder=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, 5.,65., 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, 5.,65., 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, 5.,65., ref_arrow, str(ref_arrow) + ' m/s', coordinates='data', 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') land.zsurf.plot.contour(ax=ax1, x='lon', y='lat', levels=np.arange(0.,2001.,1000.), 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 mse: msestr = '_mse' if video: vidstr='video/' plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/' + run +'/' + vidstr + windtype + msestr + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) if video: plt.savefig(plot_dir + 'wind_and_slp_zanom_' + str(int(data.xofyear[i])) + levstr + windtypestr + msestr + '.png', format='png') else: plt.savefig(plot_dir + 'wind_and_slp_zanom_' + str(int(data.xofyear[i])) + levstr + windtypestr + msestr + '.pdf', format='pdf') plt.close()
def h_w_mass_flux_monthly(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/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) month_lengths = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31] month = [] for i in range(1, 13): month = month + [i] * month_lengths[i - 1] data.coords['month'] = data.day_of_yr * 0. + np.array(month) data = data.groupby('month').mean('day_of_yr') data_u.coords['month'] = data_u.day_of_yr * 0. + np.array(month) data_u = data_u.groupby('month').mean('day_of_yr') # Create a VectorWind instance to handle the computation w = VectorWind(data_u.u.sel(pfull=np.arange(50., 950., 50.)), data.v.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 land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc' land = xr.open_dataset(land_mask) # Set figure parameters rcParams['figure.figsize'] = 15, 11 rcParams['font.size'] = 18 # 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') land.lsm[0, :, :].plot.contour(x='longitude', y='latitude', ax=ax, levels=np.arange(-1., 2., 1.), add_labels=False, colors='k') ax.fill_between([0, 360], [-26, -26], [-7, -7], alpha=0.2, color='k') ax.fill_between([0, 360], [7, 7], [26, 26], alpha=0.2, color='k') #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 + 'hadley_era.pdf' plt.savefig(figname, format='pdf') plt.close() # 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_zon.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') land.lsm[0, :, :].plot.contour(x='longitude', y='latitude', ax=ax, levels=np.arange(-1., 2., 1.), add_labels=False, colors='k') ax.fill_between([0, 360], [-26, -26], [-7, -7], alpha=0.2, color='k') ax.fill_between([0, 360], [7, 7], [26, 26], alpha=0.2, color='k') 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 zonal circulation, kgm$^{-2}$s$^{-1}$' ) figname = plot_dir + 'walker_era.pdf' plt.savefig(figname, format='pdf') plt.close()
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 plot_gill_dev(run, land_mask=None, lev=850, qscale=100., windtype='', ref_arrow=5): plot_dir = '/scratch/rg419/plots/zonal_asym_runs/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') # Take monthly averages data.coords['month'] = (data.xofyear - 1) //6 + 1 data = data.groupby('month').mean(('xofyear')) 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') data['precipitation'] = (data.precipitation*86400.) # Start figure with 12 subplots rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 14 fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row') axes = [ax1, ax2, ax3, ax4, ax5, ax6] titles = ['May', 'June', 'July', 'August', 'September', 'October'] blist=[] for i in range(6): f1 = data.precipitation[i+4,:,:].plot.contourf(x='lon', y='lat', ax=axes[i], levels = np.arange(2.,21.,2.), add_labels=False, add_colorbar=False, extend='both', cmap='Blues', zorder=1) #data_zanom.slp[i+4,:,:].plot.contour(x='lon', y='lat', ax=axes[i], levels = np.arange(-15.,16.,3.), add_labels=False, colors='0.5', alpha=0.5) axes[i].contour(data_zanom.lon, data_zanom.lat, data_zanom.slp[i+4,:,:], levels = np.arange(0.,16.,3.), colors='0.4', alpha=0.5, zorder=2) axes[i].contour(data_zanom.lon, data_zanom.lat, data_zanom.slp[i+4,:,:], levels = np.arange(-15.,0.,3.), colors='0.4', alpha=0.5, linestyle='--', zorder=2) if windtype=='div': b = axes[i].quiver(data.lon[::6], data.lat[::3], uchi_zanom[i+4,::3,::6], vchi_zanom[i+4,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) elif windtype=='rot': b = axes[i].quiver(data.lon[::6], data.lat[::3], upsi_zanom[i+4,::3,::6], vpsi_zanom[i+4,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) else: b = axes[i].quiver(data.lon[::6], data.lat[::3], data_zanom.ucomp.sel(pfull=lev)[i+4,::3,::6], data_zanom.vcomp.sel(pfull=lev)[i+4,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3) blist.append(b) axes[i].grid(True,linestyle=':') axes[i].set_ylim(-60.,60.) axes[i].set_yticks(np.arange(-60.,61.,30.)) axes[i].set_xticks(np.arange(0.,361.,90.)) axes[i].set_title(titles[i]) if not land_mask==None: land = xr.open_dataset(land_mask) land.land_mask.plot.contour(x='lon', y='lat', ax=axes[i], levels=np.arange(-1.,2.,1.), add_labels=False, colors='k') for ax in [ax1, ax4]: ax.set_ylabel('Latitude') for ax in [ax4, ax5, ax6]: ax.set_xlabel('Longitude') ax4.quiverkey(blist[3], 0.,-0.5, ref_arrow, str(ref_arrow) + ' m/s', fontproperties={'weight': 'bold', 'size': 10}, color='k', labelcolor='k', labelsep=0.03, zorder=10) 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=axes, use_gridspec=True, orientation = 'horizontal',fraction=0.05, pad=0.15, aspect=60, shrink=0.5) levstr='' windtypestr='' if lev != 850: levstr = '_' + str(lev) if windtype != '': windtypestr = '_' + windtype plt.savefig(plot_dir + 'wind_and_slp_zanom_' + run + levstr + windtypestr + '.pdf', format='pdf') plt.close()
def horiz_streamfun_monthly(run, land_mask=None): data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc') 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=150.), data.vcomp.sel(pfull=150.)) # Compute variables streamfun, vel_pot = w.sfvp() uchi_150, vchi_150, upsi_150, vpsi_150 = w.helmholtz() vel_pot_150 = (vel_pot - vel_pot.mean('lon')) / 10.**6 streamfun_150 = (streamfun - streamfun.mean('lon')) / 10.**6 w = VectorWind(data.ucomp.sel(pfull=850.), data.vcomp.sel(pfull=850.)) # Compute variables streamfun, vel_pot = w.sfvp() uchi_850, vchi_850, upsi_850, vpsi_850 = w.helmholtz() vel_pot_850 = (vel_pot - vel_pot.mean('lon')) / 10.**6 streamfun_850 = (streamfun - streamfun.mean('lon')) / 10.**6 # Set figure parameters rcParams['figure.figsize'] = 15, 8 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 = streamfun_150[i, :, :].plot.contourf(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, levels=np.arange( -30., 30.1, 5.), extend='both') streamfun_850[i, :, :].plot.contour(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, colors='k', levels=np.arange(0, 12.1, 2.)) ax.contour(streamfun_850.lon, streamfun_850.lat, streamfun_850[i, :, :], colors='k', ls='--', levels=np.arange(-12., 0., 2.)) if not land_mask == None: land = xr.open_dataset(land_mask) land.land_mask.plot.contour(x='lon', y='lat', ax=axes[i], levels=np.arange(-1., 2., 1.), add_labels=False, colors='k', alpha=0.5) #streamfun_850[i,:,:].plot.contour(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, cmap='PRGn', levels=np.arange(-12.,12.1,2.)) i = i + 1 ax.set_ylim(-35, 35) ax.set_yticks(np.arange(-30, 31, 15)) ax.set_xlim(0, 360) ax.set_xticks(np.arange(0, 361, 60)) ax.grid(True, linestyle=':') for ax in [ax1, ax5, ax9]: ax.set_ylabel('Latitude') for ax in [ax9, ax10, ax11, ax12]: ax.set_xlabel('Longitude') plt.subplots_adjust(left=0.08, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) 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('Horizontal streamfunction') plt.savefig(plot_dir + 'streamfun_' + run + '.pdf', format='pdf') plt.close() # 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 = vel_pot_150[i, :, :].plot.contourf(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, levels=np.arange( -30., 30.1, 5.), extend='both') vel_pot_850[i, :, :].plot.contour(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, colors='k', levels=np.arange(0, 12.1, 2.)) ax.contour(vel_pot_850.lon, vel_pot_850.lat, vel_pot_850[i, :, :], colors='k', ls='--', levels=np.arange(-12., 0., 2.)) if not land_mask == None: land = xr.open_dataset(land_mask) land.land_mask.plot.contour(x='lon', y='lat', ax=axes[i], levels=np.arange(-1., 2., 1.), add_labels=False, colors='k', alpha=0.5) #vel_pot_850[i,:,:].plot.contour(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, cmap='PRGn', levels=np.arange(-12.,12.1,2.)) i = i + 1 ax.set_ylim(-35, 35) ax.set_yticks(np.arange(-30, 31, 15)) ax.set_xlim(0, 360) ax.set_xticks(np.arange(0, 361, 60)) ax.grid(True, linestyle=':') for ax in [ax1, ax5, ax9]: ax.set_ylabel('Latitude') for ax in [ax9, ax10, ax11, ax12]: ax.set_xlabel('Longitude') plt.subplots_adjust(left=0.08, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) 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('Velocity potential') plt.savefig(plot_dir + 'vel_pot_' + run + '.pdf', format='pdf') plt.close() # 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 = vel_pot_150[i, :, :].plot.contourf(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, levels=np.arange( -30., 30.1, 5.), extend='both') b = axes[i].quiver(data.lon[::5], data.lat[::2], uchi_150[i, ::2, ::5], vchi_150[i, ::2, ::5], scale=100., angles='xy') if not land_mask == None: land = xr.open_dataset(land_mask) land.land_mask.plot.contour(x='lon', y='lat', ax=axes[i], levels=np.arange(-1., 2., 1.), add_labels=False, colors='k', alpha=0.5) i = i + 1 ax.set_ylim(-35, 35) ax.set_yticks(np.arange(-30, 31, 15)) ax.set_xlim(0, 360) ax.set_xticks(np.arange(0, 361, 60)) ax.grid(True, linestyle=':') for ax in [ax1, ax5, ax9]: ax.set_ylabel('Latitude') for ax in [ax9, ax10, ax11, ax12]: ax.set_xlabel('Longitude') plt.subplots_adjust(left=0.08, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) 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('Velocity potential') plt.savefig(plot_dir + 'vel_pot_vchi_' + run + '.pdf', format='pdf') plt.close()
def overturning_monthly(lonin=[-1.,361.]): 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/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) def get_lons(lonin, data): 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]] return lons lons = get_lons(lonin,data) month_lengths = [31,28,31,30,31,30,31,31,30,31,30,31] month = [] for i in range(1,13): month = month + [i]*month_lengths[i-1] data.coords['month'] = data.day_of_yr*0. + np.array(month) data = data.groupby('month').mean('day_of_yr') data_u.coords['month'] = data_u.day_of_yr*0. + np.array(month) data_u = data_u.groupby('month').mean('day_of_yr') data = data.rename(name_dict={'v': 'vcomp'}) data_u = data_u.rename(name_dict={'u': 'ucomp'}) # Create a VectorWind instance to handle the computation w = VectorWind(data_u.ucomp.sel(pfull=np.arange(50.,950.,50.)), data.vcomp.sel(pfull=np.arange(50.,950.,50.))) streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() ds_chi = xr.Dataset({'vcomp': (vchi)}, coords={'month': ('month', vchi.month), 'pfull': ('pfull', vchi.pfull), 'lat': ('lat', vchi.lat), 'lon': ('lon', vchi.lon)}) psi = mass_streamfunction(data, lons=lons, dp_in=-50., intdown=False) psi /= 1.e9 psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=-50., intdown=False) psi_chi /= 1.e9 # print(psi) # Set figure parameters rcParams['figure.figsize'] = 10, 7 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) axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12] i=0 for ax in axes: psi[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0.,601,100.), colors='k', add_labels=False) psi[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600.,0.,100.), colors='k', linestyles='dashed', add_labels=False) f1 = data_u.ucomp.sel(month=i+1).sel(lon=lons).mean('lon').plot.contourf(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-50.,50.1,5.), extend='both', add_labels=False, add_colorbar=False) m = mc.omega * mc.a**2. * np.cos(psi.lat*np.pi/180.)**2. + data_u.ucomp.sel(lon=lons).mean('lon') * mc.a * np.cos(psi.lat*np.pi/180.) m_levs = mc.omega * mc.a**2. * np.cos(np.arange(-60.,1.,5.)*np.pi/180.)**2. m.sel(month=i+1).plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=m_levs, colors='0.7', add_labels=False) i=i+1 ax.set_xlim(-35,35) ax.set_xticks(np.arange(-30,31,15)) ax.grid(True,linestyle=':') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) #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('Zonal wind speed, m/s') if lonin == [-1.,361.]: plt.savefig(plot_dir + 'psi_u_era.pdf', format='pdf') else: figname = plot_dir + 'psi_u_era_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf' plt.savefig(figname, format='pdf') plt.close() # Start figure with 12 subplots fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8), (ax9, ax10, ax11, ax12)) = plt.subplots(3, 4) axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12] i=0 for ax in axes: psi_chi[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0.,601,100.), colors='k', add_labels=False) psi_chi[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600.,0.,100.), colors='k', linestyles='dashed', add_labels=False) i=i+1 ax.set_xlim(-35,35) ax.set_xticks(np.arange(-30,31,15)) ax.grid(True,linestyle=':') plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3) 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('Zonal wind speed, m/s') if lonin == [-1.,361.]: plt.savefig(plot_dir + 'psi_chi_era.pdf', format='pdf') else: figname = plot_dir + 'psi_chi_era_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf' plt.savefig(figname, format='pdf') plt.close()