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