def energy_eq(run, month, filename='plev_daily', lev=150.):
#Load in dataset
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc'
name = name_temp % month
#read data into xarray
data = xr.open_dataset(name, decode_times=False)
# Calculate planetary vorticity
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
uchi = data.ucomp - data.upsi
vchi = data.vcomp - data.vpsi
pe_to_ke = -1. * (uchi * gr.ddx(data.height * 9.8) +
vchi * gr.ddy(data.height * 9.8, vector=False))
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
conv_1 = vor * (data.upsi * vchi - uchi * data.vpsi)
conv_2 = -1. * omega * (data.upsi * gr.ddp(uchi) +
data.vpsi * gr.ddp(vchi))
conv_3 = -1. * (gr.ddx(uchi) + gr.ddy(vchi)) * (data.upsi**2. +
data.vpsi**2.) / 2.
ds = xr.Dataset(
{
'pe_to_ke': (['time', 'pfull', 'lat', 'lon'], pe_to_ke),
'conv_1': (['time', 'pfull', 'lat', 'lon'], conv_1),
'conv_2': (['time', 'pfull', 'lat', 'lon'], conv_2),
'conv_3': (['time', 'pfull', 'lat', 'lon'], conv_3),
'upsi': (['time', 'pfull', 'lat', 'lon'], data.upsi),
'vpsi': (['time', 'pfull', 'lat', 'lon'], data.vpsi),
'uchi': (['time', 'pfull', 'lat', 'lon'], uchi),
'vchi': (['time', 'pfull', 'lat', 'lon'], vchi)
},
coords={
'time': ('time', data.time),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat),
'lon': ('lon', data.lon)
})
ds.coords['xofyear'] = np.mod(ds.time - 1., 360.) // 5 + 1.
dsout = ds.groupby('xofyear').mean(('time'))
fileout = '/scratch/rg419/Data_moist/' + run + '/run%03d/energy_eq.nc'
fileout = fileout % month
dsout.to_netcdf(path=fileout)
print 'data written to ', fileout
return dsout
def partition_advection(data, lons, lev=150):
#First do uu terms
uu_trans_dx = -86400. * gr.ddx(
(data.ucomp_sq -
data.ucomp**2).sel(pfull=lev)) # <u'u'> = <uu> - <u><u>
u = data.ucomp.sel(pfull=lev) # u
u_dx = -86400. * gr.ddx(u) # dudx
u_dudx_zav = u.sel(lon=lons).mean('lon') * u_dx.sel(lon=lons).mean(
'lon') # [u][dudx]
u_dudx_stat = (u * u_dx).sel(
lon=lons).mean('lon') - u_dudx_zav # u*dudx* = [ududx] - [u][dudx]
data['uu_trans_dx'] = (('xofyear', 'lat'),
uu_trans_dx.sel(lon=lons).mean('lon'))
data['u_dudx_stat'] = (('xofyear', 'lat'), u_dudx_stat)
data['u_dudx_zav'] = (('xofyear', 'lat'), u_dudx_zav)
print 'uu terms done'
#Next do uv terms
uv_trans_dy = -86400. * gr.ddy(
(data.ucomp_vcomp - data.ucomp * data.vcomp).sel(pfull=lev), uv=True)
v = data.vcomp.sel(pfull=lev).load() # v
u_dy = -86400. * gr.ddy(u) # dudy
v_dudy_zav = v.sel(lon=lons).mean('lon') * u_dy.sel(lon=lons).mean(
'lon') # [v][dudy]
v_dudy_stat = (v * u_dy).sel(
lon=lons).mean('lon') - v_dudy_zav # v*dudy* = [vdudy] - [v][dudy]
data['uv_trans_dy'] = (('xofyear', 'lat'),
uv_trans_dy.sel(lon=lons).mean('lon'))
data['v_dudy_stat'] = (('xofyear', 'lat'), v_dudy_stat)
data['v_dudy_zav'] = (('xofyear', 'lat'), v_dudy_zav)
print 'uv terms done'
#Finally do uw terms
uw_trans_dp = -86400. * gr.ddp(
(data.ucomp_omega - data.ucomp * data.omega).sel(lon=lons).mean('lon'))
w = data.omega.sel(pfull=lev).load() # w
u_dp = -86400. * (gr.ddp(data.ucomp)).sel(pfull=lev) # dudp
w_dudp_zav = w.sel(lon=lons).mean('lon') * u_dp.sel(lon=lons).mean('lon')
w_dudp_stat = (w * u_dp).sel(lon=lons).mean('lon') - w_dudp_zav
data['uw_trans_dp'] = (('xofyear', 'lat'), uw_trans_dp.sel(pfull=lev))
data['w_dudp_stat'] = (('xofyear', 'lat'), w_dudp_stat)
data['w_dudp_zav'] = (('xofyear', 'lat'), w_dudp_zav)
print 'uw terms done'
def transient_plot(run, ax_in, pentad, plot_land=True):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
data_vort = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc')
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp.sel(pfull=150)) # dvdx
u_dy = gr.ddy(data.ucomp.sel(pfull=150)) # dudy
vor = v_dx - u_dy + f
div = gr.ddx(data.ucomp.sel(pfull=150)) + gr.ddy(data.vcomp.sel(pfull=150))
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp.sel(pfull=150) * dvordx +
data.vcomp.sel(pfull=150) * dvordy)
stretching_mean = -86400.**2. * vor * div
transient = (data_vort.horiz_md.sel(pfull=150).values +
data_vort.stretching.sel(pfull=150).values
) * 86400.**2. - horiz_md_mean - stretching_mean
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
data['land'] = (('lat', 'lon'), land.land_mask)
f1 = transient[pentad, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-5., 5.1, 1.),
ax=ax_in,
add_labels=False,
extend='both',
add_colorbar=False)
if plot_land:
data.land.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_in,
colors='k',
add_colorbar=False,
add_labels=False)
(vor * 86400.)[pentad, :, :].plot.contour(x='lon',
y='lat',
levels=np.arange(-14., 15., 2.),
ax=ax_in,
colors='0.7',
add_labels=False)
ax_in.set_xlim(0, 180)
ax_in.set_ylim(-30, 60)
ax_in.set_xticks(np.arange(0., 185., 30.))
ax_in.set_yticks(np.arange(-30., 65., 30.))
return f1
def vort_eq(run, month, filename='plev_daily', period_fac=1.):
#Load in dataset
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/'+filename+'.nc'
name = name_temp % month
#read data into xarray
data = xr.open_dataset( name, decode_times=False)
# Calculate planetary vorticity
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat *np.pi/180)
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
#vor_rel = v_dx - u_dy
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
dvordp = gr.ddp(vor)
# Calculate horizontal material derivative
horiz_md = -1. * (data.ucomp * dvordx + data.vcomp * dvordy)
# Calculate vertical material derivative
vert_md = -1. * data.omega * dvordp
# Now do the terms involving gradients of windspeed
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching = -1. * vor * div
tilting = gr.ddp(data.ucomp) * gr.ddy(data.omega, vector=False) - gr.ddp(data.vcomp) * gr.ddx(data.omega)
total = horiz_md + vert_md + stretching + tilting
ds = xr.Dataset({'horiz_md': (['time', 'pfull', 'lat', 'lon'], horiz_md),
'vert_md': (['time', 'pfull', 'lat', 'lon'], vert_md),
'stretching': (['time', 'pfull', 'lat', 'lon'], stretching),
'tilting': (['time', 'pfull', 'lat', 'lon'], tilting),
'ucomp': (['time', 'pfull', 'lat', 'lon'], data.ucomp),
'vcomp': (['time', 'pfull', 'lat', 'lon'], data.vcomp),
'total': (['time', 'pfull', 'lat', 'lon'], total)},
coords={'time': ('time', data.time),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat),
'lon': ('lon', data.lon)})
dsout = ds#.mean(('time'))
fileout = '/scratch/rg419/Data_moist/' + run + '/run%03d/vort_eq.nc'
fileout = fileout % month
dsout.to_netcdf(path=fileout)
print 'data written to ', fileout
return dsout
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 get_stretching(run):
#Load data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
f = 2 * mc.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 - v_dx + u_dy
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching_mean = -86400.**2. * (vor * div).sel(pfull=150.)
return stretching_mean
def get_horiz_adv(run):
#Load data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
f = 2 * mc.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 - v_dx + u_dy
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_adv_mean = -86400.**2. * (data.ucomp * dvordx +
data.vcomp * dvordy).sel(pfull=150.)
return horiz_adv_mean
def calc_vort(data):
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
f = 2 * mc.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(pfull=150) * 86400.
return abs_vort
def get_abs_vort(data):
# 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
data['abs_vort'] = vor.mean('lon') * 86400.
def vort_eq_hm(run, bef_aft, lev=150, lonin=[-1., 361.]):
rcParams['figure.figsize'] = 10, 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)
#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.
print 'starting plotting'
np.abs(
vor.sel(lon=lons, xofyear=range(bef_aft[0], bef_aft[0] + 4)).mean(
('lon', 'xofyear')) * 86400.).plot(color='k')
np.abs(
vor.sel(lon=lons, xofyear=range(bef_aft[1], bef_aft[1] + 4)).mean(
('lon', 'xofyear')) * 86400.).plot(color='r')
plt.xlim([-45, 45])
plt.ylim([-12, 12])
plt.grid(True, linestyle=':')
if lonin == [-1., 361.]:
figname = 'vort_line_' + run + '.pdf'
else:
figname = 'vort_line_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def get_transient(run):
#Load data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
f = 2 * mc.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)
horiz_adv_mean = -86400.**2. * (data.ucomp * dvordx +
data.vcomp * dvordy).sel(pfull=150.)
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching_mean = -86400.**2. * (vor * div).sel(pfull=150.)
transient = data.horiz_md.sel(pfull=150.) + data.stretching.sel(
pfull=150.) - horiz_adv_mean - stretching_mean
return transient
def get_abs_vort(run):
#Load data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_'+run+'.nc')
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
f = 2 * mc.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*86400.
abs_vort = abs_vort.sel(pfull=150)
return abs_vort
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 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 vort_budg(run, ax_s, ax_h, pentad):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
data['rain'] = (('xofyear', 'lat', 'lon'),
(data.convection_rain + data.condensation_rain) * 86400.)
lons = [
data.lon[i] for i in range(len(data.lon))
if data.lon[i] >= 60. and data.lon[i] < 150.
]
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -30. and data.lat[i] < 60.
]
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp.sel(pfull=150)) # dvdx
u_dy = gr.ddy(data.ucomp.sel(pfull=150)) # dudy
vor = v_dx - u_dy + f
div = gr.ddx(data.ucomp.sel(pfull=150)) + gr.ddy(data.vcomp.sel(pfull=150))
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp.sel(pfull=150) * dvordx +
data.vcomp.sel(pfull=150) * dvordy)
stretching_mean = -86400.**2. * vor * div
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
f1 = stretching_mean[pentad:pentad +
4, :, :].mean('xofyear').plot.contourf(
x='lon',
y='lat',
levels=np.arange(-5., 5.1, 1.),
ax=ax_s,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_s,
colors='k',
add_colorbar=False,
add_labels=False)
cs = (vor * 86400.)[pentad:pentad + 4, :, :].sel(
lon=lons,
lat=lats).mean('xofyear').plot.contour(x='lon',
y='lat',
levels=np.arange(-14., 15., 2.),
ax=ax_s,
colors='0.6',
add_labels=False,
add_colorbar=False)
plt.clabel(cs, fontsize=12, inline_spacing=-1, fmt='%1.0f')
f1 = horiz_md_mean[pentad:pentad + 4, :, :].mean('xofyear').plot.contourf(
x='lon',
y='lat',
levels=np.arange(-5., 5.1, 1.),
ax=ax_h,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_h,
colors='k',
add_colorbar=False,
add_labels=False)
cs = (vor * 86400.)[pentad:pentad + 4, :, :].sel(
lon=lons,
lat=lats).mean('xofyear').plot.contour(x='lon',
y='lat',
levels=np.arange(-14., 15., 2.),
ax=ax_h,
colors='0.6',
add_labels=False,
add_colorbar=False)
plt.clabel(cs, fontsize=12, inline_spacing=-1, fmt='%1.0f')
for ax in [ax_s, ax_h]:
ax.set_xlim(60, 150)
ax.set_ylim(-30, 60)
ax.set_xticks(np.arange(60., 155., 30.))
ax.set_yticks(np.arange(-30., 65., 30.))
return f1
def vort_eq_ss(run, lonin=[-1., 361.]):
rcParams['figure.figsize'] = 10, 7.5
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/steady_state_runs/'
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'
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
f = 2 * mc.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.horiz_md * 86400.**2 + data.stretching * 86400.**2 -
horiz_md_mean - stretching_mean).values
data['transient'] = (('pfull', '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.transient.sel(lon=lons).mean('lon')
total_hm = data.total.sel(lon=lons).mean('lon') * 86400.**2.
print 'starting plotting'
levels = np.arange(-1.5, 1.6, 0.25)
# 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='lat',
y='pfull',
levels=levels,
ax=ax1,
extend='both',
add_labels=False,
add_colorbar=False,
yincrease=False)
ax1.contour(data.lat,
data.pfull,
abs_vort,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax1.set_xlabel('Latitude')
ax1.set_xlim(-60, 60)
ax1.set_xticks(np.arange(-60, 61, 30))
ax1.grid(True, linestyle=':')
#Second plot
stretching_hm.plot.contourf(x='lat',
y='pfull',
levels=levels,
ax=ax2,
extend='both',
add_labels=False,
add_colorbar=False,
yincrease=False)
ax2.contour(data.lat,
data.pfull,
abs_vort,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax2.set_xlim(-60, 60)
ax2.set_xticks(np.arange(-60, 61, 30))
ax2.grid(True, linestyle=':')
#Third plot
transient_hm.plot.contourf(x='lat',
y='pfull',
levels=levels,
ax=ax3,
extend='both',
add_labels=False,
add_colorbar=False,
yincrease=False)
ax3.contour(data.lat,
data.pfull,
abs_vort,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax3.set_xlabel('Latitude')
ax3.set_xlim(-60, 60)
ax3.set_xticks(np.arange(-60, 61, 30))
ax3.grid(True, linestyle=':')
#Fourth plot
total_hm.plot.contourf(x='lat',
y='pfull',
levels=levels,
ax=ax4,
extend='both',
add_labels=False,
add_colorbar=False,
yincrease=False)
ax4.contour(data.lat,
data.pfull,
abs_vort,
levels=np.arange(-12., 13., 2.),
colors='k',
linewidths=2,
alpha=0.25)
ax4.set_xlim(-60, 60)
ax4.set_xticks(np.arange(-60, 61, 30))
ax4.grid(True, linestyle=':')
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_ss_' + run + '.pdf'
else:
figname = 'vort_budg_ss_' + run + '_' + str(int(lonin[0])) + '_' + str(
int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def get_edge_psi_nh(run,
lonin=[-1., 361.],
sanity_check=False,
lev=500.,
thresh=0.,
do_month_av=False,
period_fac=1.,
intdown=True,
do_ageostrophic=False,
do_geostrophic=False):
"""Front end so same function can be used for both climatological and steady state data"""
#Load data
if run == 'era':
data = xr.open_dataset(
'/scratch/rg419/obs_and_reanalysis/era_v_alllevs.nc')
# Take pentad means
data.coords['xofyear'] = data.day_of_yr // 5 + 1.
data = data.groupby('xofyear').mean(('day_of_yr'))
else:
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' +
run + '.nc')
if do_ageostrophic:
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
dphidx = gr.ddx(data.height)
data = data
data['vcomp'] = data.vcomp - data.vcomp.mean('lon') - dphidx * 9.8 / f
if do_geostrophic:
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
dphidx = gr.ddx(data.height)
data['vcomp'] = dphidx * 9.8 / f
# Choose lons to average over
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 data has a time dimension, use time series fn. Otherwise use steady state fn
if 'xofyear' in data.coords:
edge_loc, psi_max = get_edge_psi_time(data,
lons,
sanity_check=sanity_check,
lev=lev,
thresh=thresh,
intdown=intdown)
# Take monthly average if wanted, return outputs
if do_month_av:
psi_max.coords['month'] = np.mod(
psi_max.xofyear - 1, 72. * period_fac) // (6 * period_fac) + 1
psi_max_month = psi_max.groupby('month').mean(('xofyear'))
edge_loc.coords['month'] = np.mod(
edge_loc.xofyear - 1, 72. * period_fac) // (6 * period_fac) + 1
edge_loc_month = edge_loc.groupby('month').mean(('xofyear'))
return edge_loc, psi_max, edge_loc_month, psi_max_month
else:
return edge_loc, psi_max
else:
edge_loc, psi_max = get_edge_psi_ss(data,
lons,
sanity_check=sanity_check,
lev=lev,
thresh=thresh,
intdown=intdown)
return edge_loc, psi_max
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/full_qflux.nc')
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
height_before = 9.8 * data.height.sel(
pfull=850.)[18:18 + 4, :, :].mean('xofyear') / 1000
height_after = 9.8 * data.height.sel(
pfull=850.)[39:39 + 4, :, :].mean('xofyear') / 1000
#Geopotential gradient
dphidx_before = -86400. * 9.8 * (gr.ddx(
data.height.sel(pfull=150.)))[18:18 + 4, :, :].mean('xofyear')
dphidx_after = -86400. * 9.8 * (gr.ddx(
data.height.sel(pfull=150.)))[39:39 + 4, :, :].mean('xofyear')
land = xr.open_dataset(
'/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc')
levels = np.arange(12., 15.1, 0.25)
levels_dphidx = np.arange(-100., 100.1, 20.)
# Two subplots
fig, (ax1, ax2) = plt.subplots(2, sharex=True)
#First plot
f1 = height_before.plot.contourf(x='lon',
y='lat',
ax=ax1,
def mom_budg(run, lev=150):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
#First do uu terms
uu_trans_dx = -86400. * gr.ddx(
(data.ucomp_sq -
data.ucomp**2).sel(pfull=lev)) # <u'u'> = <uu> - <u><u>
u = data.ucomp.sel(pfull=lev) # u
u_dx = -86400. * gr.ddx(u) # dudx
u_ed = u - u.mean('lon')
u_dx_ed = u_dx - u_dx.mean('lon')
u_dudx_zav = u.mean('lon') * u_dx.mean(
'lon') # [u][dudx], where brackets denote mean over all longitudes
u_dudx_cross1 = u.mean('lon') * u_dx_ed # [u]dudx*
u_dudx_cross2 = u_ed * u_dx.mean('lon') # u*[dudx]
u_dudx_stat = u_ed * u_dx_ed # u*dudx*
data['uu_trans_dx'] = (('xofyear', 'lat', 'lon'), uu_trans_dx)
data['u_dudx_cross1'] = (('xofyear', 'lat', 'lon'), u_dudx_cross1)
data['u_dudx_cross2'] = (('xofyear', 'lat', 'lon'), u_dudx_cross2)
data['u_dudx_stat'] = (('xofyear', 'lat', 'lon'), u_dudx_stat)
data['u_dudx_zav'] = (('xofyear', 'lat'), u_dudx_zav)
#Next do uv terms
uv_trans_dy = -86400. * gr.ddy(
(data.ucomp_vcomp - data.ucomp * data.vcomp).sel(pfull=lev), uv=True)
v = data.vcomp.sel(pfull=lev).load() # v
u_dy = -86400. * gr.ddy(u) # dudy
v_ed = v - v.mean('lon')
u_dy_ed = u_dy - u_dy.mean('lon')
v_dudy_zav = v.mean('lon') * u_dy.mean('lon') # [v][dudy]
v_dudy_cross1 = v.mean('lon') * u_dy_ed # [v]dudy*
v_dudy_cross2 = v_ed * u_dy.mean('lon') # v*[dudy]
v_dudy_stat = v_ed * u_dy_ed # v*dudy*
data['uv_trans_dy'] = (('xofyear', 'lat', 'lon'), uv_trans_dy)
data['v_dudy_cross1'] = (('xofyear', 'lat', 'lon'), v_dudy_cross1)
data['v_dudy_cross2'] = (('xofyear', 'lat', 'lon'), v_dudy_cross2)
data['v_dudy_stat'] = (('xofyear', 'lat', 'lon'), v_dudy_stat)
data['v_dudy_zav'] = (('xofyear', 'lat'), v_dudy_zav)
#Finally do uw terms
uw_trans_dp = -86400. * gr.ddp(
(data.ucomp_omega - data.ucomp * data.omega))
w = data.omega.sel(pfull=lev).load() # w
u_dp = -86400. * (gr.ddp(data.ucomp)).sel(pfull=lev) # dudp
w_ed = w - w.mean('lon')
u_dp_ed = u_dp - u_dp.mean('lon')
w_dudp_zav = w.mean('lon') * u_dp.mean('lon') # [w][dudp]
w_dudp_cross1 = w.mean('lon') * u_dp_ed # [w]dudp*
w_dudp_cross2 = w_ed * u_dp.mean('lon') # w*[dudp]
w_dudp_stat = w_ed * u_dp_ed # w*dudp*
data['uw_trans_dp'] = (('xofyear', 'lat', 'lon'),
uw_trans_dp.sel(pfull=lev))
data['w_dudp_cross1'] = (('xofyear', 'lat', 'lon'), w_dudp_cross1)
data['w_dudp_cross2'] = (('xofyear', 'lat', 'lon'), w_dudp_cross2)
data['w_dudp_stat'] = (('xofyear', 'lat', 'lon'), w_dudp_stat)
data['w_dudp_zav'] = (('xofyear', 'lat'), w_dudp_zav)
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
fv = data.vcomp.sel(pfull=lev) * f * 86400.
fv_mean = fv.mean('lon')
fv_local = fv - fv_mean
totp = (data.convection_rain + data.condensation_rain) * 86400.
#Geopotential gradient
dphidx = gr.ddx(data.height.sel(pfull=lev))
dphidx = -86400. * 9.8 * dphidx
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_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_local + fv_mean + dphidx + mom_mean + mom_trans + mom_stat + mom_cross
data['mom_mean'] = (('xofyear', 'lat'), mom_mean)
data['mom_cross'] = (('xofyear', 'lat', 'lon'), mom_cross)
data['mom_trans'] = (('xofyear', 'lat', 'lon'), mom_trans)
data['mom_stat'] = (('xofyear', 'lat', 'lon'), mom_stat)
data['fv_local'] = (('xofyear', 'lat', 'lon'), fv_local)
data['fv_ageo'] = (('xofyear', 'lat', 'lon'), fv_ageo)
data['fv_mean'] = (('xofyear', 'lat'), fv_mean)
data['dphidx'] = (('xofyear', 'lat', 'lon'), dphidx)
data['mom_sum'] = (('xofyear', 'lat', 'lon'), mom_sum)
return data
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 stretching_plot(run, ax_h, ax_s, pentad):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
lons = [
data.lon[i] for i in range(len(data.lon))
if data.lon[i] >= 60. and data.lon[i] < 150.
]
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -30. and data.lat[i] < 60.
]
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp.sel(pfull=150)) # dvdx
u_dy = gr.ddy(data.ucomp.sel(pfull=150)) # dudy
vor = v_dx - u_dy + f
div = gr.ddx(data.ucomp.sel(pfull=150)) + gr.ddy(data.vcomp.sel(pfull=150))
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp.sel(pfull=150) * dvordx +
data.vcomp.sel(pfull=150) * dvordy)
stretching_mean = -86400.**2. * vor * div
levels = np.arange(-1.5, 1.6, 0.5)
fh = horiz_md_mean[pentad, :, :].plot.contourf(x='lon',
y='lat',
levels=levels,
ax=ax_h,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_h,
colors='k',
add_colorbar=False,
add_labels=False)
fs = stretching_mean[pentad, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(
-1.5, 1.6, 0.5),
ax=ax_s,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_s,
colors='k',
add_colorbar=False,
add_labels=False)
for ax in [ax_h, ax_s]:
ax.set_xlim(60, 150)
ax.set_ylim(-30, 60)
ax.set_xticks(np.arange(60., 151., 30.))
ax.set_yticks(np.arange(-30., 61., 30.))
ax.grid(True, linestyle=':')
ax_h.set_xticklabels('')
ax_s.set_xticklabels(['60', '90', '120', '150'])
return fh, fs
filename_u = '/scratch/rg419/obs_and_reanalysis/sep_levs_u/era_u_150_dt.nc'
filename_v = '/scratch/rg419/obs_and_reanalysis/sep_levs_v/era_v_150_dt.nc'
# Load vorticity
data_vo = xr.open_dataset(filename_vo, decode_times=False)
omega = 7.2921150e-5
f = 2. * omega * np.sin(data_vo.lat * np.pi / 180.)
# Choose lon range
#lons = [data_vo.lon[i] for i in range(len(data_vo.lon)) if data_vo.lon[i] >= 60. and data_vo.lon[i] < 150.]
vor = data_vo['vo_150'] + f
# Take x and y gradients of vorticity
dvodx = gr.ddx(vor)
dvody = gr.ddy(vor, vector=False)
print 'gradients done'
# Load u and v
data_u = xr.open_dataset(filename_u, decode_times=False)
data_v = xr.open_dataset(filename_v, decode_times=False)
udvodx = (dvodx * data_u['u_150']).mean('year_no')
vdvody = (dvody * data_v['v_150']).mean('year_no')
#udvodx.sel(lon=lons).mean('lon').plot.contourf()
#plt.figure(2)
#vdvody.sel(lon=lons).mean('lon').plot.contourf()
#plt.figure(3)
def max_div(run, lats=np.arange(-90, 90.1, 0.1), rotfac=1.0):
#Open data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
#Calculate vorticity, divergence, and pentad mean vortex stretching
f = 2 * mc.omega * rotfac * 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
# Select upper level and take zonal mean
# Take a 20 day window rolling mean, and interpolate onto a 0.1 deg lat grid.
# Find max values of divergence and stretching, and locate lat of max divergence
div_hm = div.mean('lon').sel(pfull=150.)
div_rm = rolling_mean(div_hm, 4, tcoord='pentad')
div_int = interp_field(div_rm)
div_max = div_int.max('lat') * 86400.
div_argmax = div_int.lat.values[div_int.argmax('lat').values]
div_argmax = xr.DataArray(div_argmax,
coords=[div_int.xofyear],
dims=['xofyear'])
stretching_hm = stretching_mean.mean('lon').sel(pfull=150.)
stretching_hm[:, 32:] = stretching_hm[:, 32:] * -1.
stretching_rm = rolling_mean(stretching_hm, 4, tcoord='pentad')
stretching_int = interp_field(stretching_rm)
stretching_max = stretching_int.max('lat')
# Plot up max divergence and vortex stretching tendency in top panel, latitude of max divergence below
#fig, ax1 = plt.subplots()
fig, (ax1, ax3) = plt.subplots(2, sharex=True)
div_max.plot(ax=ax1, color='k')
ax1.set_xlabel('')
ax1.set_title('')
# Make the y-axis label, ticks and tick labels match the line color.
ax1.set_ylabel('Divergence, day$^{-1}$')
ax1.tick_params('y')
ax1.grid(True, linestyle=':')
ax1.set_xlim(0, max(div_max.xofyear) + 1)
ax2 = ax1.twinx()
stretching_max.plot(color='b', ax=ax2)
ax2.set_title('')
ax2.set_ylabel('Vortex stretching, day$^{-2}$', color='b')
ax2.tick_params('y', colors='b')
ax2.set_xlim(0, max(div_max.xofyear) + 1)
div_argmax.plot(ax=ax3, color='k')
ax3.set_title('')
ax3.set_xlabel('Pentad')
ax3.set_ylabel('Lat of max divergence')
ax3.grid(True, linestyle=':')
ax3.set_xlim(0, max(div_max.xofyear) + 1)
fig.tight_layout()
plt.savefig(plot_dir + 'div_tseries_' + run + '.pdf', 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 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 vort_eq_ss(run, lonin=[-1., 361.]):
rcParams['figure.figsize'] = 15, 6.25
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/steady_state_runs/'
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'
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
f = 2 * mc.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')
print 'starting plotting'
levels = np.arange(-1.5, 1.6, 0.25)
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='lat',
y='pfull',
levels=levels,
ax=ax1,
extend='both',
add_labels=False,
yincrease=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_xlabel('Latitude')
ax1.set_ylabel('Pressure, hPa')
ax1.set_xlim(-60, 60)
ax1.set_xticks(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='lat',
y='pfull',
levels=np.arange(
-1.e-10, 1.05e-10,
0.1e-10),
ax=ax2,
extend='both',
add_labels=False,
yincrease=False)
#ax2.set_ylim(-60,60)
ax2.set_xticks(np.arange(-60, 61, 30))
ax2.grid(True, linestyle=':')
#ax2.text(-7, 60, 'b)')
#Third plot
data.vcomp.sel(lon=lons).mean('lon').plot.contourf(x='lat',
y='pfull',
levels=np.arange(
-12., 13, 2.),
ax=ax3,
extend='both',
add_labels=False,
yincrease=False)
#ax3.set_ylim(-60,60)
ax3.set_xticks(np.arange(-60, 61, 30))
ax3.grid(True, linestyle=':')
#ax3.text(-7, 60, 'c)')
#Fourth plot
stretching_hm.plot.contourf(x='lat',
y='pfull',
levels=levels,
ax=ax4,
extend='both',
add_labels=False,
yincrease=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_xlabel('Latitude')
ax4.set_ylabel('Pressure, hPa')
ax4.set_xlim(-60, 60)
ax4.set_xticks(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
(div.sel(lon=lons).mean('lon') * 86400.).plot.contourf(x='lat',
y='pfull',
levels=np.arange(
-1., 1.1, 0.1),
ax=ax5,
extend='both',
add_labels=False,
yincrease=False)
ax5.set_xlim(-60, 60)
ax4.set_xlabel('Latitude')
ax5.set_xticks(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
(vor.sel(lon=lons).mean('lon') * 86400.).plot.contourf(x='lat',
y='pfull',
levels=np.arange(
-12., 13., 2.),
ax=ax6,
extend='both',
add_labels=False,
yincrease=False)
#ax6.set_ylim(-60,60)
ax4.set_xlabel('Latitude')
ax6.set_xticks(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_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_budg(run, ax_v, ax_h, ax_s):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
lons = [
data.lon[i] for i in range(len(data.lon))
if data.lon[i] >= 60. and data.lon[i] < 150.
]
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -30. and data.lat[i] < 60.
]
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp.sel(pfull=150)) # dvdx
u_dy = gr.ddy(data.ucomp.sel(pfull=150)) # dudy
vor = v_dx - u_dy + f
div = gr.ddx(data.ucomp.sel(pfull=150)) + gr.ddy(data.vcomp.sel(pfull=150))
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp.sel(pfull=150) * dvordx +
data.vcomp.sel(pfull=150) * dvordy)
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')
abs_vort = vor.sel(lon=lons).mean('lon') * 86400.
abs_vort_amb = (vor[39:39 + 4, :, :].mean('pentad') -
vor[18:18 + 4, :, :].mean('pentad')) * 86400.
abs_vort_after = vor[39:39 + 4, :, :].mean('pentad') * 86400.
levels = np.arange(-1.5, 1.6, 0.5)
fv = abs_vort_amb.plot.contourf(x='lon',
y='lat',
extend='both',
ax=ax_v,
levels=np.arange(-5., 6., 1.),
add_colorbar=False,
add_labels=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(-5., 7., 5.),
ax=ax_v,
add_colorbar=False,
add_labels=False,
colors='0.5')
ax_v.contour(data.lon,
data.lat,
abs_vort_after,
levels=np.arange(-20., 20., 2.),
colors='k')
ax_v.set_yticks(np.arange(-30., 61., 30.))
ax_v.set_xticks(np.arange(60., 151., 30.))
ax_v.set_ylim(-30, 60)
ax_v.set_xlim(60, 150)
ax_v.grid(True, linestyle=':')
fh = horiz_md_mean[22:39, :, :].mean('pentad').plot.contourf(
x='lon',
y='lat',
levels=levels,
ax=ax_h,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(-5., 7., 5.),
ax=ax_h,
add_colorbar=False,
add_labels=False,
colors='0.5')
#abs_vort.plot.contour(x='pentad', y='lat', levels = np.arange(-20.,20.,2.), ax=ax_h, add_labels = False, colors = 'k')
fs = stretching_mean[22:39, :, :].mean('pentad').plot.contourf(
x='lon',
y='lat',
levels=levels,
ax=ax_s,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(-5., 7., 5.),
ax=ax_s,
add_colorbar=False,
add_labels=False,
colors='0.5')
#abs_vort.plot.contour(x='pentad', y='lat', levels = np.arange(-20.,20.,2.), ax=ax_s, add_labels = False, colors = 'k')
for ax in [ax_h, ax_s]:
ax.set_ylim(-30, 60)
ax.set_xlim(60, 150)
ax.set_xticks(np.arange(60., 151., 30.))
ax.set_yticks(np.arange(-30., 61., 30.))
ax.grid(True, linestyle=':')
ax_h.set_xticklabels('')
ax_v.set_xticklabels('')
return fv, fh, fs
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 vort_eq(run, months, filename='plev_daily', lev=150.):
#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})
print 'files opened'
# Calculate planetary vorticity
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
# Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
vor = v_dx - u_dy + f
#vor_rel = v_dx - u_dy
print 'absolute vorticity calculated'
dvordx = gr.ddx(vor.sel(pfull=lev))
dvordy = gr.ddy(vor.sel(pfull=lev), vector=False)
dvordp = gr.ddp(vor)
# Calculate horizontal material derivative
horiz_md = -1. * (data.ucomp * dvordx + data.vcomp * dvordy)
# Calculate vertical material derivative
vert_md = -1. * data.omega.sel(pfull=lev) * dvordp.sel(pfull=lev)
print 'advective terms done'
# Now do the terms involving gradients of windspeed
div = gr.ddx(data.ucomp.sel(pfull=lev)) + gr.ddy(data.vcomp.sel(pfull=lev))
stretching = -1. * vor * div
tilting = (gr.ddp(data.ucomp)).sel(pfull=lev) * gr.ddy(
data.omega.sel(pfull=lev), vector=False) - (gr.ddp(
data.vcomp)).sel(pfull=lev) * gr.ddx(data.omega.sel(pfull=lev))
print 'stretching and tilting done'
total = horiz_md + vert_md + stretching + tilting
ds = xr.Dataset(
{
'horiz_md': (['time', 'lat', 'lon'], horiz_md),
'vert_md': (['time', 'lat', 'lon'], vert_md),
'stretching': (['time', 'lat', 'lon'], stretching),
'tilting': (['time', 'lat', 'lon'], tilting),
'total': (['time', 'lat', 'lon'], total)
},
coords={
'time': ('time', data.time),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat),
'lon': ('lon', data.lon)
})
ds.coords['xofyear'] = np.mod(ds.time - 1., 360.) // 5 + 1.
dsout = ds.groupby('xofyear').mean(('time'))
GFDL_DATA = os.environ['GFDL_DATA']
fileout = GFDL_DATA + 'climatologies/' + run + '_' + str(
int(lev)) + '_vort_eq.nc'
dsout.to_netcdf(path=fileout)
print 'data written to ', fileout
return dsout