def plot_precip(data, lons=range(0, 128)):
print 'precip'
data['totp'] = (('xofyear', 'lat', 'lon'),
(data.convection_rain + data.condensation_rain) * 86400.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
ax = data.totp[:, :, lons].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
6., 25., 3.),
add_label=False,
add_colorbar=False,
extend='min')
cb1 = plt.colorbar(ax)
cb1.set_label('Precip, mm/day')
cs = data.t_surf[:, :,
lons].mean('lon').plot.contour(x='xofyear',
y='lat',
levels=range(250, 321, 10),
add_label=False,
colors='w',
add_colorbar=False)
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
plt.ylim((-45, 45))
plt.xlim((1, 73))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.ylabel('Latitude')
#plt.title('Zonal mean precipitation evolution, mm/day')
plt.tight_layout()
def mse_adv_hm(run, lev=850., filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.]):
rcParams['figure.figsize'] = 7, 5
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/'+run+'/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/'+run+'.nc')
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]]
cp = 287.04/2*7
L = 2.500e6
data['mse_v'] = (cp*data.vcomp_temp + L*data.sphum_v).sel(lon=lons).mean('lon')
data['mse_v_mean'] = cp*data.vcomp.sel(lon=lons).mean('lon') * data.temp.sel(lon=lons).mean('lon') + L*data.sphum.sel(lon=lons).mean('lon') * data.vcomp.sel(lon=lons).mean('lon')
data['mse'] = (cp*data.temp + L*data.sphum).sel(lon=lons).mean('lon')
mse_max_loc = [data.lat[i] for i in data.mse.sel(pfull=lev).argmax('lat').values]
levels = np.arange(-8000,8000.1,1000.)
mn_dic = month_dic(1)
tickspace = range(13,72,18)
labels = [mn_dic[(k+5)/6 ] for k in tickspace]
data.mse_v.sel(pfull=lev).plot.contourf(x='xofyear', y='lat', levels=np.arange(-10.e5,10.1e5,20.e4), extend='both', add_labels=False)
data.mse.sel(pfull=lev).plot.contour(x='xofyear', y='lat', colors='k', levels=np.arange(200000.,400000.,20000.), add_colorbar=False, add_labels=False)
plt.plot(data.xofyear,mse_max_loc)
plt.grid(True,linestyle=':')
plt.xticks(tickspace,labels,rotation=25)
plt.tight_layout()
plt.ylim(-60,60)
plt.ylabel('Latitude')
if lonin == [-1.,361.]:
figname = 'mse_v.pdf'
else:
figname = 'mse_v_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
data.mse_v_mean.sel(pfull=lev).plot.contourf(x='xofyear', y='lat', levels=np.arange(-10.e5,10.1e5,20.e4), extend='both', add_labels=False)
data.mse.sel(pfull=lev).plot.contour(x='xofyear', y='lat', colors='k', levels=np.arange(200000.,400000.,20000.), add_colorbar=False, add_labels=False)
plt.plot(data.xofyear,mse_max_loc)
plt.grid(True,linestyle=':')
plt.xticks(tickspace,labels,rotation=25)
plt.tight_layout()
plt.ylim(-60,60)
plt.ylabel('Latitude')
if lonin == [-1.,361.]:
figname = 'mse_v_mean.pdf'
else:
figname = 'mse_v_mean_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def plot_u(data, lons=range(0, 128)):
print 'u'
data['ucomp_850'] = (('xofyear', 'lat', 'lon'), data.ucomp[:,
2, :, :].values)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
ax = data.ucomp_850[:, :,
lons].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-20., 21., 4.),
add_label=False,
add_colorbar=False,
extend='neither')
cb1 = plt.colorbar(ax)
cb1.set_label('Zonal wind speed, m/s (850 hPa)')
cs = data.t_surf[:, :,
lons].mean('lon').plot.contour(x='xofyear',
y='lat',
levels=range(250, 321, 10),
add_label=False,
colors='w',
add_colorbar=False)
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
plt.ylim((-45, 45))
plt.xlim((1, 73))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.ylabel('Latitude')
#plt.title('Zonal mean precipitation evolution, mm/day')
plt.tight_layout()
def bs_plot(inp_fol, years):
data = load_bs_vars(inp_fol, years)
#pd_dic = pentad_dic(1)
mn_dic = month_dic(1)
#tickspace = range(9,80,18)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
ax = data.totp.plot.contourf(x='pentad',
y='lat',
levels=np.arange(6., 31., 3.),
add_label=False,
add_colorbar=False)
cb1 = plt.colorbar(ax)
cb1.set_label('Precip, mm/day')
cs = data.t_surf.plot.contour(x='pentad',
y='lat',
levels=range(250, 321, 10),
add_label=False,
colors='w',
add_colorbar=False)
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
plt.ylim((-45, 45))
plt.xlim((1, 73))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.ylabel('Latitude')
#plt.title('Zonal mean precipitation evolution, mm/day')
plt.tight_layout()
def p_and_t(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.,
lonin=[-1., 361.]):
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
totp = (data.convection_rain + data.condensation_rain) * 86400.
lons = pick_lons(data, lonin)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
# ax=totp[:,:,lons].mean('lon').plot.contourf(x='xofyear', y='lat',levels=np.arange(6.,31.,3.), add_label = False, add_colorbar=False)
ax = totp[:, :,
lons].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(2., 15., 2.),
add_label=False,
add_colorbar=False,
extend='both')
totp[:, :,
lons].mean('lon').plot.contour(x='xofyear',
y='lat',
levels=np.arange(-92., 109., 100.),
add_label=False,
add_colorbar=False,
colors='k')
plt.grid(True, linestyle=':')
cb1 = plt.colorbar(ax)
cb1.set_label('Precip, mm/day')
cs = data.t_surf[:, :,
lons].mean('lon').plot.contour(x='xofyear',
y='lat',
levels=range(250, 321, 10),
add_label=False,
colors='w',
add_colorbar=False)
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
plt.ylim((-45, 45))
plt.xlim((1, 73))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.ylabel('Latitude')
plt.tight_layout()
plt.savefig(plot_dir + 'p_and_t.png')
plt.close()
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 vort_eq_hm(run, lev=150, rotfac=1.0, period_fac=1.):
rcParams['figure.figsize'] = 15, 6.25
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/crit_lat_test/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#Load in vorticity budget term means
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
omega = 7.2921150e-5 * rotfac
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching_mean = -86400.**2. * vor * div
stretching_hm = stretching_mean.mean('lon')
mn_dic = month_dic(1)
tickspace = np.arange(13, 72, 18) * period_fac
labels = [mn_dic[(k + 5) / 6] for k in range(13, 72, 18)]
levels = np.arange(-3., 3.1, 0.25)
stretching_max = np.unravel_index(
stretching_hm.sel(pfull=lev).isel(lat=range(32, 64)).argmin(),
(data.pentad.size, 32))
print stretching_max
print data.lat[stretching_max[1] + 32]
stretching_hm.sel(pfull=lev).plot.contourf(x='pentad',
y='lat',
levels=levels,
extend='both',
add_labels=False)
#ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
plt.plot(data.pentad[stretching_max[0]], data.lat[stretching_max[1] + 32],
'kx')
plt.ylabel('Latitude')
plt.xlabel('')
plt.ylim(-60, 60)
plt.yticks(np.arange(-60, 61, 30))
plt.xticks(tickspace, labels, rotation=25)
plt.title('Vortex stretching', fontsize=17)
plt.grid(True, linestyle=':')
plt.tight_layout()
figname = 'vort_stretching_' + run + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def plot_pd_mom_fn(inp_fol):
data_p = load_bs_vars(inp_fol, range(11, 41))
data_file = '/scratch/rg419/plots/monsoon_analysis/' + inp_fol + '/mombudg_data.nc'
data = xr.open_dataset(data_file)
print 'data loaded, starting plotting'
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/exp/topo_10m/input/land.nc'
land = xr.open_dataset(land_file)
land_plot = xr.DataArray(land.land_mask.values, [('lat', data.lat),
('lon', data.lon)])
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
def plot_mom_var(var, levels):
if var == 'fv_mn_imb':
plot_dataa = data.data_vars['fv_av'] + data.data_vars['mom_mean']
else:
plot_dataa = data.data_vars[var]
plot_data = plot_dataa[:, :, 22:65].mean(('lon')) * 10000.
g = plot_data.plot.contourf(x='pentad',
y='lat',
levels=levels,
add_label=False,
add_colorbar=False,
extend='both')
cb1 = plt.colorbar(g)
cb1.set_label('$\displaystyle10^{-4}m/s^2$')
cs = data_p.totp.plot.contour(x='pentad',
y='lat',
levels=np.arange(6., 31., 6.),
add_label=False,
add_colorbar=False,
colors='k')
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
plt.ylim(0, 90)
plt.xlim(1, 73)
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.ylabel('Latitude')
#plt.plot([float(data.pentad[36]),float(data.pentad[36])],[0.,90.],'k--')
plt.tight_layout()
plt.title('')
#plt.title(var_dic[var] + ', $\displaystyle10^{-4}m/s^2$')
plt.savefig('/scratch/rg419/plots/monsoon_analysis/' + inp_fol + '/' +
var + '_lattime.png')
plt.close()
vars = [
'dphidx_av', 'fv_av', 'mom_eddy', 'mom_mean', 'mom_sum', 'fv_mn_imb'
]
for var in vars:
plot_mom_var(var, np.arange(-2., 2.1, 0.4))
def plot_pd_mom_fn(inp_fol, years):
data_p = load_bs_vars(inp_fol, years)
data = mombudg_pd_fn(inp_fol, years)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
def plot_mom_var(var, lev, levels):
var_dic = {
'fv_av': 'fv',
'mom_eddy': 'Eddy advective terms',
'mom_mean': 'Mean state advective terms',
'mom_sum': 'Residual'
}
plot_data = data.data_vars[var]
plot_data = plot_data * 10000.
g = plot_data[:, lev, :].plot.contourf(x='pentad',
y='lat',
levels=levels,
add_label=False,
add_colorbar=False,
extend='neither')
cb1 = plt.colorbar(g)
cb1.set_label('$\displaystyle10^{-4}m/s^2$')
cs = data_p.totp.plot.contour(x='pentad',
y='lat',
levels=np.arange(6., 31., 6.),
add_label=False,
add_colorbar=False,
colors='k')
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
plt.ylim(0, 90)
plt.xlim(1, 73)
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.ylabel('Latitude')
#plt.plot([float(data.pentad[36]),float(data.pentad[36])],[0.,90.],'k--')
plt.tight_layout()
plt.title('')
#plt.title(var_dic[var] + ', $\displaystyle10^{-4}m/s^2$')
plt.savefig('/scratch/rg419/plots/mom_budg_work/' + inp_fol + '/' +
var + '_lattime.png')
plt.close()
vars = ['fv_av', 'mom_eddy', 'mom_mean', 'mom_sum']
for var in vars:
plot_mom_var(var, 9, np.arange(-2., 2.1, 0.2))
def plot_aht_eq(run):
ahteq, swabs, olr, shf, stor = aht_eq(run)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
ahteq.plot(color='m')
(-1. * stor).plot(color='c')
(-1. * olr).plot(color='g')
shf.plot(color='b')
swabs.plot(color='r')
#plt.xlim((1,72))
#plt.xlabel('')
#plt.xticks(tickspace,labels,rotation=25)
plt.ylabel('Atmospheric heat transport at the equator, PW')
plt.savefig('/scratch/rg419/plots/aht_work/aht_eq_' + run + '.pdf',
format='pdf')
plt.close()
def lapse_rate(run, pentad, period_fac=1.0):
#Load in data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
g = 9.8
Ra = 287.04
mn_dic = month_dic(1)
tickspace = np.arange(13, 72, 18) * period_fac
labels = [mn_dic[(k + 5) / 6] for k in range(13, 72, 18)]
levels = np.arange(-1.5, 1.6, 0.25)
dlnTdp = gr.ddp(np.log(data.temp.mean('lon')))
dTdz = data.pfull * 100. * g / Ra * dlnTdp * 1000.
#dTdz.mean('xofyear').plot.contourf(x='lat', y='pfull', extend = 'both', add_labels=False, yincrease=False, levels=np.arange(-10.,10.5,2.))
#plt.xlabel('Latitude')
#plt.ylabel('Pressure, hPa')
#plt.grid(True,linestyle=':')
#plt.tight_layout()
#figname = 'dTdz_' + run + '.pdf'
#plt.savefig(plot_dir + figname, format='pdf')
#plt.close()
#ind = abs(dTdz.isel(pfull=range(10,20)) -2.).mean('xofyear').argmin('pfull')
ind = abs(dTdz.isel(pfull=range(10, 20)) -
2.).isel(xofyear=pentad).argmin('pfull')
trop_height = np.zeros([64, 1])
for i in range(0, 64):
#trop_height[i] = data.height.mean(('lon','xofyear'))[ind.values[i]+10,i]
trop_height[i] = data.height.mean(
('lon')).isel(xofyear=pentad)[ind.values[i] + 10, i]
print run, trop_height[32]
return trop_height
def v_fft(run, months, before_after, filename='plev_daily', period_fac=1.):
rcParams['figure.figsize'] = 10, 15
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/'+run+'/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#Load in dataset
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/'+filename+'.nc'
names = [name_temp % m for m in range( months[0], months[1]) ]
#read data into xarray
data = xr.open_mfdataset( names,
decode_times=False, # no calendar so tell netcdf lib
# choose how data will be broken down into manageable chunks.
chunks={'time': 30})
data.coords['xofyear'] = np.mod( data.time, 360.*period_fac)
vwnd = data.vcomp.load()
# Take fourier transform and average over years
N = 72*period_fac # Number of samplepoints
data.coords['wavenumber'] = np.arange(0,len(data.lon))
data['v_fft'] = (('time','lat','wavenumber'), 2.0/N * np.abs(fft(vwnd[:,17,:,:]))) #150 hPa
data = data.groupby('xofyear').mean('time')
fft_before = data.v_fft[before_after[0]:before_after[0]+20,:,:16].mean('xofyear')
fft_after = data.v_fft[before_after[1]:before_after[1]+20,:,:16].mean('xofyear')
fft_5 = data.v_fft[:,33,:16]
fft_45 = data.v_fft[:,48,:16]
mn_dic = month_dic(1)
tickspace = range(60,360,90)
labels = [mn_dic[(k+30)/30 ] for k in tickspace]
# Two subplots
f, axarr = plt.subplots(2, sharex=True)
plt.set_cmap('inferno_r')
#First plot
f1 = axarr[0].contourf(data.wavenumber[:16], data.lat, fft_before, levels = np.arange(3.,19.,3.), extend = 'max')
axarr[0].set_ylabel('Latitude')
axarr[0].set_yticks(np.arange(-60.,61.,30.))
axarr[0].grid(True,linestyle=':')
#Second plot
f2 = axarr[1].contourf(data.wavenumber[:16], data.lat, fft_after, levels = np.arange(3.,19.,3.), extend = 'max')
axarr[1].set_ylabel('Latitude')
axarr[1].set_yticks(np.arange(-60.,61.,30.))
axarr[1].set_xlabel('Wavenumber')
axarr[1].grid(True,linestyle=':')
plt.subplots_adjust(right=0.95, top=0.95, bottom=0., hspace=0.1)
#Colorbar
cb1=f.colorbar(f2, ax=axarr.ravel().tolist(), use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.07, aspect=30)
cb1.set_label('Variance, $(m/s)^2$')
plt.savefig(plot_dir+'fft_before_after.pdf', format='pdf')
plt.close()
rcParams['figure.figsize'] = 20, 10
# Two subplots
f, axarr = plt.subplots(2, sharex=True)
#First plot
f1 = fft_5.plot.contourf(x='xofyear', y='wavenumber', ax=axarr[0], levels = np.arange(3.,19.,3.), extend = 'max', add_colorbar=False, add_labels=False, cmap='inferno_r')
axarr[0].set_ylabel('Wavenumber')
axarr[0].grid(True,linestyle=':')
#Second plot
f2 = fft_45.plot.contourf(x='xofyear', y='wavenumber', ax=axarr[1], levels = np.arange(3.,19.,3.), extend = 'max', add_colorbar=False, add_labels=False, cmap='inferno_r')
axarr[1].set_ylabel('Wavenumber')
axarr[1].set_xticks(tickspace)
axarr[1].set_xticklabels(labels,rotation=25)
axarr[1].grid(True,linestyle=':')
plt.subplots_adjust(right=0.95, top=0.95, bottom=0., hspace=0.1)
#Colorbar
cb1=f.colorbar(f2, ax=axarr.ravel().tolist(), use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.07, aspect=30)
cb1.set_label('Variance, $(m/s)^2$')
plt.savefig(plot_dir+'fft_5_45.pdf', 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()
]
return lons
plot_dir = '/scratch/rg419/plots/paper_1_figs/revisions/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset(
'/scratch/rg419/obs_and_reanalysis/gpcp_detrendedpentads.nc')
lons = pick_lons(data, [60., 150.])
data['totp'] = data.precip_clim.sel(lon=lons).mean('lon')
mn_dic = month_dic(1)
tickspace = [13, 31, 49, 68]
labels = ['Mar', 'Jun', 'Sep', 'Dec']
plevels = np.arange(2., 15., 2.)
# Begin plotting: 3 subplots
#First plot
f1 = data.totp.plot.contourf(x='xofyear',
y='lat',
levels=plevels,
add_labels=False,
extend='max',
cmap='Blues')
#data.totp.plot.contour(x='xofyear', y='lat',levels=np.arange(-92.,109.,100.), add_labels = False, add_colorbar=False, colors='k')
def vort_eq_hm(run, lev=150, lonin=[-1., 361.]):
rcParams['figure.figsize'] = 15, 6.25
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/paper_1_figs/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#Load in vorticity budget term means
data_vort = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc')
data_vort = data_vort * 86400.**2. #Convert to day^-2
print 'vorticity budget data loaded'
#Also load climatological data so that transient eddies can be calculated
data = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
print 'climatology loaded'
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]
]
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
abs_vort = vor.sel(lon=lons).mean('lon') * 86400.
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching_mean = -86400.**2. * vor * div
horiz_md_hm = horiz_md_mean.sel(lon=lons).mean('lon')
stretching_hm = stretching_mean.sel(lon=lons).mean('lon')
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
levels = np.arange(-1.5, 1.6, 0.25)
print 'starting plotting'
# Four subplots
f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
3,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f2 = horiz_md_hm.plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax1,
extend='both',
add_labels=False)
#ax1.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.set_title('Horizontal advection', fontsize=17)
ax1.set_yticks(np.arange(-60, 61, 30))
ax1.grid(True, linestyle=':')
ax1.text(-15, 60, 'a)')
#Second plot
dvordy.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-1.e-10, 1.05e-10,
0.1e-10),
ax=ax2,
extend='both',
add_labels=False)
#ax2.set_ylim(-60,60)
ax2.set_yticks(np.arange(-60, 61, 30))
ax2.set_title('$\partial (\overline{\zeta} + f) /\partial y$', fontsize=17)
ax2.grid(True, linestyle=':')
ax2.text(-7, 60, 'b)')
#Third plot
data.vcomp.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-12., 13, 2.),
ax=ax3,
extend='both',
add_labels=False)
ax3.set_ylim(-60, 60)
ax3.set_yticks(np.arange(-60, 61, 30))
ax3.set_title('$\overline{v}$', fontsize=17)
ax3.grid(True, linestyle=':')
ax3.text(-7, 60, 'c)')
#Fourth plot
stretching_hm.plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax4,
extend='both',
add_labels=False)
#ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
ax4.set_ylabel('Latitude')
ax4.set_ylim(-60, 60)
ax4.set_yticks(np.arange(-60, 61, 30))
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels, rotation=25)
ax4.set_title('Vortex stretching', fontsize=17)
ax4.grid(True, linestyle=':')
ax4.text(-15, 60, 'd)')
#Fith plot
(div.sel(lon=lons).mean('lon') * 86400.).plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-0.6, 0.7, 0.1),
ax=ax5,
extend='both',
add_labels=False)
ax5.set_ylim(-60, 60)
ax5.set_yticks(np.arange(-60, 61, 30))
ax5.set_xticks(tickspace)
ax5.set_xticklabels(labels, rotation=25)
ax5.grid(True, linestyle=':')
ax5.set_title('Divergence', fontsize=17)
ax5.text(-7, 60, 'e)')
#Sixth plot
(vor.sel(lon=lons).mean('lon') * 86400.).plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-12., 13., 2.),
ax=ax6,
extend='both',
add_labels=False)
ax6.set_ylim(-60, 60)
ax6.set_yticks(np.arange(-60, 61, 30))
ax6.set_xticks(tickspace)
ax6.set_xticklabels(labels, rotation=25)
ax6.grid(True, linestyle=':')
ax6.set_title('Absolute vorticity', fontsize=17)
ax6.text(-7, 60, 'f)')
plt.subplots_adjust(right=0.97,
left=0.08,
top=0.93,
bottom=0.1,
hspace=0.25,
wspace=0.15)
if lonin == [-1., 361.]:
figname = 'vort_breakdown_' + run + '.pdf'
else:
figname = 'vort_breakdown_' + run + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def vort_eq_hm(run, lev=150, lonin=[-1., 361.]):
rcParams['figure.figsize'] = 10, 7.5
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/vort_schematic/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#Load in vorticity budget term means
data_vort = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc')
data_vort = data_vort * 86400.**2. #Convert to day^-2
print 'vorticity budget data loaded'
#Also load climatological data so that transient eddies can be calculated
data = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
print 'climatology loaded'
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]
]
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
abs_vort = vor.sel(lon=lons).mean('lon') * 86400.
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching_mean = -86400.**2. * vor * div
transient = data_vort.horiz_md.sel(
pfull=lev).values + data_vort.stretching.sel(
pfull=lev).values - horiz_md_mean.values - stretching_mean.values
data_vort['transient'] = (('pentad', 'lat', 'lon'), transient)
horiz_md_hm = horiz_md_mean.sel(lon=lons).mean('lon')
stretching_hm = stretching_mean.sel(lon=lons).mean('lon')
transient_hm = data_vort.transient.sel(lon=lons).mean('lon')
total_hm = data_vort.total.sel(pfull=lev, lon=lons).mean('lon')
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
levels = np.arange(-1.5, 1.6, 0.1)
print 'starting plotting'
# Four subplots
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,
2,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f2 = horiz_md_hm.plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax1,
extend='both',
add_labels=False,
add_colorbar=False)
ax1.contour(data.xofyear,
data.lat,
abs_vort.T,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.set_yticks(np.arange(-60, 61, 30))
ax1.grid(True, linestyle=':')
ax1.text(-10, 60, 'a)')
#Second plot
stretching_hm.plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax2,
extend='both',
add_labels=False,
add_colorbar=False)
ax2.contour(data.xofyear,
data.lat,
abs_vort.T,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax2.set_ylim(-60, 60)
ax2.set_yticks(np.arange(-60, 61, 30))
ax2.grid(True, linestyle=':')
ax2.text(-5, 60, 'b)')
#Third plot
transient_hm.plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax3,
extend='both',
add_labels=False,
add_colorbar=False)
ax3.contour(data.xofyear,
data.lat,
abs_vort.T,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax3.set_ylabel('Latitude')
ax3.set_ylim(-60, 60)
ax3.set_yticks(np.arange(-60, 61, 30))
ax3.set_xticks(tickspace)
ax3.set_xticklabels(labels, rotation=25)
ax3.grid(True, linestyle=':')
ax3.text(-10, 60, 'c)')
#Fourth plot
total_hm.plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax4,
extend='both',
add_labels=False,
add_colorbar=False)
ax4.contour(data.xofyear,
data.lat,
abs_vort.T,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax4.set_ylim(-60, 60)
ax4.set_yticks(np.arange(-60, 61, 30))
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels, rotation=25)
ax4.grid(True, linestyle=':')
ax4.text(-5, 60, 'd)')
plt.subplots_adjust(right=0.97,
left=0.1,
top=0.95,
bottom=0.,
hspace=0.2,
wspace=0.1)
#Colorbar
cb1 = f.colorbar(f2,
ax=[ax1, ax2, ax3, ax4],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label('Vorticity tendency, day$^{-2}$')
if lonin == [-1., 361.]:
figname = 'vort_budg_hm_' + run + '.pdf'
else:
figname = 'vort_budg_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def mom_budg_hm(run,
months,
lev=150,
filename='plev_pentad',
timeav='pentad',
period_fac=1.,
lonin=[-1., 361.]):
rcParams['figure.figsize'] = 12, 10
#rcParams['figure.figsize'] = 18, 10
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#data = time_means(run, months, filename=filename, timeav=timeav,period_fac=period_fac)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
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]
]
#advective terms
partition_advection(data, lons, lev=150)
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
fv = data.vcomp.sel(pfull=lev) * f * 86400.
fv = fv.sel(lon=lons).mean('lon')
#Geopotential gradient
dphidx = gr.ddx(data.height.sel(pfull=lev))
dphidx = -86400. * 9.8 * dphidx.sel(lon=lons).mean('lon')
mom_mean = data.u_dudx_zav + data.v_dudy_zav + data.w_dudp_zav
mom_trans = data.uu_trans_dx + data.uv_trans_dy + data.uw_trans_dp
mom_stat = data.u_dudx_stat + data.v_dudy_stat + data.w_dudp_stat
mom_sum = fv + dphidx + mom_mean + mom_trans + mom_stat
levels = np.arange(-10, 10.1, 2.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
# Six subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
3,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1 = fv.plot.contourf(ax=ax1,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.grid(True, linestyle=':')
#Second plot
mom_mean.plot.contourf(ax=ax2,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax2.grid(True, linestyle=':')
ax2.set_ylim(-60, 60)
#Third plot
dphidx.plot.contourf(ax=ax3,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax3.grid(True, linestyle=':')
ax3.set_ylim(-60, 60)
#Fourth plot
mom_trans.plot.contourf(ax=ax4,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax4.grid(True, linestyle=':')
ax4.set_ylabel('Latitude')
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels, rotation=25)
ax4.set_ylim(-60, 60)
#Fifth plot
mom_stat.plot.contourf(ax=ax5,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax5.grid(True, linestyle=':')
ax5.set_xticks(tickspace)
ax5.set_xticklabels(labels, rotation=25)
ax5.set_ylim(-60, 60)
#Sixth plot
mom_sum.plot.contourf(ax=ax6,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax6.grid(True, linestyle=':')
ax6.set_xticks(tickspace)
ax6.set_xticklabels(labels, rotation=25)
ax6.set_ylim(-60, 60)
plt.subplots_adjust(right=0.97,
left=0.1,
top=0.95,
bottom=0.,
hspace=0.3,
wspace=0.1)
#Colorbar
cb1 = fig.colorbar(f1,
ax=[ax1, ax2, ax3, ax4, ax5, ax6],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30)
#cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True,fraction=0.15, aspect=30)
cb1.set_label('$ms^{-1}day^{-1}$')
if lonin == [-1., 361.]:
figname = 'zon_mom_budg.pdf'
else:
figname = 'zon_mom_budg_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def mom_budg_hm(run, lev=150, lonin=[-1., 361.]):
#NB run is a dummy variable
rcParams['figure.figsize'] = 12, 8
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/paper_1_figs/revisions/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/era_mom_vars.nc')
# Take pentad means
data.coords['pentad'] = data.day_of_yr // 5 + 1.
data = data.groupby('pentad').mean(('day_of_yr'))
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]
]
#advective terms
partition_advection(data, lons, lev=150)
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
fv = data.vcomp * f * 86400.
fv_mean = fv.mean('lon')
fv_local = (fv - fv_mean).sel(lon=lons).mean('lon')
#Geopotential gradient
dphidx = gr.ddx(data.phi)
dphidx = -86400. * dphidx.sel(lon=lons).mean('lon')
fv_ageo = fv_local + dphidx
mom_mean = data.u_dudx_zav + data.v_dudy_zav + data.w_dudp_zav
mom_cross = data.u_dudx_cross1 + data.v_dudy_cross1 + data.w_dudp_cross1 + data.u_dudx_cross2 + data.v_dudy_cross2 + data.w_dudp_cross2
mom_cross1 = data.u_dudx_cross1 + data.v_dudy_cross1 + data.w_dudp_cross1
mom_cross2 = data.u_dudx_cross2 + data.v_dudy_cross2 + data.w_dudp_cross2
mom_trans = data.uu_trans_dx + data.uv_trans_dy + data.uw_trans_dp
mom_stat = data.u_dudx_stat + data.v_dudy_stat + data.w_dudp_stat
mom_crossu = data.u_dudx_cross1 + data.u_dudx_cross2
mom_crossv = data.v_dudy_cross1 + data.v_dudy_cross2
mom_crossw = data.w_dudp_cross1 + data.w_dudp_cross2
#data.w_dudp_cross1.plot.contourf(x='pentad', y='lat', extend='both',levels = np.arange(-20,21.1,2.))
#plt.yticks(np.arange(-60.,61.,30.))
#plt.ylim(-60,60)
#plt.figure(2)
#data.w_dudp_cross2.plot.contourf(x='pentad', y='lat', extend='both',levels = np.arange(-20,21.1,2.))
#plt.ylim(-60,60)
#plt.yticks(np.arange(-60.,61.,30.))
#plt.show()
#print mom_stat
mom_sum = fv_local + fv_mean + dphidx + mom_mean + mom_trans + mom_stat + mom_cross
levels = np.arange(-20, 21.1, 2.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
# Six subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8,
ax9)) = plt.subplots(3,
3,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1 = fv_mean.plot.contourf(ax=ax1,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax1.set_ylabel('Latitude')
ax1.set_title('Zonal mean Coriolis', fontsize=17)
ax1.set_ylim(-60, 60)
ax1.grid(True, linestyle=':')
ax1.set_yticks(np.arange(-60., 61., 30.))
ax1.text(-15, 60, 'a)')
#Second plot
mom_mean.plot.contourf(ax=ax2,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax2.grid(True, linestyle=':')
ax2.set_title('Mean state advection', fontsize=17)
ax2.set_ylim(-60, 60)
ax2.text(-5, 60, 'b)')
#Third plot
mom_sum.plot.contourf(ax=ax3,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax3.grid(True, linestyle=':')
ax3.set_ylim(-60, 60)
ax3.set_title('Residual', fontsize=17)
ax3.text(-5, 60, 'c)')
#Fourth plot
mom_trans.plot.contourf(ax=ax4,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax4.grid(True, linestyle=':')
ax4.set_ylabel('Latitude')
ax4.set_ylim(-60, 60)
ax4.set_title('Transient eddy flux conv.', fontsize=17)
ax4.set_yticks(np.arange(-60., 61., 30.))
ax4.text(-15, 60, 'd)')
#Fifth plot
mom_stat.plot.contourf(ax=ax5,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax5.grid(True, linestyle=':')
ax5.set_title('Stat. eddy flux conv.', fontsize=17)
ax5.set_ylim(-60, 60)
ax5.text(-5, 60, 'e)')
#Sixth plot
mom_cross.plot.contourf(ax=ax6,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax6.grid(True, linestyle=':')
ax6.set_ylim(-60, 60)
ax6.set_title('Stat. eddy cross terms', fontsize=17)
ax6.text(-5, 60, 'f)')
#Fourth plot
dphidx.plot.contourf(ax=ax7,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax7.grid(True, linestyle=':')
ax7.set_ylabel('Latitude')
ax7.set_xticks(tickspace)
ax7.set_xticklabels(labels, rotation=25)
ax7.set_ylim(-60, 60)
ax7.set_title('Geopotential gradient', fontsize=17)
ax7.set_yticks(np.arange(-60., 61., 30.))
ax7.text(-15, 60, 'g)')
#Fifth plot
fv_local.plot.contourf(ax=ax8,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax8.grid(True, linestyle=':')
ax8.set_xticks(tickspace)
ax8.set_xticklabels(labels, rotation=25)
ax8.set_title('Stat. eddy Coriolis', fontsize=17)
ax8.set_ylim(-60, 60)
ax8.text(-5, 60, 'h)')
#Sixth plot
fv_ageo.plot.contourf(ax=ax9,
x='pentad',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax9.grid(True, linestyle=':')
ax9.set_xticks(tickspace)
ax9.set_xticklabels(labels, rotation=25)
ax9.set_ylim(-60, 60)
ax9.set_title('Ageostrophic stat. eddy Coriolis', fontsize=17)
ax9.text(-5, 60, 'i)')
plt.subplots_adjust(right=0.97,
left=0.1,
top=0.95,
bottom=0.,
hspace=0.25,
wspace=0.12)
#Colorbar
cb1 = fig.colorbar(f1,
ax=[ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.15,
aspect=30,
shrink=0.5)
#cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True,fraction=0.15, aspect=30)
#cb1.set_label('$ms^{-1}day^{-1}$')
if lonin == [-1., 361.]:
figname = 'zon_mom_budg_era.pdf'
else:
figname = 'zon_mom_budg_era_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def energy_eq_hm(run,
lev=150,
lonin=[-1., 361.],
levels=np.arange(-0.01, 0.011, 0.001)):
rcParams['figure.figsize'] = 10, 9
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#Load in energy budget term means
data = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/energy_eq_' + run + '.nc')
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]
]
pe_to_ke_hm = data.pe_to_ke.sel(pfull=lev, lon=lons).mean('lon')
conv_1_hm = data.conv_1.sel(pfull=lev, lon=lons).mean('lon')
conv_2_hm = data.conv_2.sel(pfull=lev, lon=lons).mean('lon')
conv_3_hm = data.conv_3.sel(pfull=lev, lon=lons).mean('lon')
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
# Four subplots
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,
2,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f2 = pe_to_ke_hm.plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax1,
extend='both',
add_labels=False,
add_colorbar=False)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.grid(True, linestyle=':')
#Second plot
conv_1_hm.plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax2,
extend='both',
add_labels=False,
add_colorbar=False)
ax2.set_ylim(-60, 60)
ax2.grid(True, linestyle=':')
#Third plot
conv_3_hm.plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax3,
extend='both',
add_labels=False,
add_colorbar=False)
ax3.set_ylabel('Latitude')
ax3.set_ylim(-60, 60)
ax3.set_xticks(tickspace)
ax3.set_xticklabels(labels, rotation=25)
ax3.grid(True, linestyle=':')
#Fourth plot
conv_2_hm.plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax4,
extend='both',
add_labels=False,
add_colorbar=False)
ax4.set_ylim(-60, 60)
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels, rotation=25)
ax4.grid(True, linestyle=':')
plt.subplots_adjust(right=0.97,
left=0.1,
top=0.95,
bottom=0.,
hspace=0.3,
wspace=0.1)
#Colorbar
cb1 = f.colorbar(f2,
ax=[ax1, ax2, ax3, ax4],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30,
shrink=0.9)
cb1.set_label('Energy conversion, J/s')
if lonin == [-1., 361.]:
figname = 'energy_budg_hm_' + str(int(lev)) + '.pdf'
else:
figname = 'energy_budg_hm_' + str(int(lev)) + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def abs_vort_hm(run, lev=150, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.]):
rcParams['figure.figsize'] = 10, 15
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/'+run+'/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/'+run+'.nc')
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]]
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat *np.pi/180)
abs_vort = (f + data.vor.sel(pfull=lev))*86400.
abs_vort_zmean = abs_vort.sel(lon=lons).mean('lon')
abs_vort_min = abs_vort
abs_vort_min[0:32,:,:] = -1.*abs_vort_min[0:32,:,:]
abs_vort_min = abs_vort_min.sel(lon=lons).min('lon')
abs_vort_min[0:32,:] = -1.*abs_vort_min[0:32,:]
levels = np.arange(-20.,20.1,2.)
mn_dic = month_dic(1)
tickspace = range(13,72,18)
labels = [mn_dic[(k+5)/6 ] for k in tickspace]
# Two subplots
fig, (ax1, ax2) = plt.subplots(2, sharex=True)
plt.set_cmap('RdBu_r')
#First plot
f1=abs_vort_zmean.plot.contourf(ax=ax1, x='xofyear', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False)
abs_vort_zmean.plot.contour(ax=ax1, x='xofyear', y='lat', extend='both', levels=[-1000.,0.,1000.], add_colorbar=False, add_labels=False)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60,60)
ax1.grid(True,linestyle=':')
#Second plot
abs_vort_min.plot.contourf(ax=ax2, x='xofyear', y='lat', extend='both', levels=levels, add_colorbar=False, add_labels=False)
abs_vort_min.plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=[-1000.,0.,1000.], add_colorbar=False, add_labels=False)
ax2.grid(True,linestyle=':')
ax2.set_xticks(tickspace)
ax2.set_xticklabels(labels,rotation=25)
ax1.set_ylabel('Latitude')
ax2.set_ylim(-60,60)
plt.subplots_adjust(right=0.95, top=0.95, bottom=0., hspace=0.1)
#Colorbar
cb1=fig.colorbar(f1, ax=[ax1,ax2], use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.1, aspect=30)
cb1.set_label('$day^{-1}$')
if lonin == [-1.,361.]:
figname = 'abs_vort.pdf'
else:
figname = 'abs_vort_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def mom_budg_hm(run,
lev=150,
filename='plev_pentad',
timeav='pentad',
period_fac=1.,
lonin=[-1., 361.]):
a = 6376.0e3
rcParams['figure.figsize'] = 12, 5.5
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/itcz_theory_testing/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
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]
]
#advective terms
partition_advection(data, lons, lev=150)
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180.)
fu = -1. * data.ucomp.sel(pfull=lev) * f * 86400.
fu = fu.sel(lon=lons).mean('lon')
metric_term = (-86400. * data.ucomp.sel(pfull=lev) *
data.ucomp.sel(pfull=lev) *
np.tan(data.lat * np.pi / 180.) / a).mean('lon')
#totp = ((data.convection_rain + data.condensation_rain)*86400.).sel(lon=lons).mean('lon')
#abs_vort = (data.vor + f).sel(lon=lons).mean('lon')*86400.
#Geopotential gradient
dphidy = gr.ddy(data.height.sel(pfull=lev), vector=False)
dphidy = -86400. * 9.8 * dphidy.sel(lon=lons).mean('lon')
mom_mean = data.v_dvdy_zav + data.w_dvdp_zav
mom_trans = data.vv_trans_dy + data.vw_trans_dp
mom_stat = data.v_dvdy_stat + data.w_dvdp_stat
mom_sum = mom_mean + mom_trans + mom_stat + fu + dphidy + metric_term
resid = -(fu + dphidy + metric_term)
levels = np.arange(-10., 10.1, 1.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
resid.plot.contourf(x='xofyear',
y='lat',
extend='both',
levels=levels,
add_labels=False)
plt.ylabel('Latitude')
plt.ylim(-60, 60)
plt.xticks(tickspace, labels, rotation=25)
plt.yticks(np.arange(-60., 61., 30.))
plt.grid(True, linestyle=':')
plt.tight_layout()
plt.savefig(plot_dir + 'residual.pdf')
plt.close()
mom_mean.plot.contourf(x='xofyear',
y='lat',
extend='both',
levels=levels,
add_labels=False)
plt.ylabel('Latitude')
plt.ylim(-60, 60)
plt.xticks(tickspace, labels, rotation=25)
plt.yticks(np.arange(-60., 61., 30.))
plt.grid(True, linestyle=':')
plt.tight_layout()
plt.savefig(plot_dir + 'mom_mean.pdf')
plt.close()
def vort_eq_hm(run, lev=150, rotfac=1.0, period_fac=1.):
rcParams['figure.figsize'] = 10, 7.5
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/crit_lat_test/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#Load in vorticity budget term means
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
print 'vorticity budget data loaded'
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
omega = 7.2921150e-5 * rotfac
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
abs_vort = vor.mean('lon') * 86400.
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching_mean = -86400.**2. * vor * div
transient = 86400.**2. * data.horiz_md + 86400.**2. * data.stretching - horiz_md_mean - stretching_mean
data['transient'] = (('pentad', 'pfull', 'lat', 'lon'), transient.values)
horiz_md_hm = horiz_md_mean.mean('lon')
stretching_hm = stretching_mean.mean('lon')
transient_hm = data.transient.mean('lon')
total_hm = data.total.sel(pfull=lev).mean('lon')
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
levels = np.arange(-1.5, 1.6, 0.25)
print 'starting plotting'
# Four subplots
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,
2,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f2 = horiz_md_hm.sel(pfull=lev).plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax1,
extend='both',
add_labels=False,
add_colorbar=False)
ax1.contour(data.pentad,
data.lat,
abs_vort.sel(pfull=lev).T,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.set_yticks(np.arange(-60, 61, 30))
ax1.grid(True, linestyle=':')
ax1.text(-10, 60, 'a)')
#Second plot
stretching_hm.sel(pfull=lev).plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax2,
extend='both',
add_labels=False,
add_colorbar=False)
ax2.contour(data.pentad,
data.lat,
abs_vort.sel(pfull=lev).T,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax2.set_ylim(-60, 60)
ax2.set_yticks(np.arange(-60, 61, 30))
ax2.grid(True, linestyle=':')
ax2.text(-5, 60, 'b)')
#Third plot
transient_hm.sel(pfull=lev).plot.contourf(x='pentad',
y='lat',
levels=levels,
ax=ax3,
extend='both',
add_labels=False,
add_colorbar=False)
ax3.contour(data.pentad,
data.lat,
abs_vort.sel(pfull=lev).T,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax3.set_ylabel('Latitude')
ax3.set_ylim(-60, 60)
ax3.set_yticks(np.arange(-60, 61, 30))
ax3.set_xticks(tickspace)
ax3.set_xticklabels(labels, rotation=25)
ax3.grid(True, linestyle=':')
ax3.text(-10, 60, 'c)')
#Fourth plot
(horiz_md_hm + stretching_hm).sel(pfull=lev).plot.contourf(
x='pentad',
y='lat',
levels=levels,
ax=ax4,
extend='both',
add_labels=False,
add_colorbar=False)
ax4.contour(data.pentad,
data.lat,
abs_vort.sel(pfull=lev).T,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax4.set_ylim(-60, 60)
ax4.set_yticks(np.arange(-60, 61, 30))
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels, rotation=25)
ax4.grid(True, linestyle=':')
ax4.text(-5, 60, 'd)')
plt.subplots_adjust(right=0.97,
left=0.1,
top=0.95,
bottom=0.,
hspace=0.2,
wspace=0.1)
#Colorbar
cb1 = f.colorbar(f2,
ax=[ax1, ax2, ax3, ax4],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label('Vorticity tendency, day$^{-2}$')
figname = 'vort_budg_hm_balance' + run + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def abs_vort_hm(run,
lev=150,
filename='plev_daily',
timeav='pentad',
period_fac=1.,
latin=25.):
rcParams['figure.figsize'] = 9, 9
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc'
names = [name_temp % m for m in range(397, 409)]
#read data into xarray
data = xr.open_mfdataset(
names,
decode_times=False, # no calendar so tell netcdf lib
# choose how data will be broken down into manageable chunks.
chunks={'time': 30})
uwnd = data.ucomp.sel(pfull=lev)
vwnd = data.vcomp.sel(pfull=lev)
# Create a VectorWind instance to handle the computation of streamfunction and
# velocity potential.
w = VectorWind(uwnd, vwnd)
# Compute the streamfunction and velocity potential.
data['vor'], data['div'] = w.vrtdiv()
#data = xr.open_dataset('/scratch/rg419/Data_moist/'+run+'climatologies/'+run+'.nc')
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
lat_hm = data.lat[np.argmin(np.abs(data.lat - latin))]
abs_vort = (f + data.vor).sel(lat=lat_hm) * 86400.
levels = np.arange(0., 14.1, 2.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
# Plot
f1 = abs_vort.plot.contourf(x='lon',
y='time',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
plt.set_cmap('inferno_r')
plt.ylabel('Time')
#plt.yticks(tickspace, labels, rotation=25)
plt.xlabel('Longitude')
plt.xlim(60, 150)
plt.ylim(240 + 33 * 360, 90 + 33 * 360)
plt.grid(True, linestyle=':')
plt.tight_layout()
#Colorbar
cb1 = plt.colorbar(f1,
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30)
cb1.set_label('$day^{-1}$')
figname = 'abs_vort_lon_hm.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def mom_budg_latlon(run, t, lev=150, land=False, lonin=[-1., 361.]):
rcParams['figure.figsize'] = 15, 10
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
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]
]
if land:
land_data = xr.open_dataset(
'/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc')
t_dates = data.xofyear[t:t + 4]
data_tmean = data.sel(xofyear=t_dates).mean('xofyear')
#advective terms
partition_advection(data_tmean, lons=lons, lev=lev)
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
fv = data_tmean.vcomp.sel(pfull=lev) * f * 86400.
#Geopotential gradient
dphidx = gr.ddx(data_tmean.height.sel(pfull=lev))
dphidx = -86400. * 9.8 * dphidx
mom_mean = data_tmean.u_dudx_zav + data_tmean.v_dudy_zav + data_tmean.w_dudp_zav
mom_trans = data_tmean.uu_trans_dx + data_tmean.uv_trans_dy + data_tmean.uw_trans_dp
mom_stat = data_tmean.u_dudx_stat + data_tmean.v_dudy_stat + data_tmean.w_dudp_stat
mom_sum = fv + dphidx + mom_mean + mom_trans + mom_stat
mom_mean = xr.DataArray(np.rollaxis(np.tile(mom_mean, [128, 1]), 1),
mom_stat.coords)
levels = np.arange(-30, 30.1, 5.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
# Six subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
3,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1 = fv.plot.contourf(ax=ax1,
x='lon',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.set_xlim(60, 150)
ax1.grid(True, linestyle=':')
#Second plot
mom_mean.plot.contourf(ax=ax2,
x='lon',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax2.grid(True, linestyle=':')
ax2.set_ylim(-60, 60)
ax2.set_xlim(60, 150)
#Third plot
dphidx.plot.contourf(ax=ax3,
x='lon',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax3.grid(True, linestyle=':')
ax3.set_ylim(-60, 60)
ax3.set_xlim(60, 150)
#Fourth plot
mom_trans.plot.contourf(ax=ax4,
x='lon',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax4.grid(True, linestyle=':')
ax4.set_ylabel('Latitude')
ax4.set_xlabel('Longitude')
ax4.set_ylim(-60, 60)
ax4.set_xlim(60, 150)
ax4.set_xticks(range(60, 151, 15))
#Fifth plot
mom_stat.plot.contourf(ax=ax5,
x='lon',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax5.grid(True, linestyle=':')
ax4.set_xlabel('Longitude')
ax5.set_ylim(-60, 60)
ax5.set_xlim(60, 150)
ax5.set_xticks(range(60, 151, 15))
#Sixth plot
mom_sum.plot.contourf(ax=ax6,
x='lon',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax6.grid(True, linestyle=':')
ax4.set_xlabel('Longitude')
ax6.set_ylim(-60, 60)
ax6.set_xlim(60, 150)
ax6.set_xticks(range(60, 151, 15))
if land:
for ax in [ax1, ax2, ax3, ax4, ax5, ax6]:
land_data.land_mask.plot.contour(x='lon',
y='lat',
ax=ax,
colors='w',
levels=range(-1000, 1002, 1000),
add_labels=False,
add_colorbar=False)
plt.subplots_adjust(right=0.95, top=0.95, bottom=0., hspace=0.1)
#Colorbar
cb1 = fig.colorbar(f1,
ax=[ax1, ax2, ax3, ax4, ax5, ax6],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30)
cb1.set_label('$ms^{-1}day^{-1}$')
figname = 'zon_mom_budg_ll_' + str(t) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def mom_budg_hm(run, lev=150, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.]):
rcParams['figure.figsize'] = 12, 5.5
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/paper_1_figs/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/'+run+'.nc')
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]]
#advective terms
partition_advection(data, lons, lev=150)
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat *np.pi/180)
fv = data.vcomp.sel(pfull=lev) * f * 86400.
fv = fv.sel(lon=lons).mean('lon')
totp = ((data.convection_rain + data.condensation_rain)*86400.).sel(lon=lons).mean('lon')
abs_vort = (data.vor + f).sel(lon=lons).mean('lon')*86400.
#Geopotential gradient
dphidx = gr.ddx(data.height.sel(pfull=lev))
dphidx = -86400. * 9.8 * dphidx.sel(lon=lons).mean('lon')
mom_mean = data.u_dudx_zav + data.v_dudy_zav + data.w_dudp_zav
mom_trans = data.uu_trans_dx + data.uv_trans_dy + data.uw_trans_dp
mom_stat = data.u_dudx_stat + data.v_dudy_stat + data.w_dudp_stat
mom_sum = fv + dphidx + mom_mean + mom_trans + mom_stat
levels = np.arange(-21,21.1,3.)
mn_dic = month_dic(1)
tickspace = range(13,72,18)
labels = [mn_dic[(k+5)/6 ] for k in tickspace]
# Six subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1=fv.plot.contourf(ax=ax1, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
totp.plot.contour(ax=ax1, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
ax1.set_ylabel('Latitude')
ax1.set_title('Coriolis', fontsize=17)
ax1.set_ylim(-60,60)
ax1.grid(True,linestyle=':')
ax1.set_yticks(np.arange(-60.,61.,30.))
ax1.text(-15, 60, 'a)')
#Second plot
mom_mean.plot.contourf(ax=ax2, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
#ax2.contour(data.xofyear, data.lat, abs_vort.sel(pfull=lev).T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2)
#abs_vort.sel(pfull=lev).plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-2.,6.,2.), add_colorbar=False, add_labels=False, colors='k', linewidths=2)
totp.plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
ax2.grid(True,linestyle=':')
ax2.set_title('Mean state advection', fontsize=17)
ax2.set_ylim(-60,60)
ax2.text(-5, 60, 'b)')
#Third plot
dphidx.plot.contourf(ax=ax3, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
totp.plot.contour(ax=ax3, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
ax3.grid(True,linestyle=':')
ax3.set_ylim(-60,60)
ax3.set_title('Geopotential gradient', fontsize=17)
ax3.text(-5, 60, 'c)')
#Fourth plot
mom_trans.plot.contourf(ax=ax4, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
totp.plot.contour(ax=ax4, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
ax4.grid(True,linestyle=':')
ax4.set_ylabel('Latitude')
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels,rotation=25)
ax4.set_ylim(-60,60)
ax4.set_title('Transient eddy flux conv.', fontsize=17)
ax4.set_yticks(np.arange(-60.,61.,30.))
ax4.text(-15, 60, 'd)')
#Fifth plot
mom_stat.plot.contourf(ax=ax5, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
totp.plot.contour(ax=ax5, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
ax5.grid(True,linestyle=':')
ax5.set_xticks(tickspace)
ax5.set_xticklabels(labels,rotation=25)
ax5.set_title('Stat. eddy flux conv.', fontsize=17)
ax5.set_ylim(-60,60)
ax5.text(-5, 60, 'e)')
#Sixth plot
mom_sum.plot.contourf(ax=ax6, x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
totp.plot.contour(ax=ax6, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
ax6.grid(True,linestyle=':')
ax6.set_xticks(tickspace)
ax6.set_xticklabels(labels,rotation=25)
ax6.set_ylim(-60,60)
ax6.set_title('Residual', fontsize=17)
ax6.text(-5, 60, 'f)')
plt.subplots_adjust(right=0.97, left=0.1, top=0.95, bottom=0., hspace=0.25, wspace=0.12)
#Colorbar
cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.15, aspect=30, shrink=0.5)
#cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True,fraction=0.15, aspect=30)
#cb1.set_label('$ms^{-1}day^{-1}$')
if lonin == [-1.,361.]:
figname = 'zon_mom_budg_' +run+ '.pdf'
else:
figname = 'zon_mom_budg_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def advection_hm(run,
months,
lev=150,
filename='plev_pentad',
timeav='pentad',
period_fac=1.,
lonin=[-1., 361.]):
rcParams['figure.figsize'] = 15, 10
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
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]
]
#advective terms
partition_advection(data, lons, lev=150)
data['mom_mean'] = (('xofyear', 'lat'),
data.u_dudx_zav + data.v_dudy_zav + data.w_dudp_zav)
data['mom_trans'] = (('xofyear', 'lat'), data.uu_trans_dx +
data.uv_trans_dy + data.uw_trans_dp)
data['mom_stat'] = (('xofyear', 'lat'),
data.u_dudx_stat + data.v_dudy_stat + data.w_dudp_stat)
levels = np.arange(-10, 10.1, 2.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
fig_dicts = [{
'p1': 'mom_mean',
'p2': 'u_dudx_zav',
'p3': 'v_dudy_zav',
'p4': 'w_dudp_zav',
'namebase': 'meanstate_adv'
}, {
'p1': 'mom_stat',
'p2': 'u_dudx_stat',
'p3': 'v_dudy_stat',
'p4': 'w_dudp_stat',
'namebase': 'stat_adv'
}, {
'p1': 'mom_trans',
'p2': 'uu_trans_dx',
'p3': 'uv_trans_dy',
'p4': 'uw_trans_dp',
'namebase': 'trans_adv'
}]
for fig_dict in fig_dicts:
# Six subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
3,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1 = data[fig_dict['p1']].plot.contourf(ax=ax1,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.grid(True, linestyle=':')
#Fourth plot
data[fig_dict['p2']].plot.contourf(ax=ax4,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax4.grid(True, linestyle=':')
ax4.set_ylabel('Latitude')
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels, rotation=25)
ax4.set_ylim(-60, 60)
#Fifth plot
data[fig_dict['p3']].plot.contourf(ax=ax5,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax5.grid(True, linestyle=':')
ax5.set_xticks(tickspace)
ax5.set_xticklabels(labels, rotation=25)
ax5.set_ylim(-60, 60)
#Sixth plot
data[fig_dict['p4']].plot.contourf(ax=ax6,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax6.grid(True, linestyle=':')
ax6.set_xticks(tickspace)
ax6.set_xticklabels(labels, rotation=25)
ax6.set_ylim(-60, 60)
plt.subplots_adjust(right=0.95, top=0.95, bottom=0., hspace=0.1)
#Colorbar
cb1 = fig.colorbar(f1,
ax=[ax1, ax2, ax3, ax4, ax5, ax6],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30)
cb1.set_label('$ms^{-1}day^{-1}$')
if lonin == [-1., 361.]:
figname = fig_dict['namebase'] + '.pdf'
else:
figname = fig_dict['namebase'] + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def vort_eq_hm(run, lev=150, lonin=[-1., 361.]):
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/paper_1_figs/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#Also load climatological data so that transient eddies can be calculated
data = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
print 'climatology loaded'
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]
]
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
print 'starting plotting'
# Easier to produce individual plots as scales are different.
vor.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
extend='both',
add_labels=False,
levels=np.arange(
-0.00024, 0.00025,
0.00002))
plt.ylabel('Latitude')
plt.ylim(-60, 60)
plt.yticks(np.arange(-60, 61, 30))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.grid(True, linestyle=':')
if lonin == [-1., 361.]:
figname = 'vor_hm_' + run + '.pdf'
else:
figname = 'vor_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
div.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
extend='both',
add_labels=False,
levels=np.arange(
-0.000012, 0.000012,
0.000001))
plt.ylabel('Latitude')
plt.ylim(-60, 60)
plt.yticks(np.arange(-60, 61, 30))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.grid(True, linestyle=':')
if lonin == [-1., 361.]:
figname = 'div_hm_' + run + '.pdf'
else:
figname = 'div_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
dvordx.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
extend='both',
add_labels=False)
plt.ylabel('Latitude')
plt.ylim(-60, 60)
plt.yticks(np.arange(-60, 61, 30))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.grid(True, linestyle=':')
if lonin == [-1., 361.]:
figname = 'dvordx_hm_' + run + '.pdf'
else:
figname = 'dvordx_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
dvordy.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
extend='both',
add_labels=False,
levels=np.arange(
-1.2e-10, 1.25e-10,
0.1e-10))
plt.ylabel('Latitude')
plt.ylim(-60, 60)
plt.yticks(np.arange(-60, 61, 30))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.grid(True, linestyle=':')
if lonin == [-1., 361.]:
figname = 'dvordy_hm_' + run + '.pdf'
else:
figname = 'dvordy_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
data.ucomp.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
extend='both',
add_labels=False)
plt.ylabel('Latitude')
plt.ylim(-60, 60)
plt.yticks(np.arange(-60, 61, 30))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.grid(True, linestyle=':')
if lonin == [-1., 361.]:
figname = 'u_hm_' + run + '.pdf'
else:
figname = 'u_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
data.vcomp.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
extend='both',
add_labels=False,
levels=np.arange(
-15., 15.5, 1.))
plt.ylabel('Latitude')
plt.ylim(-60, 60)
plt.yticks(np.arange(-60, 61, 30))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.grid(True, linestyle=':')
if lonin == [-1., 361.]:
figname = 'v_hm_' + run + '.pdf'
else:
figname = 'v_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.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 vort_eq_hm(run, lev=150, lonin=[-1., 361.]):
rcParams['figure.figsize'] = 15, 5.5
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/paper_1_figs/revisions/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#Load in vorticity budget term means
data_vort = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc')
data_vort = data_vort * 86400.**2. #Convert to day^-2
print 'vorticity budget data loaded'
#Also load climatological data so that transient eddies can be calculated
data = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
print 'climatology loaded'
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]
]
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
abs_vort = vor.sel(lon=lons).mean('lon') * 86400.
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching_mean = -86400.**2. * vor * div
# Zonal means
horiz_zmean = -86400.**2. * (data.ucomp.mean('lon') * dvordx.mean('lon') +
data.vcomp.mean('lon') * dvordy.mean('lon'))
stretching_zmean = -86400.**2. * vor.mean('lon') * div.mean('lon')
# Transient eddies
transient_s_hm = (data_vort.stretching.sel(pfull=lev).values -
stretching_mean).sel(lon=lons).mean('lon')
transient_h_hm = (data_vort.horiz_md.sel(pfull=lev).values -
horiz_md_mean).sel(lon=lons).mean('lon')
transient_hm = transient_s_hm + transient_h_hm
# Stationary eddies
u_ed = data.ucomp - data.ucomp.mean('lon')
v_ed = data.vcomp - data.vcomp.mean('lon')
dvordx_ed = dvordx - dvordx.mean('lon')
dvordy_ed = dvordy - dvordy.mean('lon')
div_ed = div - div.mean('lon')
vor_ed = vor - vor.mean('lon')
stat_s = -86400.**2. * (div_ed * vor_ed)
stat_h = -86400.**2. * (u_ed * dvordx_ed + v_ed * dvordy_ed)
# Stationary eddy cross terms
cross_s = -86400.**2. * (vor.mean('lon') * div_ed +
vor_ed * div.mean('lon'))
cross_h = -86400.**2. * (u_ed * dvordx.mean('lon') +
v_ed * dvordy.mean('lon') +
data.ucomp.mean('lon') * dvordx_ed +
data.vcomp.mean('lon') * dvordy_ed)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
levels = np.arange(-1.5, 1.6, 0.25)
print 'starting plotting'
# Four subplots
f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
3,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f2 = horiz_zmean.plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax1,
extend='both',
add_labels=False)
#ax1.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.set_yticks(np.arange(-60, 61, 30))
ax1.grid(True, linestyle=':')
ax1.text(-15, 60, 'a)')
#Second plot
stat_h.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax2,
extend='both',
add_labels=False)
ax2.set_ylim(-60, 60)
ax2.set_yticks(np.arange(-60, 61, 30))
ax2.grid(True, linestyle=':')
ax2.text(-7, 60, 'b)')
#Third plot
cross_h.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax3,
extend='both',
add_labels=False)
ax3.set_ylim(-60, 60)
ax3.set_yticks(np.arange(-60, 61, 30))
ax3.grid(True, linestyle=':')
ax3.text(-7, 60, 'c)')
#Fourth plot
stretching_zmean.plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax4,
extend='both',
add_labels=False)
#ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
ax4.set_ylabel('Latitude')
ax4.set_ylim(-60, 60)
ax4.set_yticks(np.arange(-60, 61, 30))
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels, rotation=25)
ax4.grid(True, linestyle=':')
ax4.text(-15, 60, 'd)')
#Fifth plot
stat_s.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax5,
extend='both',
add_labels=False)
ax5.set_ylim(-60, 60)
ax5.set_yticks(np.arange(-60, 61, 30))
ax5.set_xticks(tickspace)
ax5.set_xticklabels(labels, rotation=25)
ax5.grid(True, linestyle=':')
ax5.text(-7, 60, 'e)')
#Sixth plot
cross_s.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=levels,
ax=ax6,
extend='both',
add_labels=False)
#ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
ax6.set_ylim(-60, 60)
ax6.set_yticks(np.arange(-60, 61, 30))
ax6.set_xticks(tickspace)
ax6.set_xticklabels(labels, rotation=25)
ax6.grid(True, linestyle=':')
ax6.text(-7, 60, 'f)')
plt.subplots_adjust(right=0.97,
left=0.08,
top=0.93,
bottom=0.1,
hspace=0.25,
wspace=0.15)
if lonin == [-1., 361.]:
figname = 'vort_eddies_' + run + '.pdf'
else:
figname = 'vort_eddies_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def abs_vort_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]
]
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
abs_vort = (data.vor + f)[:, :, :, lons].mean('lon')
abs_vort *= 1.e4
uwnd = data.ucomp[:, level, :, :].load()
ax = abs_vort[:, level, :].plot.contourf(x='xofyear',
y='lat',
levels=np.arange(-1., 1.1, 0.1),
extend='both',
add_colorbar=False,
add_label=False)
cs = uwnd[:, :, lons].mean('lon').plot.contour(x='xofyear',
y='lat',
levels=np.arange(
-1000., 1001., 1000.),
extend='both',
colors='k',
add_colorbar=False,
add_label=False)
plt.grid(True, linestyle=':')
cb1 = plt.colorbar(ax)
cb1.set_label('Absolute vorticity')
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.yticks(range(-60, 61, 30))
plt.ylabel('Latitude')
plt.tight_layout()
plt.savefig(plot_dir + 'abs_vort.png')
plt.close()