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()
'/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, title, levels_clim=np.arange(-50., 51., 5.),
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()
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_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 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 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()
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) ax.coastlines() ax.gridlines() plt.colorbar(fill_sf, orientation='horizontal')
def sf_vp(run, months, filename='plev_pentad', timeav='month', period_fac=1., land_mask=False, level=9): plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/sf_vp/' + 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) sf, vp = w.sfvp() sf *= 1e-6 vp *= 1e-6 # Plot streamfunction. for i in range(0, 12): ax = sf[i, :, :].plot.contourf(x='lon', y='lat', levels=np.arange(-180., 181., 20.), extend='both', add_colorbar=False, add_label=False) cb1 = plt.colorbar(ax) cb1.set_label('Streamfunction') # Plot velocity potential. 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 = vp[i, :, :].plot.contour(x='lon', y='lat', levels=np.arange(-16., 16.1, 4.), 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 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 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_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 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_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 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_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_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()
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_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 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 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()
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(),
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, transform=ccrs.PlateCarree(), extend='both', add_colorbar=False)
def plot_sf_vp(run, land_mask=None): rcParams['figure.figsize'] = 10, 5 rcParams['font.size'] = 16 plot_dir = '/scratch/rg419/plots/climatology/' mkdir = sh.mkdir.bake('-p') mkdir(plot_dir) data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc') uwnd = data.ucomp.sel(pfull=200.) vwnd = data.vcomp.sel(pfull=200.) # Create a VectorWind instance to handle the computation w = VectorWind(uwnd, vwnd) # Compute variables streamfun, vel_pot = w.sfvp() def mn_sn_av(da): #Take seasonal and monthly averages da = da/10.**6 da.coords['season'] = np.mod(da.xofyear + 5., 72.) // 18. da.coords['month'] = (da.xofyear-1) //6 + 1 da_sn = da.groupby('season').mean(('xofyear')) da_mn = da.groupby('month').mean(('xofyear')) return da_sn, da_mn streamfun_sn, streamfun_mn = mn_sn_av(streamfun) vel_pot_sn, vel_pot_mn = mn_sn_av(vel_pot) lats= [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= 5.] lons= [data.lon[i] for i in range(len(data.lon)) if data.lon[i] <= 175. and data.lon[i] >=25.] sf_max_lat = np.zeros((12,)) vp_min_lat = np.zeros((12,)) sf_max_lon = np.zeros((12,)) vp_min_lon = np.zeros((12,)) for i in range(12): streamfun_asia_i = streamfun_mn[i,:,:].sel(lat=lats, lon=lons) sf_max_i = streamfun_asia_i.where(streamfun_asia_i==streamfun_asia_i.max(), drop=True) sf_max_lon[i] = sf_max_i.lon.values sf_max_lat[i] = sf_max_i.lat.values vp_asia_i = vel_pot_mn[i,:,:] vp_min_i = vp_asia_i.where(vp_asia_i==vp_asia_i.min(), drop=True) vp_min_lon[i] = vp_min_i.lon.values vp_min_lat[i] = vp_min_i.lat.values streamfun_asia_i.plot.contourf(x='lon', y='lat', levels = np.arange(-140.,141.,5.)) plt.plot(sf_max_lon[i], sf_max_lat[i], 'kx') plt.show() f1 = streamfun_sn[2,:,:].plot.contourf(x='lon', y='lat', 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)) plt.plot(sf_max_lon, sf_max_lat, 'kx-', mew=1.5) plt.grid(True,linestyle=':') plt.colorbar(f1) if not land_mask==None: land = xr.open_dataset(land_mask) land.land_mask.plot.contour(x='lon', y='lat', levels=np.arange(-1.,2.,1.), add_labels=False, colors='k') plt.savefig(plot_dir + 'streamfun_' + run + '.pdf', format='pdf') plt.close() f1 = vel_pot_sn[2,:,:].plot.contourf(x='lon', y='lat', levels = np.arange(-20.,21.,2.), add_labels=False, add_colorbar=False, extend='both') for i in range(12): plt.text(vp_min_lon[i]+3, vp_min_lat[i], str(i+1), fontsize=10) plt.plot(vp_min_lon, vp_min_lat, 'kx-', mew=1.5) plt.grid(True,linestyle=':') plt.colorbar(f1) if not land_mask==None: land = xr.open_dataset(land_mask) land.land_mask.plot.contour(x='lon', y='lat', levels=np.arange(-1.,2.,1.), add_labels=False, colors='k') plt.savefig(plot_dir + 'vel_pot_' + run + '.pdf', format='pdf') plt.close()