def streamfun_vid(run, ax_in, pentad, plot_land=True, tibet=True): data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') uwnd = data.ucomp.sel(pfull=150.) vwnd = data.vcomp.sel(pfull=150.) # Create a VectorWind instance to handle the computation w = VectorWind(uwnd, vwnd) # Compute variables streamfun, vel_pot = w.sfvp() 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. ] land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc' land = xr.open_dataset(land_file) streamfun_zanom = (streamfun - streamfun.mean('lon')) / 10.**6 f1 = streamfun_zanom[pentad, :, :].plot.contourf(x='lon', y='lat', ax=ax_in, levels=np.arange( -36., 37., 3.), add_labels=False, add_colorbar=False, extend='both') if plot_land: land.land_mask.plot.contour(x='lon', y='lat', levels=np.arange(0., 2., 1.), ax=ax_in, colors='k', add_colorbar=False, add_labels=False) if tibet: land.zsurf.plot.contourf(x='lon', y='lat', ax=ax_in, levels=np.arange(2500., 100001., 97500.), add_labels=False, extend='neither', add_colorbar=False, alpha=0.5, cmap='Greys_r') ax_in.set_xlim(60, 150) ax_in.set_ylim(-30, 60) ax_in.set_xticks(np.arange(60., 155., 30.)) ax_in.set_yticks(np.arange(-30., 65., 30.)) ax_in.grid(True, linestyle=':') return f1
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 test_truncation(da_urho, da_vrho, trunc_vals=[10, 20, 30, 40, 50, 100, 200]): for idx, trunc in enumerate(trunc_vals): print(f'Done {int(100*idx/da_urho.time.values.shape[0])} %') da_urho_ = da_urho.isel(time=slice(0, 8)) da_vrho_ = da_vrho.isel(time=slice(0, 8)) w = VectorWind(da_urho_, da_vrho_) da_urho_ = w.truncate(da_urho_, truncation=trunc) da_vrho_ = w.truncate(da_vrho_, truncation=trunc) lcs = LCS(timestep=-6 * 3600, timedim='time', SETTLS_order=4) ftle = lcs(u=da_urho_, v=da_vrho_, isglobal=True, s=4e6, interp_to_common_grid=True, traj_interp_order=3) ftle = .5 * np.log(ftle) toplot = ftle fig, ax = plt.subplots(1, 1, subplot_kw={'projection': ccrs.Robinson()}) toplot.plot(ax=ax, transform=ccrs.PlateCarree(), robust=True, cbar_kwargs={'shrink': .6}, vmin=.5, vmax=2) ax.coastlines() ax.set_title(f'N96 truncated at t{trunc}') plt.savefig(f'upscale/figs/trunc_test/{trunc:03d}.png', dpi=600, boundary='trim', bbox_inches='tight') 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 calculate_streamfunction(uwnd, vwnd): """Calculate the Streamfunction """ uwnd_ds = xr.open_dataarray(uwnd) vwnd_ds = xr.open_dataarray(vwnd) wind = VectorWind(uwnd_ds, vwnd_ds) psi = wind.streamfunction() return psi
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 calc_ftle(da_urho_, da_vrho_, truncation=20): print('********************************** \n' 'Truncating wind and computing FTLE\n' '**********************************') w = VectorWind(da_urho_, da_vrho_) da_urho_ = w.truncate(da_urho_, truncation=20) da_vrho_ = w.truncate(da_vrho_, truncation=20) lcs = LCS(timestep=-6 * 3600, timedim='time', SETTLS_order=4) ftle = lcs(u=da_urho_, v=da_vrho_, isglobal=True, s=4e6, interp_to_common_grid=True, traj_interp_order=3) ftle = .5 * np.log(ftle) return ftle
def calc_Koc(ds, detrend=False, plot=False): if detrend: ds['s'].data = signal.detrend(ds['s'].data, axis=0) + ds['s'].mean(dim='time').data s_mean = ds['s'].mean(dim='time') s_anom = ds['s'] - s_mean w = VectorWind(ds['u'], ds['v']) grad_s_mean = w.gradient(s_mean) grad_s_mean = VectorWind(grad_s_mean[0], grad_s_mean[1]) abs_grad_s_mean = grad_s_mean.magnitude() sq_abs_grad_s_mean = abs_grad_s_mean * abs_grad_s_mean grad_s_anom = w.gradient(s_anom) grad_s_anom = VectorWind(grad_s_anom[0], grad_s_anom[1]) abs_grad_s_anom = grad_s_anom.magnitude() sq_abs_grad_s_anom = abs_grad_s_anom * abs_grad_s_anom sq_abs_grad_s_anom = sq_abs_grad_s_anom.mean(dim='time') sq_abs_grad_s_mean = sq_abs_grad_s_mean.mean(dim='longitude') sq_abs_grad_s_anom = sq_abs_grad_s_anom.mean(dim='longitude') Koc = sq_abs_grad_s_anom / sq_abs_grad_s_mean if plot: plt.ion() fig1, ax1 = plt.subplots() ax1.plot(lats, sq_abs_grad_s_mean.data, label='mean') ax1.plot(lats, sq_abs_grad_s_anom.data, label='anom') ax11 = ax1.twinx() ax11.plot(lats, ds['q'].mean(dim=['time', 'longitude']).data) fig2, ax2 = plt.subplots() ax2.plot(lats, Koc, label='Koc') ax21 = ax2.twinx() ax21.plot(lats, ds['q'].mean(dim=['time', 'longitude']).data) plt.show() return Koc
def potential_vorticity_baroclinic(uwnd, vwnd, theta, coord, **kwargs): ''' Calculate potential vorticity on isobaric levels. Requires input uwnd, vwnd and tmp arrays to have (lat, lon, ...) format for Windspharm. Input ----- uwnd : zonal winds, array-like vwnd : meridional winds, array-like tmp : temperature, array-like theta : potential temperature, array-like coord : dimension name for pressure axis (eg. 'pfull') omega : planetary rotation rate, optional g : planetary gravitational acceleration, optional rsphere : planetary radius, in metres, optional ''' omega = kwargs.pop('omega', 7.08822e-05) g = kwargs.pop('g', 3.72076) rsphere = kwargs.pop('rsphere', 3.3962e6) w = VectorWind(uwnd.fillna(0), vwnd.fillna(0), rsphere=rsphere) relvort = w.vorticity() relvort = relvort.where(relvort != 0, other=np.nan) planvort = w.planetaryvorticity(omega=omega) absvort = relvort + planvort dthtady, dthtadx = w.gradient(theta.fillna(0)) dthtadx = dthtadx.where(dthtadx != 0, other=np.nan) dthtady = dthtady.where(dthtady != 0, other=np.nan) dthtadp = wrapped_gradient(theta, coord) dudp = wrapped_gradient(uwnd, coord) dvdp = wrapped_gradient(vwnd, coord) s = -dthtadp f = dvdp * dthtadx - dudp * dthtady ret = g * (absvort + f / s) * s return ret
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 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()
# try: # ws_ccmp.to_netcdf('datasets/CCMP_windspeed.nc') # print('saved') # except: # pass # %% #ws_ccmp1=xr.open_dataset('datasets/CCMP_windspeed.nc') #wu=xr.open_dataset('datasets/uwnd.10m.mon.mean.nc').sel(level=10).uwnd #wv=xr.open_dataset('datasets/vwnd.10m.mon.mean.nc').sel(level=10).vwnd ws_ccmp=xr.open_dataset('processed/CCMP_ws_1deg_global.nc') wu=ws_ccmp.uwnd wv=ws_ccmp.vwnd # %% Test Horizontal Divergence w = VectorWind(wu, wv) #spd = w.magnitude() divergence = w.divergence().sel(lat=slice(20,-20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01')) #div.mean(dim='time').plot() wu=wu.sel(lat=slice(-20,20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01')) wv=wv.sel(lat=slice(-20,20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01')) # %% Prepare Figure lanina=pd.read_csv('processed/indexes/la_nina_events.csv') cp_nino=pd.read_csv('processed/indexes/cp_events.csv') ep_nino=pd.read_csv('processed/indexes/ep_events.csv') fp='processed/combined_dataset/month_data_exports.nc' info=xr.open_mfdataset(fp).sel(Mooring=195).to_dataframe()
from windspharm.xarray import VectorWind from windspharm.examples import example_data_path mpl.rcParams['mathtext.default'] = 'regular' # Read zonal and meridional wind components from file using the xarray module. # The components are in separate files. ds = xr.open_mfdataset([example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')]) uwnd = ds['uwnd'] vwnd = ds['vwnd'] # Create a VectorWind instance to handle the computation of streamfunction and # velocity potential. w = VectorWind(uwnd, vwnd) # Compute the streamfunction and velocity potential. sf, vp = w.sfvp() # Pick out the field for December. sf_dec = sf[sf['time.month'] == 12] vp_dec = vp[vp['time.month'] == 12] # Plot streamfunction. clevs = [-120, -100, -80, -60, -40, -20, 0, 20, 40, 60, 80, 100, 120] ax = plt.subplot(111, projection=ccrs.PlateCarree(central_longitude=180)) sf_dec *= 1e-6 fill_sf = sf_dec[0].plot.contourf(ax=ax, levels=clevs, cmap=plt.cm.RdBu_r, transform=ccrs.PlateCarree(), extend='both', add_colorbar=False)
def abs_vort_hm(run, lev=150, filename='plev_daily', timeav='pentad', period_fac=1., latin=25.): rcParams['figure.figsize'] = 9, 9 rcParams['font.size'] = 18 rcParams['text.usetex'] = True plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc' names = [name_temp % m for m in range(397, 409)] #read data into xarray data = xr.open_mfdataset( names, decode_times=False, # no calendar so tell netcdf lib # choose how data will be broken down into manageable chunks. chunks={'time': 30}) uwnd = data.ucomp.sel(pfull=lev) vwnd = data.vcomp.sel(pfull=lev) # Create a VectorWind instance to handle the computation of streamfunction and # velocity potential. w = VectorWind(uwnd, vwnd) # Compute the streamfunction and velocity potential. data['vor'], data['div'] = w.vrtdiv() #data = xr.open_dataset('/scratch/rg419/Data_moist/'+run+'climatologies/'+run+'.nc') #Coriolis omega = 7.2921150e-5 f = 2 * omega * np.sin(data.lat * np.pi / 180) lat_hm = data.lat[np.argmin(np.abs(data.lat - latin))] abs_vort = (f + data.vor).sel(lat=lat_hm) * 86400. levels = np.arange(0., 14.1, 2.) mn_dic = month_dic(1) tickspace = range(13, 72, 18) labels = [mn_dic[(k + 5) / 6] for k in tickspace] # Plot f1 = abs_vort.plot.contourf(x='lon', y='time', extend='both', levels=levels, add_colorbar=False, add_labels=False) plt.set_cmap('inferno_r') plt.ylabel('Time') #plt.yticks(tickspace, labels, rotation=25) plt.xlabel('Longitude') plt.xlim(60, 150) plt.ylim(240 + 33 * 360, 90 + 33 * 360) plt.grid(True, linestyle=':') plt.tight_layout() #Colorbar cb1 = plt.colorbar(f1, use_gridspec=True, orientation='horizontal', fraction=0.15, pad=0.1, aspect=30) cb1.set_label('$day^{-1}$') figname = 'abs_vort_lon_hm.pdf' plt.savefig(plot_dir + 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()
from windspharm.xarray import VectorWind from windspharm.examples import example_data_path mpl.rcParams['mathtext.default'] = 'regular' # Read zonal and meridional wind components from file using the xarray module. # The components are in separate files. ds = xr.open_mfdataset([example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')]) uwnd = ds['uwnd'] vwnd = ds['vwnd'] # Create a VectorWind instance to handle the computations. w = VectorWind(uwnd, vwnd) # Compute components of rossby wave source: absolute vorticity, divergence, # irrotational (divergent) wind components, gradients of absolute vorticity. eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) etax.attrs['units'] = 'm**-1 s**-1' etay.attrs['units'] = 'm**-1 s**-1' # Combine the components to form the Rossby wave source term. S = eta * -1. * div - (uchi * etax + vchi * etay) # Pick out the field for December at 200 hPa. S_dec = S[S['time.month'] == 12]
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 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()
import xarray as xr from windspharm.xarray import VectorWind f = xr.open_dataset( "/Users/brianpm/Documents/www.ncl.ucar.edu/Applications/Data/cdf/uv300.nc") u = f["U"] v = f["V"] w = VectorWind(u, v) ## VERY IMPORTANT: VectorWind apparently reverses latitude to be decreasing (90 to -90) vort, div = w.vrtdiv() # Relative vorticity and horizontal divergence. sf, vp = w.sfvp() # The streamfunction and velocity potential respectively. uchi, vchi, upsi, vpsi = w.helmholtz() # plot the results import matplotlib as mpl import matplotlib.pyplot as plt import cartopy.crs as ccrs import numpy as np fig, ax = plt.subplots(figsize=(12, 12), nrows=2, subplot_kw={"projection": ccrs.PlateCarree()}, constrained_layout=True) N = mpl.colors.Normalize(vmin=-8e6, vmax=8e6) x, y = np.meshgrid(f['lon'], f['lat']) im0 = ax[0].contourf(x, y, vp[0, ::-1, :], norm=N, transform=ccrs.PlateCarree(),
import xarray as xr from windspharm.xarray import VectorWind from windspharm.examples import example_data_path mpl.rcParams['mathtext.default'] = 'regular' # Read zonal and meridional wind components from file using the xarray module. # The components are in separate files. ds = xr.open_mfdataset( [example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')]) uwnd = ds['uwnd'] vwnd = ds['vwnd'] # Create a VectorWind instance to handle the computations. w = VectorWind(uwnd, vwnd) # Compute components of rossby wave source: absolute vorticity, divergence, # irrotational (divergent) wind components, gradients of absolute vorticity. eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) etax.attrs['units'] = 'm**-1 s**-1' etay.attrs['units'] = 'm**-1 s**-1' # Combine the components to form the Rossby wave source term. S = eta * -1. * div - (uchi * etax + vchi * etay) # Pick out the field for December at 200 hPa. S_dec = S[S['time.month'] == 12]
#S.to_netcdf(RUTA + 'monthly_RWS.nc4') #uchi_monthly = xr.concat(uchi_monthly, dim='month') #uchi_monthly['month'] = np.array([8, 9, 10, 11, 12, 1, 2]) #uchi_monthly.to_netcdf(RUTA + 'monthly_uchi.nc4') #vchi_monthly = xr.concat(vchi_monthly, dim='month') #vchi_monthly['month'] = np.array([8, 9, 10, 11, 12, 1, 2]) #vchi_monthly.to_netcdf(RUTA + 'monthly_vchi.nc4') S_seasonal = [] uchi_seasonal = [] vchi_seasonal = [] for i in np.arange(5): ds1 = ds.isel(month=range(i, i + 3)).mean(dim='month') #create vectorwind instance w = VectorWind(ds1['u'][:, :, :], ds1['v'][:, :, :]) eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) etax.attrs['units'] = 'm**-1 s**-1' etay.attrs['units'] = 'm**-1 s**-1' # Combine the components to form the Rossby wave source term. S = eta * -1. * div - (uchi * etax + vchi * etay) S *= 1e11 S_seasonal.append(S) uchi_seasonal.append(uchi) vchi_seasonal.append(vchi) S_seasonal = xr.concat(S_seasonal, dim='season') S_seasonal['season'] = np.array(seas)
data_vo = (dvdx - dudy) * 86400. data_t = xr.open_dataset( '/disca/share/reanalysis_links/jra_55/1958_2016/temp_monthly/atmos_monthly_together.nc' ) data_q = xr.open_dataset( '/disca/share/rg419/JRA_55/sphum_monthly/atmos_monthly_together.nc') data_z = xr.open_dataset( '/disca/share/reanalysis_links/jra_55/1958_2016/height_monthly/atmos_monthly_together.nc' ) data_mse = (mc.cp_air * data_t.var11 + mc.L * data_q.var51 + 9.81 * data_z.var7) / 1000. # Create a VectorWind instance to handle the computation w = VectorWind(data_u.sel(lev=np.arange(5000., 100001., 5000.)), data_v.sel(lev=np.arange(5000., 100001., 5000.))) # Compute variables streamfun, vel_pot = w.sfvp() uchi, vchi, upsi, vpsi = w.helmholtz() coslat = np.cos(data_u.lat * np.pi / 180) dp = 5000. # Evaluate mass fluxes for the zonal and meridional components (Walker and Hadley) following Schwendike et al. 2014 mass_flux_zon = (gr.ddx(uchi)).cumsum('lev') * dp * coslat / 9.81 mass_flux_merid = (gr.ddy(vchi)).cumsum('lev') * dp * coslat / 9.81 land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc' land = xr.open_dataset(land_mask) def plot_winter_climate(data,
from windspharm.xarray import VectorWind GFDL_DATA = os.environ['GFDL_DATA'] filename = GFDL_DATA + 'full_qflux/run121/atmos_pentad.nc' fileout = GFDL_DATA + 'full_qflux/run121/atmos_test.nc' dsin= Dataset(filename, 'r', format='NETCDF3_CLASSIC') data = xr.open_dataset(filename,decode_times=False) uwnd = data.ucomp vwnd = data.vcomp # Create a VectorWind instance to handle the computation of streamfunction and # velocity potential. w = VectorWind(uwnd, vwnd) # Compute the streamfunction and velocity potential. streamfun, vel_pot = w.sfvp() dsout= Dataset(fileout, 'w', format='NETCDF3_CLASSIC') dsout = copy_netcdf_attrs(dsin, dsout) sf_out = dsout.createVariable('streamfun', 'f4', ('time', 'pfull', 'lat', 'lon',)) sf_out.setncatts({k: dsout.variables['ps'].getncattr(k) for k in dsout.variables['ps'].ncattrs()}) sf_out.setncattr('long_name', 'streamfunction') sf_out.setncattr('units', 'm**2 s**-1') sf_out[:] = streamfun.load().data
def plot_sf_vp(land_mask=None): rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 16 plot_dir = '/scratch/rg419/plots/era_wn2/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) land = xr.open_dataset(land_mask) data = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/sep_levs_u/era_u_200_mm.nc') print data.time.units uwnd = data.u.load().squeeze('level') uwnd = uwnd[0:456,:,:] data_sn = data.resample(time='Q-NOV').mean() uwnd_sn = data_sn.u.load().squeeze('level') data = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/sep_levs_v/era_v_200_mm.nc') vwnd = data.v.load() data_sn = data.resample(time='Q-NOV').mean() vwnd_sn = data_sn.v.load() # Create a VectorWind instance to handle the computation w = VectorWind(uwnd, vwnd) # Compute variables streamfun, vel_pot = w.sfvp() streamfun = streamfun/10.**6 # Create a VectorWind instance to handle the computation w = VectorWind(uwnd_sn, vwnd_sn) # Compute variables streamfun_sn, vel_pot_sn = w.sfvp() streamfun_sn = streamfun_sn/10.**6 lats = [data.latitude[i] for i in range(len(data.latitude)) if data.latitude[i] >= 5.] for year in range(1979,2017): sf_max_lat = np.zeros((12,)) sf_max_lon = np.zeros((12,)) for month in range(1,13): streamfun_asia_i = streamfun.sel(time=str(year)+'-%02d' % month , latitude=lats) sf_max_i = streamfun_asia_i.where(streamfun_asia_i==streamfun_asia_i.max(), drop=True) sf_max_lon[month-1] = sf_max_i.longitude.values sf_max_lat[month-1] = sf_max_i.latitude.values f1 = streamfun_sn.sel(time=str(year) + '-08').squeeze('time').plot.contourf(x='longitude', y='latitude', levels = np.arange(-140.,141.,10.), add_labels=False, add_colorbar=False, extend='both') for i in range(12): plt.text(sf_max_lon[i]+3, sf_max_lat[i], str(i+1), fontsize=10) plt.plot(sf_max_lon, sf_max_lat, 'kx-', mew=1.5) plt.grid(True,linestyle=':') plt.colorbar(f1) land.lsm[0,:,:].plot.contour(x='longitude', y='latitude', levels=np.arange(-1.,2.,1.), add_labels=False, colors='k') plt.savefig(plot_dir + 'streamfun_era_' + str(year) + '.pdf', format='pdf') plt.close()
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 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 plot_sf_clim(run, land_mask=None): rcParams['figure.figsize'] = 10, 8 rcParams['font.size'] = 16 plot_dir = '/scratch/rg419/plots/egu_2018_talk_plots/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') uwnd = data.ucomp.sel(pfull=150.) vwnd = data.vcomp.sel(pfull=150.) # Create a VectorWind instance to handle the computation w = VectorWind(uwnd, vwnd) # Compute variables streamfun, vel_pot = w.sfvp() def sn_av(da): #Take seasonal and monthly averages da = da / 10.**6 da.coords['season'] = np.mod(da.xofyear + 5., 72.) // 18. da_sn = da.groupby('season').mean(('xofyear')) return da_sn streamfun_sn = sn_av(streamfun) streamfun_sn = streamfun_sn - streamfun_sn.mean('lon') f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex='col', sharey='row') axes = [ax1, ax2, ax3, ax4] title = ['DJF', 'MAM', 'JJA', 'SON'] for i in range(4): f1 = streamfun_sn[i, :, :].plot.contourf(x='lon', y='lat', ax=axes[i], levels=np.arange( -36., 37., 3.), add_labels=False, add_colorbar=False, extend='both') axes[i].grid(True, linestyle=':') 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, ax3]: ax.set_ylabel('Latitude') for ax in [ax3, ax4]: ax.set_xlabel('Longitude') cb1 = f.colorbar(f1, ax=axes, use_gridspec=True, orientation='horizontal', fraction=0.15, pad=0.15, aspect=30, shrink=0.5) plt.savefig(plot_dir + 'streamfun_' + run + '.pdf', format='pdf') plt.close()
import numpy as np import xarray as xr from windspharm.xarray import VectorWind from windspharm.examples import example_data_path # Read zonal and meridional wind components from file using the xarray module. # The components are in separate files. ds = xr.open_mfdataset( [example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')]) uwnd = ds['uwnd'] vwnd = ds['vwnd'] # Create a VectorWind instance to handle the computation of streamfunction and # velocity potential. w = VectorWind(uwnd, vwnd) # Compute the streamfunction and velocity potential. sf, vp = w.sfvp() # Pick out the field for December. sf_dec = sf[sf['time.month'] == 12] vp_dec = vp[vp['time.month'] == 12] # Plot streamfunction. clevs = [-120, -100, -80, -60, -40, -20, 0, 20, 40, 60, 80, 100, 120] ax = plt.subplot(111, projection=ccrs.PlateCarree(central_longitude=180)) sf_dec *= 1e-6 fill_sf = sf_dec[0].plot.contourf(ax=ax, levels=clevs, cmap=plt.cm.RdBu_r,
def plot_sf_vp(land_mask=None): rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 16 plot_dir = '/scratch/rg419/plots/era_wn2/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) land = xr.open_dataset(land_mask) data = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/sep_levs_u/era_u_200_mm.nc') uwnd = data.u.load().squeeze('level') uwnd = uwnd[0:456, :, :] data_sn = data.resample(time='Q-NOV').mean() uwnd_sn = data_sn.u.load().squeeze('level') data = xr.open_dataset( '/scratch/rg419/obs_and_reanalysis/sep_levs_v/era_v_200_mm.nc') vwnd = data.v.load() data_sn = data.resample(time='Q-NOV').mean() vwnd_sn = data_sn.v.load() # Create a VectorWind instance to handle the computation w = VectorWind(uwnd, vwnd) # Compute variables streamfun, vel_pot = w.sfvp() streamfun = streamfun / 10.**6 # Create a VectorWind instance to handle the computation w = VectorWind(uwnd_sn, vwnd_sn) # Compute variables streamfun_sn, vel_pot_sn = w.sfvp() streamfun_sn = streamfun_sn / 10.**6 streamfun_sn_mean = streamfun_sn.groupby('time.month').mean('time') #time_list = [str(year) + '-08' for year in range(1979,2017)] #print time_list[0] #streamfun_JJA_mean = streamfun_sn.sel(time=time_list[2]).mean('time') #print streamfun_JJA_mean #streamfun_JJA_anom = streamfun_sn.sel(time=time_list[0]) - streamfun_JJA_mean #print streamfun_JJA_anom for year in range(1979, 2017): f1 = (streamfun_sn.sel(time=str(year) + '-08') - streamfun_sn_mean.sel(month=8)).squeeze('time').plot.contourf( x='longitude', y='latitude', levels=np.arange(-20., 21., 2.), add_labels=False, add_colorbar=False, extend='both') plt.grid(True, linestyle=':') plt.colorbar(f1) land.lsm[0, :, :].plot.contour(x='longitude', y='latitude', levels=np.arange(-1., 2., 1.), add_labels=False, colors='k') plt.savefig(plot_dir + 'streamfun_era_anom_' + str(year) + '.pdf', format='pdf') plt.close() f1 = (streamfun_sn.sel(time=str(year) + '-08') - streamfun_sn.sel(time=str(year) + '-08').mean('longitude') ).squeeze('time').plot.contourf(x='longitude', y='latitude', levels=np.arange(-50., 51., 5.), add_labels=False, add_colorbar=False, extend='both') plt.grid(True, linestyle=':') plt.colorbar(f1) land.lsm[0, :, :].plot.contour(x='longitude', y='latitude', levels=np.arange(-1., 2., 1.), add_labels=False, colors='k') plt.savefig(plot_dir + 'streamfun_era_zanom_' + str(year) + '.pdf', format='pdf') plt.close()