def ap_monsoon_hm(run, title):
#Produce identical plot to above, but averaged over all latitudes for a uniform mld run
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
mn_dic = month_dic(1)
tickspace = [1, 19, 37, 55]
labels = ['1st Jan', '1st Apr', '1st Jul', '1st Oct']
# Set figure parameters
rcParams['figure.figsize'] = 5, 3.5
rcParams['font.size'] = 14
fig, ax1 = plt.subplots()
try:
f1 = (data.precipitation * 86400.).mean('lon').plot.contourf(
ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(0., 21., 2.),
add_labels=False,
extend='max',
cmap='Blues',
add_colorbar=False)
except:
f1 = ((data.condensation_rain + data.convection_rain) *
86400.).mean('lon').plot.contourf(ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(0., 21., 2.),
add_labels=False,
extend='max',
cmap='Blues',
add_colorbar=False)
ax1.set_ylim(-45, 45)
ax1.grid(True, linestyle=':')
ax1.set_yticks(np.arange(-30., 31., 30.))
ax1.set_title(title)
ax1.set_ylabel('Latitude')
ax1.set_xticks(tickspace)
ax1.set_xticklabels(labels, rotation=25)
plt.subplots_adjust(left=0.18,
right=0.97,
top=0.9,
bottom=0.2,
hspace=0.3,
wspace=0.2)
#cb1=fig.colorbar(f1, ax=ax1, use_gridspec=True, orientation = 'horizontal',fraction=0.05, pad=0.15, aspect=60, shrink=0.5)
#cb1.set_label('Precipitation, mm/day')
# Save as a pdf
plt.savefig(plot_dir + 'precip_monsoons_' + run + '.pdf', format='pdf')
plt.close()
data.close()
def overturning_plot(data,
ax_in,
lonin=[-1., 361.],
do_xlabels=False,
month_labels=True,
thresh=0.,
nh=False):
lons = pick_lons(data, lonin)
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons)
psi /= 1.e9
edge_loc, psi_max, psi_max_loc = get_edge_psi(data,
lonin=lonin,
lev=500.,
thresh=thresh,
nh=nh)
f1 = psi.sel(pfull=500.).plot.contourf(ax=ax_in,
x='xofyear',
y='lat',
levels=np.arange(-500., 510., 50.),
add_colorbar=False,
add_labels=False,
extend='both')
edge_loc.plot(color='k', ax=ax_in)
psi_max_loc.plot(color='k', ax=ax_in)
ax_in.set_ylabel('Latitude')
ax_in.set_ylim(-60, 60)
ax_in.set_yticks(np.arange(-60., 61., 30.))
ax_in.grid(True, linestyle=':')
if month_labels:
mn_dic = month_dic(1)
tickspace = list(range(13, 72, 18))
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
ax_in.set_xlim((1, 72))
ax_in.set_xticks(tickspace)
if do_xlabels:
ax_in.set_xlabel('')
ax_in.set_xticklabels(labels, rotation=25)
return f1
def isca_monsoon_hm(run):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
mn_dic = month_dic(1)
tickspace = [1, 19, 37, 55]
labels = ['1st Jan', '1st Apr', '1st Jul', '1st Oct']
# Set figure parameters
rcParams['figure.figsize'] = 15, 6
rcParams['font.size'] = 14
# Start figure with 8 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7,
ax8)) = plt.subplots(2,
4,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8]
def precip_plot(data, ax, lons, title, levels=np.arange(0., 21., 2.)):
lons = pick_lons(data, lons)
f1 = (data.precipitation * 86400.).sel(
lon=lons).mean('lon').plot.contourf(ax=ax,
x='xofyear',
y='lat',
levels=levels,
add_labels=False,
extend='max',
cmap='Blues',
add_colorbar=False)
ax.set_ylim(-45, 45)
ax.grid(True, linestyle=':')
ax.set_yticks(np.arange(-30., 31., 30.))
ax.set_title(title)
return f1
f1 = precip_plot(data, ax1, [0., 10.], '0-10 E')
precip_plot(data, ax2, [80., 100.], '80-100 E')
#precip_plot(data, ax3, [125.,145.], '125-145 E')
precip_plot(data, ax3, [170., 180.], '170-180 E')
precip_plot(data, ax4, [180., 190.], '180-190 E')
#precip_plot(data, ax5, [215.,235.], '215-235 E')
precip_plot(data, ax5, [260., 280.], '260-280 E')
#precip_plot(data, ax7, [305.,325.], '305-325 E')
precip_plot(data, ax6, [350., 0.], '350-0 E')
ax1.set_ylabel('Latitude')
ax5.set_ylabel('Latitude')
for ax in axes:
ax.set_xticks(tickspace)
ax.set_xticklabels(labels, rotation=25)
plt.subplots_adjust(left=0.07,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.2)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.15,
aspect=60,
shrink=0.5)
cb1.set_label('Precipitation, mm/day')
# Save as a pdf
plt.savefig(plot_dir + 'precip_monsoons_' + run + '_diffsplit.pdf',
format='pdf')
plt.close()
data.close()
def precip_mse_plot(data,
ax_in,
lonin=[-1., 361.],
do_xlabels=False,
plot_type=None,
precip_contour=8.,
p_cent=True,
mse_max=True,
month_labels=True,
lat_bound=45.):
lons = pick_lons(data, lonin)
try:
precip_plot = (data.precipitation * 86400.).sel(lon=lons).mean('lon')
except:
precip_plot = ((data.convection_rain + data.condensation_rain) *
86400.).sel(lon=lons).mean('lon')
mse_plot = (data.temp * cp + data.sphum * L +
data.height * g).mean('lon') / 1000.
if plot_type == None:
# Default case, plot precip with mse overlaid. Choice of whether or not to highlight a specific precip contour
f1 = precip_plot.plot.contourf(ax=ax_in,
x='xofyear',
y='lat',
levels=np.arange(2., 15., 2.),
add_colorbar=False,
add_labels=False,
extend='max',
cmap='Blues')
if not precip_contour == None:
precip_plot.plot.contour(ax=ax_in,
x='xofyear',
y='lat',
levels=np.arange(precip_contour - 100.,
200., 100.),
add_labels=False,
add_colorbar=False,
colors='k',
linewidth=2)
cs = mse_plot.sel(pfull=850.).plot.contour(ax=ax_in,
x='xofyear',
y='lat',
levels=np.arange(
200., 401., 10.),
add_labels=False,
colors='0.7',
add_colorbar=False,
linewidths=2)
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
if p_cent:
data = precip_centroid(data, lonin=lonin, lat_bound=lat_bound)
data.p_cent.plot.line(color='w', ax=ax_in)
ax_in.set_xlabel('')
elif plot_type == 'precip':
# No mse, plot precip and precip centroid
f1 = precip_plot.plot.contourf(ax=ax_in,
x='xofyear',
y='lat',
levels=np.arange(3., 16., 3.),
add_colorbar=False,
add_labels=False,
extend='max',
cmap='Blues',
linewidth=2)
if p_cent:
data = precip_centroid(data, lonin=lonin, lat_bound=lat_bound)
data.p_cent.plot.line(color='k', ax=ax_in, linewidth=2)
ax_in.set_xlabel('')
elif plot_type == 'mse':
# Plot mse in colour, overplot max mse
f1 = mse_plot.sel(pfull=850.).plot.contourf(ax=ax_in,
x='xofyear',
y='lat',
levels=np.arange(
200., 401., 10.),
add_labels=False,
extend='both',
add_colorbar=False)
if mse_max:
data_mse = peak_mse(data, lonin=lonin)
data_mse.mse_max_loc.plot.line('k', ax=ax_in)
ax_in.set_xlabel('')
ax_in.set_ylabel('Latitude')
ax_in.set_ylim(-60, 60)
ax_in.set_yticks(np.arange(-60., 61., 30.))
ax_in.grid(True, linestyle=':')
if month_labels:
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
ax_in.set_xlim((1, 72))
ax_in.set_xticks(tickspace)
if do_xlabels:
ax_in.set_xlabel('')
ax_in.set_xticklabels(labels, rotation=25)
return f1
def vort_eq_hm(run, lev=150., lonin=[-1.,361.], planetary_only=False, month_labels=True, rot_fac=1., no_eddies=False, add_psi=True):
'''Plot vorticity budget Hovmoller. RG 3/11/2017
Imputs: run = run_name
lev = level to plot
lonin = longitude range to average over
planetary_only = only plot planetary vorticity terms
month_labels = label x axis with month of year (no good for non earth year length)
rot_fac = scale factor for Earth's rotation rate
no_eddies = don't plot transient eddies too'''
# Call the above function to get fields to plot. For a Hovmoller, run is by definition not steady state. Pass other inputs on.
ds = vort_budg_terms(run, lonin=lonin, rot_fac=rot_fac, planetary_only=planetary_only, no_eddies=no_eddies)
if add_psi:
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]]
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons)
psi /= 1.e9
# Determine figure size based on number of panels
if planetary_only or no_eddies:
rcParams['figure.figsize'] = 15, 6.25
else:
rcParams['figure.figsize'] = 15, 8
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
levels = np.arange(-1.5,1.6,0.25)
print('starting plotting')
# Set number of panels and declare which are the left column, bottom row.
if planetary_only or no_eddies:
f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row')
left_column = [ax1,ax4]
bottom_row = [ax4,ax5,ax6]
all_plots = [ax1,ax2,ax3,ax4,ax5,ax6]
else:
f, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row')
left_column = [ax1,ax4,ax7]
bottom_row = [ax7,ax8,ax9]
all_plots = [ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9]
plt.set_cmap('RdBu_r')
# Select the pressure level, if needed
if run in ['ap_2','full_qflux']:
ds = ds
else:
ds = ds.sel(pfull=lev)
# Plot!
f2 = ds.horiz_md.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax1, extend = 'both', add_labels=False)
if add_psi:
ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax1.set_title('Horizontal advection', fontsize=17)
ax1.text(-15, 60, 'a)')
ds.dvordy.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)
if planetary_only:
ax2.set_title('$\partial f /\partial y$', fontsize=17)
else:
ax2.set_title('$\partial (\overline{\zeta} + f) /\partial y$', fontsize=17)
ax2.text(-7, 60, 'b)')
ds.v.plot.contourf(x='xofyear', y='lat', levels=np.arange(-12.,13,2.), ax=ax3, extend = 'both', add_labels=False)
ax3.set_title('$\overline{v}$', fontsize=17)
ax3.text(-7, 60, 'c)')
ds.stretching.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax4, extend = 'both', add_labels=False)
if add_psi:
ax4.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax4.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax4.set_title('Vortex stretching', fontsize=17)
ax4.text(-15, 60, 'd)')
ds.div.plot.contourf(x='xofyear', y='lat', levels=np.arange(-0.6,0.7,0.1), ax=ax5, extend = 'both', add_labels=False)
ax5.set_title('Divergence', fontsize=17)
ax5.text(-7, 60, 'e)')
ds.vor.plot.contourf(x='xofyear', y='lat', levels=np.arange(-12.,13.,2.), ax=ax6, extend = 'both', add_labels=False)
if planetary_only:
ax6.set_title('Planetary vorticity', fontsize=17)
else:
ax6.set_title('Absolute vorticity', fontsize=17)
ax6.text(-7, 60, 'f)')
# Plot transients if needed
if not (planetary_only or no_eddies):
ds.transient_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax7, extend = 'both', add_labels=False)
if add_psi:
ax7.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax7.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax7.set_title('Transient eddy vorticity tendency', fontsize=17)
ax7.text(-15, 60, 'g)')
ds.transient_h_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax8, extend = 'both', add_labels=False)
ax8.set_title('Transient horizontal adv.', fontsize=17)
ax8.text(-7, 60, 'h)')
ds.transient_s_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax9, extend = 'both', add_labels=False)
ax9.set_title('Transient vortex stretching', fontsize=17)
ax9.text(-7, 60, 'i)')
# Add grid, set y limits and ticks
for ax in all_plots:
ax.grid(True,linestyle=':')
ax.set_ylim(-60,60)
ax.set_yticks(np.arange(-60,61,30))
# Label left column
for ax in left_column:
ax.set_ylabel('Latitude')
# If month labels are being used, load in the month dictionary and set up labelling, then label bottom row
if month_labels:
mn_dic = month_dic(1)
tickspace = list(range(13,72,18))
labels = [mn_dic[(k+5)/6 ] for k in tickspace]
for ax in bottom_row:
ax.set_xticks(tickspace)
ax.set_xticklabels(labels,rotation=25)
plt.subplots_adjust(right=0.97, left=0.08, top=0.93, bottom=0.1, hspace=0.25, wspace=0.15)
# Set plot name based on type of plot.
if lonin == [-1.,361.]:
lon_tag = ''
else:
lon_tag = '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1]))
if planetary_only:
f_tag = '_fonly'
else:
f_tag = ''
figname = 'vort_breakdown_' + run + lon_tag + f_tag + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def vort_eq_hm(run, lev=150., lonin=[-1.,361.], month_labels=True, rot_fac=1., add_psi=True):
'''Plot vorticity budget Hovmoller. RG 3/11/2017
Imputs: run = run_name
lev = level to plot
lonin = longitude range to average over
planetary_only = only plot planetary vorticity terms
month_labels = label x axis with month of year (no good for non earth year length)
rot_fac = scale factor for Earth's rotation rate
no_eddies = don't plot transient eddies too'''
# Call the above function to get fields to plot. For a Hovmoller, run is by definition not steady state. Pass other inputs on.
ds = vort_budg_terms(run, lonin=lonin, rot_fac=rot_fac)
if add_psi:
data = xr.open_dataset('/disca/share/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]]
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons)
psi /= 1.e9
# Determine figure size based on number of panels
rcParams['figure.figsize'] = 15, 4
rcParams['font.size'] = 14
rcParams['text.usetex'] = True
levels = np.arange(-1.5,1.6,0.25)
print('starting plotting')
# Set number of panels and declare which are the left column, bottom row.
f, (ax1, ax2, ax3) = plt.subplots(1, 3, sharey='row')
left_column = [ax1]
all_plots = [ax1,ax2,ax3]
plt.set_cmap('RdBu_r')
# Select the pressure level, if needed
if run in ['ap_2','full_qflux']:
ds = ds
else:
ds = ds.sel(pfull=lev)
# Plot!
f2 = ds.div_uvor.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax1, extend = 'both', add_labels=False)
if add_psi:
ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax1.set_title('Relative vorticty part', fontsize=17)
ax1.text(-15, 60, 'a)')
ds.div_uf.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax2, extend = 'both', add_labels=False)
if add_psi:
ax2.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax2.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax2.set_title('Planetary vorticity part', fontsize=17)
ax2.text(-7, 60, 'b)')
ds.transient_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax3, extend = 'both', add_labels=False)
if add_psi:
ax3.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax3.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax3.set_title('Transient eddy vorticity tendency', fontsize=17)
ax3.text(-7, 60, 'c)')
# Add grid, set y limits and ticks
for ax in all_plots:
ax.grid(True,linestyle=':')
ax.set_ylim(-60,60)
ax.set_yticks(np.arange(-60,61,30))
if month_labels:
mn_dic = month_dic(1)
tickspace = list(range(13,72,18))
labels = [mn_dic[(k+5)/6 ] for k in tickspace]
ax.set_xticks(tickspace)
ax.set_xticklabels(labels,rotation=25)
# Label left column
for ax in left_column:
ax.set_ylabel('Latitude')
plt.subplots_adjust(right=0.97, left=0.08, top=0.93, bottom=0.1, hspace=0.25, wspace=0.15)
# Set plot name based on type of plot.
if lonin == [-1.,361.]:
lon_tag = ''
else:
lon_tag = '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1]))
figname = 'alt_breakdown_' + run + lon_tag + '.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=lev)
#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 = list(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 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('/disca/share/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
levels = np.arange(-300.,300.1,30.)
mn_dic = month_dic(1)
tickspace = list(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=fu.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
dphidy.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
metric_term.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('Metric term', 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 = 'merid_mom_budg_' +run+ '.pdf'
else:
figname = 'merid_mom_budg_' + run + '_' + 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.]):
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/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_mean = fv.mean('lon')
fv_local = (fv - fv_mean).sel(lon=lons).mean('lon')
try:
totp = ((data.precipitation)*86400.).sel(lon=lons).mean('lon')
except:
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')
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
#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='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('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='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
mom_sum.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('Residual', 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_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_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='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_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='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
totp.plot.contour(ax=ax7, 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)
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='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
totp.plot.contour(ax=ax8, 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)
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='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
totp.plot.contour(ax=ax9, 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)
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_' +run+ '_rev.pdf'
else:
figname = 'zon_mom_budg_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '_rev.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
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.]) #'Asia'
#lons = pick_lons(data, [15.,60.]) # South Africa
lons = pick_lons(data, [100., 150.]) # East Asia
lons = pick_lons(data, [70., 100.]) # India
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 heat_budg_hm(run, lev=850, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.], plot_precip=True, do_theta=False):
rcParams['figure.figsize'] = 12, 8
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/heat_budg/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/disca/share/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=lev)
try:
precip_centroid(data)
except:
data['precipitation'] = data.convection_rain + data.condensation_rain
precip_centroid(data)
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
heat_mean = data.u_dTdx_zav #+ data.w_dTdp_zav #+data.v_dTdy_zav #
heat_cross = data.u_dTdx_cross1 + data.v_dTdy_cross1 + data.w_dTdp_cross1 + data.u_dTdx_cross2 + data.v_dTdy_cross2 + data.w_dTdp_cross2
heat_cross1 = data.u_dTdx_cross1 + data.v_dTdy_cross1 + data.w_dTdp_cross1
heat_cross2 = data.u_dTdx_cross2 + data.v_dTdy_cross2 + data.w_dTdp_cross2
heat_trans = data.uT_trans_dx + data.vT_trans_dy + data.wT_trans_dp
heat_stat = data.u_dTdx_stat + data.v_dTdy_stat + data.w_dTdp_stat
heat_crossu = data.u_dTdx_cross1 + data.u_dTdx_cross2
heat_crossv = data.v_dTdy_cross1 + data.v_dTdy_cross2
heat_crossw = data.w_dTdp_cross1 + data.w_dTdp_cross2
tdt_rad = data.tdt_rad.sel(pfull=lev).sel(lon=lons).mean('lon')*86400.
tdt_conv = data.dt_tg_convection.sel(pfull=lev).sel(lon=lons).mean('lon')*86400.
tdt_cond = data.dt_tg_condensation.sel(pfull=lev).sel(lon=lons).mean('lon')*86400.
tdt_diff = data.dt_tg_diffusion.sel(pfull=lev).sel(lon=lons).mean('lon')*86400.
if do_theta:
convTtotheta=(1000./data.pfull)**(2./7.)
diabatic = (tdt_rad + tdt_conv + tdt_cond + tdt_diff)*convTtotheta.sel(pfull=lev)
heat_sum = heat_mean + heat_trans + heat_stat + heat_cross + diabatic
asc_cool = heat_sum * 0.
Tdt = gr.ddt(data.temp.sel(pfull=lev)).sel(lon=lons).mean('lon') *86400.
else:
diabatic = tdt_rad + tdt_conv + tdt_cond + tdt_diff
rho = data.pfull*100./data.temp/287.04
asc_cool = (data.omega /(1004.64 * rho)).sel(pfull=lev).sel(lon=lons).mean('lon') *86400.
heat_sum = heat_mean + heat_trans + heat_stat + heat_cross + diabatic + asc_cool
Tdt = gr.ddt(data.temp.sel(pfull=lev)).sel(lon=lons).mean('lon') *86400.
#levels = np.arange(-20,21.1,2.)
#levels = np.arange(-10,10.1,1.)
levels = np.arange(-1,1.1,0.1)
mn_dic = month_dic(1)
tickspace = list(range(13,72,18))
ticklabels = [mn_dic[(k+5)/6 ] for k in tickspace]
# Nine subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row')
plt.set_cmap('RdBu_r')
plot_vars = [heat_mean, heat_trans, heat_sum, asc_cool, diabatic,
tdt_rad, tdt_diff, tdt_conv, tdt_cond]
axes = [ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9]
labels = ['a)','b)','c)','d)','e)','f)','g)','h)','i)']
for i in range(9):
f1=plot_vars[i].plot.contourf(ax=axes[i], x='xofyear', y='lat', extend='both', add_labels=False, cmap='RdBu_r', levels = levels, add_colorbar=False)
if plot_precip:
psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(-500.,0.,100.), add_labels=False, colors='0.7', linewidths=2, linestyles='--')
psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(0.,510.,100.), add_labels=False, colors='0.7', linewidths=2)
psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(-1000.,1010.,1000.), add_labels=False, colors='0.5', linewidths=2)
data.p_cent.plot.line(ax=axes[i],color='k', linewidth=2)
axes[i].set_xlabel(' ')
axes[i].set_ylabel(' ')
#totp.plot.contour(ax=axes[i], 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)
axes[i].set_ylim(-60,60)
axes[i].grid(True,linestyle=':')
axes[i].set_yticks(np.arange(-60.,61.,30.))
axes[i].set_xticks(tickspace)
for i in [0,3,6]:
axes[i].set_ylabel('Latitude')
axes[i].text(-15, 60, labels[i])
for i in [1,2,4,5,7,8]:
axes[i].text(-5, 60, labels[i])
for i in [6,7,8]:
axes[i].set_xticklabels(ticklabels,rotation=25)
# set titles
ax1.set_title('Mean state advection', fontsize=17)
ax2.set_title('Transient eddy flux conv.', fontsize=17)
ax3.set_title('Residual', fontsize=17)
ax4.set_title('w cooling', fontsize=17)
ax5.set_title('Total diabatic heating', fontsize=17)
ax6.set_title('Radiative heating', fontsize=17)
ax7.set_title('Diffusive heating', fontsize=17)
ax8.set_title('Convective heating', fontsize=17)
ax9.set_title('Condensational heating', fontsize=17)
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=axes, use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.15, aspect=30, shrink=0.5)
if lonin == [-1.,361.]:
figname = 'zon_heat_budg_' +run+ '_u_850.pdf'
else:
figname = 'zon_heat_budg_' + run + '_' + 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.], plot_precip=True, rot_fac=1.):
rcParams['figure.figsize'] = 12, 6
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/mom_budg_merid/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/disca/share/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 * rot_fac
f = 2 * omega * np.sin(data.lat *np.pi/180)
fu = -1. * data.ucomp.sel(pfull=lev) * f * 86400.
fu_mean = fu.mean('lon')
fu_local = (fu - fu_mean).sel(lon=lons).mean('lon')
if plot_precip:
try:
totp = ((data.precipitation)*86400.).sel(lon=lons).mean('lon')
except:
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')
fu_ageo = fu_local + fu_mean + dphidy
dvdt = gr.ddt(data.vcomp.sel(pfull=lev)).sel(lon=lons).mean('lon')*86400.
metric = (-86400.* data.ucomp_sq * np.tan(data.lat * np.pi/180.)/6376.0e3).sel(pfull=lev).sel(lon=lons).mean('lon')
vert_term = (-86400.*data.vcomp_omega/6376.0e3).sel(pfull=lev).sel(lon=lons).mean('lon')
mom_mean = data.u_dvdx_zav + data.v_dvdy_zav + data.w_dvdp_zav
mom_cross = data.u_dvdx_cross1 + data.v_dvdy_cross1 + data.w_dvdp_cross1 + data.u_dvdx_cross2 + data.v_dvdy_cross2 + data.w_dvdp_cross2
mom_cross1 = data.u_dvdx_cross1 + data.v_dvdy_cross1 + data.w_dvdp_cross1
mom_cross2 = data.u_dvdx_cross2 + data.v_dvdy_cross2 + data.w_dvdp_cross2
mom_trans = data.uv_trans_dx + data.vv_trans_dy + data.vw_trans_dp
mom_stat = data.u_dvdx_stat + data.v_dvdy_stat + data.w_dvdp_stat
mom_crossu = data.u_dvdx_cross1 + data.u_dvdx_cross2
mom_crossv = data.v_dvdy_cross1 + data.v_dvdy_cross2
mom_crossw = data.w_dvdp_cross1 + data.w_dvdp_cross2
mom_sum = fu_local + fu_mean + dphidy + mom_mean + mom_trans + mom_stat + mom_cross + metric
#levels = np.arange(-40,41.1,4.)
levels = np.arange(-300.,300.1,30.)
mn_dic = month_dic(1)
tickspace = list(range(13,72,18))
ticklabels = [mn_dic[(k+5)/6 ] for k in tickspace]
# Nine subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row')
plt.set_cmap('RdBu_r')
plot_vars = [fu_mean, mom_mean, mom_sum, mom_trans, dphidy, metric]
axes = [ax1,ax2,ax3,ax4,ax5,ax6]
labels = ['a)','b)','c)','d)','e)','f)']
for i in range(len(plot_vars)):
f1=plot_vars[i].plot.contourf(ax=axes[i], x='xofyear', y='lat', extend='both', levels = levels, add_colorbar=False, add_labels=False)
if plot_precip:
totp.plot.contour(ax=axes[i], 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)
axes[i].set_ylim(-60,60)
axes[i].grid(True,linestyle=':')
axes[i].set_yticks(np.arange(-60.,61.,30.))
axes[i].set_xticks(tickspace)
for i in [0,3]:
axes[i].set_ylabel('Latitude')
axes[i].text(-15, 60, labels[i])
for i in [1,2,4,5]:
axes[i].text(-5, 60, labels[i])
for i in [3,4,5]:
axes[i].set_xticklabels(ticklabels,rotation=25)
# set titles
ax1.set_title('Zonal mean Coriolis', fontsize=17)
ax2.set_title('Mean state advection', fontsize=17)
ax3.set_title('Residual', fontsize=17)
ax4.set_title('Transient eddy flux conv.', fontsize=17)
ax5.set_title('Geopotential gradient', fontsize=17)
ax6.set_title('Metric term', fontsize=17)
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=axes, use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.15, aspect=30, shrink=0.5)
if lonin == [-1.,361.]:
figname = 'merid_mom_budg_' +run+ '.pdf'
else:
figname = 'merid_mom_budg_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()