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()
